From b40cacf5ce8f6a190a002fa44e973f1aa5bac2b0 Mon Sep 17 00:00:00 2001
From: WeidiDeng <weidi_deng@icloud.com>
Date: Tue, 16 Apr 2024 01:37:37 +0800
Subject: [PATCH] reverseproxy: Wait for both ends of websocket to close
 (#6175)

---
 modules/caddyhttp/reverseproxy/reverseproxy.go | 4 +++-
 modules/caddyhttp/reverseproxy/streaming.go    | 8 ++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/modules/caddyhttp/reverseproxy/reverseproxy.go b/modules/caddyhttp/reverseproxy/reverseproxy.go
index 1648b779..deba304a 100644
--- a/modules/caddyhttp/reverseproxy/reverseproxy.go
+++ b/modules/caddyhttp/reverseproxy/reverseproxy.go
@@ -922,7 +922,9 @@ func (h *Handler) finalizeResponse(
 ) error {
 	// deal with 101 Switching Protocols responses: (WebSocket, h2c, etc)
 	if res.StatusCode == http.StatusSwitchingProtocols {
-		h.handleUpgradeResponse(logger, rw, req, res)
+		var wg sync.WaitGroup
+		h.handleUpgradeResponse(logger, &wg, rw, req, res)
+		wg.Wait()
 		return nil
 	}
 
diff --git a/modules/caddyhttp/reverseproxy/streaming.go b/modules/caddyhttp/reverseproxy/streaming.go
index 2a5b5286..6d899057 100644
--- a/modules/caddyhttp/reverseproxy/streaming.go
+++ b/modules/caddyhttp/reverseproxy/streaming.go
@@ -34,7 +34,7 @@ import (
 	"golang.org/x/net/http/httpguts"
 )
 
-func (h *Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWriter, req *http.Request, res *http.Response) {
+func (h *Handler) handleUpgradeResponse(logger *zap.Logger, wg *sync.WaitGroup, rw http.ResponseWriter, req *http.Request, res *http.Response) {
 	reqUpType := upgradeType(req.Header)
 	resUpType := upgradeType(res.Header)
 
@@ -121,7 +121,7 @@ func (h *Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWrit
 	defer deleteFrontConn()
 	defer deleteBackConn()
 
-	spc := switchProtocolCopier{user: conn, backend: backConn}
+	spc := switchProtocolCopier{user: conn, backend: backConn, wg: wg}
 
 	// setup the timeout if requested
 	var timeoutc <-chan time.Time
@@ -132,6 +132,7 @@ func (h *Handler) handleUpgradeResponse(logger *zap.Logger, rw http.ResponseWrit
 	}
 
 	errc := make(chan error, 1)
+	wg.Add(2)
 	go spc.copyToBackend(errc)
 	go spc.copyFromBackend(errc)
 	select {
@@ -529,16 +530,19 @@ func (m *maxLatencyWriter) stop() {
 // forth have nice names in stacks.
 type switchProtocolCopier struct {
 	user, backend io.ReadWriteCloser
+	wg            *sync.WaitGroup
 }
 
 func (c switchProtocolCopier) copyFromBackend(errc chan<- error) {
 	_, err := io.Copy(c.user, c.backend)
 	errc <- err
+	c.wg.Done()
 }
 
 func (c switchProtocolCopier) copyToBackend(errc chan<- error) {
 	_, err := io.Copy(c.backend, c.user)
 	errc <- err
+	c.wg.Done()
 }
 
 var streamingBufPool = sync.Pool{