diff --git a/v2/response.go b/v2/response.go index 33c2c20..5256126 100644 --- a/v2/response.go +++ b/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 diff --git a/v2/response_test.go b/v2/response_test.go index 47e1357..c5da31b 100644 --- a/v2/response_test.go +++ b/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) +}