diff --git a/args.go b/args.go index 47a97c6ac2..92e5d5aadc 100644 --- a/args.go +++ b/args.go @@ -633,3 +633,11 @@ func peekAllArgBytesToDst(dst [][]byte, h []argsKV, k []byte) [][]byte { } return dst } + +func peekArgsKeys(dst [][]byte, h []argsKV) [][]byte { + for i, n := 0, len(h); i < n; i++ { + kv := &h[i] + dst = append(dst, kv.key) + } + return dst +} diff --git a/header.go b/header.go index fd3defdcd9..72bdc7d29b 100644 --- a/header.go +++ b/header.go @@ -43,6 +43,7 @@ type ResponseHeader struct { contentEncoding []byte server []byte mulHeader [][]byte + keys [][]byte h []argsKV trailer []argsKV @@ -81,6 +82,7 @@ type RequestHeader struct { contentType []byte userAgent []byte mulHeader [][]byte + keys [][]byte h []argsKV trailer []argsKV @@ -977,6 +979,7 @@ func (h *ResponseHeader) resetSkipNormalize() { h.cookies = h.cookies[:0] h.trailer = h.trailer[:0] h.mulHeader = h.mulHeader[:0] + h.keys = h.keys[:0] } // SetNoDefaultContentType allows you to control if a default Content-Type header will be set (false) or not (true). @@ -1006,6 +1009,7 @@ func (h *RequestHeader) resetSkipNormalize() { h.userAgent = h.userAgent[:0] h.trailer = h.trailer[:0] h.mulHeader = h.mulHeader[:0] + h.keys = h.keys[:0] h.h = h.h[:0] h.cookies = h.cookies[:0] @@ -1847,7 +1851,7 @@ func (h *RequestHeader) peekAll(key []byte) [][]byte { // PeekAll returns all header value for the given key. // // The returned value is valid until the request is released, -// either though ReleaseRequest or your request handler returning. +// either though ReleaseResponse or your request handler returning. // Do not store references to returned value. Make copies instead. func (h *ResponseHeader) PeekAll(key string) [][]byte { k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing) @@ -1887,6 +1891,50 @@ func (h *ResponseHeader) peekAll(key []byte) [][]byte { return h.mulHeader } +// PeekKeys return all header keys. +// +// The returned value is valid until the request is released, +// either though ReleaseRequest or your request handler returning. +// Do not store references to returned value. Make copies instead. +func (h *RequestHeader) PeekKeys() [][]byte { + h.keys = h.keys[:0] + h.keys = peekArgsKeys(h.keys, h.h) + return h.keys +} + +// PeekTrailerKeys return all trailer keys. +// +// The returned value is valid until the request is released, +// either though ReleaseRequest or your request handler returning. +// Do not store references to returned value. Make copies instead. +func (h *RequestHeader) PeekTrailerKeys() [][]byte { + h.keys = h.keys[:0] + h.keys = peekArgsKeys(h.keys, h.trailer) + return h.keys +} + +// PeekKeys return all header keys. +// +// The returned value is valid until the request is released, +// either though ReleaseResponse or your request handler returning. +// Do not store references to returned value. Make copies instead. +func (h *ResponseHeader) PeekKeys() [][]byte { + h.keys = h.keys[:0] + h.keys = peekArgsKeys(h.keys, h.h) + return h.keys +} + +// PeekTrailerKeys return all trailer keys. +// +// The returned value is valid until the request is released, +// either though ReleaseResponse or your request handler returning. +// Do not store references to returned value. Make copies instead. +func (h *ResponseHeader) PeekTrailerKeys() [][]byte { + h.keys = h.keys[:0] + h.keys = peekArgsKeys(h.keys, h.trailer) + return h.keys +} + // Cookie returns cookie for the given key. func (h *RequestHeader) Cookie(key string) []byte { h.collectCookies() diff --git a/header_test.go b/header_test.go index fdefc3c9e8..7d2942551d 100644 --- a/header_test.go +++ b/header_test.go @@ -2935,3 +2935,43 @@ func expectResponseHeaderAll(t *testing.T, h *ResponseHeader, key string, expect t.Fatalf("Unexpected value for key %q: %q. Expected %q", key, h.PeekAll(key), expectedValue) } } + +func TestRequestHeader_Keys(t *testing.T) { + h := &RequestHeader{} + h.Add(HeaderConnection, "keep-alive") + h.Add("Content-Type", "aaa") + err := h.SetTrailer("aaa,bbb,ccc") + if err != nil { + t.Fatal(err) + } + actualKeys := h.PeekKeys() + expectedKeys := [][]byte{s2b("keep-alive"), s2b("aaa")} + if reflect.DeepEqual(actualKeys, expectedKeys) { + t.Fatalf("Unexpected value %q. Expected %q", actualKeys, expectedKeys) + } + actualTrailerKeys := h.PeekTrailerKeys() + expectedTrailerKeys := [][]byte{s2b("aaa"), s2b("bbb"), s2b("ccc")} + if reflect.DeepEqual(actualTrailerKeys, expectedTrailerKeys) { + t.Fatalf("Unexpected value %q. Expected %q", actualTrailerKeys, expectedTrailerKeys) + } +} + +func TestResponseHeader_Keys(t *testing.T) { + h := &ResponseHeader{} + h.Add(HeaderConnection, "keep-alive") + h.Add("Content-Type", "aaa") + err := h.SetTrailer("aaa,bbb,ccc") + if err != nil { + t.Fatal(err) + } + actualKeys := h.PeekKeys() + expectedKeys := [][]byte{s2b("keep-alive"), s2b("aaa")} + if reflect.DeepEqual(actualKeys, expectedKeys) { + t.Fatalf("Unexpected value %q. Expected %q", actualKeys, expectedKeys) + } + actualTrailerKeys := h.PeekTrailerKeys() + expectedTrailerKeys := [][]byte{s2b("aaa"), s2b("bbb"), s2b("ccc")} + if reflect.DeepEqual(actualTrailerKeys, expectedTrailerKeys) { + t.Fatalf("Unexpected value %q. Expected %q", actualTrailerKeys, expectedTrailerKeys) + } +}