From 86115595abbc3a17a9b8d5f67fca588ffdc2f09c Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Thu, 12 Nov 2015 13:30:27 +0200
Subject: [PATCH] Allow setting body length for Response.BodyStream via
 ResponseHeader.ContentLength

---
 http.go      | 73 +++++++++++++++++++++++++++++++++++++++++-----------
 http_test.go | 27 ++++++++++++-------
 2 files changed, 76 insertions(+), 24 deletions(-)

diff --git a/http.go b/http.go
index 78e83a6..262fcea 100644
--- a/http.go
+++ b/http.go
@@ -47,9 +47,14 @@ type Response struct {
 	//
 	// BodyStream may be set instead of Body for performance reasons only.
 	//
-	// BodyStream is read by Response.Write() until error or io.EOF.
-	// Response.Write() calls BodyStream.Close() if such method is present
-	// after BodyStream.Read() returns error or io.EOF.
+	// BodyStream is read by Response.Write() until one of the following
+	// events occur:
+	// - Response.ContentLength bytes read if it is greater than 0.
+	// - error or io.EOF is returned from BodyStream if Response.ContentLength
+	// is 0 or negative.
+	//
+	// Response.Write() calls BodyStream.Close() (io.Closer) if such method
+	// is present after finishing reading BodyStream.
 	//
 	// Either BodyStream or Body may be set, but not both.
 	//
@@ -209,23 +214,33 @@ func (req *Request) Write(w *bufio.Writer) error {
 func (resp *Response) Write(w *bufio.Writer) error {
 	var err error
 	if resp.BodyStream != nil {
-		resp.Header.ContentLength = -1
-		if err = resp.Header.Write(w); err != nil {
-			return err
-		}
-		if err = writeBodyChunked(w, resp.BodyStream); err != nil {
-			return err
+		if resp.Header.ContentLength > 0 {
+			if err = resp.Header.Write(w); err != nil {
+				return err
+			}
+			if err = writeBodyFixedSize(w, resp.BodyStream, resp.Header.ContentLength); err != nil {
+				return err
+			}
+		} else {
+			resp.Header.ContentLength = -1
+			if err = resp.Header.Write(w); err != nil {
+				return err
+			}
+			if err = writeBodyChunked(w, resp.BodyStream); err != nil {
+				return err
+			}
 		}
 		if bsc, ok := resp.BodyStream.(io.Closer); ok {
 			err = bsc.Close()
 		}
-	} else {
-		resp.Header.ContentLength = len(resp.Body)
-		if err = resp.Header.Write(w); err != nil {
-			return err
-		}
-		_, err = w.Write(resp.Body)
+		return err
 	}
+
+	resp.Header.ContentLength = len(resp.Body)
+	if err = resp.Header.Write(w); err != nil {
+		return err
+	}
+	_, err = w.Write(resp.Body)
 	return err
 }
 
@@ -261,6 +276,34 @@ func writeBodyChunked(w *bufio.Writer, r io.Reader) error {
 	return err
 }
 
+var limitReaderPool sync.Pool
+
+func writeBodyFixedSize(w *bufio.Writer, r io.Reader, size int) error {
+	vbuf := copyBufPool.Get()
+	if vbuf == nil {
+		vbuf = make([]byte, 4096)
+	}
+	buf := vbuf.([]byte)
+
+	vlr := limitReaderPool.Get()
+	if vlr == nil {
+		vlr = &io.LimitedReader{}
+	}
+	lr := vlr.(*io.LimitedReader)
+	lr.R = r
+	lr.N = int64(size)
+
+	n, err := io.CopyBuffer(w, lr, buf)
+
+	limitReaderPool.Put(vlr)
+	copyBufPool.Put(vbuf)
+
+	if n != int64(size) && err == nil {
+		err = fmt.Errorf("read %d bytes from BodyStream instead of %d bytes", n, size)
+	}
+	return err
+}
+
 func writeChunk(w *bufio.Writer, b []byte) error {
 	n := len(b)
 	writeHexInt(w, n)
diff --git a/http_test.go b/http_test.go
index 33213b2..a1fd7f2 100644
--- a/http_test.go
+++ b/http_test.go
@@ -8,19 +8,28 @@ import (
 	"testing"
 )
 
-func TestResponseBodyStream(t *testing.T) {
-	testResponseBodyStream(t, "")
-
-	body := "foobar baz aaa bbb ccc"
-	testResponseBodyStream(t, body)
-
-	body = string(createFixedBody(10001))
-	testResponseBodyStream(t, body)
+func TestResponseBodyStreamFixedSize(t *testing.T) {
+	testResponseBodyStream(t, "a", false)
+	testResponseBodyStream(t, string(createFixedBody(4097)), false)
+	testResponseBodyStream(t, string(createFixedBody(100500)), false)
 }
 
-func testResponseBodyStream(t *testing.T, body string) {
+func TestResponseBodyStreamChunked(t *testing.T) {
+	testResponseBodyStream(t, "", true)
+
+	body := "foobar baz aaa bbb ccc"
+	testResponseBodyStream(t, body, true)
+
+	body = string(createFixedBody(10001))
+	testResponseBodyStream(t, body, true)
+}
+
+func testResponseBodyStream(t *testing.T, body string, chunked bool) {
 	var resp Response
 	resp.BodyStream = bytes.NewBufferString(body)
+	if !chunked {
+		resp.Header.ContentLength = len(body)
+	}
 
 	var w bytes.Buffer
 	bw := bufio.NewWriter(&w)