From 61ffdd448099f94ec3dc9aaf93705399b97a2bb2 Mon Sep 17 00:00:00 2001
From: Michel Hollands <42814411+MichelHollands@users.noreply.github.com>
Date: Mon, 14 Sep 2020 09:32:18 +0100
Subject: [PATCH] Make HTTP middleware optional (#194)

* Make HTTP middleware optional
* Allow the Router to be set by caller

Signed-off-by: Michel Hollands <michel.hollands@grafana.com>
---
 server/server.go      | 28 +++++++++++++++++++++-------
 server/server_test.go | 15 +++++++++------
 2 files changed, 30 insertions(+), 13 deletions(-)

diff --git a/server/server.go b/server/server.go
index ae04022c..a4bedb80 100644
--- a/server/server.go
+++ b/server/server.go
@@ -62,10 +62,12 @@ type Config struct {
 	HTTPServerWriteTimeout        time.Duration `yaml:"http_server_write_timeout"`
 	HTTPServerIdleTimeout         time.Duration `yaml:"http_server_idle_timeout"`
 
-	GRPCOptions          []grpc.ServerOption            `yaml:"-"`
-	GRPCMiddleware       []grpc.UnaryServerInterceptor  `yaml:"-"`
-	GRPCStreamMiddleware []grpc.StreamServerInterceptor `yaml:"-"`
-	HTTPMiddleware       []middleware.Interface         `yaml:"-"`
+	GRPCOptions                   []grpc.ServerOption            `yaml:"-"`
+	GRPCMiddleware                []grpc.UnaryServerInterceptor  `yaml:"-"`
+	GRPCStreamMiddleware          []grpc.StreamServerInterceptor `yaml:"-"`
+	HTTPMiddleware                []middleware.Interface         `yaml:"-"`
+	Router                        *mux.Router                    `yaml:"-"`
+	DoNotAddDefaultHTTPMiddleware bool                           `yaml:"-"`
 
 	GPRCServerMaxRecvMsgSize        int           `yaml:"grpc_server_max_recv_msg_size"`
 	GRPCServerMaxSendMsgSize        int           `yaml:"grpc_server_max_send_msg_size"`
@@ -270,7 +272,12 @@ func New(cfg Config) (*Server, error) {
 	grpcServer := grpc.NewServer(grpcOptions...)
 
 	// Setup HTTP server
-	router := mux.NewRouter()
+	var router *mux.Router
+	if cfg.Router != nil {
+		router = cfg.Router
+	} else {
+		router = mux.NewRouter()
+	}
 	if cfg.PathPrefix != "" {
 		// Expect metrics and pprof handlers to be prefixed with server's path prefix.
 		// e.g. /loki/metrics or /loki/debug/pprof
@@ -279,6 +286,7 @@ func New(cfg Config) (*Server, error) {
 	if cfg.RegisterInstrumentation {
 		RegisterInstrumentation(router)
 	}
+
 	var sourceIPs *middleware.SourceIPExtractor
 	if cfg.LogSourceIPs {
 		sourceIPs, err = middleware.NewSourceIPs(cfg.LogSourceIPsHeader, cfg.LogSourceIPsRegex)
@@ -286,7 +294,8 @@ func New(cfg Config) (*Server, error) {
 			return nil, fmt.Errorf("error setting up source IP extraction: %v", err)
 		}
 	}
-	httpMiddleware := []middleware.Interface{
+
+	defaultHTTPMiddleware := []middleware.Interface{
 		middleware.Tracer{
 			RouteMatcher: router,
 			SourceIPs:    sourceIPs,
@@ -303,8 +312,13 @@ func New(cfg Config) (*Server, error) {
 			InflightRequests: inflightRequests,
 		},
 	}
+	httpMiddleware := []middleware.Interface{}
+	if cfg.DoNotAddDefaultHTTPMiddleware {
+		httpMiddleware = cfg.HTTPMiddleware
+	} else {
+		httpMiddleware = append(defaultHTTPMiddleware, cfg.HTTPMiddleware...)
+	}
 
-	httpMiddleware = append(httpMiddleware, cfg.HTTPMiddleware...)
 	httpServer := &http.Server{
 		ReadTimeout:  cfg.HTTPServerReadTimeout,
 		WriteTimeout: cfg.HTTPServerWriteTimeout,
diff --git a/server/server_test.go b/server/server_test.go
index d13f6d9f..7f85a00d 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -20,6 +20,7 @@ import (
 	"google.golang.org/grpc/status"
 
 	google_protobuf "github.com/golang/protobuf/ptypes/empty"
+	"github.com/gorilla/mux"
 	"github.com/prometheus/client_golang/prometheus"
 	node_https "github.com/prometheus/node_exporter/https"
 	"github.com/stretchr/testify/require"
@@ -425,12 +426,14 @@ func TestMiddlewareLogging(t *testing.T) {
 	var level logging.Level
 	level.Set("info")
 	cfg := Config{
-		HTTPListenAddress: "localhost",
-		HTTPListenPort:    9192,
-		GRPCListenAddress: "localhost",
-		HTTPMiddleware:    []middleware.Interface{middleware.Logging},
-		MetricsNamespace:  "testing_logging",
-		LogLevel:          level,
+		HTTPListenAddress:             "localhost",
+		HTTPListenPort:                9192,
+		GRPCListenAddress:             "localhost",
+		HTTPMiddleware:                []middleware.Interface{middleware.Logging},
+		MetricsNamespace:              "testing_logging",
+		LogLevel:                      level,
+		DoNotAddDefaultHTTPMiddleware: true,
+		Router:                        &mux.Router{},
 	}
 	server, err := New(cfg)
 	require.NoError(t, err)