net/http: allow upgrading non keepalive connections

Fixes #36381

Author: nhooyr

src/net/http/response.go

@@ -352,10 +352,10 @@ func (r *Response) bodyIsWritable() bool {
 	return ok
 }
 
-// isProtocolSwitch reports whether r is a response to a successful
-// protocol upgrade.
-func (r *Response) isProtocolSwitch() bool {
-	return r.StatusCode == StatusSwitchingProtocols &&
-		r.Header.Get("Upgrade") != "" &&
-		httpguts.HeaderValuesContainsToken(r.Header["Connection"], "Upgrade")
+// isProtocolSwitch reports whether the response code and header
+// indicate a successful protocol upgrade response.
+func isProtocolSwitchResponse(code int, h Header) bool {
+	return code == StatusSwitchingProtocols &&
+		h.Get("Upgrade") != "" &&
+		httpguts.HeaderValuesContainsToken(h["Connection"], "Upgrade")
 }

src/net/http/serve_test.go

@@ -6443,3 +6443,44 @@ func BenchmarkResponseStatusLine(b *testing.B) {
 		}
 	})
 }
+func TestDisableKeepAliveUpgrade(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping in short mode")
+	}
+
+	setParallel(t)
+	defer afterTest(t)
+
+	s := httptest.NewServer(HandlerFunc(func(w ResponseWriter, r *Request) {
+		c, _, err := w.(Hijacker).Hijack()
+		if err != nil {
+			return
+		}
+		defer c.Close()
+
+		c.Write([]byte("hello"))
+	}))
+	defer s.Close()
+
+	s.Client().Transport.(*Transport).DisableKeepAlives = true
+
+	resp, err := s.Client().Get("/")
+	if err != nil {
+		t.Fatal("failed to perform request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	rwc, ok := resp.Body.(io.ReadWriteCloser)
+	if !ok {
+		t.Fatal("body is not a rwc: %T", resp.Body)
+	}
+
+	b, err := ioutil.ReadAll(rwc)
+	if err != nil {
+		t.Fatal("failed to read rwc: %v", err)
+	}
+
+	if string(b) != "hello" {
+		t.Fatal("unexpected value read from rwc: %s", b)
+	}
+}

src/net/http/server.go

@@ -1290,7 +1290,10 @@ func (cw *chunkWriter) writeHeader(p []byte) {
 		if !connectionHeaderSet {
 			setHeader.connection = "keep-alive"
 		}
-	} else if !w.req.ProtoAtLeast(1, 1) || w.wantsClose {
+	} else if !w.req.ProtoAtLeast(1, 1) ||
+		// Only close if the request indicates the client wants to close
+		// and we are not upgrading the connection.
+		(w.wantsClose && !isProtocolSwitchResponse(w.status, header)) {
 		w.closeAfterReply = true
 	}
 

src/net/http/transport.go

@@ -2195,7 +2195,7 @@ func (pc *persistConn) readResponse(rc requestAndChan, trace *httptrace.ClientTr
 		}
 		break
 	}
-	if resp.isProtocolSwitch() {
+	if isProtocolSwitchResponse(resp.StatusCode, resp.Header) {
 		resp.Body = newReadWriteCloserBody(pc.br, pc.conn)
 	}