From 5563f2f98f94dff7ec7c698d01620aad96451312 Mon Sep 17 00:00:00 2001
From: Alexander Zobnin <alexanderzobnin@gmail.com>
Date: Thu, 19 Oct 2023 17:09:53 +0200
Subject: [PATCH] Team LBAC: Fix multiple header values (#76821)

* Team LBAC: Fix multiple header values

Co-authored-by: Eric Leijonmarck <eleijonmarck@users.noreply.github.com>

* Fix linter error

---------

Co-authored-by: Eric Leijonmarck <eleijonmarck@users.noreply.github.com>
---
 pkg/util/proxyutil/proxyutil.go      |  12 ++-
 pkg/util/proxyutil/proxyutil_test.go | 111 ++++++++++++++++++++-------
 2 files changed, 91 insertions(+), 32 deletions(-)

diff --git a/pkg/util/proxyutil/proxyutil.go b/pkg/util/proxyutil/proxyutil.go
index 2946cee9d76..43183d07800 100644
--- a/pkg/util/proxyutil/proxyutil.go
+++ b/pkg/util/proxyutil/proxyutil.go
@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"net/url"
 	"sort"
 	"strconv"
 	"strings"
@@ -170,9 +171,14 @@ func GetTeamHTTPHeaders(ds *datasources.DataSource, teams []int64) (map[string]s
 		}
 
 		for _, header := range headers {
-			// TODO: handle multiple header values
-			// add tests for these cases
-			teamHTTPHeadersMap[header.Header] = header.Value
+			// Header values should be properly escaped.
+			if value, ok := teamHTTPHeadersMap[header.Header]; ok {
+				// Add multiple header values as a comma-separated strings according to RFC 7230
+				// https://datatracker.ietf.org/doc/html/rfc7230#section-3.2.6
+				teamHTTPHeadersMap[header.Header] = fmt.Sprintf("%s,%s", value, url.PathEscape(header.Value))
+			} else {
+				teamHTTPHeadersMap[header.Header] = url.PathEscape(header.Value)
+			}
 		}
 	}
 
diff --git a/pkg/util/proxyutil/proxyutil_test.go b/pkg/util/proxyutil/proxyutil_test.go
index 5cdf76a6b5b..e05b295df26 100644
--- a/pkg/util/proxyutil/proxyutil_test.go
+++ b/pkg/util/proxyutil/proxyutil_test.go
@@ -207,40 +207,93 @@ func TestApplyUserHeader(t *testing.T) {
 }
 
 func TestApplyteamHTTPHeaders(t *testing.T) {
-	t.Run("Should apply team headers for users teams", func(t *testing.T) {
-		req, err := http.NewRequest(http.MethodGet, "/", nil)
-		require.NoError(t, err)
-		ds := &datasources.DataSource{
-			JsonData: simplejson.New(),
-		}
-		userTeams := []int64{1, 2}
-		// add team headers
-		ds.JsonData.Set("teamHttpHeaders", map[string]interface{}{
-			"1": []map[string]interface{}{
-				{
-					"header": "X-Team-Header",
-					"value":  "1",
+	testCases := []struct {
+		desc      string
+		jsonData  any
+		userTeams []int64
+		want      map[string]string
+	}{
+		{
+			desc: "Should apply team headers for users teams",
+			jsonData: map[string]interface{}{
+				"1": []map[string]interface{}{
+					{
+						"header": "X-Team-Header",
+						"value":  "1",
+					},
+				},
+				"2": []map[string]interface{}{
+					{
+						"header": "X-Prom-Label-Policy",
+						"value":  "2",
+					},
+				},
+				// user is not part of this team
+				"3": []map[string]interface{}{
+					{
+						"header": "X-Custom-Label-Policy",
+						"value":  "3",
+					},
 				},
 			},
-			"2": []map[string]interface{}{
-				{
-					"header": "X-Prom-Label-Policy",
-					"value":  "2",
+			userTeams: []int64{1, 2},
+			want: map[string]string{
+				"X-Team-Header":       "1",
+				"X-Prom-Label-Policy": "2",
+			},
+		},
+		{
+			desc: "Should be able to parse header values with commas",
+			jsonData: map[string]interface{}{
+				"101": []map[string]interface{}{
+					{
+						"header": "X-Prom-Label-Policy",
+						"value":  `1234:{ foo="bar", bar="baz" }`,
+					},
 				},
 			},
-			// user is not part of this team
-			"3": []map[string]interface{}{
-				{
-					"header": "X-Custom-Label-Policy",
-					"value":  "3",
+			userTeams: []int64{101},
+			want: map[string]string{
+				"X-Prom-Label-Policy": "1234:%7B%20foo=%22bar%22%2C%20bar=%22baz%22%20%7D",
+			},
+		}, {
+			desc: "Should be able to handle multiple header values",
+			jsonData: map[string]interface{}{
+				"101": []map[string]interface{}{
+					{
+						"header": "X-Prom-Label-Policy",
+						"value":  `1234:{ foo="bar" }`,
+					},
+					{
+						"header": "X-Prom-Label-Policy",
+						"value":  `1234:{ bar="baz" }`,
+					},
 				},
 			},
-		})
+			userTeams: []int64{101},
+			want: map[string]string{
+				"X-Prom-Label-Policy": "1234:%7B%20foo=%22bar%22%20%7D,1234:%7B%20bar=%22baz%22%20%7D",
+			},
+		},
+	}
 
-		err = ApplyTeamHTTPHeaders(req, ds, userTeams)
-		require.NoError(t, err)
-		require.Contains(t, req.Header, "X-Team-Header")
-		require.Contains(t, req.Header, "X-Prom-Label-Policy")
-		require.NotContains(t, req.Header, "X-Custom-Label-Policy")
-	})
+	for _, testCase := range testCases {
+		t.Run("Should apply team headers for users teams", func(t *testing.T) {
+			req, err := http.NewRequest(http.MethodGet, "/", nil)
+			require.NoError(t, err)
+			ds := &datasources.DataSource{
+				JsonData: simplejson.New(),
+			}
+
+			// add team headers
+			ds.JsonData.Set("teamHttpHeaders", testCase.jsonData)
+
+			err = ApplyTeamHTTPHeaders(req, ds, testCase.userTeams)
+			require.NoError(t, err)
+			for header, value := range testCase.want {
+				require.Contains(t, req.Header, header)
+				require.Equal(t, value, req.Header.Get(header))
+			}
+		})
+	}
 }