From e6a7b537038123d00af6e9961b324e5d7d5c729e Mon Sep 17 00:00:00 2001
From: Sven Grossmann <sven.grossmann@grafana.com>
Date: Wed, 23 Nov 2022 12:52:35 +0100
Subject: [PATCH] [9.2.x] Fix #58598 X-ID-Token header missing on Loki
 Datasource (#58784) (#59196)

* Fix #58598 X-ID-Token header missing on Loki Datasource (#58784)

* Fix #58598 X-ID-Token header missing on Loki Datasource

* Remove unecessary continue statements

* Add getAuthHeadersForCallResource unit tests

* Fix test and switch statement issues introduced during merge

(cherry picked from commit f1ef63791ad17d4c3b16cd296148c403515ad37d)

* update test

* missed linting

Co-authored-by: Yann Vigara <yvigara@users.noreply.github.com>
---
 pkg/tsdb/loki/loki.go      | 25 +++++++------
 pkg/tsdb/loki/loki_test.go | 72 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 87 insertions(+), 10 deletions(-)
 create mode 100644 pkg/tsdb/loki/loki_test.go

diff --git a/pkg/tsdb/loki/loki.go b/pkg/tsdb/loki/loki.go
index 6913c817ca0..c86dff06159 100644
--- a/pkg/tsdb/loki/loki.go
+++ b/pkg/tsdb/loki/loki.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net/http"
+	"net/textproto"
 	"regexp"
 	"strings"
 	"sync"
@@ -122,18 +123,22 @@ func (s *Service) CallResource(ctx context.Context, req *backend.CallResourceReq
 func getAuthHeadersForCallResource(headers map[string][]string) map[string]string {
 	data := make(map[string]string)
 
-	if auth := arrayHeaderFirstValue(headers["Authorization"]); auth != "" {
-		data["Authorization"] = auth
-	}
+	for k, values := range headers {
+		k = textproto.CanonicalMIMEHeaderKey(k)
+		firstValue := arrayHeaderFirstValue(values)
 
-	if cookie := arrayHeaderFirstValue(headers["Cookie"]); cookie != "" {
-		data["Cookie"] = cookie
+		if firstValue == "" {
+			continue
+		}
+		switch k {
+		case "Authorization":
+			data["Authorization"] = firstValue
+		case "X-Id-Token":
+			data["X-ID-Token"] = firstValue
+		case "Cookie":
+			data["Cookie"] = firstValue
+		}
 	}
-
-	if idToken := arrayHeaderFirstValue(headers["X-ID-Token"]); idToken != "" {
-		data["X-ID-Token"] = idToken
-	}
-
 	return data
 }
 
diff --git a/pkg/tsdb/loki/loki_test.go b/pkg/tsdb/loki/loki_test.go
new file mode 100644
index 00000000000..a377f6011c5
--- /dev/null
+++ b/pkg/tsdb/loki/loki_test.go
@@ -0,0 +1,72 @@
+package loki
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetHeadersForCallResource(t *testing.T) {
+	const idTokn1 = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+	const idTokn2 = "eyJhbGciOiJIUzI1NiJ9.eyJuYW1lIjoiSm9obiBEb2UiLCJleHAiOjE2Njg2MjExODQsImlhdCI6MTY2ODYyMTE4NH0.bg0Y0S245DeANhNnnLBCfGYBseTld29O0xynhQwZZlU"
+	const authTokn1 = "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
+	const authTokn2 = "Bearer eyJhbGciOiJIUzI1NiJ9.eyJuYW1lIjoiSm9obiBEb2UiLCJleHAiOjE2Njg2MjExODQsImlhdCI6MTY2ODYyMTE4NH0.bg0Y0S245DeANhNnnLBCfGYBseTld29O0xynhQwZZlU"
+
+	testCases := map[string]struct {
+		headers         map[string][]string
+		expectedHeaders map[string]string
+	}{
+		"Headers with empty value": {
+			headers: map[string][]string{
+				"X-Grafana-Org-Id": {"1"},
+				"Cookie":           {""},
+				"X-Id-Token":       {""},
+				"Accept-Encoding":  {""},
+				"Authorization":    {""},
+			},
+			expectedHeaders: map[string]string{},
+		},
+		"Headers with multiple values": {
+			headers: map[string][]string{
+				"Authorization":    {authTokn1, authTokn2},
+				"Cookie":           {"a=1"},
+				"X-Grafana-Org-Id": {"1"},
+				"Accept-Encoding":  {"gzip", "compress"},
+				"X-Id-Token":       {idTokn1, idTokn2},
+			},
+			expectedHeaders: map[string]string{
+				"Authorization": authTokn1,
+				"Cookie":        "a=1",
+				"X-ID-Token":    idTokn1,
+			},
+		},
+		"Headers with single value": {
+			headers: map[string][]string{
+				"Authorization":    {authTokn1},
+				"X-Grafana-Org-Id": {"1"},
+				"Cookie":           {"a=1"},
+				"Accept-Encoding":  {"gzip"},
+				"X-Id-Token":       {idTokn1},
+			},
+			expectedHeaders: map[string]string{
+				"Authorization": authTokn1,
+				"Cookie":        "a=1",
+				"X-ID-Token":    idTokn1,
+			},
+		},
+		"Non Canonical 'X-Id-Token' header key": {
+			headers: map[string][]string{
+				"X-ID-TOKEN": {idTokn1},
+			},
+			expectedHeaders: map[string]string{
+				"X-ID-Token": idTokn1,
+			},
+		},
+	}
+	for name, test := range testCases {
+		t.Run(name, func(t *testing.T) {
+			headers := getAuthHeadersForCallResource(test.headers)
+			assert.Equal(t, test.expectedHeaders, headers)
+		})
+	}
+}