From e377eeff50e6785605d15670eaf8740bad889e8b Mon Sep 17 00:00:00 2001
From: Tw <tw19881113@gmail.com>
Date: Sat, 23 Sep 2017 08:10:48 +0800
Subject: [PATCH] proxy: websocket proxy exits immediately if backend is
 shutdown (#1869)

Signed-off-by: Tw <tw19881113@gmail.com>
---
 caddyhttp/proxy/proxy_test.go   | 34 +++++++++++++++++++++++++++++++++
 caddyhttp/proxy/reverseproxy.go | 15 +++++++++++++--
 2 files changed, 47 insertions(+), 2 deletions(-)

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