STORY-25143 - Add prometheus metrics to smokescreen

README.md

@@ -59,6 +59,10 @@ Here are the options you can give Smokescreen:
    --deny-address value                        Add IP[:PORT] to list of blocked IPs.  Repeatable.
    --allow-address value                       Add IP[:PORT] to list of allowed IPs.  Repeatable.
    --egress-acl-file FILE                      Validate egress traffic against FILE
+   --prometheus-endpoint ENDPOINT              Expose Prometheus metrics on the given endpoint.
+                                                 If --prometheus-port isn't set, by default listens on 9810
+   --prometheus-port PORT                      Expose Prometheus metrics on the given port.
+                                                 Requires --prometheus-endpoint to be set.
    --resolver-address ADDRESS                  Make DNS requests to ADDRESS (IP:port).  Repeatable.
    --statsd-address ADDRESS                    Send metrics to statsd at ADDRESS (IP:port). (default: "127.0.0.1:8200")
    --tls-server-bundle-file FILE               Authenticate to clients using key and certs from FILE

cmd/smokescreen.go

@@ -85,6 +85,15 @@ func NewConfiguration(args []string, logger *log.Logger) (*smokescreen.Config, e
 			Name:  "egress-acl-file",
 			Usage: "Validate egress traffic against `FILE`",
 		},
+		cli.StringFlag{
+			Name:  "prometheus-endpoint",
+			Usage: "Expose metrics via prometheus on `ENDPOINT`.",
+		},
+		cli.StringFlag{
+			Name:  "prometheus-port",
+			Value: "9810",
+			Usage: "Expose metrics via prometheus on `PORT`.",
+		},
 		cli.StringSliceFlag{
 			Name:  "resolver-address",
 			Usage: "Make DNS requests to `ADDRESS` (IP:port).  Repeatable.",
@@ -229,6 +238,12 @@ func NewConfiguration(args []string, logger *log.Logger) (*smokescreen.Config, e
 			}
 		}
 
+		if c.IsSet("prometheus-endpoint") {
+			if err := conf.SetupPrometheus(c.String("prometheus-endpoint"), c.String("prometheus-port")); err != nil {
+				return err
+			}
+		}
+
 		if c.IsSet("egress-acl-file") {
 			if err := conf.SetupEgressAcl(c.String("egress-acl-file")); err != nil {
 				return err

go.mod

@@ -1,27 +1,36 @@
 module github.com/stripe/smokescreen
 
-go 1.17
+go 1.18
 
 require (
 	github.com/DataDog/datadog-go v4.5.1+incompatible
 	github.com/armon/go-proxyproto v0.0.0-20170620220930-48572f11356f
 	github.com/carlmjohnson/versioninfo v0.22.4
 	github.com/hashicorp/go-cleanhttp v0.0.0-20171218145408-d5fe4b57a186
 	github.com/patrickmn/go-cache v2.1.0+incompatible
+	github.com/prometheus/client_golang v1.13.0
 	github.com/rs/xid v1.2.1
 	github.com/sirupsen/logrus v1.7.0
 	github.com/stretchr/testify v1.8.0
 	github.com/stripe/goproxy v0.0.0-20220308202309-3f1dfba6d1a4
 	golang.org/x/net v0.0.0-20220812174116-3211cb980234
 	gopkg.in/urfave/cli.v1 v1.20.0
-	gopkg.in/yaml.v2 v2.2.8
+	gopkg.in/yaml.v2 v2.4.0
 )
 
 require (
 	github.com/Microsoft/go-winio v0.4.17 // indirect
+	github.com/beorn7/perks v1.0.1 // indirect
+	github.com/cespare/xxhash/v2 v2.1.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
+	github.com/prometheus/client_model v0.2.0 // indirect
+	github.com/prometheus/common v0.37.0 // indirect
+	github.com/prometheus/procfs v0.8.0 // indirect
 	golang.org/x/sys v0.0.0-20220728004956-3c1f35247d10 // indirect
 	golang.org/x/text v0.3.7 // indirect
+	google.golang.org/protobuf v1.28.1 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

pkg/smokescreen/config.go

@@ -306,7 +306,7 @@ func (config *Config) SetupStatsdWithNamespace(addr, namespace string) error {
 		return nil
 	}
 
-	mc, err := metrics.NewMetricsClient(addr, namespace)
+	mc, err := metrics.NewStatsdMetricsClient(addr, namespace)
 	if err != nil {
 		return err
 	}
@@ -318,6 +318,15 @@ func (config *Config) SetupStatsd(addr string) error {
 	return config.SetupStatsdWithNamespace(addr, DefaultStatsdNamespace)
 }
 
+func (config *Config) SetupPrometheus(endpoint string, port string) error {
+	metricsClient, err := metrics.NewPrometheusMetricsClient(endpoint, port)
+	if err != nil {
+		return err
+	}
+	config.MetricsClient = metricsClient
+	return nil
+}
+
 func (config *Config) SetupEgressAcl(aclFile string) error {
 	if aclFile == "" {
 		config.EgressACL = nil

pkg/smokescreen/config_loader.go

@@ -48,71 +48,71 @@ type yamlConfig struct {
 	Tls *yamlConfigTls
 	// Currently not configurable via YAML: RoleFromRequest, Log, DisabledAclPolicyActions
 
-	UnsafeAllowPrivateRanges bool	`yaml:"unsafe_allow_private_ranges"`
+	UnsafeAllowPrivateRanges bool `yaml:"unsafe_allow_private_ranges"`
 }
 
-func (c *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
+func (config *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
 	var yc yamlConfig
-	*c = *NewConfig()
+	*config = *NewConfig()
 
 	err := unmarshal(&yc)
 	if err != nil {
 		return err
 	}
 
-	c.Ip = yc.Ip
+	config.Ip = yc.Ip
 
 	if yc.Port != nil {
-		c.Port = *yc.Port
+		config.Port = *yc.Port
 	}
 
-	err = c.SetDenyRanges(yc.DenyRanges)
+	err = config.SetDenyRanges(yc.DenyRanges)
 	if err != nil {
 		return err
 	}
 
-	err = c.SetAllowRanges(yc.AllowRanges)
+	err = config.SetAllowRanges(yc.AllowRanges)
 	if err != nil {
 		return err
 	}
 
-	err = c.SetResolverAddresses(yc.Resolvers)
+	err = config.SetResolverAddresses(yc.Resolvers)
 	if err != nil {
 		return err
 	}
 
-	c.IdleTimeout = yc.IdleTimeout
-	c.ConnectTimeout = yc.ConnectTimeout
+	config.IdleTimeout = yc.IdleTimeout
+	config.ConnectTimeout = yc.ConnectTimeout
 	if yc.ExitTimeout != nil {
-		c.ExitTimeout = *yc.ExitTimeout
+		config.ExitTimeout = *yc.ExitTimeout
 	}
 
-	err = c.SetupStatsd(yc.StatsdAddress)
+	err = config.SetupStatsd(yc.StatsdAddress)
 	if err != nil {
 		return err
 	}
 
 	if yc.EgressAclFile != "" {
-		err = c.SetupEgressAcl(yc.EgressAclFile)
+		err = config.SetupEgressAcl(yc.EgressAclFile)
 		if err != nil {
 			return err
 		}
 	}
 
-	c.SupportProxyProtocol = yc.SupportProxyProtocol
+	config.SupportProxyProtocol = yc.SupportProxyProtocol
 
 	if yc.StatsSocketDir != "" {
-		c.StatsSocketDir = yc.StatsSocketDir
+		config.StatsSocketDir = yc.StatsSocketDir
 	}
 
 	if yc.StatsSocketFileMode != "" {
 		filemode, err := strconv.ParseInt(yc.StatsSocketFileMode, 8, 9)
 
 		if err != nil {
-			c.Log.Fatal(err)
+			config.Log.Fatal(err)
 		}
 
-		c.StatsSocketFileMode = os.FileMode(filemode)
+		config.StatsSocketFileMode = os.FileMode(filemode)
 	}
 
 	if yc.Tls != nil {
@@ -126,12 +126,12 @@ func (c *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
 			key_file = yc.Tls.CertFile
 		}
 
-		err = c.SetupTls(yc.Tls.CertFile, key_file, yc.Tls.ClientCAFiles)
+		err = config.SetupTls(yc.Tls.CertFile, key_file, yc.Tls.ClientCAFiles)
 		if err != nil {
 			return err
 		}
 
-		err = c.SetupCrls(yc.Tls.CRLFiles)
+		err = config.SetupCrls(yc.Tls.CRLFiles)
 		if err != nil {
 			return err
 		}
@@ -143,13 +143,13 @@ func (c *Config) UnmarshalYAML(unmarshal func(interface{}) error) error {
 		default:
 			return fmt.Errorf("invalid network type: %v", yc.Network)
 		}
-		c.Network = yc.Network
+		config.Network = yc.Network
 	}
 
-	c.AllowMissingRole = yc.AllowMissingRole
-	c.AdditionalErrorMessageOnDeny = yc.DenyMessageExtra
-	c.TimeConnect = yc.TimeConnect
-	c.UnsafeAllowPrivateRanges = yc.UnsafeAllowPrivateRanges
+	config.AllowMissingRole = yc.AllowMissingRole
+	config.AdditionalErrorMessageOnDeny = yc.DenyMessageExtra
+	config.TimeConnect = yc.TimeConnect
+	config.UnsafeAllowPrivateRanges = yc.UnsafeAllowPrivateRanges
 
 	return nil
 }

pkg/smokescreen/conntrack/conn_tracker_test.go

@@ -101,11 +101,11 @@ func TestConnSuccessRateTracker(t *testing.T) {
 			assert.InDelta(tc.expectedRate, stats.ConnSuccessRate, 0.01)
 			assert.Equal(tc.totalConns, stats.TotalConns)
 
-			v, err := mockMetricsClient.GetValues("cn.atpt.distinct_domains_success_rate")
+			v, err := mockMetricsClient.GetValues("cn.atpt.distinct_domains_success_rate", map[string]string{})
 			assert.NoError(err)
 			assert.Equal(tc.expectedRate, v[len(v)-1])
 
-			v, err = mockMetricsClient.GetValues("cn.atpt.distinct_domains")
+			v, err = mockMetricsClient.GetValues("cn.atpt.distinct_domains", map[string]string{})
 			assert.NoError(err)
 			assert.Equal(tc.totalConns, int(v[len(v)-1]))
 

pkg/smokescreen/conntrack/instrumented_conn.go

@@ -2,7 +2,6 @@ package conntrack
 
 import (
 	"encoding/json"
-	"fmt"
 	"net"
 	"sync"
 	"sync/atomic"
@@ -93,8 +92,8 @@ func (ic *InstrumentedConn) Close() error {
 	end := time.Now()
 	duration := end.Sub(ic.Start).Seconds()
 
-	tags := []string{
-		fmt.Sprintf("role:%s", ic.Role),
+	tags := map[string]string{
+		"role": ic.Role,
 	}
 
 	ic.tracker.statsc.IncrWithTags("cn.close", tags, 1)