Repository: valyala/fasthttp
Pull request #1405 (2022-10-22 17:10:43)

feat: add PeekKeys and PeekTrailerKeys
pr for https://github.com/valyala/fasthttp/issues/1400


Patch:
diff --git a/args.go b/args.go
index 47a97c6..92e5d5a 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 fd3defd..fe4f669 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()
@@ -2520,8 +2570,8 @@ func (h *ResponseHeader) parseTrailer(buf []byte) (int, error) {
 				err = fmt.Errorf("forbidden trailer key %q", s.key)
 				continue
 			}
+			h.h = appendArgBytes(h.h, s.key, s.value, argsHasValue)
 		}
-		h.h = appendArgBytes(h.h, s.key, s.value, argsHasValue)
 	}
 	if s.err != nil {
 		return 0, s.err
diff --git a/http.go b/http.go
index 97bf5d6..62ba30f 100644
--- a/http.go
+++ b/http.go
@@ -1287,7 +1287,11 @@ func (req *Request) ContinueReadBodyStream(r *bufio.Reader, maxBodySize int, pre
 		// the end of body is determined by connection close.
 		// So just ignore request body for requests without
 		// 'Content-Length' and 'Transfer-Encoding' headers.
-		req.Header.SetContentLength(0)
+
+		// refer to https://tools.ietf.org/html/rfc7230#section-3.3.2
+		if !req.Header.ignoreBody() {
+			req.Header.SetContentLength(0)
+		}
 		return nil
 	}
 


Test patch:
diff --git a/header_test.go b/header_test.go
index fdefc3c..7d29425 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)
+	}
+}
diff --git a/http_test.go b/http_test.go
index 58bc8cb..60e7891 100644
--- a/http_test.go
+++ b/http_test.go
@@ -1053,6 +1053,25 @@ func TestRequestReadNoBody(t *testing.T) {
 	}
 }
 
+func TestRequestReadNoBodyStreaming(t *testing.T) {
+	t.Parallel()
+
+	var r Request
+
+	r.Header.contentLength = -2
+
+	br := bufio.NewReader(bytes.NewBufferString("GET / HTTP/1.1\r\n\r\n"))
+	err := r.ContinueReadBodyStream(br, 0)
+	r.SetHost("foobar")
+	if err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	s := r.String()
+	if strings.Contains(s, "Content-Length: ") {
+		t.Fatalf("unexpected Content-Length")
+	}
+}
+
 func TestResponseWriteTo(t *testing.T) {
 	t.Parallel()
 
@@ -1641,10 +1660,10 @@ func TestRequestReadLimitBody(t *testing.T) {
 }
 
 func testResponseReadLimitBodyError(t *testing.T, s string, maxBodySize int, expectedErr error) {
-	var req Response
+	var resp Response
 	r := bytes.NewBufferString(s)
 	br := bufio.NewReader(r)
-	err := req.ReadLimitBody(br, maxBodySize)
+	err := resp.ReadLimitBody(br, maxBodySize)
 	if err == nil {
 		t.Fatalf("expecting error. s=%q, maxBodySize=%d", s, maxBodySize)
 	}
@@ -1654,10 +1673,10 @@ func testResponseReadLimitBodyError(t *testing.T, s string, maxBodySize int, exp
 }
 
 func testResponseReadLimitBodySuccess(t *testing.T, s string, maxBodySize int) {
-	var req Response
+	var resp Response
 	r := bytes.NewBufferString(s)
 	br := bufio.NewReader(r)
-	if err := req.ReadLimitBody(br, maxBodySize); err != nil {
+	if err := resp.ReadLimitBody(br, maxBodySize); err != nil {
 		t.Fatalf("unexpected error: %v. s=%q, maxBodySize=%d", err, s, maxBodySize)
 	}
 }