From 7bab90c006b85b0257f2bd6c6735029237d4cc8f 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 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