From 0dfac9c3aa3401e11ef279321175aba118b2baf6 Mon Sep 17 00:00:00 2001
From: Serge Zaitsev <serge.zaitsev@grafana.com>
Date: Tue, 10 Aug 2021 09:03:22 +0200
Subject: [PATCH] Macaron: convert CSP middleware (#37672)

---
 pkg/api/http_server.go            |  2 +-
 pkg/middleware/csp.go             | 61 +++++++++++++++----------------
 pkg/middleware/middleware_test.go |  2 +-
 3 files changed, 32 insertions(+), 33 deletions(-)

diff --git a/pkg/api/http_server.go b/pkg/api/http_server.go
index 2fd0a608118..6dc9d5ef994 100644
--- a/pkg/api/http_server.go
+++ b/pkg/api/http_server.go
@@ -375,7 +375,7 @@ func (hs *HTTPServer) addMiddlewaresAndStaticRoutes() {
 	}
 
 	m.Use(middleware.HandleNoCacheHeader)
-	m.Use(middleware.AddCSPHeader(hs.Cfg, hs.log))
+	m.UseMiddleware(middleware.AddCSPHeader(hs.Cfg, hs.log))
 
 	for _, mw := range hs.middlewares {
 		m.Use(mw)
diff --git a/pkg/middleware/csp.go b/pkg/middleware/csp.go
index 3fff0cbe04f..2ea9614dfe0 100644
--- a/pkg/middleware/csp.go
+++ b/pkg/middleware/csp.go
@@ -10,45 +10,44 @@ import (
 	"strings"
 
 	"github.com/grafana/grafana/pkg/infra/log"
-	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/contexthandler"
 	"github.com/grafana/grafana/pkg/setting"
-	macaron "gopkg.in/macaron.v1"
 )
 
 // AddCSPHeader adds the Content Security Policy header.
-func AddCSPHeader(cfg *setting.Cfg, logger log.Logger) macaron.Handler {
-	return func(w http.ResponseWriter, req *http.Request, c *macaron.Context) {
-		if !cfg.CSPEnabled {
-			return
-		}
+func AddCSPHeader(cfg *setting.Cfg, logger log.Logger) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+			if !cfg.CSPEnabled {
+				next.ServeHTTP(rw, req)
+				return
+			}
 
-		logger.Debug("Adding CSP header to response", "cfg", fmt.Sprintf("%p", cfg))
+			logger.Debug("Adding CSP header to response", "cfg", fmt.Sprintf("%p", cfg))
 
-		ctx, ok := c.Data["ctx"].(*models.ReqContext)
-		if !ok {
-			panic("Failed to convert context into models.ReqContext")
-		}
+			ctx := contexthandler.FromContext(req.Context())
+			if cfg.CSPTemplate == "" {
+				logger.Debug("CSP template not configured, so returning 500")
+				ctx.JsonApiErr(500, "CSP template has to be configured", nil)
+				return
+			}
 
-		if cfg.CSPTemplate == "" {
-			logger.Debug("CSP template not configured, so returning 500")
-			ctx.JsonApiErr(500, "CSP template has to be configured", nil)
-			return
-		}
+			var buf [16]byte
+			if _, err := io.ReadFull(rand.Reader, buf[:]); err != nil {
+				logger.Error("Failed to generate CSP nonce", "err", err)
+				ctx.JsonApiErr(500, "Failed to generate CSP nonce", err)
+			}
 
-		var buf [16]byte
-		if _, err := io.ReadFull(rand.Reader, buf[:]); err != nil {
-			logger.Error("Failed to generate CSP nonce", "err", err)
-			ctx.JsonApiErr(500, "Failed to generate CSP nonce", err)
-		}
+			nonce := base64.RawStdEncoding.EncodeToString(buf[:])
+			val := strings.ReplaceAll(cfg.CSPTemplate, "$NONCE", fmt.Sprintf("'nonce-%s'", nonce))
 
-		nonce := base64.RawStdEncoding.EncodeToString(buf[:])
-		val := strings.ReplaceAll(cfg.CSPTemplate, "$NONCE", fmt.Sprintf("'nonce-%s'", nonce))
-
-		re := regexp.MustCompile(`^\w+:(//)?`)
-		rootPath := re.ReplaceAllString(cfg.AppURL, "")
-		val = strings.ReplaceAll(val, "$ROOT_PATH", rootPath)
-		w.Header().Set("Content-Security-Policy", val)
-		ctx.RequestNonce = nonce
-		logger.Debug("Successfully generated CSP nonce", "nonce", nonce)
+			re := regexp.MustCompile(`^\w+:(//)?`)
+			rootPath := re.ReplaceAllString(cfg.AppURL, "")
+			val = strings.ReplaceAll(val, "$ROOT_PATH", rootPath)
+			rw.Header().Set("Content-Security-Policy", val)
+			ctx.RequestNonce = nonce
+			logger.Debug("Successfully generated CSP nonce", "nonce", nonce)
+			next.ServeHTTP(rw, req)
+		})
 	}
 }
diff --git a/pkg/middleware/middleware_test.go b/pkg/middleware/middleware_test.go
index 30ea0b1aff1..95defe3376c 100644
--- a/pkg/middleware/middleware_test.go
+++ b/pkg/middleware/middleware_test.go
@@ -632,7 +632,7 @@ func middlewareScenario(t *testing.T, desc string, fn scenarioFunc, cbs ...func(
 
 		sc.m = macaron.New()
 		sc.m.Use(AddDefaultResponseHeaders(cfg))
-		sc.m.Use(AddCSPHeader(cfg, logger))
+		sc.m.UseMiddleware(AddCSPHeader(cfg, logger))
 		sc.m.Use(macaron.Renderer(macaron.RenderOptions{
 			Directory: viewsPath,
 			Delims:    macaron.Delims{Left: "[[", Right: "]]"},