diff --git a/header.go b/header.go index b6d8ffec92..1b11779d7c 100644 --- a/header.go +++ b/header.go @@ -1382,7 +1382,7 @@ func (h *RequestHeader) del(key []byte) { } // setSpecialHeader handles special headers and return true when a header is processed. -func (h *ResponseHeader) setSpecialHeader(key, value []byte) bool { +func (h *ResponseHeader) setSpecialHeader(key, value []byte, add bool) bool { if len(key) == 0 { return false } @@ -1417,6 +1417,9 @@ func (h *ResponseHeader) setSpecialHeader(key, value []byte) bool { return true } else if caseInsensitiveCompare(strSetCookie, key) { var kv *argsKV + if !add { + h.DelAllCookies() + } h.cookies, kv = allocArg(h.cookies) kv.key = getCookieKey(kv.key, value) kv.value = append(kv.value[:0], value...) @@ -1446,7 +1449,7 @@ func (h *ResponseHeader) setNonSpecial(key, value []byte) { } // setSpecialHeader handles special headers and return true when a header is processed. -func (h *RequestHeader) setSpecialHeader(key, value []byte) bool { +func (h *RequestHeader) setSpecialHeader(key, value []byte, add bool) bool { if len(key) == 0 || h.disableSpecialHeader { return false } @@ -1473,6 +1476,9 @@ func (h *RequestHeader) setSpecialHeader(key, value []byte) bool { return true case caseInsensitiveCompare(strCookie, key): h.collectCookies() + if !add { + h.DelAllCookies() + } h.cookies = parseRequestCookies(h.cookies, value) return true } @@ -1561,7 +1567,7 @@ func (h *ResponseHeader) AddBytesV(key string, value []byte) { // If the header is set as a Trailer (forbidden trailers will not be set, see AddTrailer for more details), // it will be sent after the chunked response body. func (h *ResponseHeader) AddBytesKV(key, value []byte) { - if h.setSpecialHeader(key, value) { + if h.setSpecialHeader(key, value, true) { return } @@ -1620,7 +1626,7 @@ func (h *ResponseHeader) SetBytesKV(key, value []byte) { // If the header is set as a Trailer (forbidden trailers will not be set, see SetTrailer for more details), // it will be sent after the chunked response body. func (h *ResponseHeader) SetCanonical(key, value []byte) { - if h.setSpecialHeader(key, value) { + if h.setSpecialHeader(key, value, false) { return } h.setNonSpecial(key, value) @@ -1772,7 +1778,7 @@ func (h *RequestHeader) AddBytesV(key string, value []byte) { // If the header is set as a Trailer (forbidden trailers will not be set, see AddTrailer for more details), // it will be sent after the chunked request body. func (h *RequestHeader) AddBytesKV(key, value []byte) { - if h.setSpecialHeader(key, value) { + if h.setSpecialHeader(key, value, true) { return } @@ -1831,7 +1837,7 @@ func (h *RequestHeader) SetBytesKV(key, value []byte) { // If the header is set as a Trailer (forbidden trailers will not be set, see SetTrailer for more details), // it will be sent after the chunked request body. func (h *RequestHeader) SetCanonical(key, value []byte) { - if h.setSpecialHeader(key, value) { + if h.setSpecialHeader(key, value, false) { return } h.setNonSpecial(key, value) diff --git a/header_test.go b/header_test.go index 37b775b26d..1eedb0dab2 100644 --- a/header_test.go +++ b/header_test.go @@ -1643,7 +1643,32 @@ func TestRequestHeaderSetCookie(t *testing.T) { var h RequestHeader h.Set("Cookie", "foo=bar; baz=aaa") + if string(h.Cookie("foo")) != "bar" { + t.Fatalf("Unexpected cookie %q. Expecting %q", h.Cookie("foo"), "bar") + } + if string(h.Cookie("baz")) != "aaa" { + t.Fatalf("Unexpected cookie %q. Expecting %q", h.Cookie("baz"), "aaa") + } + h.Set("cOOkie", "xx=yyy") + if len(h.Cookie("foo")) != 0 { + t.Fatalf("Unexpected cookie %q. Expecting empty string", h.Cookie("foo")) + } + if len(h.Cookie("baz")) != 0 { + t.Fatalf("Unexpected cookie %q. Expecting empty string", h.Cookie("baz")) + } + if string(h.Cookie("xx")) != "yyy" { + t.Fatalf("unexpected cookie %q. Expecting %q", h.Cookie("xx"), "yyy") + } +} + +func TestRequestHeaderAddCookie(t *testing.T) { + t.Parallel() + + var h RequestHeader + + h.Add("Cookie", "foo=bar; baz=aaa") + h.Add("cOOkie", "xx=yyy") if string(h.Cookie("foo")) != "bar" { t.Fatalf("Unexpected cookie %q. Expecting %q", h.Cookie("foo"), "bar") @@ -1662,8 +1687,44 @@ func TestResponseHeaderSetCookie(t *testing.T) { var h ResponseHeader h.Set("set-cookie", "foo=bar; path=/aa/bb; domain=aaa.com") + + var c Cookie + c.SetKey("foo") + if !h.Cookie(&c) { + t.Fatalf("cannot obtain %q cookie", c.Key()) + } + if string(c.Value()) != "bar" { + t.Fatalf("unexpected cookie value %q. Expected %q", c.Value(), "bar") + } + if string(c.Path()) != "/aa/bb" { + t.Fatalf("unexpected cookie path %q. Expected %q", c.Path(), "/aa/bb") + } + if string(c.Domain()) != "aaa.com" { + t.Fatalf("unexpected cookie domain %q. Expected %q", c.Domain(), "aaa.com") + } + h.Set(HeaderSetCookie, "aaaaa=bxx") + if h.Cookie(&c) { + t.Fatalf("obtain %q cookie", c.Key()) + } + c.SetKey("aaaaa") + if !h.Cookie(&c) { + t.Fatalf("cannot obtain %q cookie", c.Key()) + } + if string(c.Value()) != "bxx" { + t.Fatalf("unexpected cookie value %q. Expecting %q", c.Value(), "bxx") + } +} + +func TestResponseHeaderAddCookie(t *testing.T) { + t.Parallel() + + var h ResponseHeader + + h.Add("set-cookie", "foo=bar; path=/aa/bb; domain=aaa.com") + h.Add(HeaderSetCookie, "aaaaa=bxx") + var c Cookie c.SetKey("foo") if !h.Cookie(&c) {