From 12ba9899844060189449e9768dc0900a56d3624b Mon Sep 17 00:00:00 2001 From: Aliaksandr Valialkin Date: Tue, 29 Mar 2016 17:11:43 +0300 Subject: [PATCH] Issue #69: Added RequestHeader.Add and ResponseHeader.Add for adding multiple headers with the same key --- args_test.go | 2 + header.go | 58 +++++++++++++++++++++++++++ header_test.go | 105 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 165 insertions(+) diff --git a/args_test.go b/args_test.go index de6e95a..914d5c8 100644 --- a/args_test.go +++ b/args_test.go @@ -48,6 +48,8 @@ func TestArgsAdd(t *testing.T) { t.Fatalf("unexpected value: %q. Expecting %q", v, "23") } baFound = true + default: + t.Fatalf("unexpected key found %q", k) } }) if !barFound || !bazFound || !oneFound || !baFound { diff --git a/header.go b/header.go index d6b6b07..3aabbbe 100644 --- a/header.go +++ b/header.go @@ -834,6 +834,35 @@ func (h *RequestHeader) del(key []byte) { h.h = delAllArgsBytes(h.h, key) } +// Add adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *ResponseHeader) Add(key, value string) { + k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing) + h.h = appendArg(h.h, b2s(k), value) +} + +// AddBytesK adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *ResponseHeader) AddBytesK(key []byte, value string) { + h.Add(b2s(key), value) +} + +// AddBytesV adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *ResponseHeader) AddBytesV(key string, value []byte) { + h.Add(key, b2s(value)) +} + +// AddBytesKV adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *ResponseHeader) AddBytesKV(key, value []byte) { + h.Add(b2s(key), b2s(value)) +} + // Set sets the given 'key: value' header. func (h *ResponseHeader) Set(key, value string) { initHeaderKV(&h.bufKV, key, value, h.disableNormalizing) @@ -917,6 +946,35 @@ func (h *RequestHeader) SetCookieBytesKV(key, value []byte) { h.cookies = setArgBytes(h.cookies, key, value) } +// Add adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *RequestHeader) Add(key, value string) { + k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing) + h.h = appendArg(h.h, b2s(k), value) +} + +// AddBytesK adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *RequestHeader) AddBytesK(key []byte, value string) { + h.Add(b2s(key), value) +} + +// AddBytesV adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *RequestHeader) AddBytesV(key string, value []byte) { + h.Add(key, b2s(value)) +} + +// AddBytesKV adds the given 'key: value' header. +// +// Multiple headers with the same key may be added. +func (h *RequestHeader) AddBytesKV(key, value []byte) { + h.Add(b2s(key), b2s(value)) +} + // Set sets the given 'key: value' header. func (h *RequestHeader) Set(key, value string) { initHeaderKV(&h.bufKV, key, value, h.disableNormalizing) diff --git a/header_test.go b/header_test.go index c39b250..8bbb055 100644 --- a/header_test.go +++ b/header_test.go @@ -10,6 +10,111 @@ import ( "testing" ) +func TestResponseHeaderAdd(t *testing.T) { + m := make(map[string]struct{}) + var h ResponseHeader + h.Add("aaa", "bbb") + m["bbb"] = struct{}{} + for i := 0; i < 10; i++ { + v := fmt.Sprintf("%d", i) + h.Add("Foo-Bar", v) + m[v] = struct{}{} + } + if h.Len() != 12 { + t.Fatalf("unexpected header len %d. Expecting 12", h.Len()) + } + + h.VisitAll(func(k, v []byte) { + switch string(k) { + case "Aaa", "Foo-Bar": + if _, ok := m[string(v)]; !ok { + t.Fatalf("unexpected value found %q. key %q", v, k) + } + delete(m, string(v)) + case "Content-Type": + default: + t.Fatalf("unexpected key found: %q", k) + } + }) + if len(m) > 0 { + t.Fatalf("%d headers are missed", len(m)) + } + + s := h.String() + br := bufio.NewReader(bytes.NewBufferString(s)) + var h1 ResponseHeader + if err := h1.Read(br); err != nil { + t.Fatalf("unexpected error: %s", err) + } + + h.VisitAll(func(k, v []byte) { + switch string(k) { + case "Aaa", "Foo-Bar": + m[string(v)] = struct{}{} + case "Content-Type": + default: + t.Fatalf("unexpected key found: %q", k) + } + }) + if len(m) != 11 { + t.Fatalf("unexpected number of headers: %d. Expecting 11", len(m)) + } +} + +func TestRequestHeaderAdd(t *testing.T) { + m := make(map[string]struct{}) + var h RequestHeader + h.Add("aaa", "bbb") + m["bbb"] = struct{}{} + for i := 0; i < 10; i++ { + v := fmt.Sprintf("%d", i) + h.Add("Foo-Bar", v) + m[v] = struct{}{} + } + if h.Len() != 11 { + t.Fatalf("unexpected header len %d. Expecting 11", h.Len()) + } + + h.VisitAll(func(k, v []byte) { + switch string(k) { + case "Aaa", "Foo-Bar": + if _, ok := m[string(v)]; !ok { + t.Fatalf("unexpected value found %q. key %q", v, k) + } + delete(m, string(v)) + default: + t.Fatalf("unexpected key found: %q", k) + } + }) + if len(m) > 0 { + t.Fatalf("%d headers are missed", len(m)) + } + + s := h.String() + br := bufio.NewReader(bytes.NewBufferString(s)) + var h1 RequestHeader + if err := h1.Read(br); err != nil { + t.Fatalf("unexpected error: %s", err) + } + + h.VisitAll(func(k, v []byte) { + switch string(k) { + case "Aaa", "Foo-Bar": + m[string(v)] = struct{}{} + case "User-Agent": + default: + t.Fatalf("unexpected key found: %q", k) + } + }) + if len(m) != 11 { + t.Fatalf("unexpected number of headers: %d. Expecting 11", len(m)) + } + s1 := h1.String() + if s != s1 { + t.Fatalf("unexpected headers %q. Expecting %q", s1, s) + } +} + func TestHasHeaderValue(t *testing.T) { testHasHeaderValue(t, "foobar", "foobar", true) testHasHeaderValue(t, "foobar", "foo", false)