Skip to content

Commit

Permalink
feat: add PeekKeys and PeekTrailerKeys
Browse files Browse the repository at this point in the history
  • Loading branch information
li-jin-gou committed Oct 23, 2022
1 parent 128e9b3 commit 276d8cd
Show file tree
Hide file tree
Showing 3 changed files with 93 additions and 1 deletion.
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
}
46 changes: 45 additions & 1 deletion header.go
Expand Up @@ -1847,7 +1847,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)
Expand Down Expand Up @@ -1887,6 +1887,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.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.
// 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.
// 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.
// 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)
}
}

0 comments on commit 276d8cd

Please sign in to comment.