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 <[email protected]>

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