diff --git a/sdk/typescript/src/client.test.ts b/sdk/typescript/src/client.test.ts index 86fdfa3972..756ff6e44d 100644 --- a/sdk/typescript/src/client.test.ts +++ b/sdk/typescript/src/client.test.ts @@ -12,6 +12,7 @@ import type { MessageInitShape } from '@bufbuild/protobuf'; import { Code, ConnectError, createRouterTransport, type ServiceImpl, type Transport } from '@connectrpc/connect'; import { afterEach, describe, expect, it } from 'vitest'; import { + __setSshCleanupTimeoutMsForTests, __setWatchLogsBackoffSleepForTests, errorCode, Pager, @@ -2236,6 +2237,81 @@ describe('forward', () => { expect(revokedToken).toBe('fwd-tok'); }); + it('revokes the token of a connection still open at close()', async () => { + const revoked: string[] = []; + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id-open'), + createSshSession: () => ({ + sandboxId: 'sb-id-open', + token: 'open-tok', + gatewayHost: 'gw', + gatewayPort: 443, + gatewayScheme: 'https', + hostKeyFingerprint: '', + expirationTime: undefined, + }), + // A real transport never delivers a call whose signal is already aborted. + revokeSshSession: (req, ctx) => { + if (!ctx.signal.aborted) revoked.push(req.token); + return { outcome: 1 }; + }, + forwardTcp: async function* (requests) { + for await (const frame of requests) { + if (frame.payload.case === 'data') yield { payload: { case: 'data', value: frame.payload.value } }; + } + }, + }); + + const handle = await sandbox.forward('sb', { targetPort: 9000 }); + const socket = net.connect(handle.localPort, handle.localHost); + socket.on('error', () => {}); + await new Promise((resolve) => { + socket.once('data', () => resolve()); + socket.write('ping'); + }); + + await handle.close(); + expect(revoked).toEqual(['open-tok']); + }); + + it('close() resolves within the cleanup bound when revoke never resolves', async () => { + __setSshCleanupTimeoutMsForTests(100); + try { + const sandbox = client({ + getSandbox: () => readySandbox('sb', 'sb-id-stall'), + createSshSession: () => ({ + sandboxId: 'sb-id-stall', + token: 'stall-tok', + gatewayHost: 'gw', + gatewayPort: 443, + gatewayScheme: 'https', + hostKeyFingerprint: '', + expirationTime: undefined, + }), + revokeSshSession: () => new Promise(() => {}), + forwardTcp: async function* (requests) { + for await (const frame of requests) { + if (frame.payload.case === 'data') yield { payload: { case: 'data', value: frame.payload.value } }; + } + }, + }); + + const handle = await sandbox.forward('sb', { targetPort: 9000 }); + const socket = net.connect(handle.localPort, handle.localHost); + socket.on('error', () => {}); + await new Promise((resolve) => { + socket.once('data', () => resolve()); + socket.write('ping'); + }); + + const start = Date.now(); + await handle.close(); + expect(Date.now() - start).toBeLessThan(2_000); + } finally { + __setSshCleanupTimeoutMsForTests(); + } + }); + it('rejects when the sandbox is not ready', async () => { const sandbox = client({ getSandbox: () => ({ diff --git a/sdk/typescript/src/client.ts b/sdk/typescript/src/client.ts index f82242d4a7..1a92768e5b 100644 --- a/sdk/typescript/src/client.ts +++ b/sdk/typescript/src/client.ts @@ -485,7 +485,8 @@ export interface ForwardOptions extends SandboxWorkspaceOptions { // A process-lifetime local listener that tunnels each accepted connection into // the sandbox. Call `close()` on teardown; `closed` resolves once the listener -// is fully torn down. An in-process forward cannot outlive the Node process. +// is fully torn down. `close()` may wait up to 5s for best-effort revokes of +// open connections' tokens. An in-process forward cannot outlive the Node process. export interface ForwardHandle { localHost: string; localPort: number; @@ -723,6 +724,15 @@ function versionPin(value: string | undefined): bigint { const FORWARD_CHUNK = 64 * 1024; +// Bound for revoking a forward's SSH session token during cleanup. +const SSH_CLEANUP_TIMEOUT_MS = 5_000; +let sshCleanupTimeoutMs = SSH_CLEANUP_TIMEOUT_MS; + +// Test hook; not re-exported from index.ts. No argument restores the default. +export function __setSshCleanupTimeoutMsForTests(ms?: number): void { + sshCleanupTimeoutMs = ms ?? SSH_CLEANUP_TIMEOUT_MS; +} + function workspaceName(options?: SandboxWorkspaceOptions | null): string { const workspace = options?.workspace ?? 'default'; if (workspace.trim() === '') throw new SdkError('invalid_config', 'workspace must be non-empty'); @@ -1894,7 +1904,11 @@ export class SandboxClient { input.end(); if (token !== undefined) { try { - await this.grpc.revokeSshSession({ token, allowMissing: true }, { signal }); + // Not the connection signal: close() aborts it before this runs. + await this.grpc.revokeSshSession( + { token, allowMissing: true }, + { signal: AbortSignal.timeout(sshCleanupTimeoutMs) }, + ); } catch { // Best-effort revoke; the token expires on its own regardless. }