diff options
author | feistel <6742251-feistel@users.noreply.gitlab.com> | 2022-02-16 23:29:14 +0300 |
---|---|---|
committer | feistel <6742251-feistel@users.noreply.gitlab.com> | 2022-02-16 23:29:14 +0300 |
commit | fba96f28a6d0d9af3810d9d996bf32654771eaca (patch) | |
tree | edfe047558f47cae8019ec5acc27da90fce33719 /internal/customheaders | |
parent | 4afc9bb98bec0c8bbb0aadb7801d1da71be93841 (diff) |
refactor: simplify parsing and remove whitespace handling
Diffstat (limited to 'internal/customheaders')
-rw-r--r-- | internal/customheaders/customheaders.go | 30 | ||||
-rw-r--r-- | internal/customheaders/customheaders_test.go | 39 |
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) } }) |