From 7d7c17c8d7bd88bee712f814bcf675dc744a5cde Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Wed, 6 Jan 2016 15:22:52 +0200
Subject: [PATCH] Do not send response body (and content-length) for 1xx, 204
 and 304 responses

---
 header.go    | 17 +++++++++++++++++
 http.go      | 19 ++++++++-----------
 http_test.go | 38 ++++++++++++++++++++++++++++++++++----
 3 files changed, 59 insertions(+), 15 deletions(-)

diff --git a/header.go b/header.go
index 1777feb..ad4bbf6 100644
--- a/header.go
+++ b/header.go
@@ -181,6 +181,9 @@ func (h *ResponseHeader) ContentLength() int {
 // -1 means Transfer-Encoding: chunked.
 // -2 means Transfer-Encoding: identity.
 func (h *ResponseHeader) SetContentLength(contentLength int) {
+	if h.mustSkipContentLength() {
+		return
+	}
 	h.contentLength = contentLength
 	if contentLength >= 0 {
 		h.contentLengthBytes = AppendUint(h.contentLengthBytes[:0], contentLength)
@@ -196,6 +199,20 @@ func (h *ResponseHeader) SetContentLength(contentLength int) {
 	}
 }
 
+func (h *ResponseHeader) mustSkipContentLength() bool {
+	// From http/1.1 specs:
+	// All 1xx (informational), 204 (no content), and 304 (not modified) responses MUST NOT include a message-body
+	statusCode := h.StatusCode()
+
+	// Fast path.
+	if statusCode < 100 || statusCode == StatusOK {
+		return false
+	}
+
+	// Slow path.
+	return statusCode == StatusNotModified || statusCode == StatusNoContent || statusCode < 200
+}
+
 // ContentLength returns Content-Length header value.
 //
 // It may be negative:
diff --git a/http.go b/http.go
index 7b2aa32..a2e2b98 100644
--- a/http.go
+++ b/http.go
@@ -674,7 +674,7 @@ func (resp *Response) ReadLimitBody(r *bufio.Reader, maxBodySize int) error {
 		}
 	}
 
-	if !isSkipResponseBody(resp.Header.StatusCode()) && !resp.SkipBody {
+	if !resp.mustSkipBody() {
 		resp.body, err = readBody(r, resp.Header.ContentLength(), maxBodySize, resp.body)
 		if err != nil {
 			resp.Reset()
@@ -685,13 +685,8 @@ func (resp *Response) ReadLimitBody(r *bufio.Reader, maxBodySize int) error {
 	return nil
 }
 
-func isSkipResponseBody(statusCode int) bool {
-	// From http/1.1 specs:
-	// All 1xx (informational), 204 (no content), and 304 (not modified) responses MUST NOT include a message-body
-	if statusCode >= 100 && statusCode < 200 {
-		return true
-	}
-	return statusCode == StatusNoContent || statusCode == StatusNotModified
+func (resp *Response) mustSkipBody() bool {
+	return resp.SkipBody || resp.Header.mustSkipContentLength()
 }
 
 var errRequestHostRequired = errors.New("Missing required Host header in request")
@@ -844,6 +839,8 @@ func (resp *Response) deflateBody(level int) error {
 // Write doesn't flush response to w for performance reasons.
 func (resp *Response) Write(w *bufio.Writer) error {
 	var err error
+	sendBody := !resp.mustSkipBody()
+
 	if resp.bodyStream != nil {
 		contentLength := resp.Header.ContentLength()
 		if contentLength < 0 {
@@ -860,7 +857,7 @@ func (resp *Response) Write(w *bufio.Writer) error {
 			if err = resp.Header.Write(w); err != nil {
 				return err
 			}
-			if !resp.SkipBody {
+			if sendBody {
 				if err = writeBodyFixedSize(w, resp.bodyStream, int64(contentLength)); err != nil {
 					return err
 				}
@@ -870,7 +867,7 @@ func (resp *Response) Write(w *bufio.Writer) error {
 			if err = resp.Header.Write(w); err != nil {
 				return err
 			}
-			if !resp.SkipBody {
+			if sendBody {
 				if err = writeBodyChunked(w, resp.bodyStream); err != nil {
 					return err
 				}
@@ -883,7 +880,7 @@ func (resp *Response) Write(w *bufio.Writer) error {
 	if err = resp.Header.Write(w); err != nil {
 		return err
 	}
-	if !resp.SkipBody {
+	if sendBody {
 		_, err = w.Write(resp.body)
 	}
 	return err
diff --git a/http_test.go b/http_test.go
index cc0f522..e7fb90b 100644
--- a/http_test.go
+++ b/http_test.go
@@ -13,17 +13,47 @@ import (
 func TestResponseSkipBody(t *testing.T) {
 	var r Response
 
-	r.SkipBody = true
+	// set StatusNotModified
+	r.Header.SetStatusCode(StatusNotModified)
 	r.SetBodyString("foobar")
 	s := r.String()
 	if strings.Contains(s, "\r\n\r\nfoobar") {
-		t.Fatalf("unexpected non-zero body in request %q", s)
+		t.Fatalf("unexpected non-zero body in response %q", s)
+	}
+	if strings.Contains(s, "Content-Length: ") {
+		t.Fatalf("unexpected content-length in response %q", s)
+	}
+	if strings.Contains(s, "Content-Type: ") {
+		t.Fatalf("unexpected content-type in response %q", s)
+	}
+
+	// set StatusNoContent
+	r.Header.SetStatusCode(StatusNoContent)
+	r.SetBodyString("foobar")
+	s = r.String()
+	if strings.Contains(s, "\r\n\r\nfoobar") {
+		t.Fatalf("unexpected non-zero body in response %q", s)
+	}
+	if strings.Contains(s, "Content-Length: ") {
+		t.Fatalf("unexpected content-length in response %q", s)
+	}
+	if strings.Contains(s, "Content-Type: ") {
+		t.Fatalf("unexpected content-type in response %q", s)
+	}
+
+	// explicitly skip body
+	r.Header.SetStatusCode(StatusOK)
+	r.SkipBody = true
+	r.SetBodyString("foobar")
+	s = r.String()
+	if strings.Contains(s, "\r\n\r\nfoobar") {
+		t.Fatalf("unexpected non-zero body in response %q", s)
 	}
 	if !strings.Contains(s, "Content-Length: 6\r\n") {
-		t.Fatalf("unexpected content-length in request %q", s)
+		t.Fatalf("unexpected content-length in response %q", s)
 	}
 	if !strings.Contains(s, "Content-Type: ") {
-		t.Fatalf("unexpected content-type in request %q", s)
+		t.Fatalf("expecting content-type in response %q", s)
 	}
 }