net/http/httputil: fix missing previous headers in response when switching protocol...

When using switching protocol, previous headers set before the reverse proxy are lost.

Fixes #29407

Change-Id: Ia2b9784022d9bccef8625519ccbabbe8a276dfc0
GitHub-Last-Rev: 79bb493dcbb9b76d9d2ff9cd0854b29d634f8b73
GitHub-Pull-Request: golang/go#29408
Reviewed-on: https://go-review.googlesource.com/c/155741
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>

diff --git a/src/net/http/httputil/reverseproxy.go b/src/net/http/httputil/reverseproxy.go
index 5d07ba3d361873d834c36a794e8ca3973271c119..c13b99ff72fedc0808bbc2fed89e69da00d0c95e 100644 (file)
--- a/src/net/http/httputil/reverseproxy.go
+++ b/src/net/http/httputil/reverseproxy.go
@@ -497,6 +497,9 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R
                p.getErrorHandler()(rw, req, fmt.Errorf("backend tried to switch protocol %q when %q was requested", resUpType, reqUpType))
                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))

diff --git a/src/net/http/httputil/reverseproxy_test.go b/src/net/http/httputil/reverseproxy_test.go
index 5caa2060660a9766caaa0b334cfbf1bb2cf3c508..bda569acc73bdc5bde1362d03659c0bd56f43d88 100644 (file)
--- a/src/net/http/httputil/reverseproxy_test.go
+++ b/src/net/http/httputil/reverseproxy_test.go
@@ -1013,7 +1013,12 @@ func TestReverseProxyWebSocket(t *testing.T) {
        rproxy := NewSingleHostReverseProxy(backURL)
        rproxy.ErrorLog = log.New(ioutil.Discard, "", 0) // quiet for tests
 
-       frontendProxy := httptest.NewServer(rproxy)
+       handler := http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
+               rw.Header().Set("X-Header", "X-Value")
+               rproxy.ServeHTTP(rw, req)
+       })
+
+       frontendProxy := httptest.NewServer(handler)
        defer frontendProxy.Close()
 
        req, _ := http.NewRequest("GET", frontendProxy.URL, nil)
@@ -1028,6 +1033,13 @@ func TestReverseProxyWebSocket(t *testing.T) {
        if res.StatusCode != 101 {
                t.Fatalf("status = %v; want 101", res.Status)
        }
+
+       got := res.Header.Get("X-Header")
+       want := "X-Value"
+       if got != want {
+               t.Errorf("Header(XHeader) = %q; want %q", got, want)
+       }
+
        if upgradeType(res.Header) != "websocket" {
                t.Fatalf("not websocket upgrade; got %#v", res.Header)
        }
@@ -1042,8 +1054,8 @@ func TestReverseProxyWebSocket(t *testing.T) {
        if !bs.Scan() {
                t.Fatalf("Scan: %v", bs.Err())
        }
-       got := bs.Text()
-       want := `backend got "Hello"`
+       got = bs.Text()
+       want = `backend got "Hello"`
        if got != want {
                t.Errorf("got %#q, want %#q", got, want)
        }