From b024c68728c02d418c0d75cd4cafb411afad5cc7 Mon Sep 17 00:00:00 2001 From: Jinjing <6427696+AmethystLiang@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:37:58 -0700 Subject: [PATCH] fix(ssh): close transports owned by cancelled connect attempts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When a connect is cancelled after its transport has opened, that cancelled attempt still owns the transport and must close it — otherwise it leaks. Add disconnectConnection() to close by identity (not by target ID) so a cancelled attempt closes only the transport it minted, without tearing down its replacement's live transport. Track priorConnection to detect whether this attempt opened a new transport or reused an existing one, and close only on abandonment if this attempt owns the session. --- src/main/ipc/ssh.test.ts | 103 ++++++++++++++++++ src/main/ipc/ssh.ts | 37 ++++++- src/main/ssh/ssh-connection-manager.test.ts | 32 ++++++ src/main/ssh/ssh-connection-manager.ts | 12 ++ ...ay-pty-consumer-owner-displacement.test.ts | 3 +- src/shared/pty-consumer-session.test.ts | 4 + 6 files changed, 189 insertions(+), 2 deletions(-) diff --git a/src/main/ipc/ssh.test.ts b/src/main/ipc/ssh.test.ts index e76aa5302b5..6609899dce1 100644 --- a/src/main/ipc/ssh.test.ts +++ b/src/main/ipc/ssh.test.ts @@ -45,6 +45,7 @@ const { mockConnectionManager: { connect: vi.fn(), disconnect: vi.fn(), + disconnectConnection: vi.fn(), reconnect: vi.fn(), getConnection: vi.fn(), getState: vi.fn(), @@ -319,6 +320,7 @@ describe('SSH IPC handlers', () => { const createConnectionManagerMock = () => ({ connect: vi.fn(), disconnect: vi.fn(), + disconnectConnection: vi.fn(), reconnect: vi.fn(), getConnection: vi.fn(), getState: vi.fn(), @@ -388,6 +390,7 @@ describe('SSH IPC handlers', () => { mockConnectionManager.connect.mockReset() mockConnectionManager.disconnect.mockReset() + mockConnectionManager.disconnectConnection.mockReset().mockResolvedValue(undefined) mockConnectionManager.reconnect.mockReset() mockConnectionManager.getConnection.mockReset() mockConnectionManager.getState.mockReset() @@ -2032,6 +2035,106 @@ describe('SSH IPC handlers', () => { expect(mockDeployAndLaunchRelay).toHaveBeenCalledTimes(1) }) + it('closes the transport a cancelled connect opened after the disconnect finished', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + const lateConn = { id: 'late-transport' } + let resolveStaleConnect!: (connection: unknown) => void + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getConnection.mockReturnValue(undefined) + mockConnectionManager.connect.mockReturnValueOnce( + new Promise((resolve) => { + resolveStaleConnect = resolve + }) + ) + mockConnectionManager.disconnect.mockResolvedValue(undefined) + + const staleConnect = handlers.get('ssh:connect')!(null, { + targetId: 'ssh-1' + }) as Promise + await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1)) + // Why await the whole disconnect: the leak only exists once its teardown has already run, so + // nothing else is left to close the transport this attempt opens afterwards. + await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' }) + + resolveStaleConnect(lateConn) + + await expect(staleConnect).rejects.toThrow('SSH connection attempt was cancelled') + expect(mockConnectionManager.disconnectConnection).toHaveBeenCalledWith('ssh-1', lateConn) + expect(getActiveMultiplexer('ssh-1')).toBeUndefined() + }) + + it('closes the transport when establish resumes after the connect was invalidated', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + const conn = { id: 'establishing-transport' } + let releaseRelayLaunch = (): void => {} + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getConnection.mockReturnValue(undefined) + mockConnectionManager.connect.mockResolvedValue(conn) + mockConnectionManager.disconnect.mockResolvedValue(undefined) + mockDeployAndLaunchRelay.mockImplementationOnce( + () => + new Promise((resolve) => { + releaseRelayLaunch = () => resolve(createRelayLaunchResult()) + }) + ) + + const connect = handlers.get('ssh:connect')!(null, { + targetId: 'ssh-1' + }) as Promise + await vi.waitFor(() => expect(mockDeployAndLaunchRelay).toHaveBeenCalledTimes(1)) + await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' }) + + releaseRelayLaunch() + + await expect(connect).rejects.toThrow('SSH connection attempt was cancelled') + expect(mockConnectionManager.disconnectConnection).toHaveBeenCalledWith('ssh-1', conn) + expect(getActiveMultiplexer('ssh-1')).toBeUndefined() + }) + + it('leaves a reused transport to its replacement when the connect is cancelled', async () => { + const target: SshTarget = { + id: 'ssh-1', + label: 'Server', + host: 'example.com', + port: 22, + username: 'deploy' + } + // Why: connect() hands back the already-open transport, so this attempt never owned it. + const sharedConn = { id: 'shared-transport' } + let resolveStaleConnect!: (connection: unknown) => void + mockSshStore.getTarget.mockReturnValue(target) + mockConnectionManager.getConnection.mockReturnValue(sharedConn) + mockConnectionManager.connect.mockReturnValueOnce( + new Promise((resolve) => { + resolveStaleConnect = resolve + }) + ) + mockConnectionManager.disconnect.mockResolvedValue(undefined) + + const staleConnect = handlers.get('ssh:connect')!(null, { + targetId: 'ssh-1' + }) as Promise + await vi.waitFor(() => expect(mockConnectionManager.connect).toHaveBeenCalledTimes(1)) + await handlers.get('ssh:disconnect')!(null, { targetId: 'ssh-1' }) + + resolveStaleConnect(sharedConn) + + await expect(staleConnect).rejects.toThrow('SSH connection attempt was cancelled') + expect(mockConnectionManager.disconnectConnection).not.toHaveBeenCalled() + }) + it('keeps reconnect behind transport disconnect when forward teardown fails', async () => { const target: SshTarget = { id: 'ssh-1', diff --git a/src/main/ipc/ssh.ts b/src/main/ipc/ssh.ts index db72b6f8928..3b362aa7ef4 100644 --- a/src/main/ipc/ssh.ts +++ b/src/main/ipc/ssh.ts @@ -7,7 +7,7 @@ import { listUserSshConfigHostSummaries, resolveUserSshConfigHost } from '../ssh/ssh-config-host-picker' -import type { SshConnectionCallbacks } from '../ssh/ssh-connection' +import type { SshConnection, SshConnectionCallbacks } from '../ssh/ssh-connection' import { SshConnectionManager } from '../ssh/ssh-connection-manager' import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer' import { SshRelaySession, type SshRelayAiVaultHostInfo } from '../ssh/ssh-relay-session' @@ -293,6 +293,33 @@ async function abandonFailedSshSession(targetId: string, session: SshRelaySessio } } +// Why: a connect cancelled after its transport already opened still owns that transport and its +// unpublished session — nothing else will reach them, so it has to close them itself. Why only the +// connection it minted, and by identity: disconnecting by target id (or closing a transport it merely +// reused) would tear down the replacement connect's live transport. +async function abandonCancelledConnectAttempt( + targetId: string, + session: SshRelaySession, + mintedConnection: SshConnection | null +): Promise { + // Why the identity guard: every path that removes a session from activeSessions detaches it first, + // so re-detaching a superseded session would only clobber the replacement's lease state. + if (activeSessions.get(targetId) === session) { + await abandonFailedSshSession(targetId, session) + } + if (!mintedConnection) { + return + } + try { + await connectionManager!.disconnectConnection(targetId, mintedConnection) + } catch (error) { + // Why: the caller is about to throw the cancellation; a teardown throw must not replace it. + console.warn( + `[ssh] Failed to disconnect cancelled connect transport for ${targetId}: ${error instanceof Error ? error.message : String(error)}` + ) + } +} + function relayGracePeriodForTarget(target: SshTarget | null | undefined): number | undefined { return target?.relayGracePeriodSeconds } @@ -1163,6 +1190,12 @@ export function registerSshHandlers( const ownsSession = (): boolean => isCurrentConnectAttempt(targetId, authority) && activeSessions.get(targetId) === session + // Why captured here and not with existingState: connect() reuses an already-connected transport, + // and only a transport this attempt opened is this attempt's to close when it loses the race. + const priorConnection = connectionManager!.getConnection(targetId) + const mintedConnection = (): SshConnection | null => + conn && conn !== priorConnection ? conn : null + try { conn = await connectionManager!.connect(target) if (!ownsSession()) { @@ -1173,6 +1206,7 @@ export function registerSshHandlers( const errObj = err instanceof Error ? err : new Error(String(err)) const status: SshConnectionStatus = isAuthError(errObj) ? 'auth-failed' : 'error' if (!ownsSession()) { + await abandonCancelledConnectAttempt(targetId, session, mintedConnection()) throw createCancelledConnectAttemptError() } // Why: clear this failed connect's flag so a later non-prompting connect isn't deferred. @@ -1213,6 +1247,7 @@ export function registerSshHandlers( }) } catch (err) { if (!ownsSession()) { + await abandonCancelledConnectAttempt(targetId, session, mintedConnection()) throw createCancelledConnectAttemptError() } await abandonFailedSshSession(targetId, session) diff --git a/src/main/ssh/ssh-connection-manager.test.ts b/src/main/ssh/ssh-connection-manager.test.ts index 3a7a73cbcbf..b2de497a006 100644 --- a/src/main/ssh/ssh-connection-manager.test.ts +++ b/src/main/ssh/ssh-connection-manager.test.ts @@ -94,4 +94,36 @@ describe('SshConnectionManager', () => { expect(await manager.connect(target)).toBe(replacement) expect(manager.getConnection(target.id)).toBe(replacement) }) + + it('closes a late-resolved connection without evicting the replacement', async () => { + let resolveFirst!: () => void + mockState.connectResults.push( + new Promise((resolve) => { + resolveFirst = resolve + }), + Promise.resolve() + ) + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + + const firstConnect = manager.connect(target) + await manager.disconnect(target.id) + const replacement = await manager.connect(target) + resolveFirst() + const late = await firstConnect + + await manager.disconnectConnection(target.id, late) + + expect(mockState.instances[0].disconnect).toHaveBeenCalledTimes(2) + expect(manager.getConnection(target.id)).toBe(replacement) + }) + + it('clears the pool entry when the closed connection is still the registered one', async () => { + const manager = new SshConnectionManager({ onStateChange: vi.fn() }) + const conn = await manager.connect(target) + + await manager.disconnectConnection(target.id, conn) + + expect(mockState.instances[0].disconnect).toHaveBeenCalledOnce() + expect(manager.getConnection(target.id)).toBeUndefined() + }) }) diff --git a/src/main/ssh/ssh-connection-manager.ts b/src/main/ssh/ssh-connection-manager.ts index 8bc17d06b70..e553bf906b5 100644 --- a/src/main/ssh/ssh-connection-manager.ts +++ b/src/main/ssh/ssh-connection-manager.ts @@ -76,6 +76,18 @@ export class SshConnectionManager { } } + /** + * Close one specific connection, clearing the pool entry only when it is still the registered one. + * Why: a cancelled connect whose transport opened late owns that exact connection — disconnecting + * by target id would tear down the replacement's live transport instead. + */ + async disconnectConnection(targetId: string, conn: SshConnection): Promise { + await conn.disconnect() + if (this.connections.get(targetId) === conn) { + this.connections.delete(targetId) + } + } + async reconnect(targetId: string): Promise { const conn = this.connections.get(targetId) if (!conn) { diff --git a/src/relay/relay-pty-consumer-owner-displacement.test.ts b/src/relay/relay-pty-consumer-owner-displacement.test.ts index f8db3a37c8f..1c7bd076b20 100644 --- a/src/relay/relay-pty-consumer-owner-displacement.test.ts +++ b/src/relay/relay-pty-consumer-owner-displacement.test.ts @@ -300,7 +300,8 @@ describe('relay PTY consumer owner displacement', () => { const firstReconnectClientId = dispatcher.attachClient( (data, settle) => { firstReconnectWrites.push(Buffer.from(data)) - firstReconnectSettlement = settle + // Why ??=: the test settles the grant response, not whatever frame the dispatcher wrote last. + firstReconnectSettlement ??= settle return true }, { supportsWriteCallback: true }, diff --git a/src/shared/pty-consumer-session.test.ts b/src/shared/pty-consumer-session.test.ts index 8f315672bd9..d1999639e6f 100644 --- a/src/shared/pty-consumer-session.test.ts +++ b/src/shared/pty-consumer-session.test.ts @@ -274,6 +274,8 @@ describe('PtyConsumerSession', () => { ownerLease: first.grant.ownerLease }) expect(retry.displacedOwner?.connectionId).toBe('connection-2') + // Why: the committed replacement already retired the incumbent, so only connection-2 is left to displace. + expect(session.activeGrant('connection-1')).toBeNull() }) it('retries overlapping recovery against the incumbent after publication rolls back', () => { @@ -299,6 +301,8 @@ describe('PtyConsumerSession', () => { ownerLease: first.grant.ownerLease }) expect(retry.displacedOwner?.connectionId).toBe('connection-1') + // Why: the rolled-back replacement never published, so it holds no grant the retry could displace. + expect(session.activeGrant('connection-2')).toBeNull() }) it('types mismatched recovery without disturbing principal or lease ownership', () => {