From 4fee0645a1d45683b01338b70d65c73ff7654cdd Mon Sep 17 00:00:00 2001
From: "Grot (@grafanabot)" <43478413+grafanabot@users.noreply.github.com>
Date: Mon, 7 Nov 2022 17:11:10 +0100
Subject: [PATCH] Datasource Loki: preserve header `X-ID-Token` (#57878)
 (#58346)

(cherry picked from commit bc280d07492f66456e47eb262f935c933312b1ed)

Co-authored-by: Si Mon <85333972+siiimooon@users.noreply.github.com>
---
 pkg/tsdb/loki/auth_test.go | 12 +++++++++++-
 pkg/tsdb/loki/loki.go      |  4 ++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/pkg/tsdb/loki/auth_test.go b/pkg/tsdb/loki/auth_test.go
index da8d95c1758..c694a6dc846 100644
--- a/pkg/tsdb/loki/auth_test.go
+++ b/pkg/tsdb/loki/auth_test.go
@@ -52,7 +52,7 @@ func TestOauthForwardIdentity(t *testing.T) {
 		auth   bool
 		cookie bool
 	}{
-		{name: "when auth header exists => add auth header", auth: true, cookie: false},
+		{name: "when auth headers exist => add auth headers", auth: true, cookie: false},
 		{name: "when cookie header exists => add cookie header", auth: false, cookie: true},
 		{name: "when cookie&auth headers exist => add cookie&auth headers", auth: true, cookie: true},
 		{name: "when no header exists => do not add headers", auth: false, cookie: false},
@@ -62,6 +62,8 @@ func TestOauthForwardIdentity(t *testing.T) {
 	authValue := "auth"
 	cookieName := "Cookie"
 	cookieValue := "a=1"
+	idTokenName := "X-ID-Token"
+	idTokenValue := "idtoken"
 
 	for _, test := range tt {
 		t.Run("QueryData: "+test.name, func(t *testing.T) {
@@ -90,10 +92,13 @@ func TestOauthForwardIdentity(t *testing.T) {
 				// as an array
 				authValues := req.Header.Values(authName)
 				cookieValues := req.Header.Values(cookieName)
+				idTokenValues := req.Header.Values(idTokenName)
 				if test.auth {
 					require.Equal(t, []string{authValue}, authValues)
+					require.Equal(t, []string{idTokenValue}, idTokenValues)
 				} else {
 					require.Len(t, authValues, 0)
+					require.Len(t, idTokenValues, 0)
 				}
 				if test.cookie {
 					require.Equal(t, []string{cookieValue}, cookieValues)
@@ -114,6 +119,7 @@ func TestOauthForwardIdentity(t *testing.T) {
 
 			if test.auth {
 				req.Headers[authName] = authValue
+				req.Headers[idTokenName] = idTokenValue
 			}
 
 			if test.cookie {
@@ -145,13 +151,16 @@ func TestOauthForwardIdentity(t *testing.T) {
 				clientUsed = true
 				authValues := req.Header.Values(authName)
 				cookieValues := req.Header.Values(cookieName)
+				idTokenValues := req.Header.Values(idTokenName)
 				// we need to check for "header does not exist",
 				// and the only way i can find is to get the values
 				// as an array
 				if test.auth {
 					require.Equal(t, []string{authValue}, authValues)
+					require.Equal(t, []string{idTokenValue}, idTokenValues)
 				} else {
 					require.Len(t, authValues, 0)
+					require.Len(t, idTokenValues, 0)
 				}
 				if test.cookie {
 					require.Equal(t, []string{cookieValue}, cookieValues)
@@ -168,6 +177,7 @@ func TestOauthForwardIdentity(t *testing.T) {
 
 			if test.auth {
 				req.Headers[authName] = []string{authValue}
+				req.Headers[idTokenName] = []string{idTokenValue}
 			}
 			if test.cookie {
 				req.Headers[cookieName] = []string{cookieValue}
diff --git a/pkg/tsdb/loki/loki.go b/pkg/tsdb/loki/loki.go
index 624cc88a796..6913c817ca0 100644
--- a/pkg/tsdb/loki/loki.go
+++ b/pkg/tsdb/loki/loki.go
@@ -130,6 +130,10 @@ func getAuthHeadersForCallResource(headers map[string][]string) map[string]strin
 		data["Cookie"] = cookie
 	}
 
+	if idToken := arrayHeaderFirstValue(headers["X-ID-Token"]); idToken != "" {
+		data["X-ID-Token"] = idToken
+	}
+
 	return data
 }