From 98242a4dfae69057b7699d455f3aa02fb20d8449 Mon Sep 17 00:00:00 2001
From: Jaz Volpert <ericvolp12@gmail.com>
Date: Sun, 1 Oct 2023 16:37:52 +0000
Subject: [PATCH] Use custom metrics middleware for Palomar to avoid 404DOSing

---
 search/metrics.go | 93 +++++++++++++++++++++++++++++++++++++++++++++++
 search/server.go  |  6 +--
 2 files changed, 96 insertions(+), 3 deletions(-)

diff --git a/search/metrics.go b/search/metrics.go
index 829c9bc1d..bd53581e3 100644
--- a/search/metrics.go
+++ b/search/metrics.go
@@ -1,6 +1,12 @@
 package search
 
 import (
+	"errors"
+	"net/http"
+	"strconv"
+	"time"
+
+	"github.com/labstack/echo/v4"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 )
@@ -49,3 +55,90 @@ var currentSeq = promauto.NewGauge(prometheus.GaugeOpts{
 	Name: "search_current_seq",
 	Help: "Current sequence number",
 })
+
+var reqSz = promauto.NewHistogramVec(prometheus.HistogramOpts{
+	Name:    "http_request_size_bytes",
+	Help:    "A histogram of request sizes for requests.",
+	Buckets: prometheus.ExponentialBuckets(100, 10, 8),
+}, []string{"code", "method", "path"})
+
+var reqDur = promauto.NewHistogramVec(prometheus.HistogramOpts{
+	Name:    "http_request_duration_seconds",
+	Help:    "A histogram of latencies for requests.",
+	Buckets: prometheus.ExponentialBuckets(0.0001, 2, 18),
+}, []string{"code", "method", "path"})
+
+var reqCnt = promauto.NewCounterVec(prometheus.CounterOpts{
+	Name: "http_requests_total",
+	Help: "A counter for requests to the wrapped handler.",
+}, []string{"code", "method", "path"})
+
+var resSz = promauto.NewHistogramVec(prometheus.HistogramOpts{
+	Name:    "http_response_size_bytes",
+	Help:    "A histogram of response sizes for requests.",
+	Buckets: prometheus.ExponentialBuckets(100, 10, 8),
+}, []string{"code", "method", "path"})
+
+// MetricsMiddleware defines handler function for metrics middleware
+func MetricsMiddleware(next echo.HandlerFunc) echo.HandlerFunc {
+	return func(c echo.Context) error {
+		path := c.Path()
+		if path == "/metrics" || path == "/_health" {
+			return next(c)
+		}
+
+		start := time.Now()
+		requestSize := computeApproximateRequestSize(c.Request())
+
+		err := next(c)
+
+		status := c.Response().Status
+		if err != nil {
+			var httpError *echo.HTTPError
+			if errors.As(err, &httpError) {
+				status = httpError.Code
+			}
+			if status == 0 || status == http.StatusOK {
+				status = http.StatusInternalServerError
+			}
+		}
+
+		elapsed := float64(time.Since(start)) / float64(time.Second)
+
+		statusStr := strconv.Itoa(status)
+		method := c.Request().Method
+
+		responseSize := float64(c.Response().Size)
+
+		reqDur.WithLabelValues(statusStr, method, path).Observe(elapsed)
+		reqCnt.WithLabelValues(statusStr, method, path).Inc()
+		reqSz.WithLabelValues(statusStr, method, path).Observe(float64(requestSize))
+		resSz.WithLabelValues(statusStr, method, path).Observe(responseSize)
+
+		return err
+	}
+}
+
+func computeApproximateRequestSize(r *http.Request) int {
+	s := 0
+	if r.URL != nil {
+		s = len(r.URL.Path)
+	}
+
+	s += len(r.Method)
+	s += len(r.Proto)
+	for name, values := range r.Header {
+		s += len(name)
+		for _, value := range values {
+			s += len(value)
+		}
+	}
+	s += len(r.Host)
+
+	// N.B. r.Form and r.MultipartForm are assumed to be included in r.URL.
+
+	if r.ContentLength != -1 {
+		s += int(r.ContentLength)
+	}
+	return s
+}
diff --git a/search/server.go b/search/server.go
index 7a80b27d2..82dbccd23 100644
--- a/search/server.go
+++ b/search/server.go
@@ -13,8 +13,8 @@ import (
 	"github.com/bluesky-social/indigo/backfill"
 	"github.com/bluesky-social/indigo/util/version"
 	"github.com/bluesky-social/indigo/xrpc"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
 
-	"github.com/labstack/echo-contrib/echoprometheus"
 	"github.com/labstack/echo/v4"
 	"github.com/labstack/echo/v4/middleware"
 	es "github.com/opensearch-project/opensearch-go/v2"
@@ -183,7 +183,7 @@ func (s *Server) RunAPI(listen string) error {
 	e.HideBanner = true
 	e.Use(slogecho.New(s.logger))
 	e.Use(middleware.Recover())
-	e.Use(echoprometheus.NewMiddleware("palomar"))
+	e.Use(MetricsMiddleware)
 	e.Use(middleware.BodyLimit("64M"))
 
 	e.HTTPErrorHandler = func(err error, ctx echo.Context) {
@@ -197,7 +197,7 @@ func (s *Server) RunAPI(listen string) error {
 
 	e.Use(middleware.CORS())
 	e.GET("/_health", s.handleHealthCheck)
-	e.GET("/metrics", echoprometheus.NewHandler())
+	e.GET("/metrics", echo.WrapHandler(promhttp.Handler()))
 	e.GET("/xrpc/app.bsky.unspecced.searchPostsSkeleton", s.handleSearchPostsSkeleton)
 	e.GET("/xrpc/app.bsky.unspecced.searchActorsSkeleton", s.handleSearchActorsSkeleton)
 	e.GET("/xrpc/app.bsky.unspecced.indexRepos", s.handleIndexRepos)