diff --git a/http.go b/http.go index 10dc465..79241db 100644 --- a/http.go +++ b/http.go @@ -67,6 +67,7 @@ type Response struct { bodyStream io.Reader w responseBodyWriter body *bytebufferpool.ByteBuffer + bodyRaw []byte // Response.Read() skips reading body if set to true. // Use it for reading HEAD responses. @@ -321,6 +322,9 @@ func (resp *Response) Body() []byte { } func (resp *Response) bodyBytes() []byte { + if resp.bodyRaw != nil { + return resp.bodyRaw + } if resp.body == nil { return nil } @@ -338,6 +342,7 @@ func (resp *Response) bodyBuffer() *bytebufferpool.ByteBuffer { if resp.body == nil { resp.body = responseBodyPool.Get() } + resp.bodyRaw = nil return resp.body } @@ -462,6 +467,7 @@ func (resp *Response) SetBodyString(body string) { // ResetBody resets response body. func (resp *Response) ResetBody() { + resp.bodyRaw = nil resp.closeBodyStream() if resp.body != nil { if resp.keepBodyBuffer { @@ -473,6 +479,14 @@ func (resp *Response) ResetBody() { } } +// SetBodyRaw sets response body, but without copying it. +// +// From this point onward the body argument must not be changed. +func (resp *Response) SetBodyRaw(body []byte) { + resp.ResetBody() + resp.bodyRaw = body +} + // ReleaseBody retires the response body if it is greater than "size" bytes. // // This permits GC to reclaim the large buffer. If used, must be before @@ -481,6 +495,7 @@ func (resp *Response) ResetBody() { // Use this method only if you really understand how it works. // The majority of workloads don't need this method. func (resp *Response) ReleaseBody(size int) { + resp.bodyRaw = nil if cap(resp.body.B) > size { resp.closeBodyStream() resp.body = nil @@ -519,6 +534,8 @@ func (resp *Response) SwapBody(body []byte) []byte { } } + resp.bodyRaw = nil + oldBody := bb.B bb.B = body return oldBody @@ -639,7 +656,12 @@ func (req *Request) copyToSkipBody(dst *Request) { // CopyTo copies resp contents to dst except of body stream. func (resp *Response) CopyTo(dst *Response) { resp.copyToSkipBody(dst) - if resp.body != nil { + if resp.bodyRaw != nil { + dst.bodyRaw = resp.bodyRaw + if dst.body != nil { + dst.body.Reset() + } + } else if resp.body != nil { dst.bodyBuffer().Set(resp.body.B) } else if dst.body != nil { dst.body.Reset() @@ -661,6 +683,7 @@ func swapRequestBody(a, b *Request) { func swapResponseBody(a, b *Response) { a.body, b.body = b.body, a.body + a.bodyRaw, b.bodyRaw = b.bodyRaw, a.bodyRaw a.bodyStream, b.bodyStream = b.bodyStream, a.bodyStream } @@ -1265,6 +1288,7 @@ func (resp *Response) gzipBody(level int) error { responseBodyPool.Put(resp.body) } resp.body = w + resp.bodyRaw = nil } resp.Header.SetCanonical(strContentEncoding, strGzip) return nil @@ -1319,6 +1343,7 @@ func (resp *Response) deflateBody(level int) error { responseBodyPool.Put(resp.body) } resp.body = w + resp.bodyRaw = nil } resp.Header.SetCanonical(strContentEncoding, strDeflate) return nil diff --git a/http_test.go b/http_test.go index 3253e95..653142c 100644 --- a/http_test.go +++ b/http_test.go @@ -1891,3 +1891,35 @@ Content-Type: application/json t.Fatalf("unexpected output %q", w.Bytes()) } } + + +func TestResponseRawBodySet(t *testing.T) { + var resp Response + + expectedS := "test" + body := []byte(expectedS) + resp.SetBodyRaw(body) + + testBodyWriteTo(t, &resp, expectedS, true) +} + +func TestResponseRawBodyReset(t *testing.T) { + var resp Response + + body := []byte("test") + resp.SetBodyRaw(body) + resp.ResetBody() + + testBodyWriteTo(t, &resp, "", true) +} + +func TestResponseRawBodyCopyTo(t *testing.T) { + var resp Response + + expectedS := "test" + body := []byte(expectedS) + resp.SetBodyRaw(body) + + testResponseCopyTo(t, &resp) +} +