Merge branch 'fix/read-mime-header' into 'master'

refactor: parse custom headers using ReadMIMEHeader

Closes #536

See merge request gitlab-org/gitlab-pages!517

2 files changed, 46 insertions, 37 deletions

diff --git a/internal/customheaders/customheaders.go b/internal/customheaders/customheaders.go
index b54df585..92f50069 100644
--- a/internal/customheaders/customheaders.go
+++ b/internal/customheaders/customheaders.go
@@ -1,12 +1,20 @@
 package customheaders
 
 import (
+	"bufio"
 	"errors"
+	"fmt"
 	"net/http"
+	"net/textproto"
 	"strings"
+
+	"github.com/hashicorp/go-multierror"
 )
 
-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) {
@@ -19,17 +27,30 @@ 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 {
-		keyValue := strings.SplitN(keyValueString, ":", 2)
-		if len(keyValue) != 2 {
-			return nil, errInvalidHeaderParameter
+	headers := make(http.Header, len(customHeaders))
+
+	var result *multierror.Error
+	for _, h := range customHeaders {
+		h = h + "\n\n"
+		tp := textproto.NewReader(bufio.NewReader(strings.NewReader(h)))
+
+		mimeHeader, err := tp.ReadMIMEHeader()
+		if err != nil {
+			result = multierror.Append(result, fmt.Errorf("parsing error %s: %w", h, errInvalidHeaderParameter))
 		}
 
-		key := strings.TrimSpace(keyValue[0])
-		value := strings.TrimSpace(keyValue[1])
+		for key, value := range mimeHeader {
+			if _, ok := headers[key]; ok {
+				result = multierror.Append(result, fmt.Errorf("%s already specified with value '%s': %w", key, value, errDuplicateHeader))
+			}
 
-		headers[key] = append(headers[key], value)
+			headers[key] = value
+		}
+	}
+
+	if result.ErrorOrNil() != nil {
+		return nil, result
 	}
+
 	return headers, nil
 }
diff --git a/internal/customheaders/customheaders_test.go b/internal/customheaders/customheaders_test.go
index a667f43a..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"},
@@ -122,12 +106,12 @@ func TestAddCustomHeaders(t *testing.T) {
 		{
 			name:          "Content security header case",
 			headerStrings: []string{"content-security-policy: default-src 'self'"},
-			wantHeaders:   map[string]string{"content-security-policy": "default-src 'self'"},
+			wantHeaders:   map[string]string{"Content-Security-Policy": "default-src 'self'"},
 		},
 		{
 			name:          "Multiple header strings",
-			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"},
+			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,11 @@ func TestAddCustomHeaders(t *testing.T) {
 			customheaders.AddCustomHeaders(w, headers)
 			rsp := w.Result()
 			for k, v := range tt.wantHeaders {
-				require.Equal(t, v, rsp.Header.Get(k), "Expected header %+v, got %+v", v, rsp.Header.Get(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)
 			}
 		})
 	}