From 1196b4a60992585bbbb8cb82e52dbc56ecbd81cc Mon Sep 17 00:00:00 2001
From: Marcus Efraimsson <marcus.efraimsson@gmail.com>
Date: Tue, 31 May 2022 17:02:58 +0200
Subject: [PATCH] Datasource: Fix allowed cookies to be forwarded as header to
 backend datasources (#49541)

Co-authored-by: Will Browne <wbrowne@users.noreply.github.com>
---
 pkg/api/dtos/models.go           |  3 +++
 pkg/api/metrics.go               |  2 ++
 pkg/services/query/query.go      | 20 +++++++++++++++++
 pkg/services/query/query_test.go | 37 ++++++++++++++++++++++++++++++++
 4 files changed, 62 insertions(+)

diff --git a/pkg/api/dtos/models.go b/pkg/api/dtos/models.go
index f0ceb354802..dbed74c01b7 100644
--- a/pkg/api/dtos/models.go
+++ b/pkg/api/dtos/models.go
@@ -3,6 +3,7 @@ package dtos
 import (
 	"crypto/md5"
 	"fmt"
+	"net/http"
 	"regexp"
 	"strings"
 
@@ -68,6 +69,8 @@ type MetricRequest struct {
 	Queries []*simplejson.Json `json:"queries"`
 	// required: false
 	Debug bool `json:"debug"`
+
+	HTTPRequest *http.Request `json:"-"`
 }
 
 func GetGravatarUrl(text string) string {
diff --git a/pkg/api/metrics.go b/pkg/api/metrics.go
index 3b6649ba7c8..46929748a92 100644
--- a/pkg/api/metrics.go
+++ b/pkg/api/metrics.go
@@ -44,6 +44,8 @@ func (hs *HTTPServer) QueryMetricsV2(c *models.ReqContext) response.Response {
 		return response.Error(http.StatusBadRequest, "bad request data", err)
 	}
 
+	reqDTO.HTTPRequest = c.Req
+
 	resp, err := hs.queryDataService.QueryData(c.Req.Context(), c.SignedInUser, c.SkipCache, reqDTO, true)
 	if err != nil {
 		return hs.handleQueryMetricsError(err)
diff --git a/pkg/services/query/query.go b/pkg/services/query/query.go
index b449648ca22..2bd50694b59 100644
--- a/pkg/services/query/query.go
+++ b/pkg/services/query/query.go
@@ -3,6 +3,7 @@ package query
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"strings"
 	"time"
 
@@ -18,6 +19,7 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/tsdb/grafanads"
 	"github.com/grafana/grafana/pkg/tsdb/legacydata"
+	"github.com/grafana/grafana/pkg/util/proxyutil"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 )
@@ -149,6 +151,19 @@ func (s *Service) handleQueryData(ctx context.Context, user *models.SignedInUser
 		req.Headers[k] = v
 	}
 
+	if parsedReq.httpRequest != nil && parsedReq.httpRequest.Header.Get("Cookie") != "" && ds.JsonData != nil {
+		keepCookieNames := []string{}
+
+		if keepCookies := ds.JsonData.Get("keepCookies"); keepCookies != nil {
+			keepCookieNames = keepCookies.MustStringArray()
+		}
+
+		proxyutil.ClearCookieHeader(parsedReq.httpRequest, keepCookieNames)
+		if cookieStr := parsedReq.httpRequest.Header.Get("Cookie"); cookieStr != "" {
+			req.Headers["Cookie"] = cookieStr
+		}
+	}
+
 	for _, q := range parsedReq.parsedQueries {
 		req.Queries = append(req.Queries, q.query)
 	}
@@ -164,6 +179,7 @@ type parsedQuery struct {
 type parsedRequest struct {
 	hasExpression bool
 	parsedQueries []parsedQuery
+	httpRequest   *http.Request
 }
 
 func customHeaders(jsonData *simplejson.Json, decryptedJsonData map[string]string) map[string]string {
@@ -243,6 +259,10 @@ func (s *Service) parseMetricRequest(ctx context.Context, user *models.SignedInU
 		}
 	}
 
+	if reqDTO.HTTPRequest != nil {
+		req.httpRequest = reqDTO.HTTPRequest
+	}
+
 	return req, nil
 }
 
diff --git a/pkg/services/query/query_test.go b/pkg/services/query/query_test.go
index d4f1d7b87ad..27714d41316 100644
--- a/pkg/services/query/query_test.go
+++ b/pkg/services/query/query_test.go
@@ -61,6 +61,43 @@ func TestQueryData(t *testing.T) {
 		}
 		require.Equal(t, expected, tc.pluginContext.req.Headers)
 	})
+
+	t.Run("it doesn't add cookie header to the request when keepCookies configured and no cookies provided", func(t *testing.T) {
+		tc := setup(t)
+		json, err := simplejson.NewJson([]byte(`{"keepCookies": [ "foo", "bar" ]}`))
+		require.NoError(t, err)
+		tc.dataSourceCache.ds.JsonData = json
+
+		metricReq := metricRequest()
+		httpReq, err := http.NewRequest(http.MethodGet, "/", nil)
+		require.NoError(t, err)
+		metricReq.HTTPRequest = httpReq
+		_, err = tc.queryService.QueryData(context.Background(), nil, true, metricReq, false)
+		require.NoError(t, err)
+
+		require.Empty(t, tc.pluginContext.req.Headers)
+	})
+
+	t.Run("it adds cookie header to the request when keepCookies configured and cookie provided", func(t *testing.T) {
+		tc := setup(t)
+		json, err := simplejson.NewJson([]byte(`{"keepCookies": [ "foo", "bar" ]}`))
+		require.NoError(t, err)
+		tc.dataSourceCache.ds.JsonData = json
+
+		metricReq := metricRequest()
+		httpReq, err := http.NewRequest(http.MethodGet, "/", nil)
+		require.NoError(t, err)
+		httpReq.AddCookie(&http.Cookie{Name: "a"})
+		httpReq.AddCookie(&http.Cookie{Name: "bar", Value: "rab"})
+		httpReq.AddCookie(&http.Cookie{Name: "b"})
+		httpReq.AddCookie(&http.Cookie{Name: "foo", Value: "oof"})
+		httpReq.AddCookie(&http.Cookie{Name: "c"})
+		metricReq.HTTPRequest = httpReq
+		_, err = tc.queryService.QueryData(context.Background(), nil, true, metricReq, false)
+		require.NoError(t, err)
+
+		require.Equal(t, map[string]string{"Cookie": "bar=rab; foo=oof"}, tc.pluginContext.req.Headers)
+	})
 }
 
 func setup(t *testing.T) *testContext {