net/http: allow upgrading non keepalive connections

Fixes #36381

If one was using http.Transport with DisableKeepAlives and trying
to upgrade a connection against net/http's Server, the Server
would not allow a "Connection: Upgrade" header to be written
and instead override it to "Connection: Close" which would
break the handshake.

This commit ensures net/http's Server does not override the
connection header for successful protocol switch responses.

Author: nhooyr

src/net/http/response.go

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

@@ -251,7 +251,9 @@ func TestHostHandlers(t *testing.T) {
 	for _, h := range handlers {
 		mux.Handle(h.pattern, stringHandler(h.msg))
 	}
-	ts := httptest.NewServer(mux)
+	ts := httptest.NewUnstartedServer(mux)
+	ts.Config.SetKeepAlivesEnabled(false)
+	ts.Start()
 	defer ts.Close()
 
 	conn, err := net.Dial("tcp", ts.Listener.Addr().String())
@@ -6443,3 +6445,50 @@ 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()
+
+		io.Copy(c, c)
+	}))
+	defer s.Close()
+
+	s.Client().Transport.(*Transport).DisableKeepAlives = true
+
+	resp, err := s.Client().Get("/")
+	if err != nil {
+		t.Fatalf("failed to perform request: %v", err)
+	}
+	defer resp.Body.Close()
+
+	rwc, ok := resp.Body.(io.ReadWriteCloser)
+	if !ok {
+		t.Fatalf("body is not a io.ReadWriteCloser: %T", resp.Body)
+	}
+
+	_, err = rwc.Write([]byte("hello"))
+	if err != nil {
+		t.Fatalf("failed to write to body: %v", err)
+	}
+
+	b := make([]byte, 5)
+	_, err = io.ReadFull(rwc, b)
+	if err != nil {
+		t.Fatalf("failed to read from body: %v", err)
+	}
+
+	if string(b) != "hello" {
+		t.Fatalf("unexpected value read from body: %q", 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
 	}