From db11bd9cdbafbbe1bfd5fc70204ca38f34a99c30 Mon Sep 17 00:00:00 2001 From: "Jason T. Greene" Date: Sat, 3 Dec 2022 01:42:08 -0600 Subject: [PATCH] Fix interrupted pending writes on socket write shutdown from eager close Some channel/connection implementations may signal EOF to parallel readers before tasks related to the CloseWrite (shutdown) have completed progressing. This creates the potential for a race with a parallel Close(), leading to a premature abort of certain activies (cancelling the send of buffered data). This change ensures that the two goroutines copying each direction of the stream wait until CloseWrite has completed in both directions before fully closing. Signed-off-by: Jason T. Greene --- pkg/sshclient/ssh_forwarder.go | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/pkg/sshclient/ssh_forwarder.go b/pkg/sshclient/ssh_forwarder.go index 79de4978e..38075c230 100644 --- a/pkg/sshclient/ssh_forwarder.go +++ b/pkg/sshclient/ssh_forwarder.go @@ -8,6 +8,7 @@ import ( "os" "runtime" "strings" + "sync" "time" "github.com/containers/gvisor-tap-vsock/pkg/fs" @@ -228,14 +229,22 @@ func acceptConnection(ctx context.Context, listener net.Listener, bastion *Basti return nil // eat } - go forward(src, dest) - go forward(dest, src) + complete := new(sync.WaitGroup) + complete.Add(2) + go forward(src, dest, complete) + go forward(dest, src, complete) + + go func() { + complete.Wait() + src.Close() + dest.Close() + }() return nil } -func forward(src io.ReadCloser, dest CloseWriteStream) { - defer src.Close() +func forward(src io.ReadCloser, dest CloseWriteStream, complete *sync.WaitGroup) { + defer complete.Done() _, _ = io.Copy(dest, src) // Trigger an EOF on the other end