diff --git a/src/net/http/httputil/reverseproxy.go b/src/net/http/httputil/reverseproxy.go index bba8f0f254093e..d7ca5c29e36cde 100644 --- a/src/net/http/httputil/reverseproxy.go +++ b/src/net/http/httputil/reverseproxy.go @@ -837,16 +837,25 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R reqUpType := upgradeType(req.Header) resUpType := upgradeType(res.Header) if !ascii.IsPrint(resUpType) { // We know reqUpType is ASCII, it's checked by the caller. + if res.Body != nil { + res.Body.Close() + } p.getErrorHandler()(rw, req, fmt.Errorf("backend tried to switch to invalid protocol %q", resUpType)) return } if !ascii.EqualFold(reqUpType, resUpType) { + if res.Body != nil { + res.Body.Close() + } p.getErrorHandler()(rw, req, fmt.Errorf("backend tried to switch protocol %q when %q was requested", resUpType, reqUpType)) return } backConn, ok := res.Body.(io.ReadWriteCloser) if !ok { + if res.Body != nil { + res.Body.Close() + } p.getErrorHandler()(rw, req, fmt.Errorf("internal error: 101 switching protocols response with non-writable body")) return } @@ -854,6 +863,7 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R rc := http.NewResponseController(rw) conn, brw, hijackErr := rc.Hijack() if errors.Is(hijackErr, http.ErrNotSupported) { + backConn.Close() p.getErrorHandler()(rw, req, fmt.Errorf("can't switch protocols using non-Hijacker ResponseWriter type %T", rw)) return } @@ -877,6 +887,9 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R defer conn.Close() copyHeader(rw.Header(), res.Header) + removeHopByHopHeaders(rw.Header()) + rw.Header().Set("Connection", "Upgrade") + rw.Header().Set("Upgrade", resUpType) res.Header = rw.Header() res.Body = nil // so res.Write only writes the headers; we have res.Body in backConn above @@ -888,8 +901,8 @@ func (p *ReverseProxy) handleUpgradeResponse(rw http.ResponseWriter, req *http.R p.getErrorHandler()(rw, req, fmt.Errorf("response flush: %v", err)) return } - errc := make(chan error, 1) - spc := switchProtocolCopier{user: conn, backend: backConn} + errc := make(chan error, 2) + spc := switchProtocolCopier{user: conn, userReader: brw.Reader, backend: backConn} go spc.copyToBackend(errc) go spc.copyFromBackend(errc) @@ -907,6 +920,7 @@ var errCopyDone = errors.New("hijacked connection copy complete") // forth have nice names in stacks. type switchProtocolCopier struct { user, backend io.ReadWriter + userReader io.Reader } func (c switchProtocolCopier) copyFromBackend(errc chan<- error) { @@ -925,7 +939,7 @@ func (c switchProtocolCopier) copyFromBackend(errc chan<- error) { } func (c switchProtocolCopier) copyToBackend(errc chan<- error) { - if _, err := io.Copy(c.backend, c.user); err != nil { + if _, err := io.Copy(c.backend, c.userReader); err != nil { errc <- err return } diff --git a/src/net/http/httputil/reverseproxy_test.go b/src/net/http/httputil/reverseproxy_test.go index 4e285e1e39a710..88fd129818f77f 100644 --- a/src/net/http/httputil/reverseproxy_test.go +++ b/src/net/http/httputil/reverseproxy_test.go @@ -2258,3 +2258,66 @@ func (rc *testReadWriteCloser) Close() error { } return nil } + + +func TestReverseProxyWebSocketDataLossAndHopByHop(t *testing.T) { + backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, _, err := w.(http.Hijacker).Hijack() + if err != nil { + t.Fatal(err) + } + defer c.Close() + + // Send 101 with a hop-by-hop header that should be removed by proxy + io.WriteString(c, "HTTP/1.1 101 Switching Protocols\r\nConnection: upgrade\r\nUpgrade: WebSocket\r\nKeep-Alive: timeout=5\r\n\r\n") + + bs := bufio.NewScanner(c) + if !bs.Scan() { + t.Errorf("backend failed to read: %v", bs.Err()) + return + } + fmt.Fprintf(c, "echo: %s\n", bs.Text()) + })) + defer backendServer.Close() + + backURL, _ := url.Parse(backendServer.URL) + rproxy := NewSingleHostReverseProxy(backURL) + + frontendProxy := httptest.NewServer(http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) { + rproxy.ServeHTTP(rw, req) + })) + defer frontendProxy.Close() + + u, _ := url.Parse(frontendProxy.URL) + conn, err := net.Dial("tcp", u.Host) + if err != nil { + t.Fatal(err) + } + defer conn.Close() + + // Write upgrade request AND immediate WebSocket payload in one batch to test bufio buffering + rawReq := fmt.Sprintf("GET / HTTP/1.1\r\nHost: %s\r\nConnection: Upgrade\r\nUpgrade: websocket\r\n\r\nearly client payload\n", u.Host) + if _, err := conn.Write([]byte(rawReq)); err != nil { + t.Fatal(err) + } + + br := bufio.NewReader(conn) + resp, err := http.ReadResponse(br, &http.Request{Method: "GET"}) + if err != nil { + t.Fatal(err) + } + if resp.StatusCode != 101 { + t.Fatalf("expected status 101, got %d", resp.StatusCode) + } + if resp.Header.Get("Keep-Alive") != "" { + t.Errorf("expected Keep-Alive to be stripped, got %q", resp.Header.Get("Keep-Alive")) + } + + line, err := br.ReadString('\n') + if err != nil { + t.Fatalf("failed reading echo: %v", err) + } + if want := "echo: early client payload\n"; line != want { + t.Errorf("got %q, want %q", line, want) + } +}