diff --git a/caddyhttp/proxy/proxy_test.go b/caddyhttp/proxy/proxy_test.go index 3166e85f..6412cc1b 100644 --- a/caddyhttp/proxy/proxy_test.go +++ b/caddyhttp/proxy/proxy_test.go @@ -304,6 +304,40 @@ func TestWebSocketReverseProxyNonHijackerPanic(t *testing.T) { p.ServeHTTP(nonHijacker, r) } +func TestWebSocketReverseProxyBackendShutDown(t *testing.T) { + shutdown := make(chan struct{}) + backend := httptest.NewServer(websocket.Handler(func(ws *websocket.Conn) { + shutdown <- struct{}{} + })) + defer backend.Close() + + go func() { + <-shutdown + backend.Close() + }() + + // Get proxy to use for the test + p := newWebSocketTestProxy(backend.URL, false) + backendProxy := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + p.ServeHTTP(w, r) + })) + defer backendProxy.Close() + + // Set up WebSocket client + url := strings.Replace(backendProxy.URL, "http://", "ws://", 1) + ws, err := websocket.Dial(url, "", backendProxy.URL) + + if err != nil { + t.Fatal(err) + } + defer ws.Close() + + var actualMsg string + if rcvErr := websocket.Message.Receive(ws, &actualMsg); rcvErr == nil { + t.Errorf("we don't get backend shutdown notification") + } +} + func TestWebSocketReverseProxyServeHTTPHandler(t *testing.T) { // No-op websocket backend simply allows the WS connection to be // accepted then it will be immediately closed. Perfect for testing. diff --git a/caddyhttp/proxy/reverseproxy.go b/caddyhttp/proxy/reverseproxy.go index 49660749..c22cf101 100644 --- a/caddyhttp/proxy/reverseproxy.go +++ b/caddyhttp/proxy/reverseproxy.go @@ -320,8 +320,13 @@ func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request, } defer backendConn.Close() + proxyDone := make(chan struct{}, 2) + // Proxy backend -> frontend. - go pooledIoCopy(conn, backendConn) + go func() { + pooledIoCopy(conn, backendConn) + proxyDone <- struct{}{} + }() // Proxy frontend -> backend. // @@ -336,7 +341,13 @@ func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request, backendConn.Write(rbuf) } } - pooledIoCopy(backendConn, conn) + go func() { + pooledIoCopy(backendConn, conn) + proxyDone <- struct{}{} + }() + + // If one side is done, we are done. + <-proxyDone } else { // NOTE: // Closing the Body involves acquiring a mutex, which is a