proxy: add country to prometheus metrics

Dockerfile

@@ -1,4 +1,4 @@
-FROM docker.io/library/golang:1.23 AS build
+FROM docker.io/library/golang:1.23-bookworm AS build
 
 # Set some labels
 # io.containers.autoupdate label will instruct podman to reach out to the corres
@@ -9,6 +9,8 @@ FROM docker.io/library/golang:1.23 AS build
 LABEL io.containers.autoupdate=registry
 LABEL org.opencontainers.image.authors="anti-censorship-team@lists.torproject.org"
 
+RUN apt-get update && apt-get install -y tor-geoipdb
+
 ADD . /app
 
 WORKDIR /app/proxy
@@ -19,6 +21,7 @@ FROM scratch
 
 COPY --from=build /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt
 COPY --from=build /usr/share/zoneinfo /usr/share/zoneinfo
+COPY --from=build /usr/share/tor/geoip* /usr/share/tor/
 COPY --from=build /app/proxy/proxy /bin/proxy
 
 ENTRYPOINT [ "/bin/proxy" ]

common/event/interface.go

@@ -77,8 +77,7 @@ func (e EventOnProxyClientConnected) String() string {
 
 type EventOnProxyConnectionOver struct {
 	SnowflakeEvent
-	InboundTraffic  int64
+	Country string
-	OutboundTraffic int64
 }
 
 func (e EventOnProxyConnectionOver) String() string {

proxy/lib/metrics.go

@@ -15,16 +15,18 @@ const (
 type Metrics struct {
 	totalInBoundTraffic  prometheus.Counter
 	totalOutBoundTraffic prometheus.Counter
-	totalConnections     prometheus.Counter
+	totalConnections     *prometheus.CounterVec
 }
 
 func NewMetrics() *Metrics {
 	return &Metrics{
-		totalConnections: prometheus.NewCounter(prometheus.CounterOpts{
+		totalConnections: prometheus.NewCounterVec(prometheus.CounterOpts{
 			Namespace: metricNamespace,
 			Name:      "connections_total",
 			Help:      "The total number of connections handled by the snowflake proxy",
-		}),
+		},
+			[]string{"country"},
+		),
 		totalInBoundTraffic: prometheus.NewCounter(prometheus.CounterOpts{
 			Namespace: metricNamespace,
 			Name:      "traffic_inbound_bytes_total",
@@ -71,6 +73,8 @@ func (m *Metrics) TrackOutBoundTraffic(value int64) {
 }
 
 // TrackNewConnection counts the new connections
-func (m *Metrics) TrackNewConnection() {
+func (m *Metrics) TrackNewConnection(country string) {
-	m.totalConnections.Inc()
+	m.totalConnections.
+		With(prometheus.Labels{"country": country}).
+		Inc()
 }

proxy/lib/pt_event_metrics.go

@@ -7,7 +7,7 @@ import (
 type EventCollector interface {
 	TrackInBoundTraffic(value int64)
 	TrackOutBoundTraffic(value int64)
-	TrackNewConnection()
+	TrackNewConnection(country string)
 }
 
 type EventMetrics struct {
@@ -25,6 +25,7 @@ func (em *EventMetrics) OnNewSnowflakeEvent(e event.SnowflakeEvent) {
 		em.collector.TrackInBoundTraffic(e.InboundBytes)
 		em.collector.TrackOutBoundTraffic(e.OutboundBytes)
 	case event.EventOnProxyConnectionOver:
-		em.collector.TrackNewConnection()
+		e := e.(event.EventOnProxyConnectionOver)
+		em.collector.TrackNewConnection(e.Country)
 	}
 }

proxy/lib/snowflake.go

@@ -35,6 +35,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"reflect"
 	"strings"
 	"sync"
 	"time"
@@ -114,6 +115,10 @@ var (
 	client http.Client
 )
 
+type GeoIP interface {
+	GetCountryByAddr(net.IP) (string, bool)
+}
+
 // SnowflakeProxy is used to configure an embedded
 // Snowflake in another Go application.
 // For some more info also see CLI parameter descriptions in README.
@@ -166,6 +171,9 @@ type SnowflakeProxy struct {
 	// SummaryInterval is the time interval at which proxy stats will be logged
 	SummaryInterval time.Duration
 
+	// GeoIP will be used to detect the country of the clients if provided
+	GeoIP GeoIP
+
 	periodicProxyStats *periodicProxyStats
 	bytesLogger        bytesLogger
 }
@@ -449,6 +457,7 @@ func (sf *SnowflakeProxy) makePeerConnectionFromOffer(
 
 		pr, pw := io.Pipe()
 		conn := newWebRTCConn(pc, dc, pr, sf.bytesLogger)
+		remoteIP := conn.RemoteIP()
 
 		dc.SetBufferedAmountLowThreshold(bufferedAmountLowThreshold)
 
@@ -479,7 +488,13 @@ func (sf *SnowflakeProxy) makePeerConnectionFromOffer(
 			conn.lock.Lock()
 			defer conn.lock.Unlock()
 			log.Printf("Data Channel %s-%d close\n", dc.Label(), dc.ID())
-			sf.EventDispatcher.OnNewSnowflakeEvent(event.EventOnProxyConnectionOver{})
+
+			country := ""
+			if sf.GeoIP != nil && !reflect.ValueOf(sf.GeoIP).IsNil() && remoteIP != nil {
+				country, _ = sf.GeoIP.GetCountryByAddr(remoteIP)
+			}
+			sf.EventDispatcher.OnNewSnowflakeEvent(event.EventOnProxyConnectionOver{Country: country})
+
 			conn.dc = nil
 			dc.Close()
 			pw.Close()
@@ -503,7 +518,7 @@ func (sf *SnowflakeProxy) makePeerConnectionFromOffer(
 			}
 		})
 
-		go handler(conn, conn.RemoteIP())
+		go handler(conn, remoteIP)
 	})
 	// As of v3.0.0, pion-webrtc uses trickle ICE by default.
 	// We have to wait for candidate gathering to complete

proxy/main.go

@@ -11,6 +11,7 @@ import (
 	"strings"
 	"time"
 
+	"gitlab.torproject.org/tpo/anti-censorship/geoip"
 	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/ptutil/safelog"
 	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/event"
 	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/version"
@@ -45,6 +46,8 @@ func main() {
 	metricsPort := flag.Int("metrics-port", 9999, "set port for the metrics service")
 	verboseLogging := flag.Bool("verbose", false, "increase log verbosity")
 	ephemeralPortsRangeFlag := flag.String("ephemeral-ports-range", "", "Set the `range` of ports used for client connections (format:\"<min>:<max>\").\nUseful in conjunction with port forwarding, in order to make the proxy NAT type \"unrestricted\".\nIf omitted, the ports will be chosen automatically from a wide range.\nWhen specifying the range, make sure it's at least 2x as wide as the amount of clients that you are hoping to serve concurrently (see the \"capacity\" flag).")
+	geoipDatabase := flag.String("geoipdb", "/usr/share/tor/geoip", "path to correctly formatted geoip database mapping IPv4 address ranges to country codes")
+	geoip6Database := flag.String("geoip6db", "/usr/share/tor/geoip6", "path to correctly formatted geoip database mapping IPv6 address ranges to country codes")
 	versionFlag := flag.Bool("version", false, "display version info to stderr and quit")
 
 	var ephemeralPortsRange []uint16 = []uint16{0, 0}
@@ -92,6 +95,12 @@ func main() {
 		}
 	}
 
+	gip, err := geoip.New(*geoipDatabase, *geoip6Database)
+	if *enableMetrics && err != nil {
+		// The geoip DB is only used for metrics, let's only report the error if enabled
+		log.Println("Error loading geoip db for country based metrics:", err)
+	}
+
 	proxy := sf.SnowflakeProxy{
 		PollInterval:       *pollInterval,
 		Capacity:           uint(*capacity),
@@ -112,6 +121,7 @@ func main() {
 		AllowNonTLSRelay:                *allowNonTLSRelay,
 
 		SummaryInterval: *summaryInterval,
+		GeoIP:           gip,
 	}
 
 	var logOutput = io.Discard
@@ -163,7 +173,7 @@ func main() {
 
 	log.Printf("snowflake-proxy %s\n", version.GetVersion())
 
-	err := proxy.Start()
+	err = proxy.Start()
 	if err != nil {
 		log.Fatal(err)
 	}