[v9.3.x] SSE: Make sure to forward headers, user and cookies/OAuth token (#59390)

SSE: Make sure to forward headers, user and cookies/OAuth token (#58897)

Fixes #58793 and Fixes https://github.com/grafana/azure-data-explorer-datasource/issues/513

Co-authored-by: Marcus Efraimsson <marcus.efraimsson@gmail.com>
(cherry picked from commit 5623b5afaf)

Co-authored-by: Kyle Brandt <kyle@grafana.com>

pkg/services/query/query_parsing.go

@@ -0,0 +1,157 @@
+package query
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
+	"github.com/grafana/grafana/pkg/components/simplejson"
+	"github.com/grafana/grafana/pkg/expr"
+	"github.com/grafana/grafana/pkg/infra/httpclient/httpclientprovider"
+	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/oauthtoken"
+	"github.com/grafana/grafana/pkg/services/user"
+	"github.com/grafana/grafana/pkg/util/proxyutil"
+	"golang.org/x/oauth2"
+)
+
+type parsedQuery struct {
+	datasource *datasources.DataSource
+	query      backend.DataQuery
+	rawQuery   *simplejson.Json
+}
+
+type parsedRequest struct {
+	hasExpression bool
+	parsedQueries map[string][]parsedQuery
+	dsTypes       map[string]bool
+	httpRequest   *http.Request
+}
+
+func (pr parsedRequest) getFlattenedQueries() []parsedQuery {
+	queries := make([]parsedQuery, 0)
+	for _, pq := range pr.parsedQueries {
+		queries = append(queries, pq...)
+	}
+	return queries
+}
+
+func (pr parsedRequest) validateRequest() error {
+	if pr.httpRequest == nil {
+		return nil
+	}
+
+	if pr.hasExpression {
+		hasExpr := pr.httpRequest.URL.Query().Get("expression")
+		if hasExpr == "" || hasExpr == "true" {
+			return nil
+		}
+		return ErrQueryParamMismatch
+	}
+
+	vals := splitHeaders(pr.httpRequest.Header.Values(HeaderDatasourceUID))
+	count := len(vals)
+	if count > 0 { // header exists
+		if count != len(pr.parsedQueries) {
+			return ErrQueryParamMismatch
+		}
+		for _, t := range vals {
+			if pr.parsedQueries[t] == nil {
+				return ErrQueryParamMismatch
+			}
+		}
+	}
+
+	vals = splitHeaders(pr.httpRequest.Header.Values(HeaderPluginID))
+	count = len(vals)
+	if count > 0 { // header exists
+		if count != len(pr.dsTypes) {
+			return ErrQueryParamMismatch
+		}
+		for _, t := range vals {
+			if !pr.dsTypes[t] {
+				return ErrQueryParamMismatch
+			}
+		}
+	}
+	return nil
+}
+
+func (pr parsedRequest) createDataSourceQueryEnrichers(ctx context.Context, signedInUser *user.SignedInUser, oAuthTokenService oauthtoken.OAuthTokenService, disallowedCookies []string) map[string]expr.QueryDataRequestEnricher {
+	datasourcesHeaderProvider := map[string]expr.QueryDataRequestEnricher{}
+
+	if pr.httpRequest == nil {
+		return datasourcesHeaderProvider
+	}
+
+	for uid, queries := range pr.parsedQueries {
+		if expr.IsDataSource(uid) {
+			continue
+		}
+
+		if len(queries) == 0 || queries[0].datasource == nil {
+			continue
+		}
+
+		if _, exists := datasourcesHeaderProvider[uid]; exists {
+			continue
+		}
+
+		ds := queries[0].datasource
+		allowedCookies := ds.AllowedCookies()
+		clonedReq := pr.httpRequest.Clone(pr.httpRequest.Context())
+
+		var token *oauth2.Token
+
+		if oAuthTokenService.IsOAuthPassThruEnabled(ds) {
+			token = oAuthTokenService.GetCurrentOAuthToken(ctx, signedInUser)
+		}
+
+		datasourcesHeaderProvider[uid] = func(ctx context.Context, req *backend.QueryDataRequest) context.Context {
+			if len(req.Headers) == 0 {
+				req.Headers = map[string]string{}
+			}
+
+			if len(allowedCookies) > 0 {
+				proxyutil.ClearCookieHeader(clonedReq, allowedCookies, disallowedCookies)
+				if cookieStr := clonedReq.Header.Get("Cookie"); cookieStr != "" {
+					req.Headers["Cookie"] = cookieStr
+				}
+
+				ctx = httpclient.WithContextualMiddleware(ctx, httpclientprovider.ForwardedCookiesMiddleware(clonedReq.Cookies(), allowedCookies, disallowedCookies))
+			}
+
+			if token != nil {
+				req.Headers["Authorization"] = fmt.Sprintf("%s %s", token.Type(), token.AccessToken)
+
+				idToken, ok := token.Extra("id_token").(string)
+				if ok && idToken != "" {
+					req.Headers["X-ID-Token"] = idToken
+				}
+
+				ctx = httpclient.WithContextualMiddleware(ctx, httpclientprovider.ForwardedOAuthIdentityMiddleware(token))
+			}
+
+			return ctx
+		}
+	}
+
+	return datasourcesHeaderProvider
+}
+
+func splitHeaders(headers []string) []string {
+	out := []string{}
+	for _, v := range headers {
+		if strings.Contains(v, ",") {
+			for _, sub := range strings.Split(v, ",") {
+				out = append(out, strings.TrimSpace(sub))
+			}
+		} else {
+			out = append(out, v)
+		}
+	}
+	return out
+}