proxy: add country to prometheus metrics

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() {
-	m.totalConnections.Inc()
+func (m *Metrics) TrackNewConnection(country string) {
+	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