Commit e377eeff authored by Tw's avatar Tw Committed by Matt Holt

proxy: websocket proxy exits immediately if backend is shutdown (#1869)

Signed-off-by: default avatarTw <tw19881113@gmail.com>
parent 84a2f8e8
...@@ -304,6 +304,40 @@ func TestWebSocketReverseProxyNonHijackerPanic(t *testing.T) { ...@@ -304,6 +304,40 @@ func TestWebSocketReverseProxyNonHijackerPanic(t *testing.T) {
p.ServeHTTP(nonHijacker, r) 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) { func TestWebSocketReverseProxyServeHTTPHandler(t *testing.T) {
// No-op websocket backend simply allows the WS connection to be // No-op websocket backend simply allows the WS connection to be
// accepted then it will be immediately closed. Perfect for testing. // accepted then it will be immediately closed. Perfect for testing.
......
...@@ -320,8 +320,13 @@ func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request, ...@@ -320,8 +320,13 @@ func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request,
} }
defer backendConn.Close() defer backendConn.Close()
proxyDone := make(chan struct{}, 2)
// Proxy backend -> frontend. // Proxy backend -> frontend.
go pooledIoCopy(conn, backendConn) go func() {
pooledIoCopy(conn, backendConn)
proxyDone <- struct{}{}
}()
// Proxy frontend -> backend. // Proxy frontend -> backend.
// //
...@@ -336,7 +341,13 @@ func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request, ...@@ -336,7 +341,13 @@ func (rp *ReverseProxy) ServeHTTP(rw http.ResponseWriter, outreq *http.Request,
backendConn.Write(rbuf) backendConn.Write(rbuf)
} }
} }
pooledIoCopy(backendConn, conn) go func() {
pooledIoCopy(backendConn, conn)
proxyDone <- struct{}{}
}()
// If one side is done, we are done.
<-proxyDone
} else { } else {
// NOTE: // NOTE:
// Closing the Body involves acquiring a mutex, which is a // Closing the Body involves acquiring a mutex, which is a
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment