From 69e39bad7dc2bbb411fa35755c46020969029fa7 Mon Sep 17 00:00:00 2001
From: Damien Neil <dneil@google.com>
Date: Fri, 5 Nov 2021 10:13:53 -0700
Subject: [PATCH] http2: close conns after use when req.Close is set

Avoid reusing connections after receiving a response to a request
for which cc.Close is set or a "Connection: close" header is present.

Adjust the tests for connection reuse to test both the case where
new conns are created by the http2 package and when they are
created by the net/http package.

Fixes golang/go#49375

Change-Id: I58594972c832a20b66a3910c17acb51a98a9f7a5
Reviewed-on: https://go-review.googlesource.com/c/net/+/361498
Trust: Damien Neil <dneil@google.com>
Run-TryBot: Damien Neil <dneil@google.com>
TryBot-Result: Go Bot <gobot@golang.org>
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
---
 http2/transport.go      |  3 +++
 http2/transport_test.go | 59 +++++++++++++++++++++++++++++------------
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/http2/transport.go b/http2/transport.go
index 9c41179f7..b5e2ac64d 100644
--- a/http2/transport.go
+++ b/http2/transport.go
@@ -1213,6 +1213,9 @@ func (cs *clientStream) writeRequest(req *http.Request) (err error) {
 		return err
 	}
 	cc.addStreamLocked(cs) // assigns stream ID
+	if isConnectionCloseRequest(req) {
+		cc.doNotReuse = true
+	}
 	cc.mu.Unlock()
 
 	// TODO(bradfitz): this is a copy of the logic in net/http. Unify somewhere?
diff --git a/http2/transport_test.go b/http2/transport_test.go
index 7ed4477f8..952422a44 100644
--- a/http2/transport_test.go
+++ b/http2/transport_test.go
@@ -190,7 +190,7 @@ func TestTransport(t *testing.T) {
 	}
 }
 
-func onSameConn(t *testing.T, modReq func(*http.Request)) bool {
+func testTransportReusesConns(t *testing.T, useClient, wantSame bool, modReq func(*http.Request)) {
 	st := newServerTester(t, func(w http.ResponseWriter, r *http.Request) {
 		io.WriteString(w, r.RemoteAddr)
 	}, optOnlyServer, func(c net.Conn, st http.ConnState) {
@@ -198,6 +198,9 @@ func onSameConn(t *testing.T, modReq func(*http.Request)) bool {
 	})
 	defer st.Close()
 	tr := &Transport{TLSClientConfig: tlsConfigInsecure}
+	if useClient {
+		tr.ConnPool = noDialClientConnPool{new(clientConnPool)}
+	}
 	defer tr.CloseIdleConnections()
 	get := func() string {
 		req, err := http.NewRequest("GET", st.ts.URL, nil)
@@ -205,7 +208,14 @@ func onSameConn(t *testing.T, modReq func(*http.Request)) bool {
 			t.Fatal(err)
 		}
 		modReq(req)
-		res, err := tr.RoundTrip(req)
+		var res *http.Response
+		if useClient {
+			c := st.ts.Client()
+			ConfigureTransports(c.Transport.(*http.Transport))
+			res, err = c.Do(req)
+		} else {
+			res, err = tr.RoundTrip(req)
+		}
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -222,24 +232,39 @@ func onSameConn(t *testing.T, modReq func(*http.Request)) bool {
 	}
 	first := get()
 	second := get()
-	return first == second
-}
-
-func TestTransportReusesConns(t *testing.T) {
-	if !onSameConn(t, func(*http.Request) {}) {
-		t.Errorf("first and second responses were on different connections")
+	if got := first == second; got != wantSame {
+		t.Errorf("first and second responses on same connection: %v; want %v", got, wantSame)
 	}
 }
 
-func TestTransportReusesConn_RequestClose(t *testing.T) {
-	if onSameConn(t, func(r *http.Request) { r.Close = true }) {
-		t.Errorf("first and second responses were not on different connections")
-	}
-}
-
-func TestTransportReusesConn_ConnClose(t *testing.T) {
-	if onSameConn(t, func(r *http.Request) { r.Header.Set("Connection", "close") }) {
-		t.Errorf("first and second responses were not on different connections")
+func TestTransportReusesConns(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		modReq   func(*http.Request)
+		wantSame bool
+	}{{
+		name:     "ReuseConn",
+		modReq:   func(*http.Request) {},
+		wantSame: true,
+	}, {
+		name:     "RequestClose",
+		modReq:   func(r *http.Request) { r.Close = true },
+		wantSame: false,
+	}, {
+		name:     "ConnClose",
+		modReq:   func(r *http.Request) { r.Header.Set("Connection", "close") },
+		wantSame: false,
+	}} {
+		t.Run(test.name, func(t *testing.T) {
+			t.Run("Transport", func(t *testing.T) {
+				const useClient = false
+				testTransportReusesConns(t, useClient, test.wantSame, test.modReq)
+			})
+			t.Run("Client", func(t *testing.T) {
+				const useClient = true
+				testTransportReusesConns(t, useClient, test.wantSame, test.modReq)
+			})
+		})
 	}
 }