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..4cbfff3a84 100644 --- a/header.go +++ b/header.go @@ -1801,6 +1801,7 @@ func (h *RequestHeader) peek(key []byte) []byte { // // The returned value is valid until the request is released, // either though ReleaseRequest or your request handler returning. +// Any future calls to the Peek* will modify the returned value. // Do not store references to returned value. Make copies instead. func (h *RequestHeader) PeekAll(key string) [][]byte { k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing) @@ -1847,7 +1848,8 @@ 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. +// Any future calls to the Peek* will modify the returned value. // 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 +1889,54 @@ 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. +// Any future calls to the Peek* will modify the returned value. +// Do not store references to returned value. Make copies instead. +func (h *RequestHeader) PeekKeys() [][]byte { + h.mulHeader = h.mulHeader[:0] + h.mulHeader = peekArgsKeys(h.mulHeader, h.h) + return h.mulHeader +} + +// PeekTrailerKeys return all trailer keys. +// +// The returned value is valid until the request is released, +// either though ReleaseRequest or your request handler returning. +// Any future calls to the Peek* will modify the returned value. +// Do not store references to returned value. Make copies instead. +func (h *RequestHeader) PeekTrailerKeys() [][]byte { + h.mulHeader = h.mulHeader[:0] + h.mulHeader = peekArgsKeys(h.mulHeader, h.trailer) + return h.mulHeader +} + +// PeekKeys return all header keys. +// +// The returned value is valid until the request is released, +// either though ReleaseResponse or your request handler returning. +// Any future calls to the Peek* will modify the returned value. +// Do not store references to returned value. Make copies instead. +func (h *ResponseHeader) PeekKeys() [][]byte { + h.mulHeader = h.mulHeader[:0] + h.mulHeader = peekArgsKeys(h.mulHeader, h.h) + return h.mulHeader +} + +// PeekTrailerKeys return all trailer keys. +// +// The returned value is valid until the request is released, +// either though ReleaseResponse or your request handler returning. +// Any future calls to the Peek* will modify the returned value. +// Do not store references to returned value. Make copies instead. +func (h *ResponseHeader) PeekTrailerKeys() [][]byte { + h.mulHeader = h.mulHeader[:0] + h.mulHeader = peekArgsKeys(h.mulHeader, h.trailer) + return h.mulHeader +} + // 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) + } +}