net/http2: reset client stream after processing response headers

When a client receives a HEADER frame with a END_STREAM flag,
clientConn.streamByID closes the stream before processing the headers
which may contain a full non-error response. This causes the request's
bodyWriter cancelation to race with the response.

Closing the stream after processing headers allows the response to be
available before the bodyWriter is canceled.

Updates golang/go#20521

Change-Id: I70740e88f75240836e922163a54a6cd89535f7b3
Reviewed-on: https://go-review.googlesource.com/70510
Run-TryBot: Emmanuel Odeke <emm.odeke@gmail.com>
Reviewed-by: Tom Bergan <tombergan@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
diff --git a/http2/transport.go b/http2/transport.go
index adb77ff..c112d22 100644
--- a/http2/transport.go
+++ b/http2/transport.go
@@ -1536,7 +1536,17 @@
 
 func (rl *clientConnReadLoop) processHeaders(f *MetaHeadersFrame) error {
 	cc := rl.cc
-	cs := cc.streamByID(f.StreamID, f.StreamEnded())
+	if f.StreamEnded() {
+		// Issue 20521: If the stream has ended, streamByID() causes
+		// clientStream.done to be closed, which causes the request's bodyWriter
+		// to be closed with an errStreamClosed, which may be received by
+		// clientConn.RoundTrip before the result of processing these headers.
+		// Deferring stream closure allows the header processing to occur first.
+		// clientConn.RoundTrip may still receive the bodyWriter error first, but
+		// the fix for issue 16102 prioritises any response.
+		defer cc.streamByID(f.StreamID, true)
+	}
+	cs := cc.streamByID(f.StreamID, false)
 	if cs == nil {
 		// We'd get here if we canceled a request while the
 		// server had its response still in flight. So if this
diff --git a/http2/transport_test.go b/http2/transport_test.go
index 0126ff4..12a9869 100644
--- a/http2/transport_test.go
+++ b/http2/transport_test.go
@@ -3678,6 +3678,34 @@
 	}
 }
 
+type infiniteReader struct{}
+
+func (r infiniteReader) Read(b []byte) (int, error) {
+	return len(b), nil
+}
+
+// Issue 20521: it is not an error to receive a response and end stream
+// from the server without the body being consumed.
+func TestTransportResponseAndResetWithoutConsumingBodyRace(t *testing.T) {
+	st := newServerTester(t, func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	}, optOnlyServer)
+	defer st.Close()
+
+	tr := &Transport{TLSClientConfig: tlsConfigInsecure}
+	defer tr.CloseIdleConnections()
+
+	// The request body needs to be big enough to trigger flow control.
+	req, _ := http.NewRequest("PUT", st.ts.URL, infiniteReader{})
+	res, err := tr.RoundTrip(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if res.StatusCode != http.StatusOK {
+		t.Fatalf("Response code = %v; want %v", res.StatusCode, http.StatusOK)
+	}
+}
+
 func BenchmarkClientRequestHeaders(b *testing.B) {
 	b.Run("   0 Headers", func(b *testing.B) { benchSimpleRoundTrip(b, 0) })
 	b.Run("  10 Headers", func(b *testing.B) { benchSimpleRoundTrip(b, 10) })