From f809cebb139df4f5560a8456973351c95a3dfa97 Mon Sep 17 00:00:00 2001
From: Anmol Sethi <hi@nhooyr.io>
Date: Tue, 8 Dec 2020 20:15:26 -0500
Subject: [PATCH] net/http: attempt deadlock fix in TestDisableKeepAliveUpgrade

1. The test now checks the response status code.
2. The transport has been changed to not set "Connection: Close" if
   DisableKeepAlive is set and the request is a HTTP/1.1 protocol
   upgrade.

Updates #43073
---
 src/net/http/response.go   | 9 +++++++--
 src/net/http/serve_test.go | 4 ++++
 src/net/http/transport.go  | 4 +++-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/src/net/http/response.go b/src/net/http/response.go
index b95abae646b10a..b8985da3c80fc4 100644
--- a/src/net/http/response.go
+++ b/src/net/http/response.go
@@ -361,7 +361,12 @@ func (r *Response) isProtocolSwitch() bool {
 // isProtocolSwitchResponse reports whether the response code and
 // response header indicate a successful protocol upgrade response.
 func isProtocolSwitchResponse(code int, h Header) bool {
-	return code == StatusSwitchingProtocols &&
-		h.Get("Upgrade") != "" &&
+	return code == StatusSwitchingProtocols && isProtocolSwitchHeader(h)
+}
+
+// isProtocolSwitchHeader reports whether the request or response header
+// is for a protocol switch.
+func isProtocolSwitchHeader(h Header) bool {
+	return h.Get("Upgrade") != "" &&
 		httpguts.HeaderValuesContainsToken(h["Connection"], "Upgrade")
 }
diff --git a/src/net/http/serve_test.go b/src/net/http/serve_test.go
index b1bf8e6c5e8088..95e6bf4adbaeec 100644
--- a/src/net/http/serve_test.go
+++ b/src/net/http/serve_test.go
@@ -6481,6 +6481,10 @@ func TestDisableKeepAliveUpgrade(t *testing.T) {
 	}
 	defer resp.Body.Close()
 
+	if resp.StatusCode != StatusSwitchingProtocols {
+		t.Fatalf("unexpected status code: %v", resp.StatusCode)
+	}
+
 	rwc, ok := resp.Body.(io.ReadWriteCloser)
 	if !ok {
 		t.Fatalf("Response.Body is not a io.ReadWriteCloser: %T", resp.Body)
diff --git a/src/net/http/transport.go b/src/net/http/transport.go
index 8de0f3a6a08369..4d1344aab22294 100644
--- a/src/net/http/transport.go
+++ b/src/net/http/transport.go
@@ -2564,7 +2564,9 @@ func (pc *persistConn) roundTrip(req *transportRequest) (resp *Response, err err
 		continueCh = make(chan struct{}, 1)
 	}
 
-	if pc.t.DisableKeepAlives && !req.wantsClose() {
+	if pc.t.DisableKeepAlives &&
+		!req.wantsClose() &&
+		!isProtocolSwitchHeader(req.Header) {
 		req.extraHeaders().Set("Connection", "close")
 	}