From 9aadbf5755dc8e3c3b0a224f513c92b804a1a3a7 Mon Sep 17 00:00:00 2001
From: Meng Zhuo <mengzhuo1203@gmail.com>
Date: Tue, 27 Nov 2018 16:16:43 +0800
Subject: [PATCH] net/http: prevent transport sends two "Connection: close"
 headers

There are three functions that do Connection header write:
1. transport.go/ persistConn.roundTrip
2. transfer.go/ transferWriter.writeHeader
3. request.go/ Request.write

The root cause is roundTrip didn't lookup into request.Close and
transferWriter
didn't take care of extraHeaders.

Fixes #28886

Change-Id: I1d131019c7cd42eb1bcc972c631b7df7511c1f39
Reviewed-on: https://go-review.googlesource.com/c/150722
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
---
 src/net/http/request.go        |  3 ++
 src/net/http/transport.go      |  2 +-
 src/net/http/transport_test.go | 54 ++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/src/net/http/request.go b/src/net/http/request.go
index 5b7e6564ae..d994e81d23 100644
--- a/src/net/http/request.go
+++ b/src/net/http/request.go
@@ -1330,6 +1330,9 @@ func (r *Request) wantsHttp10KeepAlive() bool {
 }
 
 func (r *Request) wantsClose() bool {
+	if r.Close {
+		return true
+	}
 	return hasToken(r.Header.get("Connection"), "close")
 }
 
diff --git a/src/net/http/transport.go b/src/net/http/transport.go
index e1cfc668ea..ad0201d554 100644
--- a/src/net/http/transport.go
+++ b/src/net/http/transport.go
@@ -2135,7 +2135,7 @@ func (pc *persistConn) roundTrip(req *transportRequest) (resp *Response, err err
 		continueCh = make(chan struct{}, 1)
 	}
 
-	if pc.t.DisableKeepAlives {
+	if pc.t.DisableKeepAlives && !req.wantsClose() {
 		req.extraHeaders().Set("Connection", "close")
 	}
 
diff --git a/src/net/http/transport_test.go b/src/net/http/transport_test.go
index 22ca3f9550..1021ce5aa2 100644
--- a/src/net/http/transport_test.go
+++ b/src/net/http/transport_test.go
@@ -41,6 +41,8 @@ import (
 	"sync/atomic"
 	"testing"
 	"time"
+
+	"golang_org/x/net/http/httpguts"
 )
 
 // TODO: test 5 pipelined requests with responses: 1) OK, 2) OK, Connection: Close
@@ -310,6 +312,58 @@ func TestTransportConnectionCloseOnRequestDisableKeepAlive(t *testing.T) {
 	}
 }
 
+// Test that Transport only sends one "Connection: close", regardless of
+// how "close" was indicated.
+func TestTransportRespectRequestWantsClose(t *testing.T) {
+	tests := []struct {
+		disableKeepAlives bool
+		close             bool
+	}{
+		{disableKeepAlives: false, close: false},
+		{disableKeepAlives: false, close: true},
+		{disableKeepAlives: true, close: false},
+		{disableKeepAlives: true, close: true},
+	}
+
+	for _, tc := range tests {
+		t.Run(fmt.Sprintf("DisableKeepAlive=%v,RequestClose=%v", tc.disableKeepAlives, tc.close),
+			func(t *testing.T) {
+				defer afterTest(t)
+				ts := httptest.NewServer(hostPortHandler)
+				defer ts.Close()
+
+				c := ts.Client()
+				c.Transport.(*Transport).DisableKeepAlives = tc.disableKeepAlives
+				req, err := NewRequest("GET", ts.URL, nil)
+				if err != nil {
+					t.Fatal(err)
+				}
+				count := 0
+				trace := &httptrace.ClientTrace{
+					WroteHeaderField: func(key string, field []string) {
+						if key != "Connection" {
+							return
+						}
+						if httpguts.HeaderValuesContainsToken(field, "close") {
+							count += 1
+						}
+					},
+				}
+				req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace))
+				req.Close = tc.close
+				res, err := c.Do(req)
+				if err != nil {
+					t.Fatal(err)
+				}
+				defer res.Body.Close()
+				if want := tc.disableKeepAlives || tc.close; count > 1 || (count == 1) != want {
+					t.Errorf("expecting want:%v, got 'Connection: close':%d", want, count)
+				}
+			})
+	}
+
+}
+
 func TestTransportIdleCacheKeys(t *testing.T) {
 	defer afterTest(t)
 	ts := httptest.NewServer(hostPortHandler)
-- 
2.48.1