Chore: Clean up the context handler (#75864)

* Chore: Clean up the context handler

Co-authored-by: Kalle Persson <kalle.persson@grafana.com>

* Better comment

Co-authored-by: Kalle Persson <kalle.persson@grafana.com>

---------

Co-authored-by: Kalle Persson <kalle.persson@grafana.com>

pkg/services/contexthandler/contexthandler.go

@@ -83,13 +83,11 @@ func CopyWithReqContext(ctx context.Context) context.Context {
 // Middleware provides a middleware to initialize the request context.
 func (h *ContextHandler) Middleware(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		ctx := r.Context()
-		mContext := web.FromContext(ctx)
-		_, span := h.tracer.Start(ctx, "Auth - Middleware")
-		defer span.End()
+		ctx, span := h.tracer.Start(r.Context(), "Auth - Middleware")
+		defer span.End() // this will span to next handlers as well
 
 		reqContext := &contextmodel.ReqContext{
-			Context: mContext,
+			Context: web.FromContext(ctx), // Extract web context from context (no knowledge of the trace)
 			SignedInUser: &user.SignedInUser{
 				Permissions: map[int64]map[string][]string{},
 			},
@@ -99,17 +97,20 @@ func (h *ContextHandler) Middleware(next http.Handler) http.Handler {
 			Logger:         log.New("context"),
 		}
 
-		// Inject ReqContext into http.Request.Context
-		*r = *r.WithContext(context.WithValue(ctx, reqContextKey{}, reqContext))
+		// inject ReqContext in the context
+		ctx = context.WithValue(ctx, reqContextKey{}, reqContext)
 		// store list of possible auth header in context
-		*reqContext.Req = *reqContext.Req.WithContext(WithAuthHTTPHeaders(reqContext.Req.Context(), h.Cfg))
+		ctx = WithAuthHTTPHeaders(ctx, h.Cfg)
+		// Set the context for the http.Request.Context
+		// This modifies both r and reqContext.Req since they point to the same value
+		*reqContext.Req = *reqContext.Req.WithContext(ctx)
 
-		traceID := tracing.TraceIDFromContext(mContext.Req.Context(), false)
+		traceID := tracing.TraceIDFromContext(reqContext.Req.Context(), false)
 		if traceID != "" {
 			reqContext.Logger = reqContext.Logger.New("traceID", traceID)
 		}
 
-		identity, err := h.authnService.Authenticate(ctx, &authn.Request{HTTPRequest: reqContext.Req, Resp: reqContext.Resp})
+		identity, err := h.authnService.Authenticate(reqContext.Req.Context(), &authn.Request{HTTPRequest: reqContext.Req, Resp: reqContext.Resp})
 		if err != nil {
 			if errors.Is(err, auth.ErrInvalidSessionToken) || errors.Is(err, authn.ErrExpiredAccessToken) {
 				// Burn the cookie in case of invalid, expired or missing token