Make easier to extend the list of known proxy types

And include iptproxy as a valid proxy type.

4 files changed, 55 insertions, 62 deletions

diff --git a/broker/ipc.go b/broker/ipc.go
index 768c0b7..9b47b90 100644
--- a/broker/ipc.go
+++ b/broker/ipc.go
@@ -25,18 +25,15 @@ type IPC struct {
 }
 
 func (i *IPC) Debug(_ interface{}, response *string) error {
-	var webexts, browsers, standalones, unknowns int
+	var unknowns int
 	var natRestricted, natUnrestricted, natUnknown int
+	proxyTypes := make(map[string]int)
 
 	i.ctx.snowflakeLock.Lock()
 	s := fmt.Sprintf("current snowflakes available: %d\n", len(i.ctx.idToSnowflake))
 	for _, snowflake := range i.ctx.idToSnowflake {
-		if snowflake.proxyType == "badge" {
-			browsers++
-		} else if snowflake.proxyType == "webext" {
-			webexts++
-		} else if snowflake.proxyType == "standalone" {
-			standalones++
+		if messages.KnownProxyTypes[snowflake.proxyType] {
+			proxyTypes[snowflake.proxyType]++
 		} else {
 			unknowns++
 		}
@@ -53,10 +50,10 @@ func (i *IPC) Debug(_ interface{}, response *string) error {
 	}
 	i.ctx.snowflakeLock.Unlock()
 
-	s += fmt.Sprintf("\tstandalone proxies: %d", standalones)
-	s += fmt.Sprintf("\n\tbrowser proxies: %d", browsers)
-	s += fmt.Sprintf("\n\twebext proxies: %d", webexts)
-	s += fmt.Sprintf("\n\tunknown proxies: %d", unknowns)
+	for pType, num := range proxyTypes {
+		s += fmt.Sprintf("\t%s proxies: %d\n", pType, num)
+	}
+	s += fmt.Sprintf("\tunknown proxies: %d", unknowns)
 
 	s += fmt.Sprintf("\nNAT Types available:")
 	s += fmt.Sprintf("\n\trestricted: %d", natRestricted)
diff --git a/broker/metrics.go b/broker/metrics.go
index 8229e0f..c642045 100644
--- a/broker/metrics.go
+++ b/broker/metrics.go
@@ -14,6 +14,7 @@ import (
 	"sync"
 	"time"
 
+	"git.torproject.org/pluggable-transports/snowflake.git/v2/common/messages"
 	"github.com/prometheus/client_golang/prometheus"
 	"gitlab.torproject.org/tpo/anti-censorship/geoip"
 )
@@ -24,10 +25,9 @@ const (
 )
 
 type CountryStats struct {
-	standalone map[string]bool
-	badge      map[string]bool
-	webext     map[string]bool
-	unknown    map[string]bool
+	// map[proxyType][address]bool
+	proxies map[string]map[string]bool
+	unknown map[string]bool
 
 	natRestricted   map[string]bool
 	natUnrestricted map[string]bool
@@ -96,22 +96,17 @@ func (m *Metrics) UpdateCountryStats(addr string, proxyType string, natType stri
 	var country string
 	var ok bool
 
-	if proxyType == "standalone" {
-		if m.countryStats.standalone[addr] {
-			return
-		}
-	} else if proxyType == "badge" {
-		if m.countryStats.badge[addr] {
-			return
-		}
-	} else if proxyType == "webext" {
-		if m.countryStats.webext[addr] {
+	addresses, ok := m.countryStats.proxies[proxyType]
+	if !ok {
+		if m.countryStats.unknown[addr] {
 			return
 		}
+		m.countryStats.unknown[addr] = true
 	} else {
-		if m.countryStats.unknown[addr] {
+		if addresses[addr] {
 			return
 		}
+		addresses[addr] = true
 	}
 
 	ip := net.ParseIP(addr)
@@ -122,18 +117,7 @@ func (m *Metrics) UpdateCountryStats(addr string, proxyType string, natType stri
 	if !ok {
 		country = "??"
 	}
-
-	//update map of unique ips and counts
 	m.countryStats.counts[country]++
-	if proxyType == "standalone" {
-		m.countryStats.standalone[addr] = true
-	} else if proxyType == "badge" {
-		m.countryStats.badge[addr] = true
-	} else if proxyType == "webext" {
-		m.countryStats.webext[addr] = true
-	} else {
-		m.countryStats.unknown[addr] = true
-	}
 
 	m.promMetrics.ProxyTotal.With(prometheus.Labels{
 		"nat":  natType,
@@ -166,14 +150,15 @@ func NewMetrics(metricsLogger *log.Logger) (*Metrics, error) {
 
 	m.countryStats = CountryStats{
 		counts:          make(map[string]int),
-		standalone:      make(map[string]bool),
-		badge:           make(map[string]bool),
-		webext:          make(map[string]bool),
+		proxies:         make(map[string]map[string]bool),
 		unknown:         make(map[string]bool),
 		natRestricted:   make(map[string]bool),
 		natUnrestricted: make(map[string]bool),
 		natUnknown:      make(map[string]bool),
 	}
+	for pType := range messages.KnownProxyTypes {
+		m.countryStats.proxies[pType] = make(map[string]bool)
+	}
 
 	m.logger = metricsLogger
 	m.promMetrics = initPrometheus()
@@ -197,11 +182,12 @@ func (m *Metrics) printMetrics() {
 	m.lock.Lock()
 	m.logger.Println("snowflake-stats-end", time.Now().UTC().Format("2006-01-02 15:04:05"), fmt.Sprintf("(%d s)", int(metricsResolution.Seconds())))
 	m.logger.Println("snowflake-ips", m.countryStats.Display())
-	m.logger.Println("snowflake-ips-total", len(m.countryStats.standalone)+
-		len(m.countryStats.badge)+len(m.countryStats.webext)+len(m.countryStats.unknown))
-	m.logger.Println("snowflake-ips-standalone", len(m.countryStats.standalone))
-	m.logger.Println("snowflake-ips-badge", len(m.countryStats.badge))
-	m.logger.Println("snowflake-ips-webext", len(m.countryStats.webext))
+	total := len(m.countryStats.unknown)
+	for pType, addresses := range m.countryStats.proxies {
+		m.logger.Printf("snowflake-ips-%s %d\n", pType, len(addresses))
+		total += len(addresses)
+	}
+	m.logger.Println("snowflake-ips-total", total)
 	m.logger.Println("snowflake-idle-count", binCount(m.proxyIdleCount))
 	m.logger.Println("client-denied-count", binCount(m.clientDeniedCount))
 	m.logger.Println("client-restricted-denied-count", binCount(m.clientRestrictedDeniedCount))
@@ -221,9 +207,9 @@ func (m *Metrics) zeroMetrics() {
 	m.clientUnrestrictedDeniedCount = 0
 	m.clientProxyMatchCount = 0
 	m.countryStats.counts = make(map[string]int)
-	m.countryStats.standalone = make(map[string]bool)
-	m.countryStats.badge = make(map[string]bool)
-	m.countryStats.webext = make(map[string]bool)
+	for pType := range m.countryStats.proxies {
+		m.countryStats.proxies[pType] = make(map[string]bool)
+	}
 	m.countryStats.unknown = make(map[string]bool)
 	m.countryStats.natRestricted = make(map[string]bool)
 	m.countryStats.natUnrestricted = make(map[string]bool)
diff --git a/broker/snowflake-broker_test.go b/broker/snowflake-broker_test.go
index 9c975eb..f7850f8 100644
--- a/broker/snowflake-broker_test.go
+++ b/broker/snowflake-broker_test.go
@@ -549,8 +549,13 @@ func TestMetrics(t *testing.T) {
 			p.offerChannel <- nil
 			<-done
 			ctx.metrics.printMetrics()
-			So(buf.String(), ShouldResemble, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips CA=4\nsnowflake-ips-total 4\nsnowflake-ips-standalone 1\nsnowflake-ips-badge 1\nsnowflake-ips-webext 1\nsnowflake-idle-count 8\nclient-denied-count 0\nclient-restricted-denied-count 0\nclient-unrestricted-denied-count 0\nclient-snowflake-match-count 0\nsnowflake-ips-nat-restricted 0\nsnowflake-ips-nat-unrestricted 0\nsnowflake-ips-nat-unknown 1\n")
 
+			metricsStr := buf.String()
+			So(metricsStr, ShouldStartWith, "snowflake-stats-end "+time.Now().UTC().Format("2006-01-02 15:04:05")+" (86400 s)\nsnowflake-ips CA=4\n")
+			So(metricsStr, ShouldContainSubstring, "\nsnowflake-ips-standalone 1\n")
+			So(metricsStr, ShouldContainSubstring, "\nsnowflake-ips-badge 1\n")
+			So(metricsStr, ShouldContainSubstring, "\nsnowflake-ips-webext 1\n")
+			So(metricsStr, ShouldEndWith, "\nsnowflake-ips-total 4\nsnowflake-idle-count 8\nclient-denied-count 0\nclient-restricted-denied-count 0\nclient-unrestricted-denied-count 0\nclient-snowflake-match-count 0\nsnowflake-ips-nat-restricted 0\nsnowflake-ips-nat-unrestricted 0\nsnowflake-ips-nat-unknown 1\n")
 		})
 
 		//Test addition of client failures
@@ -570,7 +575,11 @@ func TestMetrics(t *testing.T) {
 			buf.Reset()
 			ctx.metrics.zeroMetrics()
 			ctx.metrics.printMetrics()
-			So(buf.String(), ShouldContainSubstring, "snowflake-ips \nsnowflake-ips-total 0\nsnowflake-ips-standalone 0\nsnowflake-ips-badge 0\nsnowflake-ips-webext 0\nsnowflake-idle-count 0\nclient-denied-count 0\nclient-restricted-denied-count 0\nclient-unrestricted-denied-count 0\nclient-snowflake-match-count 0\nsnowflake-ips-nat-restricted 0\nsnowflake-ips-nat-unrestricted 0\nsnowflake-ips-nat-unknown 0\n")
+			So(buf.String(), ShouldContainSubstring, "\nsnowflake-ips \n")
+			So(buf.String(), ShouldContainSubstring, "\nsnowflake-ips-standalone 0\n")
+			So(buf.String(), ShouldContainSubstring, "\nsnowflake-ips-badge 0\n")
+			So(buf.String(), ShouldContainSubstring, "\nsnowflake-ips-webext 0\n")
+			So(buf.String(), ShouldContainSubstring, "\nsnowflake-ips-total 0\nsnowflake-idle-count 0\nclient-denied-count 0\nclient-restricted-denied-count 0\nclient-unrestricted-denied-count 0\nclient-snowflake-match-count 0\nsnowflake-ips-nat-restricted 0\nsnowflake-ips-nat-unrestricted 0\nsnowflake-ips-nat-unknown 0\n")
 		})
 		//Test addition of client matches
 		Convey("for client-proxy match", func() {
@@ -690,7 +699,9 @@ func TestMetrics(t *testing.T) {
 			<-done
 
 			ctx.metrics.printMetrics()
-			So(buf.String(), ShouldContainSubstring, "snowflake-ips CA=1\nsnowflake-ips-total 1")
+			metricsStr := buf.String()
+			So(metricsStr, ShouldContainSubstring, "snowflake-ips CA=1\n")
+			So(metricsStr, ShouldContainSubstring, "snowflake-ips-total 1\n")
 		})
 		//Test NAT types
 		Convey("proxy counts by NAT type", func() {
diff --git a/common/messages/proxy.go b/common/messages/proxy.go
index 64f139b..dcfe0ab 100644
--- a/common/messages/proxy.go
+++ b/common/messages/proxy.go
@@ -12,14 +12,17 @@ import (
 )
 
 const (
-	version = "1.2"
-
-	ProxyStandalone = "standalone"
-	ProxyWebext     = "webext"
-	ProxyBadge      = "badge"
-	ProxyUnknown    = "unknown"
+	version      = "1.2"
+	ProxyUnknown = "unknown"
 )
 
+var KnownProxyTypes = map[string]bool{
+	"standalone": true,
+	"webext":     true,
+	"badge":      true,
+	"iptproxy":   true,
+}
+
 /* Version 1.2 specification:
 
 == ProxyPollRequest ==
@@ -138,11 +141,7 @@ func DecodeProxyPollRequest(data []byte) (sid string, proxyType string, natType
 
 	// we don't reject polls with an unknown proxy type because we encourage
 	// projects that embed proxy code to include their own type
-	switch message.Type {
-	case ProxyStandalone:
-	case ProxyWebext:
-	case ProxyBadge:
-	default:
+	if !KnownProxyTypes[message.Type] {
 		message.Type = ProxyUnknown
 	}