diff --git a/sdk/go/docs/src/api/tcp.md b/sdk/go/docs/src/api/tcp.md index 025ac1e987..a0b602ec4b 100644 --- a/sdk/go/docs/src/api/tcp.md +++ b/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 { diff --git a/sdk/go/openshell/v1/tcp.go b/sdk/go/openshell/v1/tcp.go index a8fcae133d..c828d5a677 100644 --- a/sdk/go/openshell/v1/tcp.go +++ b/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) } diff --git a/sdk/go/openshell/v1/tcp_client.go b/sdk/go/openshell/v1/tcp_client.go index 24f8062062..d747b5068f 100644 --- a/sdk/go/openshell/v1/tcp_client.go +++ b/sdk/go/openshell/v1/tcp_client.go @@ -10,6 +10,7 @@ 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" @@ -17,6 +18,8 @@ import ( "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() diff --git a/sdk/go/openshell/v1/tcp_client_test.go b/sdk/go/openshell/v1/tcp_client_test.go index 46b0f586d9..563926d40a 100644 --- a/sdk/go/openshell/v1/tcp_client_test.go +++ b/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)