Repository: cloudwego/hertz
Pull request #569 (2023-01-31 04:31:18)

feat: add PeekAll function
#### What type of PR is this?
feat
<!--
Add one of the following kinds:

build: Changes that affect the build system or external dependencies (example scopes: gulp, broccoli, npm)
ci: Changes to our CI configuration files and scripts (example scopes: Travis, Circle, BrowserStack, SauceLabs)
docs: Documentation only changes
feat: A new feature
optimize: A new optimization
fix: A bug fix
perf: A code change that improves performance
refactor: A code change that neither fixes a bug nor adds a feature
style: Changes that do not affect the meaning of the code (white space, formatting, missing semi-colons, etc)
test: Adding missing tests or correcting existing tests
chore: Changes to the build process or auxiliary tools and libraries such as documentation generation
-->

#### Check the PR title.
<!--
The description of the title will be attached in Release Notes, 
so please describe it from user-oriented, what this PR does / why we need it.
Please check your PR title with the below requirements:
-->
- [x] This PR title match the format: \<type\>(optional scope): \<description\>
- [x] The description of this PR title is user-oriented and clear enough for others to understand.


#### (Optional) Translate the PR title into Chinese.

添加 PeekAll 函数



#### Which issue(s) this PR fixes:
<!--
Automatically closes linked issue when PR is merged.
Eg: `Fixes #<issue number>`, or `Fixes (paste link of issue)`.
-->


codecov[bot]: # [Codecov](https://codecov.io/gh/cloudwego/hertz/pull/569?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) Report
Base: **71.23**% // Head: **71.52**% // Increases project coverage by **`+0.29%`** :tada:
> Coverage data is based on head [(`6844d61`)](https://codecov.io/gh/cloudwego/hertz/pull/569?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) compared to base [(`75ab005`)](https://codecov.io/gh/cloudwego/hertz/commit/75ab005e05850981648a244311b6c6c0f3e90e28?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).
> Patch coverage: 89.23% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##           develop     #569      +/-   ##
===========================================
+ Coverage    71.23%   71.52%   +0.29%     
===========================================
  Files           93       93              
  Lines         8923     8988      +65     
===========================================
+ Hits          6356     6429      +73     
+ Misses        2181     2169      -12     
- Partials       386      390       +4     
```


| [Impacted Files](https://codecov.io/gh/cloudwego/hertz/pull/569?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego) | Coverage Δ | |
|---|---|---|
| [pkg/protocol/header.go](https://codecov.io/gh/cloudwego/hertz/pull/569?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2hlYWRlci5nbw==) | `77.62% <88.13%> (+2.02%)` | :arrow_up: |
| [pkg/protocol/args.go](https://codecov.io/gh/cloudwego/hertz/pull/569?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2FyZ3MuZ28=) | `81.86% <100.00%> (+0.54%)` | :arrow_up: |
| [pkg/protocol/cookie.go](https://codecov.io/gh/cloudwego/hertz/pull/569?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego#diff-cGtnL3Byb3RvY29sL2Nvb2tpZS5nbw==) | `90.53% <0.00%> (+2.05%)` | :arrow_up: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/cloudwego/hertz/pull/569?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cloudwego).


Patch:
diff --git a/pkg/protocol/args.go b/pkg/protocol/args.go
index 5731db5..ebf9ae5 100644
--- a/pkg/protocol/args.go
+++ b/pkg/protocol/args.go
@@ -224,6 +224,16 @@ func peekArgBytes(h []argsKV, k []byte) []byte {
 	return nil
 }
 
+func peekAllArgBytesToDst(dst [][]byte, h []argsKV, k []byte) [][]byte {
+	for i, n := 0, len(h); i < n; i++ {
+		kv := &h[i]
+		if bytes.Equal(kv.key, k) {
+			dst = append(dst, kv.value)
+		}
+	}
+	return dst
+}
+
 func delAllArgsBytes(args []argsKV, key []byte) []argsKV {
 	return delAllArgs(args, bytesconv.B2s(key))
 }
diff --git a/pkg/protocol/header.go b/pkg/protocol/header.go
index b447cce..dfc5c8d 100644
--- a/pkg/protocol/header.go
+++ b/pkg/protocol/header.go
@@ -83,6 +83,7 @@ type RequestHeader struct {
 	host        []byte
 	contentType []byte
 	userAgent   []byte
+	mulHeader   [][]byte
 
 	h     []argsKV
 	bufKV argsKV
@@ -121,6 +122,7 @@ type ResponseHeader struct {
 
 	contentType []byte
 	server      []byte
+	mulHeader   [][]byte
 
 	h     []argsKV
 	bufKV argsKV
@@ -486,6 +488,7 @@ func (h *ResponseHeader) ResetSkipNormalize() {
 
 	h.h = h.h[:0]
 	h.cookies = h.cookies[:0]
+	h.mulHeader = h.mulHeader[:0]
 }
 
 // ContentLength returns Content-Length header value.
@@ -687,6 +690,94 @@ func (h *ResponseHeader) peek(key []byte) []byte {
 	}
 }
 
+// PeekAll returns all header value for the given key.
+//
+// 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. Use ResponseHeader.GetAll(key) instead.
+func (h *ResponseHeader) PeekAll(key string) [][]byte {
+	k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing)
+	return h.peekAll(k)
+}
+
+func (h *ResponseHeader) peekAll(key []byte) [][]byte {
+	h.mulHeader = h.mulHeader[:0]
+	switch string(key) {
+	case consts.HeaderContentType:
+		if contentType := h.ContentType(); len(contentType) > 0 {
+			h.mulHeader = append(h.mulHeader, contentType)
+		}
+	case consts.HeaderContentEncoding:
+		if contentEncoding := h.ContentEncoding(); len(contentEncoding) > 0 {
+			h.mulHeader = append(h.mulHeader, contentEncoding)
+		}
+	case consts.HeaderServer:
+		if server := h.Server(); len(server) > 0 {
+			h.mulHeader = append(h.mulHeader, server)
+		}
+	case consts.HeaderConnection:
+		if h.ConnectionClose() {
+			h.mulHeader = append(h.mulHeader, bytestr.StrClose)
+		} else {
+			h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
+		}
+	case consts.HeaderContentLength:
+		h.mulHeader = append(h.mulHeader, h.contentLengthBytes)
+	case consts.HeaderSetCookie:
+		h.mulHeader = append(h.mulHeader, appendResponseCookieBytes(nil, h.cookies))
+	default:
+		h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
+	}
+	return h.mulHeader
+}
+
+// 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.
+// Any future calls to the Peek* will modify the returned value.
+// Do not store references to returned value. Use RequestHeader.GetAll(key) instead.
+func (h *RequestHeader) PeekAll(key string) [][]byte {
+	k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing)
+	return h.peekAll(k)
+}
+
+func (h *RequestHeader) peekAll(key []byte) [][]byte {
+	h.mulHeader = h.mulHeader[:0]
+	switch string(key) {
+	case consts.HeaderHost:
+		if host := h.Host(); len(host) > 0 {
+			h.mulHeader = append(h.mulHeader, host)
+		}
+	case consts.HeaderContentType:
+		if contentType := h.ContentType(); len(contentType) > 0 {
+			h.mulHeader = append(h.mulHeader, contentType)
+		}
+	case consts.HeaderUserAgent:
+		if ua := h.UserAgent(); len(ua) > 0 {
+			h.mulHeader = append(h.mulHeader, ua)
+		}
+	case consts.HeaderConnection:
+		if h.ConnectionClose() {
+			h.mulHeader = append(h.mulHeader, bytestr.StrClose)
+		} else {
+			h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
+		}
+	case consts.HeaderContentLength:
+		h.mulHeader = append(h.mulHeader, h.contentLengthBytes)
+	case consts.HeaderCookie:
+		if h.cookiesCollected {
+			h.mulHeader = append(h.mulHeader, appendRequestCookieBytes(nil, h.cookies))
+		} else {
+			h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
+		}
+	default:
+		h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
+	}
+	return h.mulHeader
+}
+
 // SetContentTypeBytes sets Content-Type header value.
 func (h *ResponseHeader) SetContentTypeBytes(contentType []byte) {
 	h.contentType = append(h.contentType[:0], contentType...)
@@ -1299,6 +1390,7 @@ func (h *RequestHeader) ResetSkipNormalize() {
 	h.cookiesCollected = false
 
 	h.rawHeaders = h.rawHeaders[:0]
+	h.mulHeader = h.mulHeader[:0]
 }
 
 func peekRawHeader(buf, key []byte) []byte {
@@ -1482,6 +1574,28 @@ func (h *ResponseHeader) Get(key string) string {
 	return string(h.Peek(key))
 }
 
+// GetAll returns all header value for the given key
+// it is concurrent safety and long lifetime.
+func (h *RequestHeader) GetAll(key string) []string {
+	res := make([]string, 0)
+	headers := h.PeekAll(key)
+	for _, header := range headers {
+		res = append(res, string(header))
+	}
+	return res
+}
+
+// GetAll returns all header value for the given key and is concurrent safety.
+// it is concurrent safety and long lifetime.
+func (h *ResponseHeader) GetAll(key string) []string {
+	res := make([]string, 0)
+	headers := h.PeekAll(key)
+	for _, header := range headers {
+		res = append(res, string(header))
+	}
+	return res
+}
+
 func appendHeaderLine(dst, key, value []byte) []byte {
 	dst = append(dst, key...)
 	dst = append(dst, bytestr.StrColonSpace...)


Test patch:
diff --git a/pkg/protocol/header_test.go b/pkg/protocol/header_test.go
index 0a7baf2..fc58a29 100644
--- a/pkg/protocol/header_test.go
+++ b/pkg/protocol/header_test.go
@@ -572,3 +572,72 @@ func TestRequestHeaderSetNoDefaultContentType(t *testing.T) {
 	b = h.AppendBytes(nil)
 	assert.DeepEqual(t, b, []byte("POST / HTTP/1.1\r\n\r\n"))
 }
+
+func TestRequestHeader_PeekAll(t *testing.T) {
+	t.Parallel()
+	h := &RequestHeader{}
+	h.Add(consts.HeaderConnection, "keep-alive")
+	h.Add("Content-Type", "aaa")
+	h.Add(consts.HeaderHost, "aaabbb")
+	h.Add("User-Agent", "asdfas")
+	h.Add("Content-Length", "1123")
+	h.Add("Cookie", "foobar=baz")
+	h.Add("aaa", "aaa")
+	h.Add("aaa", "bbb")
+
+	expectRequestHeaderAll(t, h, consts.HeaderConnection, [][]byte{[]byte("keep-alive")})
+	expectRequestHeaderAll(t, h, "Content-Type", [][]byte{[]byte("aaa")})
+	expectRequestHeaderAll(t, h, consts.HeaderHost, [][]byte{[]byte("aaabbb")})
+	expectRequestHeaderAll(t, h, "User-Agent", [][]byte{[]byte("asdfas")})
+	expectRequestHeaderAll(t, h, "Content-Length", [][]byte{[]byte("1123")})
+	expectRequestHeaderAll(t, h, "Cookie", [][]byte{[]byte("foobar=baz")})
+	expectRequestHeaderAll(t, h, "aaa", [][]byte{[]byte("aaa"), []byte("bbb")})
+
+	h.DelBytes([]byte("Content-Type"))
+	h.DelBytes([]byte((consts.HeaderHost)))
+	h.DelBytes([]byte("aaa"))
+	expectRequestHeaderAll(t, h, "Content-Type", [][]byte{})
+	expectRequestHeaderAll(t, h, consts.HeaderHost, [][]byte{})
+	expectRequestHeaderAll(t, h, "aaa", [][]byte{})
+}
+
+func expectRequestHeaderAll(t *testing.T, h *RequestHeader, key string, expectedValue [][]byte) {
+	if len(h.PeekAll(key)) != len(expectedValue) {
+		t.Fatalf("Unexpected size for key %q: %d. Expected %d", key, len(h.PeekAll(key)), len(expectedValue))
+	}
+	assert.DeepEqual(t, h.PeekAll(key), expectedValue)
+}
+
+func TestResponseHeader_PeekAll(t *testing.T) {
+	t.Parallel()
+
+	h := &ResponseHeader{}
+	h.Add(consts.HeaderContentType, "aaa/bbb")
+	h.Add(consts.HeaderContentEncoding, "gzip")
+	h.Add(consts.HeaderConnection, "close")
+	h.Add(consts.HeaderContentLength, "1234")
+	h.Add(consts.HeaderServer, "aaaa")
+	h.Add(consts.HeaderSetCookie, "cccc")
+	h.Add("aaa", "aaa")
+	h.Add("aaa", "bbb")
+
+	expectResponseHeaderAll(t, h, consts.HeaderContentType, [][]byte{[]byte("aaa/bbb")})
+	expectResponseHeaderAll(t, h, consts.HeaderContentEncoding, [][]byte{[]byte("gzip")})
+	expectResponseHeaderAll(t, h, consts.HeaderConnection, [][]byte{[]byte("close")})
+	expectResponseHeaderAll(t, h, consts.HeaderContentLength, [][]byte{[]byte("1234")})
+	expectResponseHeaderAll(t, h, consts.HeaderServer, [][]byte{[]byte("aaaa")})
+	expectResponseHeaderAll(t, h, consts.HeaderSetCookie, [][]byte{[]byte("cccc")})
+	expectResponseHeaderAll(t, h, "aaa", [][]byte{[]byte("aaa"), []byte("bbb")})
+
+	h.Del(consts.HeaderContentType)
+	h.Del(consts.HeaderContentEncoding)
+	expectResponseHeaderAll(t, h, consts.HeaderContentType, [][]byte{bytestr.DefaultContentType})
+	expectResponseHeaderAll(t, h, consts.HeaderContentEncoding, [][]byte{})
+}
+
+func expectResponseHeaderAll(t *testing.T, h *ResponseHeader, key string, expectedValue [][]byte) {
+	if len(h.PeekAll(key)) != len(expectedValue) {
+		t.Fatalf("Unexpected size for key %q: %d. Expected %d", key, len(h.PeekAll(key)), len(expectedValue))
+	}
+	assert.DeepEqual(t, h.PeekAll(key), expectedValue)
+}