net/http/httputil: make ReverseProxy automatically proxy WebSocket requests

Fixes #26937

Change-Id: I6cdc1bad4cf476cd2ea1462b53444eccd8841e14
Reviewed-on: https://go-review.googlesource.com/c/146437
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Dmitri Shuralyov <dmitshur@golang.org>

1 files changed, 78 insertions, 10 deletions

diff --git a/src/net/http/httputil/reverseproxy_test.go b/src/net/http/httputil/reverseproxy_test.go
index ddae11b168..039273e7c5 100644
--- a/src/net/http/httputil/reverseproxy_test.go
+++ b/src/net/http/httputil/reverseproxy_test.go
@@ -153,15 +153,20 @@ func TestReverseProxy(t *testing.T) {
 func TestReverseProxyStripHeadersPresentInConnection(t *testing.T) {
 	const fakeConnectionToken = "X-Fake-Connection-Token"
 	const backendResponse = "I am the backend"
+
+	// someConnHeader is some arbitrary header to be declared as a hop-by-hop header
+	// in the Request's Connection header.
+	const someConnHeader = "X-Some-Conn-Header"
+
 	backend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if c := r.Header.Get(fakeConnectionToken); c != "" {
 			t.Errorf("handler got header %q = %q; want empty", fakeConnectionToken, c)
 		}
-		if c := r.Header.Get("Upgrade"); c != "" {
-			t.Errorf("handler got header %q = %q; want empty", "Upgrade", c)
+		if c := r.Header.Get(someConnHeader); c != "" {
+			t.Errorf("handler got header %q = %q; want empty", someConnHeader, c)
 		}
-		w.Header().Set("Connection", "Upgrade, "+fakeConnectionToken)
-		w.Header().Set("Upgrade", "should be deleted")
+		w.Header().Set("Connection", someConnHeader+", "+fakeConnectionToken)
+		w.Header().Set(someConnHeader, "should be deleted")
 		w.Header().Set(fakeConnectionToken, "should be deleted")
 		io.WriteString(w, backendResponse)
 	}))
@@ -173,15 +178,15 @@ func TestReverseProxyStripHeadersPresentInConnection(t *testing.T) {
 	proxyHandler := NewSingleHostReverseProxy(backendURL)
 	frontend := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		proxyHandler.ServeHTTP(w, r)
-		if c := r.Header.Get("Upgrade"); c != "original value" {
-			t.Errorf("handler modified header %q = %q; want %q", "Upgrade", c, "original value")
+		if c := r.Header.Get(someConnHeader); c != "original value" {
+			t.Errorf("handler modified header %q = %q; want %q", someConnHeader, c, "original value")
 		}
 	}))
 	defer frontend.Close()
 
 	getReq, _ := http.NewRequest("GET", frontend.URL, nil)
-	getReq.Header.Set("Connection", "Upgrade, "+fakeConnectionToken)
-	getReq.Header.Set("Upgrade", "original value")
+	getReq.Header.Set("Connection", someConnHeader+", "+fakeConnectionToken)
+	getReq.Header.Set(someConnHeader, "original value")
 	getReq.Header.Set(fakeConnectionToken, "should be deleted")
 	res, err := frontend.Client().Do(getReq)
 	if err != nil {
@@ -195,8 +200,8 @@ func TestReverseProxyStripHeadersPresentInConnection(t *testing.T) {
 	if got, want := string(bodyBytes), backendResponse; got != want {
 		t.Errorf("got body %q; want %q", got, want)
 	}
-	if c := res.Header.Get("Upgrade"); c != "" {
-		t.Errorf("handler got header %q = %q; want empty", "Upgrade", c)
+	if c := res.Header.Get(someConnHeader); c != "" {
+		t.Errorf("handler got header %q = %q; want empty", someConnHeader, c)
 	}
 	if c := res.Header.Get(fakeConnectionToken); c != "" {
 		t.Errorf("handler got header %q = %q; want empty", fakeConnectionToken, c)
@@ -980,3 +985,66 @@ func TestSelectFlushInterval(t *testing.T) {
 		})
 	}
 }
+
+func TestReverseProxyWebSocket(t *testing.T) {
+	backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if upgradeType(r.Header) != "websocket" {
+			t.Error("unexpected backend request")
+			http.Error(w, "unexpected request", 400)
+			return
+		}
+		c, _, err := w.(http.Hijacker).Hijack()
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer c.Close()
+		io.WriteString(c, "HTTP/1.1 101 Switching Protocols\r\nConnection: upgrade\r\nUpgrade: WebSocket\r\n\r\n")
+		bs := bufio.NewScanner(c)
+		if !bs.Scan() {
+			t.Errorf("backend failed to read line from client: %v", bs.Err())
+			return
+		}
+		fmt.Fprintf(c, "backend got %q\n", bs.Text())
+	}))
+	defer backendServer.Close()
+
+	backURL, _ := url.Parse(backendServer.URL)
+	rproxy := NewSingleHostReverseProxy(backURL)
+	rproxy.ErrorLog = log.New(ioutil.Discard, "", 0) // quiet for tests
+
+	frontendProxy := httptest.NewServer(rproxy)
+	defer frontendProxy.Close()
+
+	req, _ := http.NewRequest("GET", frontendProxy.URL, nil)
+	req.Header.Set("Connection", "Upgrade")
+	req.Header.Set("Upgrade", "websocket")
+
+	c := frontendProxy.Client()
+	res, err := c.Do(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if res.StatusCode != 101 {
+		t.Fatalf("status = %v; want 101", res.Status)
+	}
+	if upgradeType(res.Header) != "websocket" {
+		t.Fatalf("not websocket upgrade; got %#v", res.Header)
+	}
+	rwc, ok := res.Body.(io.ReadWriteCloser)
+	if !ok {
+		t.Fatalf("response body is of type %T; does not implement ReadWriteCloser", res.Body)
+	}
+	defer rwc.Close()
+
+	io.WriteString(rwc, "Hello\n")
+	bs := bufio.NewScanner(rwc)
+	if !bs.Scan() {
+		t.Fatalf("Scan: %v", bs.Err())
+	}
+	got := bs.Text()
+	want := `backend got "Hello"`
+	if got != want {
+		t.Errorf("got %#q, want %#q", got, want)
+	}
+}