better prometheus metrics

cmd/root.go

@@ -44,6 +44,11 @@ var rootCmd = &cobra.Command{
 			log.Panic(err)
 		}
 
+		// start metrics
+		if err := pkg.StartMetrics(config); err != nil {
+			log.Panic(fmt.Errorf("failed to start metrics: %w", err))
+		}
+
 		// start the broker
 		teardown, err := StartNetworkBroker(config)
 		if err != nil {

go.mod

@@ -7,13 +7,13 @@ require (
 	github.com/gin-gonic/gin v1.10.0
 	github.com/mcuadros/go-defaults v1.2.0
 	github.com/mitchellh/mapstructure v1.5.0
+	github.com/prometheus/client_golang v1.14.0
 	github.com/sirupsen/logrus v1.9.3
 	github.com/spf13/cobra v1.8.1
 	github.com/spf13/viper v1.19.0
 	github.com/stretchr/testify v1.9.0
 	github.com/ucarion/urlpath v0.0.0-20200424170820-7ccc79b76bbb
 	github.com/whuang8/redactrus v1.0.2
-	github.com/zsais/go-gin-prometheus v0.1.0
 	golang.zx2c4.com/wireguard v0.0.0-20231010133717-42ec952eadc2
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
 	google.golang.org/protobuf v1.35.1
@@ -51,7 +51,6 @@ require (
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/pelletier/go-toml/v2 v2.2.2 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
-	github.com/prometheus/client_golang v1.14.0 // indirect
 	github.com/prometheus/client_model v0.3.0 // indirect
 	github.com/prometheus/common v0.37.0 // indirect
 	github.com/prometheus/procfs v0.8.0 // indirect

go.sum

@@ -307,8 +307,6 @@ github.com/whuang8/redactrus v1.0.2/go.mod h1:/QqU95wNV2zWg3nD5/uatl9Uz0cJUROT4S
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 github.com/yuin/goldmark v1.1.32/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
-github.com/zsais/go-gin-prometheus v0.1.0 h1:bkLv1XCdzqVgQ36ScgRi09MA2UC1t3tAB6nsfErsGO4=
-github.com/zsais/go-gin-prometheus v0.1.0/go.mod h1:Slirjzuz8uM8Cw0jmPNqbneoqcUtY2GGjn2bEd4NRLY=
 go.opencensus.io v0.21.0/go.mod h1:mSImk1erAIZhrmZN+AvHh14ztQfjbGwt4TtuofqLduU=
 go.opencensus.io v0.22.0/go.mod h1:+kGneAE2xo2IficOXnaByMWTGM9T73dGwxeWcUqIpI8=
 go.opencensus.io v0.22.2/go.mod h1:yxeiOL68Rb0Xd1ddK5vPZ/oVn4vY4Ynel7k9FzqtOIw=

pkg/config.go

@@ -262,9 +262,15 @@ type OutboundProxyConfig struct {
 	ListenPort int                            `mapstructure:"listenPort" json:"listenPort" validate:"gte=0" default:"8080"`
 }
 
+type MetricsConfig struct {
+	Disabled bool   `mapstructure:"disabled" json:"disabled"`
+	Addr     string `mapstructure:"addr" json:"addr" default:":9000"`
+}
+
 type Config struct {
 	Inbound  InboundProxyConfig  `mapstructure:"inbound" json:"inbound"`
 	Outbound OutboundProxyConfig `mapstructure:"outbound" json:"outbound"`
+	Metrics  MetricsConfig       `mapstructure:"metrics" json:"metrics"`
 }
 
 func LoadConfig(configFiles []string, deploymentId int) (*Config, error) {

pkg/heartbeat.go

@@ -41,13 +41,16 @@ func (config *HeartbeatConfig) Start(tnet *netstack.Net, userAgent string) (func
 			} else {
 				log.WithField("failure_count", failures).WithField("status_code", resp.StatusCode).Warn("heartbeat.failure")
 			}
+			heartbeatFailureCounter.Inc()
 			return false
 		} else {
 			if failures != 0 {
 				log.WithField("message", "Established connectivity with Semgrep").Info("heartbeat.success")
 			}
 			log.Debug("heartbeat.success")
 			failures = 0
+			heartbeatSuccessCounter.Inc()
+			heartbeatLastSuccessTimestamp.SetToCurrentTime()
 			return true
 		}
 	}

pkg/inbound_proxy.go

@@ -10,15 +10,16 @@ import (
 	"net/url"
 
 	"github.com/gin-gonic/gin"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 	log "github.com/sirupsen/logrus"
-	ginprometheus "github.com/zsais/go-gin-prometheus"
 	"golang.zx2c4.com/wireguard/tun/netstack"
 	"gopkg.in/dealancer/validate.v2"
 )
 
 const errorResponseHeader = "X-Semgrep-Private-Link-Error"
 const proxyResponseHeader = "X-Semgrep-Private-Link"
 const healthcheckPath = "/healthcheck"
+const metricsPath = "/metrics"
 const destinationUrlParam = "destinationUrl"
 const proxyPath = "/proxy/*" + destinationUrlParam
 
@@ -49,9 +50,9 @@ func (config *InboundProxyConfig) Start(tnet *netstack.Net) error {
 	log.WithField("path", healthcheckPath).Info("healthcheck.configured")
 
 	// setup metrics
-	p := ginprometheus.NewPrometheus("gin")
-	p.Use(r)
-	log.WithField("path", p.MetricsPath).Info("metrics.configured")
+	promHandler := promhttp.Handler()
+	r.GET(metricsPath, gin.WrapH(promHandler))
+	log.WithField("path", metricsPath).Info("internal_metrics.configured")
 
 	// setup http proxy
 	r.Any(proxyPath, func(c *gin.Context) {
@@ -80,6 +81,12 @@ func (config *InboundProxyConfig) Start(tnet *netstack.Net) error {
 
 		logger = logger.WithField("allowlist_match", allowlistMatch.URL)
 
+		instrumentedTransport, err := BuildInstrumentedRoundTripper(transport, allowlistMatch.URL)
+		if err != nil {
+			logger.WithError(err).Warn("roundtripper.instrument_error")
+			instrumentedTransport = transport
+		}
+
 		reqLogger := logger
 		if config.Logging.LogRequestBody || allowlistMatch.LogRequestBody {
 			reqBody := &bytes.Buffer{}
@@ -96,7 +103,7 @@ func (config *InboundProxyConfig) Start(tnet *netstack.Net) error {
 		reqLogger.Info("proxy.request")
 
 		proxy := httputil.ReverseProxy{
-			Transport: transport,
+			Transport: instrumentedTransport,
 			Director: func(req *http.Request) {
 				req.URL = destinationUrl
 				req.Host = destinationUrl.Host

pkg/logger.go

@@ -6,6 +6,7 @@ import (
 	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/prometheus/client_golang/prometheus"
 	log "github.com/sirupsen/logrus"
 )
 
@@ -68,3 +69,25 @@ func LoggerWithConfig(logger *log.Logger, notlogged []string) gin.HandlerFunc {
 		}
 	}
 }
+
+type MetricsHook struct {
+	CounterFunc func(labelValues ...string) (prometheus.Counter, error)
+}
+
+func (mh *MetricsHook) Levels() []log.Level {
+	return log.AllLevels
+}
+
+func (mh *MetricsHook) Fire(entry *log.Entry) error {
+	if metric, err := mh.CounterFunc(entry.Level.String(), entry.Message); err == nil {
+		metric.Inc()
+	}
+	return nil
+}
+
+func init() {
+	h := &MetricsHook{
+		CounterFunc: logEventsCounter.GetMetricWithLabelValues,
+	}
+	log.AddHook(h)
+}

pkg/metrics.go

@@ -0,0 +1,71 @@
+package pkg
+
+import (
+	"fmt"
+	"net"
+	"net/http"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+	log "github.com/sirupsen/logrus"
+)
+
+var logEventsCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
+	Name: "network_broker_log_events_total",
+	Help: "Total number of log events",
+}, []string{"level", "event"})
+
+var heartbeatCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
+	Name: "network_broker_heartbeat_total",
+	Help: "Total number of heartbeat attempts",
+}, []string{"result"})
+
+var heartbeatSuccessCounter = heartbeatCounter.WithLabelValues("success")
+var heartbeatFailureCounter = heartbeatCounter.WithLabelValues("failure")
+
+var heartbeatLastSuccessTimestamp = prometheus.NewGauge(prometheus.GaugeOpts{
+	Name: "network_broker_heartbeat_last_success_timestamp_seconds",
+	Help: "Timestamp of last successful heartbeat attempt in seconds",
+}) // TODO: refactor heartbeat.go to be per-endpoint, and then include peer_endpoint as a label
+
+var proxyInFlightGauge = prometheus.NewGauge(prometheus.GaugeOpts{
+	Name: "network_broker_proxy_in_flight_requests",
+	Help: "Number of in-flight proxy requests",
+})
+
+var proxyCounter = prometheus.NewCounterVec(prometheus.CounterOpts{
+	Name: "network_broker_proxy_requests_total",
+	Help: " Total number of proxy requests",
+}, []string{"allowlist", "method", "code"})
+
+func StartMetrics(config *Config) error {
+	if config.Metrics.Disabled {
+		log.WithField("addr", config.Metrics.Addr).Info("external_metrics.disabled")
+		return nil
+	}
+
+	prometheus.MustRegister(logEventsCounter, heartbeatCounter, heartbeatLastSuccessTimestamp, proxyInFlightGauge, proxyCounter)
+
+	promHandler := promhttp.Handler()
+	httpServer := &http.Server{Addr: config.Metrics.Addr, Handler: promHandler}
+	listener, err := net.Listen("tcp", httpServer.Addr)
+	if err != nil {
+		return fmt.Errorf("failed to start external metrics server: %w", err)
+	}
+	go httpServer.Serve(listener)
+	log.WithField("addr", config.Metrics.Addr).Info("external_metrics.started")
+
+	return nil
+}
+
+func BuildInstrumentedRoundTripper(transport http.RoundTripper, allowlist string) (http.RoundTripper, error) {
+	labels := prometheus.Labels{"allowlist": allowlist}
+	counter, err := proxyCounter.CurryWith(labels)
+	if err != nil {
+		return nil, err
+	}
+	instrumentedTransport := promhttp.InstrumentRoundTripperInFlight(proxyInFlightGauge,
+		promhttp.InstrumentRoundTripperCounter(counter, transport),
+	)
+	return instrumentedTransport, nil
+}

pkg/relay.go

@@ -14,7 +14,6 @@ import (
 	"github.com/PaesslerAG/jsonpath"
 	"github.com/gin-gonic/gin"
 	log "github.com/sirupsen/logrus"
-	ginprometheus "github.com/zsais/go-gin-prometheus"
 	"gopkg.in/dealancer/validate.v2"
 )
 
@@ -129,10 +128,6 @@ func (config *OutboundProxyConfig) Start() error {
 	r.GET(healthcheckPath, func(c *gin.Context) { c.JSON(http.StatusOK, "OK") })
 	log.WithField("path", healthcheckPath).Info("healthcheck.configured")
 
-	// setup metrics
-	p := ginprometheus.NewPrometheus("gin")
-	p.Use(r)
-
 	// setup http proxy
 	r.Any("/relay/:name", func(c *gin.Context) {
 		relayName := c.Param("name")