diff --git a/handler/tunnel/entrypoint.go b/handler/tunnel/entrypoint.go index 074ed2d6..0ec16549 100644 --- a/handler/tunnel/entrypoint.go +++ b/handler/tunnel/entrypoint.go @@ -440,8 +440,8 @@ func (ep *entrypoint) handleUpgradeResponse(ctx context.Context, rw io.ReadWrite return xnet.Pipe(ctx, rw, backConn) } -func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriter, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { - errc := make(chan error, 1) +func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.ReadWriteCloser, ro *xrecorder.HandlerRecorderObject, log logger.Logger) error { + errc := make(chan error, 2) sampleRate := ep.websocketSampleRate if sampleRate == 0 { @@ -503,6 +503,9 @@ func (ep *entrypoint) sniffingWebsocketFrame(ctx context.Context, rw, cc io.Read } }() + <-errc + rw.Close() + cc.Close() <-errc return nil } diff --git a/handler/tunnel/tunnel.go b/handler/tunnel/tunnel.go index d1e407ef..3c2e12b1 100644 --- a/handler/tunnel/tunnel.go +++ b/handler/tunnel/tunnel.go @@ -230,8 +230,8 @@ func (t *Tunnel) Close() error { } func (t *Tunnel) CloseOnIdle() bool { - t.mu.RLock() - defer t.mu.RUnlock() + t.mu.Lock() + defer t.mu.Unlock() select { case <-t.close: