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>
diff --git a/http2/transport.go b/http2/transport.go
index 9c41179..b5e2ac6 100644
--- a/http2/transport.go
+++ b/http2/transport.go
@@ -1213,6 +1213,9 @@
 		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 7ed4477..952422a 100644
--- a/http2/transport_test.go
+++ b/http2/transport_test.go
@@ -190,7 +190,7 @@
 	}
 }
 
-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 @@
 	})
 	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 @@
 			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 @@
 	}
 	first := get()
 	second := get()
-	return first == second
+	if got := first == second; got != wantSame {
+		t.Errorf("first and second responses on same connection: %v; want %v", got, wantSame)
+	}
 }
 
 func TestTransportReusesConns(t *testing.T) {
-	if !onSameConn(t, func(*http.Request) {}) {
-		t.Errorf("first and second responses were on different connections")
-	}
-}
-
-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")
+	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)
+			})
+		})
 	}
 }