Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add PeekKeys and PeekTrailerKeys #1405

Merged
merged 2 commits into from Oct 29, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
8 changes: 8 additions & 0 deletions args.go
Expand Up @@ -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
}
52 changes: 51 additions & 1 deletion header.go
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand Down
40 changes: 40 additions & 0 deletions header_test.go
Expand Up @@ -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)
}
}