From 81ebee8c79ae193cb2e30ed2b326ece3681278a8 Mon Sep 17 00:00:00 2001
From: Erik Dubbelboer <erik@dubbelboer.com>
Date: Sat, 9 Aug 2025 13:50:07 +0200
Subject: [PATCH] Fix PeekKeys()

RequestHeader.PeekKeys() and ResponseHeader.PeekKeys() were both
implemented wrong. The tests were also wrong causing this to never be
noticed. They both never actually returned all header keys, this has
been fixed now.

While this is a backwards incompatible change, I'm still going to
release it. Anyone using these functions would have noticed they
didn't work as documented and probably would not have continued using
them.

Fixes https://github.com/valyala/fasthttp/issues/2044
---
 args.go        |  8 --------
 header.go      | 20 ++++++++++++++++++--
 header_test.go | 16 ++++++++--------
 3 files changed, 26 insertions(+), 18 deletions(-)

diff --git a/args.go b/args.go
index 938e62e..a2290a8 100644
--- a/args.go
+++ b/args.go
@@ -644,11 +644,3 @@ 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 b6f9a8f..9971bb5 100644
--- a/header.go
+++ b/header.go
@@ -1992,9 +1992,25 @@ func (h *ResponseHeader) peekAll(key []byte) [][]byte {
 // 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 *header) PeekKeys() [][]byte {
+func (h *RequestHeader) PeekKeys() [][]byte {
 	h.mulHeader = h.mulHeader[:0]
-	h.mulHeader = peekArgsKeys(h.mulHeader, h.h)
+	for key := range h.All() {
+		h.mulHeader = append(h.mulHeader, key)
+	}
+	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 *ResponseHeader) PeekKeys() [][]byte {
+	h.mulHeader = h.mulHeader[:0]
+	for key := range h.All() {
+		h.mulHeader = append(h.mulHeader, key)
+	}
 	return h.mulHeader
 }
 
diff --git a/header_test.go b/header_test.go
index 430f292..da0b75e 100644
--- a/header_test.go
+++ b/header_test.go
@@ -3426,13 +3426,13 @@ func TestRequestHeader_Keys(t *testing.T) {
 		t.Fatal(err)
 	}
 	actualKeys := h.PeekKeys()
-	expectedKeys := [][]byte{s2b("keep-alive"), s2b("aaa")}
-	if reflect.DeepEqual(actualKeys, expectedKeys) {
+	expectedKeys := [][]byte{[]byte("Content-Type"), []byte("Trailer"), []byte("Connection")}
+	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) {
+	expectedTrailerKeys := [][]byte{s2b("Aaa"), s2b("Bbb"), s2b("Ccc")}
+	if !reflect.DeepEqual(actualTrailerKeys, expectedTrailerKeys) {
 		t.Fatalf("Unexpected value %q. Expected %q", actualTrailerKeys, expectedTrailerKeys)
 	}
 }
@@ -3446,13 +3446,13 @@ func TestResponseHeader_Keys(t *testing.T) {
 		t.Fatal(err)
 	}
 	actualKeys := h.PeekKeys()
-	expectedKeys := [][]byte{s2b("keep-alive"), s2b("aaa")}
-	if reflect.DeepEqual(actualKeys, expectedKeys) {
+	expectedKeys := [][]byte{[]byte("Content-Type"), []byte("Trailer"), []byte("Connection")}
+	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) {
+	expectedTrailerKeys := [][]byte{s2b("Aaa"), s2b("Bbb"), s2b("Ccc")}
+	if !reflect.DeepEqual(actualTrailerKeys, expectedTrailerKeys) {
 		t.Fatalf("Unexpected value %q. Expected %q", actualTrailerKeys, expectedTrailerKeys)
 	}
 }