reverseproxy: Wait for both ends of websocket to close (#6175)

This commit is contained in:
WeidiDeng 2024-04-16 01:37:37 +08:00 committed by Matthew Holt
parent 93ec641d5d
commit b1fb7d813e
No known key found for this signature in database
GPG key ID: 2A349DD577D586A5
2 changed files with 8 additions and 2 deletions

View file

@ -947,7 +947,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
}

View file

@ -182,7 +182,7 @@ func (h *Handler) handleUpgradeResponse(logger *zap.Logger, wg *sync.WaitGroup,
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
@ -193,6 +193,7 @@ func (h *Handler) handleUpgradeResponse(logger *zap.Logger, wg *sync.WaitGroup,
}
errc := make(chan error, 1)
wg.Add(2)
go spc.copyToBackend(errc)
go spc.copyFromBackend(errc)
select {
@ -590,16 +591,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{