refactor: simplify parsing and remove whitespace handling

2 files changed, 33 insertions, 36 deletions

diff --git a/internal/customheaders/customheaders.go b/internal/customheaders/customheaders.go
index 4c2289e0..96662fe1 100644
--- a/internal/customheaders/customheaders.go
+++ b/internal/customheaders/customheaders.go
@@ -3,12 +3,16 @@ package customheaders
 import (
 	"bufio"
 	"errors"
+	"fmt"
 	"net/http"
 	"net/textproto"
 	"strings"
 )
 
-var errInvalidHeaderParameter = errors.New("invalid syntax specified as header parameter")
+var (
+	errInvalidHeaderParameter = errors.New("invalid syntax specified as header parameter")
+	errDuplicateHeader        = errors.New("duplicate header")
+)
 
 // AddCustomHeaders adds a map of Headers to a Response
 func AddCustomHeaders(w http.ResponseWriter, headers http.Header) {
@@ -21,19 +25,25 @@ func AddCustomHeaders(w http.ResponseWriter, headers http.Header) {
 
 // ParseHeaderString parses a string of key values into a map
 func ParseHeaderString(customHeaders []string) (http.Header, error) {
-	headers := http.Header{}
-	for _, keyValueString := range customHeaders {
-		keyValueString = strings.TrimSpace(keyValueString) + "\n\n"
-		tp := textproto.NewReader(bufio.NewReader(strings.NewReader(keyValueString)))
-		keyValue, err := tp.ReadMIMEHeader()
+	headers := make(http.Header, len(customHeaders))
+
+	for _, h := range customHeaders {
+		h = h + "\n\n"
+		tp := textproto.NewReader(bufio.NewReader(strings.NewReader(h)))
+
+		mimeHeader, err := tp.ReadMIMEHeader()
 		if err != nil {
-			return nil, errInvalidHeaderParameter
+			return nil, fmt.Errorf("parsing error %s: %w", h, errInvalidHeaderParameter)
 		}
 
-		for k, v := range keyValue {
-			k = textproto.CanonicalMIMEHeaderKey(strings.TrimSpace(k))
-			headers[k] = append(headers[k], v...)
+		for key, value := range mimeHeader {
+			if _, ok := headers[key]; ok {
+				return nil, fmt.Errorf("%s already specified with value '%s': %w", key, value, errDuplicateHeader)
+			}
+
+			headers[key] = value
 		}
 	}
+
 	return headers, nil
 }
diff --git a/internal/customheaders/customheaders_test.go b/internal/customheaders/customheaders_test.go
index 5b2214fc..857c45e0 100644
--- a/internal/customheaders/customheaders_test.go
+++ b/internal/customheaders/customheaders_test.go
@@ -23,18 +23,6 @@ func TestParseHeaderString(t *testing.T) {
 			expectedLen:   1,
 		},
 		{
-			name:          "Whitespace trim case",
-			headerStrings: []string{"   X-Test-String   :   Test  "},
-			valid:         true,
-			expectedLen:   1,
-		},
-		{
-			name:          "Whitespace in key, value case",
-			headerStrings: []string{"My amazing header: This is a test"},
-			valid:         true,
-			expectedLen:   1,
-		},
-		{
 			name:          "Non-tracking header case",
 			headerStrings: []string{"Tk: N"},
 			valid:         true,
@@ -63,6 +51,11 @@ func TestParseHeaderString(t *testing.T) {
 			valid:         false,
 		},
 		{
+			name:          "duplicate headers",
+			headerStrings: []string{"Tk: N", "Tk: M"},
+			valid:         false,
+		},
+		{
 			name:          "Not valid case",
 			headerStrings: []string{"X-Test-String Some-Test"},
 			valid:         false,
@@ -99,22 +92,13 @@ func TestAddCustomHeaders(t *testing.T) {
 		name          string
 		headerStrings []string
 		wantHeaders   map[string]string
-	}{{
-		name:          "Normal case",
-		headerStrings: []string{"X-Test-String: Test"},
-		wantHeaders:   map[string]string{"X-Test-String": "Test"},
-	},
+	}{
 		{
-			name:          "Whitespace trim case",
-			headerStrings: []string{"   X-Test-String   :   Test  "},
+			name:          "Normal case",
+			headerStrings: []string{"X-Test-String: Test"},
 			wantHeaders:   map[string]string{"X-Test-String": "Test"},
 		},
 		{
-			name:          "Whitespace in key, value case",
-			headerStrings: []string{"My amazing header: This is a test"},
-			wantHeaders:   map[string]string{"My amazing header": "This is a test"},
-		},
-		{
 			name:          "Non-tracking header case",
 			headerStrings: []string{"Tk: N"},
 			wantHeaders:   map[string]string{"Tk": "N"},
@@ -126,7 +110,7 @@ func TestAddCustomHeaders(t *testing.T) {
 		},
 		{
 			name:          "Multiple header strings",
-			headerStrings: []string{"content-security-policy: default-src 'self'", "X-Test-String: Test", "My amazing header : Amazing"},
+			headerStrings: []string{"content-security-policy: default-src 'self'", "X-Test-String: Test", "My amazing header: Amazing"},
 			wantHeaders:   map[string]string{"Content-Security-Policy": "default-src 'self'", "X-Test-String": "Test", "My amazing header": "Amazing"},
 		},
 	}
@@ -139,7 +123,10 @@ func TestAddCustomHeaders(t *testing.T) {
 			customheaders.AddCustomHeaders(w, headers)
 			rsp := w.Result()
 			for k, v := range tt.wantHeaders {
-				got := rsp.Header[k]
+				require.Len(t, rsp.Header[k], 1)
+
+				// use the map directly to make sure ParseHeaderString is adding the canonical keys
+				got := rsp.Header[k][0]
 				require.Equal(t, v, got, "Expected header %+v, got %+v", v, got)
 			}
 		})