diff --git a/proxy/proxy.go b/proxy/proxy.go index 182dc0e..62606a9 100644 --- a/proxy/proxy.go +++ b/proxy/proxy.go @@ -47,6 +47,7 @@ func (p *Proxy) accept(conn net.Conn) { ctx, cancel := context.WithCancel(context.Background()) defer func() { + cancel() conn.Close() // nolint: errcheck if err := recover(); err != nil { @@ -74,6 +75,12 @@ func (p *Proxy) accept(conn net.Conn) { } defer serverConn.(io.Closer).Close() // nolint: errcheck + go func() { + <-ctx.Done() + serverConn.(io.Closer).Close() + clientConn.(io.Closer).Close() + }() + wait := &sync.WaitGroup{} wait.Add(2)