Merge middlewares the other way round (#4249)

Author: simonswine

pkg/api/api.go

@@ -15,7 +15,6 @@ import (
 	"connectrpc.com/connect"
 	"github.com/felixge/fgprof"
 	"github.com/go-kit/log"
-	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/kv/memberlist"
 	"github.com/grafana/dskit/middleware"
 	"github.com/grafana/dskit/server"
@@ -89,27 +88,7 @@ func New(cfg Config, s *server.Server, grpcGatewayMux *grpcgw.ServeMux, logger l
 //
 // Register Options allow to filter the HTTP methods and apply middlewares.
 func (a *API) RegisterRoute(path string, handler http.Handler, registerOpts ...RegisterOption) {
-	opts := applyRegisterOptions(registerOpts...)
-
-	level.Debug(a.logger).Log(append([]interface{}{
-		"msg", "api: registering route"}, opts.logFields(path)...)...)
-
-	// handle path prefixing
-	route := a.server.HTTP.Path(path)
-	if opts.isPrefix {
-		route = a.server.HTTP.PathPrefix(path)
-	}
-
-	// limit the route to the given methods
-	if len(opts.methods) > 0 {
-		route = route.Methods(opts.methods...)
-	}
-
-	for _, middleware := range opts.middlewares {
-		handler = middleware.Wrap(handler)
-	}
-
-	route.Handler(handler)
+	registerRoute(a.logger, a.server.HTTP, path, handler, registerOpts...)
 }
 
 // RegisterAPI registers the standard endpoints associated with a running Pyroscope.

pkg/api/register_options.go

@@ -1,8 +1,12 @@
 package api
 
 import (
+	"net/http"
 	"strings"
 
+	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
+	"github.com/gorilla/mux"
 	"github.com/grafana/dskit/middleware"
 
 	"github.com/grafana/pyroscope/pkg/util/gziphandler"
@@ -124,3 +128,29 @@ func (a *API) registerOptionsRingPage() []RegisterOption {
 		WithMethod("POST"),
 	}
 }
+
+func registerRoute(logger log.Logger, mux *mux.Router, path string, handler http.Handler, registerOpts ...RegisterOption) {
+	opts := applyRegisterOptions(registerOpts...)
+
+	level.Debug(logger).Log(append([]interface{}{
+		"msg", "api: registering route"}, opts.logFields(path)...)...)
+
+	// handle path prefixing
+	route := mux.Path(path)
+	if opts.isPrefix {
+		route = mux.PathPrefix(path)
+	}
+
+	// limit the route to the given methods
+	if len(opts.methods) > 0 {
+		route = route.Methods(opts.methods...)
+	}
+
+	// registering the middlewares in reverse order (similar like middleware.Merge)
+	for idx := range opts.middlewares {
+		mw := opts.middlewares[len(opts.middlewares)-idx-1]
+		handler = mw.Wrap(handler)
+	}
+
+	route.Handler(handler)
+}

pkg/api/register_options_test.go

@@ -0,0 +1,70 @@
+package api
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/go-kit/log"
+	"github.com/gorilla/mux"
+	"github.com/grafana/dskit/middleware"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type contextKey uint8
+
+const (
+	contextKeyTest contextKey = iota
+)
+
+func newTestMiddleware(name string) middleware.Interface {
+	return middleware.Func(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+			middlewares, ok := ctx.Value(contextKeyTest).([]string)
+			if !ok {
+				middlewares = []string{}
+			}
+			middlewares = append(middlewares, name)
+			ctx = context.WithValue(ctx, contextKeyTest, middlewares)
+			next.ServeHTTP(w, r.WithContext(ctx))
+		})
+	})
+
+}
+
+func Test_registerRoute(t *testing.T) {
+	router := mux.NewRouter()
+	registerRoute(
+		log.NewNopLogger(),
+		router,
+		"/test",
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			middlewares := r.Context().Value(contextKeyTest).([]string)
+			assert.Equal(t, []string{"outer", "middle", "inner"}, middlewares)
+
+			w.WriteHeader(http.StatusOK)
+		}),
+		func(r *registerParams) {
+			r.middlewares = append(r.middlewares, registerMiddleware{newTestMiddleware("outer"), "outer"})
+		},
+		func(r *registerParams) {
+			r.middlewares = append(r.middlewares, registerMiddleware{newTestMiddleware("middle"), "middle"})
+		},
+		func(r *registerParams) {
+			r.middlewares = append(r.middlewares, registerMiddleware{newTestMiddleware("inner"), "inner"})
+		},
+	)
+
+	testServer := httptest.NewServer(router)
+	defer testServer.Close()
+
+	req, err := http.NewRequest("GET", testServer.URL+"/test", nil)
+	require.NoError(t, err)
+
+	resp, err := testServer.Client().Do(req)
+	require.NoError(t, err)
+	assert.Equal(t, http.StatusOK, resp.StatusCode)
+}