From 827a5b702128edfe0b7a98c74acce929a660f753 Mon Sep 17 00:00:00 2001
From: Serge Zaitsev <serge.zaitsev@grafana.com>
Date: Mon, 26 Jul 2021 06:11:43 +0200
Subject: [PATCH] add FromContext API for macaron and reqcontext, add API for
 standard middleware in macaron (#37057)

---
 pkg/macaron/macaron.go                        | 39 ++++++++++++++++++-
 pkg/services/contexthandler/contexthandler.go | 14 +++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/pkg/macaron/macaron.go b/pkg/macaron/macaron.go
index c3957cc47cd..5ae6c5def6b 100644
--- a/pkg/macaron/macaron.go
+++ b/pkg/macaron/macaron.go
@@ -18,6 +18,7 @@
 package macaron // import "gopkg.in/macaron.v1"
 
 import (
+	"context"
 	"log"
 	"net/http"
 	"os"
@@ -151,6 +152,41 @@ func (m *Macaron) Handlers(handlers ...Handler) {
 // Macaron stops future process when it returns true.
 type BeforeHandler func(rw http.ResponseWriter, req *http.Request) bool
 
+// macaronContextKey is used to store/fetch macaron.Context inside context.Context
+type macaronContextKey struct{}
+
+// FromContext returns the macaron context stored in a context.Context, if any.
+func FromContext(c context.Context) *Context {
+	if mc, ok := c.Value(macaronContextKey{}).(*Context); ok {
+		return mc
+	}
+	return nil
+}
+
+// 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 macaron.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.Request = req
+		if mrw, ok := rw.(*responseWriter); ok {
+			c.Resp = mrw
+		} else {
+			c.Resp = NewResponseWriter(req.Method, rw)
+		}
+		c.Map(req)
+		c.MapTo(rw, (*http.ResponseWriter)(nil))
+		c.Next()
+	})
+	m.handlers = append(m.handlers, Handler(middleware(next)))
+}
+
 // 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.
@@ -165,15 +201,16 @@ func (m *Macaron) createContext(rw http.ResponseWriter, req *http.Request) *Cont
 		handlers: m.handlers,
 		index:    0,
 		Router:   m.Router,
-		Req:      Request{req},
 		Resp:     NewResponseWriter(req.Method, rw),
 		Render:   &DummyRender{rw},
 		Data:     make(map[string]interface{}),
 	}
+	req = req.WithContext(context.WithValue(req.Context(), macaronContextKey{}, c))
 	c.SetParent(m)
 	c.Map(c)
 	c.MapTo(c.Resp, (*http.ResponseWriter)(nil))
 	c.Map(req)
+	c.Req = Request{req}
 	return c
 }
 
diff --git a/pkg/services/contexthandler/contexthandler.go b/pkg/services/contexthandler/contexthandler.go
index 33628a7ab5e..236f122226f 100644
--- a/pkg/services/contexthandler/contexthandler.go
+++ b/pkg/services/contexthandler/contexthandler.go
@@ -62,6 +62,16 @@ func (h *ContextHandler) Init() error {
 	return nil
 }
 
+type reqContextKey struct{}
+
+// FromContext returns the ReqContext value stored in a context.Context, if any.
+func FromContext(c context.Context) *models.ReqContext {
+	if reqCtx, ok := c.Value(reqContextKey{}).(*models.ReqContext); ok {
+		return reqCtx
+	}
+	return nil
+}
+
 // Middleware provides a middleware to initialize the Macaron context.
 func (h *ContextHandler) Middleware(mContext *macaron.Context) {
 	span, _ := opentracing.StartSpanFromContext(mContext.Req.Context(), "Auth - Middleware")
@@ -76,6 +86,10 @@ func (h *ContextHandler) Middleware(mContext *macaron.Context) {
 		Logger:         log.New("context"),
 	}
 
+	// Inject ReqContext into a request context and replace the request instance in the macaron context
+	mContext.Req.Request = mContext.Req.WithContext(context.WithValue(mContext.Req.Context(), reqContextKey{}, reqContext))
+	mContext.Map(mContext.Req.Request)
+
 	traceID, exists := cw.ExtractTraceID(mContext.Req.Request.Context())
 	if exists {
 		reqContext.Logger = reqContext.Logger.New("traceID", traceID)