feat: add PeekKeys and PeekTrailerKeys (#1405)

* feat: add PeekKeys and PeekTrailerKeys * Improve warning Co-authored-by: Erik Dubbelboer <[email protected]>
valyala · Oct 29, 2022 · 3963a79 · 3963a79
1 parent eca86de
commit 3963a79
Show file tree

Hide file tree

Showing 3 changed files with 99 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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)
+	}
+}