From 0cd73496863721ab5b3f7bf82551731bc4b4f27d Mon Sep 17 00:00:00 2001
From: Vladimir Shteinman <vladimir.s@cloudinary.com>
Date: Mon, 15 Mar 2021 19:10:51 +0200
Subject: [PATCH] ImmediateHeaderFlush when no body (#995)

---
 http.go      |  8 +++---
 http_test.go | 80 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+), 4 deletions(-)

diff --git a/http.go b/http.go
index 5ea7847556..7be092b4f6 100644
--- a/http.go
+++ b/http.go
@@ -1716,21 +1716,21 @@ func (resp *Response) writeBodyStream(w *bufio.Writer, sendBody bool) (err error
 		}
 	}
 	if contentLength >= 0 {
-		if err = resp.Header.Write(w); err == nil && sendBody {
+		if err = resp.Header.Write(w); err == nil {
 			if resp.ImmediateHeaderFlush {
 				err = w.Flush()
 			}
-			if err == nil {
+			if err == nil && sendBody {
 				err = writeBodyFixedSize(w, resp.bodyStream, int64(contentLength))
 			}
 		}
 	} else {
 		resp.Header.SetContentLength(-1)
-		if err = resp.Header.Write(w); err == nil && sendBody {
+		if err = resp.Header.Write(w); err == nil {
 			if resp.ImmediateHeaderFlush {
 				err = w.Flush()
 			}
-			if err == nil {
+			if err == nil && sendBody {
 				err = writeBodyChunked(w, resp.bodyStream)
 			}
 		}
diff --git a/http_test.go b/http_test.go
index 7d02bb5cc2..b6ac168e5c 100644
--- a/http_test.go
+++ b/http_test.go
@@ -2212,6 +2212,7 @@ func TestRequestRawBodyCopyTo(t *testing.T) {
 type testReader struct {
 	read chan (int)
 	cb   chan (struct{})
+	onClose func() error 
 }
 
 func (r *testReader) Read(b []byte) (int, error) {
@@ -2230,6 +2231,13 @@ func (r *testReader) Read(b []byte) (int, error) {
 	return read, nil
 }
 
+func (r *testReader) Close() error {
+	if r.onClose != nil {
+		return r.onClose()
+	}
+	return nil
+}
+
 func TestResponseImmediateHeaderFlushRegressionFixedLength(t *testing.T) {
 	t.Parallel()
 
@@ -2301,6 +2309,42 @@ func TestResponseImmediateHeaderFlushFixedLength(t *testing.T) {
 	<-waitForIt
 }
 
+func TestResponseImmediateHeaderFlushFixedLengthSkipBody(t *testing.T) {
+	t.Parallel()
+
+	var r Response
+
+	r.ImmediateHeaderFlush = true
+	r.SkipBody = true
+
+	ch := make(chan int)
+	cb := make(chan struct{})
+
+	buf := &testReader{read: ch, cb: cb}
+
+	r.SetBodyStream(buf, 0)
+
+	b := []byte{}
+	w := bytes.NewBuffer(b)
+	bb := bufio.NewWriter(w)
+
+	var headersOnClose string
+	buf.onClose = func() error {
+		headersOnClose = w.String()
+		return nil
+	}
+
+	bw := &r
+
+	if err := bw.Write(bb); err != nil {
+		t.Errorf("unexpected error: %s", err)
+	}
+
+	if !strings.Contains(headersOnClose, "Content-Length: 0") {
+		t.Fatalf("Expected headers to be eagerly flushed")
+	}
+}
+
 func TestResponseImmediateHeaderFlushChunked(t *testing.T) {
 	t.Parallel()
 
@@ -2347,6 +2391,42 @@ func TestResponseImmediateHeaderFlushChunked(t *testing.T) {
 	<-waitForIt
 }
 
+func TestResponseImmediateHeaderFlushChunkedNoBody(t *testing.T) {
+	t.Parallel()
+
+	var r Response
+
+	r.ImmediateHeaderFlush = true
+	r.SkipBody = true
+
+	ch := make(chan int)
+	cb := make(chan struct{})
+
+	buf := &testReader{read: ch, cb: cb}
+
+	r.SetBodyStream(buf, -1)
+
+	b := []byte{}
+	w := bytes.NewBuffer(b)
+	bb := bufio.NewWriter(w)
+
+	var headersOnClose string
+	buf.onClose = func() error {
+		headersOnClose = w.String()
+		return nil
+	}
+
+	bw := &r
+
+	if err := bw.Write(bb); err != nil {
+		t.Errorf("unexpected error: %s", err)
+	}
+
+	if !strings.Contains(headersOnClose, "Transfer-Encoding: chunked") {
+		t.Fatalf("Expected headers to be eagerly flushed")
+	}
+}
+
 type ErroneousBodyStream struct {
 	errOnRead  bool
 	errOnClose bool