From 9bc9d1f33339e329c49d325a8752bd435bdd8503 Mon Sep 17 00:00:00 2001
From: Steven Lee <steven.lee@kynrai.com>
Date: Sat, 15 Jun 2024 12:44:57 +0100
Subject: [PATCH] FIX: fix headers for multi strings (#51)

- Return the native http.Headers object
---
 checks/headers.go      | 11 ++---------
 checks/headers_test.go |  6 +++---
 2 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/checks/headers.go b/checks/headers.go
index 242ea06..ecbd827 100644
--- a/checks/headers.go
+++ b/checks/headers.go
@@ -13,7 +13,7 @@ func NewHeaders(client *http.Client) *Headers {
 	return &Headers{client: client}
 }
 
-func (h *Headers) List(ctx context.Context, url string) (map[string]string, error) {
+func (h *Headers) List(ctx context.Context, url string) (http.Header, error) {
 	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
 	if err != nil {
 		return nil, err
@@ -25,12 +25,5 @@ func (h *Headers) List(ctx context.Context, url string) (map[string]string, erro
 	}
 	defer resp.Body.Close()
 
-	responseHeaders := make(map[string]string)
-	for k, v := range resp.Header {
-		for _, s := range v {
-			responseHeaders[k] = s
-		}
-	}
-
-	return responseHeaders, nil
+	return resp.Header, nil
 }
diff --git a/checks/headers_test.go b/checks/headers_test.go
index 57725aa..e0a8b40 100644
--- a/checks/headers_test.go
+++ b/checks/headers_test.go
@@ -14,7 +14,7 @@ func TestList(t *testing.T) {
 
 	c := testutils.MockClient(&http.Response{
 		Header: http.Header{
-			"Cache-Control":    {"private, max-age=0"},
+			"Cache-Control":    {"private", "max-age=0"},
 			"X-Xss-Protection": {"0"},
 		},
 	})
@@ -23,6 +23,6 @@ func TestList(t *testing.T) {
 	actual, err := h.List(context.Background(), "example.com")
 	assert.NoError(t, err)
 
-	assert.Equal(t, "private, max-age=0", actual["Cache-Control"])
-	assert.Equal(t, "0", actual["X-Xss-Protection"])
+	assert.Equal(t, []string{"private", "max-age=0"}, actual["Cache-Control"])
+	assert.Equal(t, []string{"0"}, actual["X-Xss-Protection"])
 }