Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add header.PeekAll #1394

Merged
merged 1 commit into from
Oct 15, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions args.go
Original file line number Diff line number Diff line change
Expand Up @@ -623,3 +623,13 @@ func decodeArgAppendNoPlus(dst, src []byte) []byte {
}
return dst
}

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
}
1 change: 1 addition & 0 deletions bytesconv_table_gen.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"bytes"
"fmt"
"log"
"os"
)

const (
Expand Down
3 changes: 2 additions & 1 deletion fasthttpadaptor/request_test.go
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
package fasthttpadaptor

import (
"github.com/valyala/fasthttp"
"net/http"
"testing"

"github.com/valyala/fasthttp"
)

func BenchmarkConvertRequest(b *testing.B) {
Expand Down
83 changes: 83 additions & 0 deletions header.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ type ResponseHeader struct {
contentType []byte
contentEncoding []byte
server []byte
mulHeader [][]byte

h []argsKV
trailer []argsKV
Expand Down Expand Up @@ -79,6 +80,7 @@ type RequestHeader struct {
host []byte
contentType []byte
userAgent []byte
mulHeader [][]byte

h []argsKV
trailer []argsKV
Expand Down Expand Up @@ -974,6 +976,7 @@ func (h *ResponseHeader) resetSkipNormalize() {
h.h = h.h[:0]
h.cookies = h.cookies[:0]
h.trailer = h.trailer[:0]
h.mulHeader = h.mulHeader[:0]
}

// SetNoDefaultContentType allows you to control if a default Content-Type header will be set (false) or not (true).
Expand Down Expand Up @@ -1002,6 +1005,7 @@ func (h *RequestHeader) resetSkipNormalize() {
h.contentType = h.contentType[:0]
h.userAgent = h.userAgent[:0]
h.trailer = h.trailer[:0]
h.mulHeader = h.mulHeader[:0]

h.h = h.h[:0]
h.cookies = h.cookies[:0]
Expand Down Expand Up @@ -1793,6 +1797,85 @@ func (h *RequestHeader) 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 ReleaseRequest or your request handler returning.
// Do not store references to returned value. Make copies instead.
func (h *RequestHeader) PeekAll(key string) [][]byte {
k := getHeaderKeyBytes(&h.bufKV, key, h.disableNormalizing)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if we should do the same trick here as with bufKV. Where we put a [][]byte in RequestHeader and reuse that instead of allocating the return value on the heap.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks,You're right

Copy link
Contributor Author

@li-jin-gou li-jin-gou Oct 13, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I modified the logic and added a mulHeader field.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add a comment like this to these functions?

// The returned value is not safe to use after the handler returns, the request is released or other calls to PeekAll.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

thanks

return h.peekAll(k)
}

func (h *RequestHeader) peekAll(key []byte) [][]byte {
h.mulHeader = h.mulHeader[:0]
switch string(key) {
case HeaderHost:
h.mulHeader = append(h.mulHeader, h.Host())
case HeaderContentType:
h.mulHeader = append(h.mulHeader, h.ContentType())
case HeaderUserAgent:
h.mulHeader = append(h.mulHeader, h.UserAgent())
case HeaderConnection:
if h.ConnectionClose() {
h.mulHeader = append(h.mulHeader, strClose)
} else {
h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
}
case HeaderContentLength:
h.mulHeader = append(h.mulHeader, h.contentLengthBytes)
case HeaderCookie:
if h.cookiesCollected {
h.mulHeader = append(h.mulHeader, appendRequestCookieBytes(nil, h.cookies))
return [][]byte{appendRequestCookieBytes(nil, h.cookies)}
} else {
h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
}
case HeaderTrailer:
h.mulHeader = append(h.mulHeader, appendArgsKeyBytes(nil, h.trailer, strCommaSpace))
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.
// Do not store references to returned value. Make copies 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 HeaderContentType:
h.mulHeader = append(h.mulHeader, h.ContentType())
case HeaderContentEncoding:
h.mulHeader = append(h.mulHeader, h.ContentEncoding())
case HeaderServer:
h.mulHeader = append(h.mulHeader, h.Server())
case HeaderConnection:
if h.ConnectionClose() {
h.mulHeader = append(h.mulHeader, strClose)
} else {
h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
}
case HeaderContentLength:
h.mulHeader = append(h.mulHeader, h.contentLengthBytes)
case HeaderSetCookie:
h.mulHeader = append(h.mulHeader, appendResponseCookieBytes(nil, h.cookies))
case HeaderTrailer:
h.mulHeader = append(h.mulHeader, appendArgsKeyBytes(nil, h.trailer, strCommaSpace))
default:
h.mulHeader = peekAllArgBytesToDst(h.mulHeader, h.h, key)
}
return h.mulHeader
}

// Cookie returns cookie for the given key.
func (h *RequestHeader) Cookie(key string) []byte {
h.collectCookies()
Expand Down
62 changes: 62 additions & 0 deletions header_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -2861,3 +2861,65 @@ func verifyTrailer(t *testing.T, r *bufio.Reader, expectedTrailers map[string]st
}
verifyResponseTrailer(t, &resp.Header, expectedTrailers)
}

func TestRequestHeader_PeekAll(t *testing.T) {
t.Parallel()
h := &RequestHeader{}
h.Add(HeaderConnection, "keep-alive")
h.Add("Content-Type", "aaa")
h.Add(HeaderHost, "aaabbb")
h.Add("User-Agent", "asdfas")
h.Add("Content-Length", "1123")
h.Add("Cookie", "foobar=baz")
h.Add(HeaderTrailer, "foo, bar")
h.Add("aaa", "aaa")
h.Add("aaa", "bbb")

expectRequestHeaderAll(t, h, HeaderConnection, [][]byte{s2b("keep-alive")})
expectRequestHeaderAll(t, h, "Content-Type", [][]byte{s2b("aaa")})
expectRequestHeaderAll(t, h, HeaderHost, [][]byte{s2b("aaabbb")})
expectRequestHeaderAll(t, h, "User-Agent", [][]byte{s2b("asdfas")})
expectRequestHeaderAll(t, h, "Content-Length", [][]byte{s2b("1123")})
expectRequestHeaderAll(t, h, "Cookie", [][]byte{s2b("foobar=baz")})
expectRequestHeaderAll(t, h, HeaderTrailer, [][]byte{s2b("Foo, Bar")})
expectRequestHeaderAll(t, h, "aaa", [][]byte{s2b("aaa"), s2b("bbb")})
}
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))
}
if !reflect.DeepEqual(h.PeekAll(key), expectedValue) {
t.Fatalf("Unexpected value for key %q: %q. Expected %q", key, h.PeekAll(key), expectedValue)
}
}

func TestResponseHeader_PeekAll(t *testing.T) {
t.Parallel()

h := &ResponseHeader{}
h.Add(HeaderContentType, "aaa/bbb")
h.Add(HeaderContentEncoding, "gzip")
h.Add(HeaderConnection, "close")
h.Add(HeaderContentLength, "1234")
h.Add(HeaderServer, "aaaa")
h.Add(HeaderSetCookie, "cccc")
h.Add("aaa", "aaa")
h.Add("aaa", "bbb")

expectResponseHeaderAll(t, h, HeaderContentType, [][]byte{s2b("aaa/bbb")})

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

does it make sense to add a test for when content type is deleted, whether the expected result is the default type or empty? see #1402 as it is ambiguous.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for reminding

expectResponseHeaderAll(t, h, HeaderContentEncoding, [][]byte{s2b("gzip")})
expectResponseHeaderAll(t, h, HeaderConnection, [][]byte{s2b("close")})
expectResponseHeaderAll(t, h, HeaderContentLength, [][]byte{s2b("1234")})
expectResponseHeaderAll(t, h, HeaderServer, [][]byte{s2b("aaaa")})
expectResponseHeaderAll(t, h, HeaderSetCookie, [][]byte{s2b("cccc")})
expectResponseHeaderAll(t, h, "aaa", [][]byte{s2b("aaa"), s2b("bbb")})
}

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))
}
if !reflect.DeepEqual(h.PeekAll(key), expectedValue) {
t.Fatalf("Unexpected value for key %q: %q. Expected %q", key, h.PeekAll(key), expectedValue)
}
}