From 49dc9da9c1e8dd6ea9ae53e1810932e5c6a5f89f Mon Sep 17 00:00:00 2001
From: Stephanie Closson <srclosson@gmail.com>
Date: Wed, 8 Jun 2022 08:25:53 -0300
Subject: [PATCH] Prometheus: Templating queries should pass on custom headers
 (#50344)

* pass on all headers except for accept headers

* touch up and testing

* add custom header values to resource queries

* remove my picture. oops

* handle gzip responses as well

* fix linting issues

* add my space

* no lint

* removed cookies from being proxied

* clean up and handle errors from io.reader.Close() calls
---
 pkg/tsdb/prometheus/resource/resource.go | 126 +++++++++++++++++++----
 1 file changed, 106 insertions(+), 20 deletions(-)

diff --git a/pkg/tsdb/prometheus/resource/resource.go b/pkg/tsdb/prometheus/resource/resource.go
index 5ff90575264..44babce41c5 100644
--- a/pkg/tsdb/prometheus/resource/resource.go
+++ b/pkg/tsdb/prometheus/resource/resource.go
@@ -1,11 +1,15 @@
 package resource
 
 import (
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
+	"net/http"
 	"net/url"
+	"strings"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	"github.com/grafana/grafana/pkg/infra/httpclient"
@@ -16,8 +20,54 @@ import (
 )
 
 type Resource struct {
-	provider *client.Provider
-	log      log.Logger
+	provider      *client.Provider
+	log           log.Logger
+	customHeaders map[string]string
+}
+
+// Hop-by-hop headers. These are removed when sent to the backend.
+// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
+var hopHeaders = []string{
+	"Connection",
+	"Keep-Alive",
+	"Proxy-Authenticate",
+	"Proxy-Authorization",
+	"Te", // canonicalized version of "TE"
+	"Trailers",
+	"Transfer-Encoding",
+	"Upgrade",
+}
+
+// The following headers will be removed from the request
+var stopHeaders = []string{
+	"cookie",
+	"Cookie",
+}
+
+func delHopHeaders(header http.Header) {
+	for _, h := range hopHeaders {
+		header.Del(h)
+	}
+}
+
+func delStopHeaders(header http.Header) {
+	for _, h := range stopHeaders {
+		header.Del(h)
+	}
+}
+
+func addHeaders(header http.Header, toAdd map[string]string) {
+	for k, v := range toAdd {
+		header.Add(k, v)
+	}
+}
+
+func normalizeReqHeaders(headers map[string][]string) map[string]string {
+	h := make(map[string]string, len(headers))
+	for k, v := range headers {
+		h[k] = strings.Join(v, ",")
+	}
+	return h
 }
 
 func New(
@@ -34,14 +84,46 @@ func New(
 
 	p := client.NewProvider(settings, jsonData, httpClientProvider, cfg, features, plog)
 
+	customHeaders := make(map[string]string)
+	var jsonDataMap map[string]interface{}
+
+	err := json.Unmarshal(settings.JSONData, &jsonDataMap)
+	if err != nil {
+		return nil, err
+	}
+
+	index := 1
+	for {
+		headerNameSuffix := fmt.Sprintf("httpHeaderName%d", index)
+		headerValueSuffix := fmt.Sprintf("httpHeaderValue%d", index)
+
+		key := jsonDataMap[headerNameSuffix]
+		if key == nil {
+			// No (more) header values are available
+			break
+		}
+
+		if val, ok := settings.DecryptedSecureJSONData[headerValueSuffix]; ok {
+			switch k := key.(type) {
+			case string:
+				customHeaders[k] = val
+			}
+		}
+		index++
+	}
+
 	return &Resource{
-		log:      plog,
-		provider: p,
+		log:           plog,
+		provider:      p,
+		customHeaders: customHeaders,
 	}, nil
 }
 
 func (r *Resource) Execute(ctx context.Context, req *backend.CallResourceRequest) (int, []byte, error) {
-	client, err := r.provider.GetClient(reqHeaders(req.Headers))
+	delHopHeaders(req.Headers)
+	delStopHeaders(req.Headers)
+	addHeaders(req.Headers, r.customHeaders)
+	client, err := r.provider.GetClient(normalizeReqHeaders(req.Headers))
 	if err != nil {
 		return 500, nil, err
 	}
@@ -65,25 +147,29 @@ func (r *Resource) fetch(ctx context.Context, client *client.Client, req *backen
 		return statusCode, nil, err
 	}
 
-	defer resp.Body.Close() //nolint (we don't care about the error being returned by resp.Body.Close())
+	defer func() {
+		err = resp.Body.Close()
+	}()
 
-	data, err := ioutil.ReadAll(resp.Body)
+	// Check that the server actually sent compressed data
+	var reader io.ReadCloser
+	switch resp.Header.Get("Content-Encoding") {
+	case "gzip":
+		reader, err = gzip.NewReader(resp.Body)
+		defer func() {
+			err = reader.Close()
+		}()
+		if err != nil {
+			return 500, nil, err
+		}
+	default:
+		reader = resp.Body
+	}
+
+	data, err := ioutil.ReadAll(reader)
 	if err != nil {
 		return 500, nil, err
 	}
 
 	return resp.StatusCode, data, err
 }
-
-func reqHeaders(headers map[string][]string) map[string]string {
-	// Keep only the authorization header, incase downstream the authorization header is required.
-	// Strip all the others out as appropriate headers will be applied to speak with prometheus.
-	h := make(map[string]string)
-	accessValues := headers["Authorization"]
-
-	if len(accessValues) > 0 {
-		h["Authorization"] = accessValues[0]
-	}
-
-	return h
-}