select by city

1 files changed, 84 insertions, 24 deletions

diff --git a/pkg/vpn/bonafide/gateways.go b/pkg/vpn/bonafide/gateways.go
index 8bfcee0..f1694a2 100644
--- a/pkg/vpn/bonafide/gateways.go
+++ b/pkg/vpn/bonafide/gateways.go
@@ -25,24 +25,40 @@ type Gateway struct {
 	Protocols    []string
 	Options      map[string]string
 	Transport    string
-	Label        string
 }
 
-/* TODO add a String method with a human representation: Label (cc) */
-/* For that, we should pass the locations to genLabels, and generate a string repr */
+// Load reflects the fullness metric that menshen returns, if available.
+type Load struct {
+	Host     string
+	Fullness string
+}
 
 type gatewayDistance struct {
 	gateway  Gateway
 	distance int
 }
 
+/* a map between locations and hostnames, to be able to select by city */
+type cityMap struct {
+	gws map[string][]string
+}
+
+func (g *cityMap) Get(key string) []string {
+	if val, ok := g.gws[key]; ok {
+		return val
+	}
+	return make([]string, 0)
+}
+
 type gatewayPool struct {
 	available  []Gateway
 	userChoice []byte
-	/* ranked is, for now, just an array of hostnames (fetched from the
-	menshen service). it should be a map in the future, to keep track of
-	quantitative metrics */
-	ranked []string
+	byCity     cityMap
+
+	/* recommended is an array of hostnames, fetched from the old geoip service.
+	 *  this should be deprecated in favor of recommendedWithLoad when new menshen is deployed */
+	recommended         []string
+	recommendedWithLoad []Load
 
 	/* TODO locations are just used to get the timezone for each gateway. I
 	* think it's easier to just merge that info into the version-agnostic
@@ -51,11 +67,52 @@ type gatewayPool struct {
 	locations map[string]Location
 }
 
+func (p *gatewayPool) populateCityList() {
+	for _, gw := range p.available {
+		loc := gw.Location
+		gws := p.cityMap.Get(loc)
+		p.cityMap[loc] = append(gws, gw.Host)
+	}
+}
+
+func (p *gatewayPool) getCities() []string {
+	c := make([]string, 0)
+	for city, _ := range p.byCity {
+		c = append(c, city)
+	}
+	return c
+}
+
+func (p *gatewayPool) isValidCity(city string) bool {
+	cities := p.getCities()
+	valid := stringInSlice(city, cities)
+	return valid
+}
+
+/* this method should only be used if we have no usable menshen list */
+func (p *gatewayPool) getRandomGatewayByCity(city string) (Gateway, error) {
+	if !p.isValidCity(city) {
+		return Gateway{}, errors.New("bonafide: BUG not a valid city: " + city)
+	}
+	gws := p.byCity.Get(city)
+	if len(gws) == 0 {
+		return Gateway{}, errors.New("bonafide: BUG no gw for city " + city)
+	}
+	s := rand.NewSource(time.Now().Unix())
+	r := rand.New(s)
+	host := gws[r.Intn(len(gws))]
+	for _, gw := range p.available {
+		if gw.Host == host {
+			return gw, nil
+		}
+	}
+	return Gateway{}, errors.New("bonafide: BUG should not reach here")
+}
+
 /* genLabels generates unique, human-readable labels for a gateway. It gives a serial
    number to each gateway in the same location (paris-1, paris-2,...). The
    current implementation will give a different label to each transport.
    An alternative (to discuss) would be to give the same label to the same hostname.
-*/
 func (p *gatewayPool) genLabels() {
 	acc := make(map[string]int)
 	for i, gw := range p.available {
@@ -67,7 +124,6 @@ func (p *gatewayPool) genLabels() {
 		gw.Label = gw.Location + "-" + strconv.Itoa(acc[gw.Location])
 		p.available[i] = gw
 	}
-	/* skip suffix if only one occurrence */
 	for i, gw := range p.available {
 		if acc[gw.Location] == 1 {
 			gw.Label = gw.Location
@@ -75,13 +131,14 @@ func (p *gatewayPool) genLabels() {
 		}
 	}
 }
+*/
 
+/*
 func (p *gatewayPool) getLabels() []string {
 	labels := make([]string, 0)
 	for _, gw := range p.available {
 		labels = append(labels, gw.Label)
 	}
-	/* TODO return error if called when no labels have been generated */
 	return labels
 }
 
@@ -99,6 +156,7 @@ func (p *gatewayPool) getGatewayByLabel(label string) (Gateway, error) {
 	}
 	return Gateway{}, errors.New("bonafide: not a valid label")
 }
+*/
 
 func (p *gatewayPool) getGatewayByIP(ip string) (Gateway, error) {
 	for _, gw := range p.available {
@@ -113,11 +171,11 @@ func (p *gatewayPool) setAutomaticChoice() {
 	p.userChoice = []byte("")
 }
 
-func (p *gatewayPool) setUserChoice(label []byte) error {
-	if !p.isValidLabel(string(label)) {
-		return errors.New("bonafide: not a valid label for gateway choice")
+func (p *gatewayPool) setUserChoice(city []byte) error {
+	if !p.isValidCity(string(city)) {
+		return errors.New("bonafide: not a valid city for gateway choice")
 	}
-	p.userChoice = label
+	p.userChoice = city
 	return nil
 }
 
@@ -129,38 +187,40 @@ func (p *gatewayPool) setRanking(hostnames []string) {
 
 	for _, host := range hostnames {
 		if !stringInSlice(host, hosts) {
-			log.Println("ERROR: invalid host in ranked hostnames", host)
+			log.Println("ERROR: invalid host in recommended list of hostnames", host)
 			return
 		}
 	}
 
-	p.ranked = hostnames
+	p.recommended = hostnames
 }
 
 func (p *gatewayPool) getBest(transport string, tz, max int) ([]Gateway, error) {
 	gws := make([]Gateway, 0)
 	if len(p.userChoice) != 0 {
-		gw, err := p.getGatewayByLabel(string(p.userChoice))
+		/* FIXME this is random because we still do not get menshen to return us load. after "new" menshen is deployed,
+		   we can just get them by the order menshen reeturned */
+		gw, err := p.getRandomGatewayByCity(string(p.userChoice))
 		gws = append(gws, gw)
 		return gws, err
-	} else if len(p.ranked) != 0 {
-		return p.getGatewaysByServiceRank(transport, max)
+	} else if len(p.recommended) != 0 {
+		return p.getGatewaysFromMenshen(transport, max)
 	} else {
 		return p.getGatewaysByTimezone(transport, tz, max)
 	}
 }
 
 func (p *gatewayPool) getAll(transport string, tz int) ([]Gateway, error) {
-	if len(p.ranked) != 0 {
-		return p.getGatewaysByServiceRank(transport, 999)
+	if len(p.recommended) != 0 {
+		return p.getGatewaysFromMenshen(transport, 999)
 	} else {
 		return p.getGatewaysByTimezone(transport, tz, 999)
 	}
 }
 
-func (p *gatewayPool) getGatewaysByServiceRank(transport string, max int) ([]Gateway, error) {
+func (p *gatewayPool) getGatewaysFromMenshen(transport string, max int) ([]Gateway, error) {
 	gws := make([]Gateway, 0)
-	for _, host := range p.ranked {
+	for _, host := range p.recommended {
 		for _, gw := range p.available {
 			if gw.Transport != transport {
 				continue
@@ -217,7 +277,7 @@ func newGatewayPool(eip *eipService) *gatewayPool {
 	p := gatewayPool{}
 	p.available = eip.getGateways()
 	p.locations = eip.Locations
-	p.genLabels()
+	p.populateCityList()
 	return &p
 }