fix(sdk): Change tcp client listener to accept connections concurrently

sdk/go/docs/src/api/tcp.md

@@ -37,6 +37,12 @@ Bind a local address that forwards every connection to a sandbox port. The
 returned `ForwardListener` owns its accept loop and all bridge goroutines; it is
 a lifecycle handle with `Addr` and `Close`, not a `net.Listener`.
 
+Each accepted connection sets up its tunnel independently, including when you
+use `WithSSHTunnel()`. Setup has a 30-second timeout. If setup fails or times out,
+the SDK closes that local connection without stopping the listener. The setup
+timeout does not limit an established tunnel's lifetime; canceling `ctx` or
+calling `Close()` still closes pending setups and active tunnels.
+
 ```go
 forward, err := client.TCP().Listen(ctx, "default", "sandbox-123", 8080, 0)
 if err != nil {

sdk/go/openshell/v1/tcp.go

@@ -72,5 +72,9 @@ func WithListenServiceID(id string) ListenOption {
 // Methods accept a sandbox name and resolve it to an ID internally.
 type TCPInterface interface {
 	Forward(ctx context.Context, workspace, sandboxName string, port uint32, opts ...ForwardOption) (io.ReadWriteCloser, error)
+
+	// Listen forwards accepted local connections concurrently. Each connection
+	// has a 30-second tunnel setup timeout, including with [WithSSHTunnel].
+	// Established tunnels remain open until closed or ctx is canceled.
 	Listen(ctx context.Context, workspace, sandboxName string, remotePort uint32, localPort uint32, opts ...ListenOption) (ForwardListener, error)
 }

sdk/go/openshell/v1/tcp_client.go

@@ -10,13 +10,16 @@ import (
 	"net"
 	"strconv"
 	"sync"
+	"time"
 
 	"github.com/NVIDIA/OpenShell/sdk/go/openshell/v1/internal/converter"
 	"github.com/NVIDIA/OpenShell/sdk/go/openshell/v1/internal/options"
 	pb "github.com/NVIDIA/OpenShell/sdk/go/proto/openshellv1"
 	"google.golang.org/grpc"
 )
 
+const tunnelSetupTimeout = time.Second * 30
+
 type tcpClient struct {
 	client    pb.OpenShellClient
 	sandboxes SandboxInterface
@@ -171,58 +174,71 @@ type tunnelListener struct {
 	closeErr    error
 }
 
-func (tl *tunnelListener) acceptAndBridge() error {
-	for {
-		conn, err := tl.inner.Accept()
-		if err != nil {
-			return err
+// setupAndBridge applies a setup-only timeout, then bridges the connection.
+// It owns the worker's WaitGroup entry and closes the local connection on exit.
+func (tl *tunnelListener) setupAndBridge(conn net.Conn) {
+	defer tl.wg.Done()
+	defer func() { _ = conn.Close() }()
+
+	// Keep the context alive for bridging; stop only its setup timer on success.
+	ctx, cancel := context.WithCancel(tl.ctx)
+	defer cancel()
+	timer := time.AfterFunc(tunnelSetupTimeout, cancel)
+	defer timer.Stop()
+
+	// Establish the tunnel to the sandbox.
+	var tunnel io.ReadWriteCloser
+	var err error
+
+	if tl.cfg.useSSHTunnel && tl.ssh != nil {
+		var tunnelOpts []TunnelOption
+		if tl.cfg.serviceID != "" {
+			tunnelOpts = append(tunnelOpts, WithTunnelServiceID(tl.cfg.serviceID))
 		}
-
-		// Establish the tunnel to the sandbox.
-		var tunnel io.ReadWriteCloser
-		if tl.cfg.useSSHTunnel && tl.ssh != nil {
-			var tunnelOpts []TunnelOption
-			if tl.cfg.serviceID != "" {
-				tunnelOpts = append(tunnelOpts, WithTunnelServiceID(tl.cfg.serviceID))
-			}
-			tunnel, err = tl.ssh.Tunnel(tl.ctx, tl.workspace, tl.sandboxName, tl.remotePort, tunnelOpts...)
-		} else {
-			var fwdOpts []ForwardOption
-			if tl.cfg.serviceID != "" {
-				fwdOpts = append(fwdOpts, WithForwardServiceID(tl.cfg.serviceID))
-			}
-			tunnel, err = tl.tcp.Forward(tl.ctx, tl.workspace, tl.sandboxName, tl.remotePort, fwdOpts...)
+		tunnel, err = tl.ssh.Tunnel(ctx, tl.workspace, tl.sandboxName, tl.remotePort, tunnelOpts...)
+	} else {
+		var fwdOpts []ForwardOption
+		if tl.cfg.serviceID != "" {
+			fwdOpts = append(fwdOpts, WithForwardServiceID(tl.cfg.serviceID))
 		}
+		tunnel, err = tl.tcp.Forward(ctx, tl.workspace, tl.sandboxName, tl.remotePort, fwdOpts...)
+	}
 
-		if err != nil {
-			_ = conn.Close()
-			select {
-			case <-tl.ctx.Done():
-				return tl.ctx.Err()
-			default:
-				continue
-			}
-		}
+	if err != nil {
+		return
+	}
 
-		tl.mu.Lock()
-		if tl.closing {
-			tl.mu.Unlock()
-			_ = conn.Close()
-			_ = tunnel.Close()
-			return net.ErrClosed
-		}
-		tl.wg.Add(1)
-		tl.mu.Unlock()
-		go tl.bridge(conn, tunnel)
+	if !timer.Stop() {
+		// Connection setup timed out
+		_ = tunnel.Close()
+		return
+	}
+
+	tl.bridge(conn, tunnel)
+}
+
+func (tl *tunnelListener) acceptAndBridge() error {
+	conn, err := tl.inner.Accept()
+	if err != nil {
+		return err
+	}
 
-		return nil
+	tl.mu.Lock()
+	if tl.closing {
+		tl.mu.Unlock()
+		_ = conn.Close()
+		return net.ErrClosed
 	}
+	tl.wg.Add(1)
+	tl.mu.Unlock()
+
+	go tl.setupAndBridge(conn)
+	return nil
 }
 
 // bridge copies data bidirectionally between the local connection and the
-// tunnel. It runs in its own goroutine and decrements the WaitGroup on exit.
+// tunnel. It runs in its own goroutine.
 func (tl *tunnelListener) bridge(local net.Conn, tunnel io.ReadWriteCloser) {
-	defer tl.wg.Done()
 	defer func() { _ = local.Close() }()
 	defer func() { _ = tunnel.Close() }()
 
@@ -247,7 +263,7 @@ func (tl *tunnelListener) bridge(local net.Conn, tunnel io.ReadWriteCloser) {
 }
 
 // Close stops the listener from accepting new connections, cancels all
-// active tunnels, and blocks until all bridge goroutines finish.
+// pending setups and active tunnels, and waits for their workers to finish.
 func (tl *tunnelListener) Close() error {
 	tl.closeOnce.Do(func() {
 		tl.mu.Lock()

sdk/go/openshell/v1/tcp_client_test.go

@@ -9,6 +9,7 @@ import (
 	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -100,6 +101,53 @@ func setupTCPTest(t *testing.T, mock *mockTCPServer) (*tcpClient, func()) {
 	}
 }
 
+type stubTunnelSetupResolver struct {
+	SandboxInterface
+	get func(context.Context) error
+}
+
+func (r *stubTunnelSetupResolver) Get(ctx context.Context, _, name string) (*Sandbox, error) {
+	if err := r.get(ctx); err != nil {
+		return nil, err
+	}
+	return &Sandbox{ID: "sb-" + name, Name: name}, nil
+}
+
+func setupTCPListenTest(t *testing.T, useSSH bool, get func(context.Context) error) (*tcpClient, []ListenOption) {
+	t.Helper()
+	client, cleanup := setupTCPTest(t, newMockTCPServer())
+	t.Cleanup(cleanup)
+	resolver := &stubTunnelSetupResolver{SandboxInterface: client.sandboxes, get: get}
+	client.sandboxes = resolver
+	if !useSSH {
+		return client, nil
+	}
+
+	ssh, sshCleanup := setupSSHTestWithSandboxes(t, newMockSSHServer(), resolver)
+	t.Cleanup(sshCleanup)
+	client.ssh = ssh
+	return client, []ListenOption{WithSSHTunnel()}
+}
+
+func dialForwardListener(t *testing.T, ln ForwardListener) net.Conn {
+	t.Helper()
+	conn, err := net.DialTimeout("tcp", ln.Addr().String(), 5*time.Second)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	require.NoError(t, conn.SetDeadline(time.Now().Add(5*time.Second)))
+	return conn
+}
+
+func assertTCPForwardEcho(t *testing.T, conn net.Conn) {
+	t.Helper()
+	_, err := conn.Write([]byte("ping"))
+	require.NoError(t, err)
+	buf := make([]byte, len("ping"))
+	_, err = io.ReadFull(conn, buf)
+	require.NoError(t, err)
+	assert.Equal(t, "ping", string(buf))
+}
+
 // --- Tests ---
 
 func TestTCPForward_InitFrame(t *testing.T) {
@@ -370,7 +418,6 @@ func TestTCPForward_ServerError(t *testing.T) {
 	defer cleanup()
 
 	rwc, err := client.Forward(context.Background(), "default", "my-sandbox", 8080)
-
 	// The stream opens successfully (gRPC bidi streams don't fail on open),
 	// but the first write or read should surface the server error.
 	if err != nil {
@@ -534,6 +581,50 @@ func TestTCPListen_ConcurrentConnections(t *testing.T) {
 	}
 }
 
+// tcp listener should allow 4 concurrent connections that do not block each other
+func TestTCPListen_ConcurrentTunnelSetup(t *testing.T) {
+	for _, transport := range []string{"tcp", "ssh"} {
+		t.Run(transport, func(t *testing.T) {
+			const numConns = 4
+			entered := make(chan struct{}, numConns)
+			release := make(chan struct{})
+			client, opts := setupTCPListenTest(t, transport == "ssh", func(ctx context.Context) error {
+				entered <- struct{}{}
+				select {
+				case <-release:
+					return nil
+				case <-ctx.Done():
+					return ctx.Err()
+				}
+			})
+
+			ln, err := client.Listen(context.Background(), "default", "my-sandbox", 8080, 0, opts...)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = ln.Close() })
+
+			conns := make([]net.Conn, numConns)
+			for i := range numConns {
+				conns[i] = dialForwardListener(t, ln)
+			}
+
+			// Every setup must start before any is allowed to finish.
+			timer := time.NewTimer(5 * time.Second)
+			defer timer.Stop()
+			for i := range numConns {
+				select {
+				case <-entered:
+				case <-timer.C:
+					t.Fatalf("only %d/%d tunnel setups started; setup is serialized", i, numConns)
+				}
+			}
+			close(release)
+			for _, conn := range conns {
+				assertTCPForwardEcho(t, conn)
+			}
+		})
+	}
+}
+
 func TestTCPListen_EphemeralPort(t *testing.T) {
 	mock := newMockTCPServer()
 	client, cleanup := setupTCPTest(t, mock)
@@ -637,6 +728,46 @@ func TestTCPListen_InvalidRemotePort(t *testing.T) {
 
 // --- Graceful shutdown tests ---
 
+func TestTCPListen_HungTunnelSetup(t *testing.T) {
+	for _, transport := range []string{"tcp", "ssh"} {
+		t.Run(transport, func(t *testing.T) {
+			var calls atomic.Int32
+			started := make(chan struct{})
+			client, opts := setupTCPListenTest(t, transport == "ssh", func(ctx context.Context) error {
+				if calls.Add(1) == 1 {
+					close(started)
+					<-ctx.Done()
+					return ctx.Err()
+				}
+				return nil
+			})
+
+			ln, err := client.Listen(context.Background(), "default", "my-sandbox", 8080, 0, opts...)
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = ln.Close() })
+			_ = dialForwardListener(t, ln)
+			select {
+			case <-started:
+			case <-time.After(5 * time.Second):
+				t.Fatal("first tunnel setup did not start")
+			}
+
+			// A hung setup must not block another connection's data flow.
+			assertTCPForwardEcho(t, dialForwardListener(t, ln))
+
+			// Close must cancel the hung setup and wait for workers to finish.
+			closed := make(chan error, 1)
+			go func() { closed <- ln.Close() }()
+			select {
+			case err := <-closed:
+				require.NoError(t, err)
+			case <-time.After(5 * time.Second):
+				t.Fatal("Close did not cancel and drain tunnel setup workers")
+			}
+		})
+	}
+}
+
 func TestTCPListen_CloseTerminatesConnections(t *testing.T) {
 	mock := newMockTCPServer()
 	client, cleanup := setupTCPTest(t, mock)