From e3241b8940870760a6a5120648f53ff991b24a83 Mon Sep 17 00:00:00 2001
From: codebien <2103732+codebien@users.noreply.github.com>
Date: Fri, 22 Sep 2023 16:54:48 +0200
Subject: [PATCH] New format for passing custom HTTP headers

---
 pkg/remotewrite/config.go      | 13 +++++++++++++
 pkg/remotewrite/config_test.go | 17 +++++++++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/pkg/remotewrite/config.go b/pkg/remotewrite/config.go
index c609913..6d9b01a 100644
--- a/pkg/remotewrite/config.go
+++ b/pkg/remotewrite/config.go
@@ -237,6 +237,19 @@ func parseEnvs(env map[string]string) (Config, error) {
 		c.Headers[k] = v
 	}
 
+	if headers, headersDefined := env["K6_PROMETHEUS_RW_HTTP_HEADERS"]; headersDefined {
+		if c.Headers == nil {
+			c.Headers = make(map[string]string)
+		}
+		for _, kvPair := range strings.Split(headers, ",") {
+			header := strings.Split(kvPair, ":")
+			if len(header) != 2 {
+				continue
+			}
+			c.Headers[header[0]] = header[1]
+		}
+	}
+
 	if b, err := getEnvBool(env, "K6_PROMETHEUS_RW_TREND_AS_NATIVE_HISTOGRAM"); err != nil {
 		return c, err
 	} else if b.Valid {
diff --git a/pkg/remotewrite/config_test.go b/pkg/remotewrite/config_test.go
index ae499ad..251585a 100644
--- a/pkg/remotewrite/config_test.go
+++ b/pkg/remotewrite/config_test.go
@@ -255,8 +255,15 @@ func TestOptionHeaders(t *testing.T) {
 		env     map[string]string
 		jsonRaw json.RawMessage
 	}{
-		"JSON": {jsonRaw: json.RawMessage(`{"headers":{"X-MY-HEADER1":"hval1","X-MY-HEADER2":"hval2"}}`)},
-		"Env":  {env: map[string]string{"K6_PROMETHEUS_RW_HEADERS_X-MY-HEADER1": "hval1", "K6_PROMETHEUS_RW_HEADERS_X-MY-HEADER2": "hval2"}},
+		"JSON": {jsonRaw: json.RawMessage(
+			`{"headers":{"X-MY-HEADER1":"hval1","X-MY-HEADER2":"hval2","X-Scope-OrgID":"my-org-id","another-header":"true"}}`)},
+		"Env": {env: map[string]string{
+			"K6_PROMETHEUS_RW_HEADERS_X-MY-HEADER1": "hval1",
+			"K6_PROMETHEUS_RW_HEADERS_X-MY-HEADER2": "hval2",
+			// it assert that the new method using HTTP_HEADERS overwrites it
+			"K6_PROMETHEUS_RW_HEADERS_X-Scope-OrgID": "my-org-id-old-method",
+			"K6_PROMETHEUS_RW_HTTP_HEADERS":          "X-Scope-OrgID:my-org-id,another-header:true",
+		}},
 		//nolint:gocritic
 		//"Arg":  {arg: "headers.X-MY-HEADER1=hval1,headers.X-MY-HEADER2=hval2"},
 	}
@@ -266,8 +273,10 @@ func TestOptionHeaders(t *testing.T) {
 		InsecureSkipTLSVerify: null.BoolFrom(false),
 		PushInterval:          types.NullDurationFrom(5 * time.Second),
 		Headers: map[string]string{
-			"X-MY-HEADER1": "hval1",
-			"X-MY-HEADER2": "hval2",
+			"X-MY-HEADER1":   "hval1",
+			"X-MY-HEADER2":   "hval2",
+			"X-Scope-OrgID":  "my-org-id",
+			"another-header": "true",
 		},
 		TrendStats:   []string{"p(99)"},
 		StaleMarkers: null.BoolFrom(false),