ssh/proxy: improve cancellation and resource cleanup (#3965)

## Changes
<!-- Brief summary of your changes that is easy to understand -->
There are edge cases where the proxy connection could hang indefinitely
during shutdown if the sending or receiving loops were blocked on I/O
(e.g. reading from stdin or the websocket).

- Introduce a dedicated cleanup goroutine in `start` that waits for
context cancellation and explicitly closes both the websocket connection
and the source. This ensures that blocked `Read` calls return
immediately with an error, allowing the `errgroup` to finish.
- Ensure that sending and receiving loops always return a non-nil error
to trigger `errgroup` cancellation for the other loop.
- Update `RunClientProxy` and `proxyConnection.start` to accept
`io.ReadCloser` instead of `io.Reader`, allowing the proxy to close the
input source.


## Tests
Existing tests with additional `-count` and `-race` flags.

<!-- If your PR needs to be included in the release notes for next
release,
add a separate entry in NEXT_CHANGELOG.md as part of your PR. -->

Author: ilia-db

experimental/ssh/internal/proxy/client.go

@@ -10,7 +10,7 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-func RunClientProxy(ctx context.Context, src io.Reader, dst io.Writer, requestHandoverTick func() <-chan time.Time, createConn createWebsocketConnectionFunc) error {
+func RunClientProxy(ctx context.Context, src io.ReadCloser, dst io.Writer, requestHandoverTick func() <-chan time.Time, createConn createWebsocketConnectionFunc) error {
 	proxy := newProxyConnection(createConn)
 	cmdio.LogString(ctx, "Establishing SSH proxy connection...")
 	g, gCtx := errgroup.WithContext(ctx)

experimental/ssh/internal/proxy/client_server_test.go

@@ -48,7 +48,7 @@ func createTestClient(t *testing.T, serverURL string, requestHandoverTick func()
 	}
 	go func() {
 		err := RunClientProxy(ctx, clientInput, clientOutput, requestHandoverTick, createConn)
-		if err != nil && !isNormalClosure(err) && !errors.Is(err, context.Canceled) {
+		if err != nil && !errors.Is(err, context.Canceled) {
 			if errChan != nil {
 				errChan <- err
 			} else {

experimental/ssh/internal/proxy/proxy.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"os"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -15,7 +16,11 @@ import (
 	"golang.org/x/sync/errgroup"
 )
 
-var errProxyEOF = errors.New("proxy EOF error")
+var (
+	errProxyEOF             = errors.New("proxy EOF error")
+	errSendingLoopStopped   = errors.New("sending loop stopped")
+	errReceivingLoopStopped = errors.New("receiving loop stopped")
+)
 
 const (
 	// Same as gorilla/websocket default read/write buffer sizes. Bigger payloads will be split into multiple ws frames.
@@ -105,13 +110,27 @@ func newProxyConnection(createConn createWebsocketConnectionFunc) *proxyConnecti
 	}
 }
 
-func (pc *proxyConnection) start(ctx context.Context, src io.Reader, dst io.Writer) error {
+func (pc *proxyConnection) start(ctx context.Context, src io.ReadCloser, dst io.Writer) error {
 	g, gCtx := errgroup.WithContext(ctx)
 	g.Go(func() error {
-		return pc.runSendingLoop(gCtx, src)
+		err := pc.runSendingLoop(gCtx, src)
+		// Always return a non nil error to cancel the errgroup context
+		return errors.Join(err, errSendingLoopStopped)
+	})
+	g.Go(func() error {
+		err := pc.runReceivingLoop(gCtx, dst)
+		// Always return a non nil error to cancel the errgroup context
+		return errors.Join(err, errReceivingLoopStopped)
 	})
 	g.Go(func() error {
-		return pc.runReceivingLoop(gCtx, dst)
+		// Wait for the context to be cancelled. There can be multiple reasons:
+		// - Sending loop finished (e.g. EOF from source)
+		// - Receiving loop finished (e.g. connection closed)
+		// - Parent context cancelled
+		// Both loops can still be stuck on conn.ReadMessage or src.Read and won't notice context cancellation,
+		// so we close the connection and the source (sshd stdout pipe or ssh client stdio) to unblock them.
+		<-gCtx.Done()
+		return errors.Join(pc.close(), pc.closeSource(src))
 	})
 	err := g.Wait()
 	if err == nil || isNormalClosure(err) {
@@ -236,6 +255,14 @@ func (pc *proxyConnection) close() error {
 	return nil
 }
 
+func (pc *proxyConnection) closeSource(src io.ReadCloser) error {
+	err := src.Close()
+	if err != nil && (errors.Is(err, os.ErrClosed) || errors.Is(err, io.ErrClosedPipe)) {
+		return nil
+	}
+	return err
+}
+
 func (pc *proxyConnection) initiateHandover(ctx context.Context) error {
 	// Blocks proxying any outgoing messages during the entire handover
 	pc.handoverMutex.Lock()

experimental/ssh/internal/proxy/proxy_test.go

@@ -115,7 +115,7 @@ func setupTestServer(ctx context.Context, t *testing.T) *TestProxy {
 		}
 		defer serverProxy.close()
 		err = serverProxy.start(ctx, serverInput, serverOutput)
-		if err != nil && !errors.Is(err, errProxyEOF) {
+		if err != nil && !errors.Is(err, context.Canceled) {
 			t.Errorf("server error: %v", err)
 			return
 		}
@@ -152,7 +152,7 @@ func setupTestClient(ctx context.Context, t *testing.T, serverURL string) *TestP
 
 	go func() {
 		err := clientProxy.start(ctx, clientInput, clientOutput)
-		if err != nil && !errors.Is(err, errProxyEOF) {
+		if err != nil && !errors.Is(err, context.Canceled) {
 			t.Errorf("proxy error: %v", err)
 		}
 	}()