pkg/web: closure-style middlewares (#51238)

* pkg/web: closure-style middlewares Switches the middleware execution model from web.Handlers in a slice to web.Middleware. Middlewares are temporarily kept in a slice to preserve ordering, but prior to execution they are applied, forming a giant call-stack, giving granular control over the execution flow. * pkg/middleware: adapt to web.Middleware * pkg/middleware/recovery: use c.Req over req c.Req gets updated by future handlers, while req stays static. The current recovery implementation needs this newer information * pkg/web: correct middleware ordering * pkg/webtest: adapt middleware * pkg/web/hack: set w and r onto web.Context By adopting std middlewares, it may happen they invoke next(w,r) without putting their modified w,r into the web.Context, leading old-style handlers to operate on outdated fields. pkg/web now takes care of this * pkg/middleware: selectively use future context * pkg/web: accept closure-style on Use() * webtest: Middleware testing adds a utility function to web/webtest to obtain a http.ResponseWriter, http.Request and http.Handler the same as a middleware that runs would receive * *: cleanup * pkg/web: don't wrap Middleware from Router * pkg/web: require chain to write response * *: remove temp files * webtest: don't require chain write * *: cleanup
2022-08-09 14:58:50 +02:00
parent 3893c46976
commit 534ece064b
17 changed files with 357 additions and 264 deletions
@@ -29,8 +29,7 @@ import (
 // Context represents the runtime context of current request of Macaron instance.
 // It is the integration of most frequently used middlewares and helper methods.
 type Context struct {
-	handlers []http.Handler
-	index    int
+	mws []Middleware

 	*Router
 	Req      *http.Request
@@ -39,30 +38,20 @@ type Context struct {
 	logger   log.Logger
 }

-func (ctx *Context) handler() http.Handler {
-	if ctx.index < len(ctx.handlers) {
-		return ctx.handlers[ctx.index]
-	}
-	if ctx.index == len(ctx.handlers) {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
-	}
-	panic("invalid index for context handler")
-}
-
-// Next runs the next handler in the context chain
-func (ctx *Context) Next() {
-	ctx.index++
-	ctx.run()
-}
-
 func (ctx *Context) run() {
-	for ctx.index <= len(ctx.handlers) {
-		ctx.handler().ServeHTTP(ctx.Resp, ctx.Req)
+	h := http.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
+	for i := len(ctx.mws) - 1; i >= 0; i-- {
+		h = ctx.mws[i](h)
+	}

-		ctx.index++
-		if ctx.Resp.Written() {
-			return
-		}
+	rw := ctx.Resp
+	h.ServeHTTP(ctx.Resp, ctx.Req)
+
+	// Prevent the handler chain from not writing anything.
+	// This indicates nearly always that a middleware is misbehaving and not calling its next.ServeHTTP().
+	// In rare cases where a blank http.StatusOK without any body is wished, explicitly state that using w.WriteStatus(http.StatusOK)
+	if !rw.Written() {
+		panic("chain did not write HTTP response")
 	}
 }

@@ -53,28 +53,16 @@ type Handler interface{}
 //go:linkname hack_wrap github.com/grafana/grafana/pkg/api/response.wrap_handler
 func hack_wrap(Handler) http.HandlerFunc

-// validateAndWrapHandler makes sure a handler is a callable function, it panics if not.
-// When the handler is also potential to be any built-in inject.FastInvoker,
-// it wraps the handler automatically to have some performance gain.
-func validateAndWrapHandler(h Handler) http.Handler {
+// wrapHandler turns any supported handler type into a http.Handler by wrapping it accordingly
+func wrapHandler(h Handler) http.Handler {
 	return hack_wrap(h)
 }

-// validateAndWrapHandlers preforms validation and wrapping for each input handler.
-// It accepts an optional wrapper function to perform custom wrapping on handlers.
-func validateAndWrapHandlers(handlers []Handler) []http.Handler {
-	wrappedHandlers := make([]http.Handler, len(handlers))
-	for i, h := range handlers {
-		wrappedHandlers[i] = validateAndWrapHandler(h)
-	}
-
-	return wrappedHandlers
-}
-
 // Macaron represents the top level web application.
 // Injector methods can be invoked to map services on a global level.
 type Macaron struct {
-	handlers []http.Handler
+	// handlers    []http.Handler
+	mws []Middleware

 	urlPrefix string // For suburl support.
 	*Router
@@ -119,43 +107,50 @@ func SetURLParams(r *http.Request, vars map[string]string) *http.Request {
 	return r.WithContext(context.WithValue(r.Context(), paramsKey{}, vars))
 }

-// UseMiddleware is a traditional approach to writing middleware in Go.
-// A middleware is a function that has a reference to the next handler in the chain
-// and returns the actual middleware handler, that may do its job and optionally
-// call next.
-// Due to how Macaron handles/injects requests and responses we patch the web.Context
-// to use the new ResponseWriter and http.Request here. The caller may only call
-// `next.ServeHTTP(rw, req)` to pass a modified response writer and/or a request to the
-// further middlewares in the chain.
-func (m *Macaron) UseMiddleware(middleware func(http.Handler) http.Handler) {
-	next := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
-		c := FromContext(req.Context())
-		c.Req = req
-		if mrw, ok := rw.(*responseWriter); ok {
-			c.Resp = mrw
-		} else {
-			c.Resp = NewResponseWriter(req.Method, rw)
-		}
-		c.Next()
-	})
-	m.handlers = append(m.handlers, middleware(next))
+type Middleware = func(next http.Handler) http.Handler
+
+// UseMiddleware registers the given Middleware
+func (m *Macaron) UseMiddleware(mw Middleware) {
+	m.mws = append(m.mws, mw)
 }

-// Use adds a middleware Handler to the stack,
-// and panics if the handler is not a callable func.
-// Middleware Handlers are invoked in the order that they are added.
-func (m *Macaron) Use(handler Handler) {
-	h := validateAndWrapHandler(handler)
-	m.handlers = append(m.handlers, h)
+// Use registers the provided Handler as a middleware.
+// The argument may be any supported handler or the Middleware type
+// Deprecated: use UseMiddleware instead
+func (m *Macaron) Use(h Handler) {
+	m.mws = append(m.mws, mwFromHandler(h))
+}
+
+func mwFromHandler(handler Handler) Middleware {
+	if mw, ok := handler.(Middleware); ok {
+		return mw
+	}
+
+	h := wrapHandler(handler)
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			mrw, ok := w.(*responseWriter)
+			if !ok {
+				mrw = NewResponseWriter(r.Method, w).(*responseWriter)
+			}
+
+			h.ServeHTTP(mrw, r)
+			if mrw.Written() {
+				return
+			}
+
+			ctx := r.Context().Value(macaronContextKey{}).(*Context)
+			next.ServeHTTP(ctx.Resp, ctx.Req)
+		})
+	}
 }

 func (m *Macaron) createContext(rw http.ResponseWriter, req *http.Request) *Context {
 	c := &Context{
-		handlers: m.handlers,
-		index:    0,
-		Router:   m.Router,
-		Resp:     NewResponseWriter(req.Method, rw),
-		logger:   log.New("macaron.context"),
+		mws:    m.mws,
+		Router: m.Router,
+		Resp:   NewResponseWriter(req.Method, rw),
+		logger: log.New("macaron.context"),
 	}

 	c.Req = req.WithContext(context.WithValue(req.Context(), macaronContextKey{}, c))
@@ -26,7 +26,7 @@ import (
 // Renderer is a Middleware that injects a template renderer into the macaron context, enabling ctx.HTML calls in the handlers.
 // If MACARON_ENV is set to "development" then templates will be recompiled on every request. For more performance, set the
 // MACARON_ENV environment variable to "production".
-func Renderer(dir, leftDelim, rightDelim string) func(http.Handler) http.Handler {
+func Renderer(dir, leftDelim, rightDelim string) Middleware {
 	fs := os.DirFS(dir)
 	t, err := compileTemplates(fs, leftDelim, rightDelim)
 	if err != nil {
@@ -46,6 +46,16 @@ func NewResponseWriter(method string, rw http.ResponseWriter) ResponseWriter {
 	return &responseWriter{method, rw, 0, 0, nil}
 }

+// Rw returns a ResponseWriter. If the argument already satisfies the interface,
+// it is returned as is, otherwise it is wrapped using NewResponseWriter
+func Rw(rw http.ResponseWriter, req *http.Request) ResponseWriter {
+	if mrw, ok := rw.(ResponseWriter); ok {
+		return mrw
+	}
+
+	return NewResponseWriter(req.Method, rw)
+}
+
 type responseWriter struct {
 	method string
 	http.ResponseWriter
@@ -146,13 +146,12 @@ func (r *Router) Handle(method string, pattern string, handlers []Handler) {
 		h = append(h, handlers...)
 		handlers = h
 	}
-	httpHandlers := validateAndWrapHandlers(handlers)

 	r.handle(method, pattern, func(resp http.ResponseWriter, req *http.Request, params map[string]string) {
 		c := r.m.createContext(resp, SetURLParams(req, params))
-		c.handlers = make([]http.Handler, 0, len(r.m.handlers)+len(handlers))
-		c.handlers = append(c.handlers, r.m.handlers...)
-		c.handlers = append(c.handlers, httpHandlers...)
+		for _, h := range handlers {
+			c.mws = append(c.mws, mwFromHandler(h))
+		}
 		c.run()
 	})
 }
@@ -194,12 +193,11 @@ func (r *Router) Any(pattern string, h ...Handler) { r.Handle("*", pattern, h) }
 // found. If it is not set, http.NotFound is used.
 // Be sure to set 404 response code in your handler.
 func (r *Router) NotFound(handlers ...Handler) {
-	httpHandlers := validateAndWrapHandlers(handlers)
 	r.notFound = func(rw http.ResponseWriter, req *http.Request) {
 		c := r.m.createContext(rw, req)
-		c.handlers = make([]http.Handler, 0, len(r.m.handlers)+len(handlers))
-		c.handlers = append(c.handlers, r.m.handlers...)
-		c.handlers = append(c.handlers, httpHandlers...)
+		for _, h := range handlers {
+			c.mws = append(c.mws, mwFromHandler(h))
+		}
 		c.run()
 	}
 }
@@ -0,0 +1,73 @@
+package webtest
+
+import (
+	"net/http"
+	"net/http/httptest"
+
+	"github.com/grafana/grafana/pkg/web"
+)
+
+type Context struct {
+	Req  *http.Request
+	Rw   http.ResponseWriter
+	Next http.Handler
+}
+
+// Middleware is a utility for testing middlewares
+type Middleware struct {
+	// Before are run ahead of the returned context
+	Before []web.Handler
+	// After are part of the http.Handler chain
+	After []web.Handler
+	// The actual handler at the end of the chain
+	Handler web.Handler
+}
+
+// MiddlewareContext returns a *http.Request, http.ResponseWriter and http.Handler
+// exactly as if it was passed to a middleware
+func MiddlewareContext(test Middleware, req *http.Request) *Context {
+	m := web.New()
+
+	// pkg/web requires the chain to write an HTTP response.
+	// While this ensures a basic amount of correctness for real handler chains,
+	// it is naturally incompatible with this package, as we terminate the chain early to pass its
+	// state to the surrounding test.
+	// By replacing the http.ResponseWriter and writing to the old one we make pkg/web happy.
+	m.Use(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusOK)
+
+			rw := web.Rw(httptest.NewRecorder(), r)
+			next.ServeHTTP(rw, r)
+		})
+	})
+
+	for _, mw := range test.Before {
+		m.Use(mw)
+	}
+
+	ch := make(chan *Context)
+	m.Use(func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ch <- &Context{
+				Req:  r,
+				Rw:   w,
+				Next: next,
+			}
+		})
+	})
+
+	for _, mw := range test.After {
+		m.Use(mw)
+	}
+
+	// set the provided (or noop) handler to exactly the queried path
+	handler := test.Handler
+	if handler == nil {
+		handler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {})
+	}
+	m.Handle(req.Method, req.URL.RequestURI(), []web.Handler{handler})
+	go m.ServeHTTP(httptest.NewRecorder(), req)
+
+	return <-ch
+}
@@ -37,7 +37,7 @@ func NewServer(t testing.TB, routeRegister routing.RouteRegister) *Server {
 		c.Req = c.Req.WithContext(ctxkey.Set(c.Req.Context(), initCtx))
 	})

-	m.Use(requestContextMiddleware())
+	m.UseMiddleware(requestContextMiddleware())

 	routeRegister.Register(m.Router)
 	testServer := httptest.NewServer(m)
@@ -126,24 +126,25 @@ func requestContextFromRequest(req *http.Request) *models.ReqContext {
 	return val
 }

-func requestContextMiddleware() web.Handler {
-	return func(res http.ResponseWriter, req *http.Request) {
-		c := ctxkey.Get(req.Context()).(*models.ReqContext)
+func requestContextMiddleware() web.Middleware {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			c := ctxkey.Get(r.Context()).(*models.ReqContext)

-		ctx := requestContextFromRequest(req)
-		if ctx == nil {
-			c.Next()
-			return
-		}
+			ctx := requestContextFromRequest(r)
+			if ctx != nil {
+				c.SignedInUser = ctx.SignedInUser
+				c.UserToken = ctx.UserToken
+				c.IsSignedIn = ctx.IsSignedIn
+				c.IsRenderCall = ctx.IsRenderCall
+				c.AllowAnonymous = ctx.AllowAnonymous
+				c.SkipCache = ctx.SkipCache
+				c.RequestNonce = ctx.RequestNonce
+				c.PerfmonTimer = ctx.PerfmonTimer
+				c.LookupTokenErr = ctx.LookupTokenErr
+			}

-		c.SignedInUser = ctx.SignedInUser
-		c.UserToken = ctx.UserToken
-		c.IsSignedIn = ctx.IsSignedIn
-		c.IsRenderCall = ctx.IsRenderCall
-		c.AllowAnonymous = ctx.AllowAnonymous
-		c.SkipCache = ctx.SkipCache
-		c.RequestNonce = ctx.RequestNonce
-		c.PerfmonTimer = ctx.PerfmonTimer
-		c.LookupTokenErr = ctx.LookupTokenErr
+			next.ServeHTTP(w, r)
+		})
 	}
 }