net/http/httputil: copy response header back to http.ResponseWriter in ReverseProxy.handleUpgradeResponse

Fixes: #41634

Change-Id: Ib78cc37a4d2ca0753d567eafb616238e4103484e
Reviewed-on: https://go-review.googlesource.com/c/go/+/257777
Reviewed-by: Damien Neil <dneil@google.com>
Trust: Damien Neil <dneil@google.com>
Trust: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Damien Neil <dneil@google.com>
TryBot-Result: Go Bot <gobot@golang.org>

Author: aofei

src/net/http/httputil/reverseproxy.go

@@ -549,8 +549,6 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R
 		return
 	}
 
-	copyHeader(res.Header, rw.Header())
-
 	hj, ok := rw.(http.Hijacker)
 	if !ok {
 		p.getErrorHandler()(rw, req, fmt.Errorf("can't switch protocols using non-Hijacker ResponseWriter type %T", rw))
@@ -581,6 +579,10 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R
 		return
 	}
 	defer conn.Close()
+
+	copyHeader(rw.Header(), res.Header)
+
+	res.Header = rw.Header()
 	res.Body = nil // so res.Write only writes the headers; we have res.Body in backConn above
 	if err := res.Write(brw); err != nil {
 		p.getErrorHandler()(rw, req, fmt.Errorf("response write: %v", err))

src/net/http/httputil/reverseproxy_test.go

@@ -1157,6 +1157,9 @@ func TestReverseProxyWebSocket(t *testing.T) {
 	handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
 		rw.Header().Set("X-Header", "X-Value")
 		rproxy.ServeHTTP(rw, req)
+		if got, want := rw.Header().Get("X-Modified"), "true"; got != want {
+			t.Errorf("response writer X-Modified header = %q; want %q", got, want)
+		}
 	})
 
 	frontendProxy := httptest.NewServer(handler)