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>
This commit is contained in:
Brad Fitzpatrick 2018-10-29 23:21:40 +00:00
parent de578dcdd6
commit ee55f0856a
4 changed files with 161 additions and 12 deletions

View file

@ -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)
}
}