feat: add option to expose the stats by using metrics

3 files changed, 124 insertions, 2 deletions

diff --git a/proxy/lib/metrics.go b/proxy/lib/metrics.go
new file mode 100644
index 0000000..adae459
--- /dev/null
+++ b/proxy/lib/metrics.go
@@ -0,0 +1,76 @@
+package snowflake_proxy
+
+import (
+	"net/http"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+const (
+	// metricNamespace represent prometheus namespace
+	metricNamespace = "tor_snowflake_proxy"
+)
+
+type Metrics struct {
+	totalInBoundTraffic  prometheus.Counter
+	totalOutBoundTraffic prometheus.Counter
+	totalConnections     prometheus.Counter
+}
+
+func NewMetrics() *Metrics {
+	return &Metrics{
+		totalConnections: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: metricNamespace,
+			Name:      "connections_total",
+			Help:      "The total number of connections handled by the snowflake proxy",
+		}),
+		totalInBoundTraffic: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: metricNamespace,
+			Name:      "traffic_inbound_bytes_total",
+			Help:      "The total in bound traffic by the snowflake proxy",
+		}),
+		totalOutBoundTraffic: prometheus.NewCounter(prometheus.CounterOpts{
+			Namespace: metricNamespace,
+			Name:      "traffic_outbound_bytes_total",
+			Help:      "The total out bound traffic by the snowflake proxy ",
+		}),
+	}
+}
+
+// Start register the metrics server and serve them on the given address
+func (m *Metrics) Start(addr string) error {
+	go func() {
+		http.Handle("/internal/metrics", promhttp.Handler())
+		if err := http.ListenAndServe(addr, nil); err != nil {
+			panic(err)
+		}
+	}()
+
+	return prometheus.Register(m)
+}
+
+func (m *Metrics) Collect(ch chan<- prometheus.Metric) {
+	m.totalConnections.Collect(ch)
+	m.totalInBoundTraffic.Collect(ch)
+	m.totalOutBoundTraffic.Collect(ch)
+}
+
+func (m *Metrics) Describe(descs chan<- *prometheus.Desc) {
+	prometheus.DescribeByCollect(m, descs)
+}
+
+// TrackInBoundTraffic counts the received traffic by the snowflake proxy
+func (m *Metrics) TrackInBoundTraffic(value int64) {
+	m.totalInBoundTraffic.Add(float64(value))
+}
+
+// TrackOutBoundTraffic counts the transmitted traffic by the snowflake proxy
+func (m *Metrics) TrackOutBoundTraffic(value int64) {
+	m.totalOutBoundTraffic.Add(float64(value))
+}
+
+// TrackNewConnection counts the new connections
+func (m *Metrics) TrackNewConnection() {
+	m.totalConnections.Inc()
+}
diff --git a/proxy/lib/pt_event_metrics.go b/proxy/lib/pt_event_metrics.go
new file mode 100644
index 0000000..64a371a
--- /dev/null
+++ b/proxy/lib/pt_event_metrics.go
@@ -0,0 +1,29 @@
+package snowflake_proxy
+
+import (
+	"gitlab.torproject.org/tpo/anti-censorship/pluggable-transports/snowflake/v2/common/event"
+)
+
+type EventCollector interface {
+	TrackInBoundTraffic(value int64)
+	TrackOutBoundTraffic(value int64)
+	TrackNewConnection()
+}
+
+type EventMetrics struct {
+	collector EventCollector
+}
+
+func NewEventMetrics(collector EventCollector) *EventMetrics {
+	return &EventMetrics{collector: collector}
+}
+
+func (em *EventMetrics) OnNewSnowflakeEvent(e event.SnowflakeEvent) {
+	switch e.(type) {
+	case event.EventOnProxyConnectionOver:
+		e := e.(event.EventOnProxyConnectionOver)
+		em.collector.TrackInBoundTraffic(e.InboundTraffic)
+		em.collector.TrackOutBoundTraffic(e.OutboundTraffic)
+		em.collector.TrackNewConnection()
+	}
+}
diff --git a/proxy/main.go b/proxy/main.go
index 87a88f9..66dbc8f 100644
--- a/proxy/main.go
+++ b/proxy/main.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"io/ioutil"
 	"log"
+	"net"
 	"os"
 	"strconv"
 	"strings"
@@ -32,6 +33,9 @@ func main() {
 		"the time interval in second before NAT type is retested, 0s disables retest. Valid time units are \"s\", \"m\", \"h\". ")
 	SummaryInterval := flag.Duration("summary-interval", time.Hour,
 		"the time interval to output summary, 0s disables summaries. Valid time units are \"s\", \"m\", \"h\". ")
+	disableStatsLogger := flag.Bool("disable-stats-logger", false, "disable the exposing mechanism for stats using logs")
+	enableMetrics := flag.Bool("metrics", false, "enable the exposing mechanism for stats using metrics")
+	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", "", "ICE UDP ephemeral ports range (format:\"<min>:<max>\")")
 	versionFlag := flag.Bool("version", false, "display version info to stderr and quit")
@@ -120,8 +124,21 @@ func main() {
 		log.SetOutput(&safelog.LogScrubber{Output: logOutput})
 	}
 
-	periodicEventLogger := sf.NewProxyEventLogger(*SummaryInterval, eventlogOutput)
-	eventLogger.AddSnowflakeEventListener(periodicEventLogger)
+	if !*disableStatsLogger {
+		periodicEventLogger := sf.NewProxyEventLogger(*SummaryInterval, eventlogOutput)
+		eventLogger.AddSnowflakeEventListener(periodicEventLogger)
+	}
+
+	if *enableMetrics {
+		metrics := sf.NewMetrics()
+
+		err := metrics.Start(net.JoinHostPort("localhost", strconv.Itoa(*metricsPort)))
+		if err != nil {
+			log.Fatalf("could not enable metrics: %v", err)
+		}
+
+		eventLogger.AddSnowflakeEventListener(sf.NewEventMetrics(metrics))
+	}
 
 	log.Printf("snowflake-proxy %s\n", version.GetVersion())