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

BodyDecoded() for request and responses #1308

Merged
merged 2 commits into from Jun 6, 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.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
15 changes: 15 additions & 0 deletions header.go
Expand Up @@ -371,6 +371,21 @@ func (h *RequestHeader) SetContentTypeBytes(contentType []byte) {
h.contentType = append(h.contentType[:0], contentType...)
}

// ContentEncoding returns Content-Encoding header value.
func (h *RequestHeader) ContentEncoding() []byte {
return peekArgBytes(h.h, strContentEncoding)
}

// SetContentEncoding sets Content-Encoding header value.
func (h *RequestHeader) SetContentEncoding(contentEncoding string) {
h.SetBytesK(strContentEncoding, contentEncoding)
}

// SetContentEncodingBytes sets Content-Encoding header value.
func (h *RequestHeader) SetContentEncodingBytes(contentEncoding []byte) {
h.setNonSpecial(strContentEncoding, contentEncoding)
}

// SetMultipartFormBoundary sets the following Content-Type:
// 'multipart/form-data; boundary=...'
// where ... is substituted by the given boundary.
Expand Down
42 changes: 42 additions & 0 deletions http.go
Expand Up @@ -486,6 +486,48 @@ func inflateData(p []byte) ([]byte, error) {
return bb.B, nil
}

var ErrContentEncodingUnsupported = errors.New("unsupported Content-Encoding")

// BodyUncompressed returns body data and if needed decompress it from gzip, deflate or Brotli.
//
// This method may be used if the response header contains
// 'Content-Encoding' for reading uncompressed request body.
// Use Body for reading the raw request body.
func (req *Request) BodyUncompressed() ([]byte, error) {
switch string(req.Header.ContentEncoding()) {
case "":
return req.Body(), nil
case "deflate":
return req.BodyInflate()
case "gzip":
return req.BodyGunzip()
case "br":
return req.BodyUnbrotli()
default:
return nil, ErrContentEncodingUnsupported
}
}

// BodyUncompressed returns body data and if needed decompress it from gzip, deflate or Brotli.
//
// This method may be used if the response header contains
// 'Content-Encoding' for reading uncompressed response body.
// Use Body for reading the raw response body.
func (resp *Response) BodyUncompressed() ([]byte, error) {
switch string(resp.Header.ContentEncoding()) {
case "":
return resp.Body(), nil
case "deflate":
return resp.BodyInflate()
case "gzip":
return resp.BodyGunzip()
case "br":
return resp.BodyUnbrotli()
default:
return nil, ErrContentEncodingUnsupported
}
}

// BodyWriteTo writes request body to w.
func (req *Request) BodyWriteTo(w io.Writer) error {
if req.bodyStream != nil {
Expand Down
56 changes: 54 additions & 2 deletions http_test.go
Expand Up @@ -347,6 +347,12 @@ func testResponseBodyStreamDeflate(t *testing.T, body []byte, bodySize int) {
if !bytes.Equal(respBody, body) {
t.Fatalf("unexpected body: %q. Expecting %q", respBody, body)
}
// check for invalid
resp.SetBodyRaw([]byte("invalid"))
_, errDeflate := resp.BodyInflate()
if errDeflate == nil || errDeflate.Error() != "zlib: invalid header" {
t.Fatalf("expected error: 'zlib: invalid header' but was %v", errDeflate)
}
}

func testResponseBodyStreamGzip(t *testing.T, body []byte, bodySize int) {
Expand Down Expand Up @@ -375,6 +381,12 @@ func testResponseBodyStreamGzip(t *testing.T, body []byte, bodySize int) {
if !bytes.Equal(respBody, body) {
t.Fatalf("unexpected body: %q. Expecting %q", respBody, body)
}
// check for invalid
resp.SetBodyRaw([]byte("invalid"))
_, errUnzip := resp.BodyGunzip()
if errUnzip == nil || errUnzip.Error() != "unexpected EOF" {
t.Fatalf("expected error: 'unexpected EOF' but was %v", errUnzip)
}
}

func TestResponseWriteGzipNilBody(t *testing.T) {
Expand Down Expand Up @@ -405,6 +417,46 @@ func TestResponseWriteDeflateNilBody(t *testing.T) {
}
}

func TestResponseBodyUncompressed(t *testing.T) {
body := "body"
var r Response
r.SetBodyStream(bytes.NewReader([]byte(body)), len(body))

w := &bytes.Buffer{}
bw := bufio.NewWriter(w)
if err := r.WriteDeflate(bw); err != nil {
t.Fatalf("unexpected error: %v", err)
}
if err := bw.Flush(); err != nil {
t.Fatalf("unexpected error: %v", err)
}

var resp Response
br := bufio.NewReader(w)
if err := resp.Read(br); err != nil {
t.Fatalf("unexpected error: %v", err)
}

ce := resp.Header.ContentEncoding()
if string(ce) != "deflate" {
t.Fatalf("unexpected Content-Encoding: %s", ce)
}
respBody, err := resp.BodyUncompressed()
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
if string(respBody) != body {
t.Fatalf("unexpected body: %q. Expecting %q", respBody, body)
}

// check for invalid encoding
resp.Header.SetContentEncoding("invalid")
_, decodeErr := resp.BodyUncompressed()
if decodeErr != ErrContentEncodingUnsupported {
t.Fatalf("unexpected error: %v", decodeErr)
}
}

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

Expand Down Expand Up @@ -1145,8 +1197,8 @@ func TestRequestReadGzippedBody(t *testing.T) {
t.Fatalf("unexpected error: %v", err)
}

if string(r.Header.Peek(HeaderContentEncoding)) != "gzip" {
t.Fatalf("unexpected content-encoding: %q. Expecting %q", r.Header.Peek(HeaderContentEncoding), "gzip")
if string(r.Header.ContentEncoding()) != "gzip" {
t.Fatalf("unexpected content-encoding: %q. Expecting %q", r.Header.ContentEncoding(), "gzip")
}
if r.Header.ContentLength() != len(body) {
t.Fatalf("unexpected content-length: %d. Expecting %d", r.Header.ContentLength(), len(body))
Expand Down