Fix duplicate set-cookie headers

v2/response.go

@@ -59,9 +59,14 @@ func (w *ResponseWriter) WriteHeader(status int) {
 
 	h := make(map[string]string)
 	mvh := make(map[string][]string)
+	var cookies []string
 
 	for k, v := range w.Header() {
-		if len(v) == 1 {
+		if http.CanonicalHeaderKey(k) == "Set-Cookie" {
+			// cookies are not returned in header maps
+			// see https://aws.amazon.com/blogs/compute/simply-serverless-using-aws-lambda-to-expose-custom-cookies-with-api-gateway/
+			cookies = append(cookies, v...)
+		} else if len(v) == 1 {
 			h[k] = v[0]
 		} else if len(v) > 1 {
 			mvh[k] = v
@@ -70,6 +75,7 @@ func (w *ResponseWriter) WriteHeader(status int) {
 
 	w.out.Headers = h
 	w.out.MultiValueHeaders = mvh
+	w.out.Cookies = cookies
 	w.wroteHeader = true
 }
 
@@ -88,10 +94,6 @@ func (w *ResponseWriter) End() events.APIGatewayV2HTTPResponse {
 		w.out.Body = w.buf.String()
 	}
 
-	// see https://aws.amazon.com/blogs/compute/simply-serverless-using-aws-lambda-to-expose-custom-cookies-with-api-gateway/
-	w.out.Cookies = w.header["Set-Cookie"]
-	w.header.Del("Set-Cookie")
-
 	// notify end
 	w.closeNotifyCh <- true
 

v2/response_test.go

@@ -2,6 +2,7 @@ package gateway
 
 import (
 	"bytes"
+	"net/http"
 	"testing"
 
 	"github.com/tj/assert"
@@ -105,3 +106,20 @@ func TestResponseWriter_WriteHeader(t *testing.T) {
 	assert.Equal(t, "Not Found\n", e.Body)
 	assert.Equal(t, "text/plain; charset=utf8", e.Headers["Content-Type"])
 }
+
+func TestResponseWriter_Cookie(t *testing.T) {
+	w := NewResponse()
+	http.SetCookie(w, &http.Cookie{
+		Name:  "foo",
+		Value: "bar",
+	})
+	w.WriteHeader(200)
+	w.Write([]byte("Have a cookie\n"))
+
+	e := w.End()
+	assert.Equal(t, 200, e.StatusCode)
+	assert.Equal(t, "Have a cookie\n", e.Body)
+	assert.Equal(t, 1, len(e.Headers))
+	assert.Equal(t, "text/plain; charset=utf8", e.Headers["Content-Type"])
+	assert.Equal(t, []string{"foo=bar"}, e.Cookies)
+}