diff --git a/src/main/providers/pty-provider-events.ts b/src/main/providers/pty-provider-events.ts index 6de435cb23f..0048b85f9dd 100644 --- a/src/main/providers/pty-provider-events.ts +++ b/src/main/providers/pty-provider-events.ts @@ -1,11 +1,18 @@ import type { TerminalGitHubPRLink } from '../../shared/terminal-github-pr-link-detector' +export type PtyDataUpstreamCredit = { + charCount: number + acknowledge(charCount: number): void +} + export type PtyDataEvent = { id: string data: string sequenceChars?: number transformed?: boolean seq?: number + /** Main-process-only credit captured from the exact provider generation that emitted this data. */ + upstreamCredit?: PtyDataUpstreamCredit } /** Notification-bearing fact a thinning transport detected while it held diff --git a/src/main/providers/ssh-agent-session-process-list.ts b/src/main/providers/ssh-agent-session-process-list.ts index cb1447aff71..f5287302ca5 100644 --- a/src/main/providers/ssh-agent-session-process-list.ts +++ b/src/main/providers/ssh-agent-session-process-list.ts @@ -1,25 +1,156 @@ import { isAgentSessionOwnerBinding } from '../../shared/agent-session-host-authority' +import { MAX_CLAIMED_AGENT_PTY_OWNER_ENTRIES } from '../../shared/claimed-agent-pty-owner' import { isPtyIncarnationId } from '../../shared/pty-incarnation' import type { PtyProcessInfo } from './types' +import { MAX_SSH_PTY_LIVE_ROSTER_ENTRIES } from './ssh-pty-live-roster' +import { admittedSshRelayPtyIdBytes } from './ssh-pty-wire-admission' + +export const MAX_SSH_PTY_PROCESS_LIST_ENTRIES = MAX_SSH_PTY_LIVE_ROSTER_ENTRIES +export const MAX_SSH_PTY_PROCESS_OWNERS_PER_ENTRY = 256 +export const MAX_SSH_PTY_PROCESS_LIST_OWNERS = MAX_CLAIMED_AGENT_PTY_OWNER_ENTRIES +export const MAX_SSH_PTY_PROCESS_LIST_BYTES = 8 * 1024 * 1024 +export const MAX_SSH_PTY_PROCESS_CWD_BYTES = 128 * 1024 +export const MAX_SSH_PTY_PROCESS_TITLE_BYTES = 16 * 1024 +const MAX_SSH_PTY_PROCESS_WORKTREE_ID_BYTES = 128 * 1024 +const MAX_SSH_PTY_PROCESS_TERMINAL_HANDLE_BYTES = 1024 +const MAX_SSH_PTY_PROCESS_WSL_DISTRO_BYTES = 4 * 1024 + +function stringBytes(value: unknown, maxBytes: number): number | null { + if (typeof value !== 'string') { + return null + } + const bytes = Buffer.byteLength(value, 'utf8') + return bytes <= maxBytes ? bytes : null +} + +function optionalStringBytes(value: unknown, maxBytes: number): number | null { + return value === undefined ? 0 : stringBytes(value, maxBytes) +} + +function ownerBytes(owner: unknown): number | null { + if (!isAgentSessionOwnerBinding(owner)) { + return null + } + const strings = [ + owner.claim.keyId, + owner.claim.identityDigest, + owner.claim.worktreeScopeDigest, + owner.claim.agent, + owner.generation, + owner.ptyId, + owner.surface.worktreeId, + owner.surface.tabId, + owner.surface.leafId, + owner.surface.terminalHandle + ] + return strings.reduce((total, value) => total + Buffer.byteLength(value, 'utf8'), 0) +} + +function assertProcessList( + sessions: unknown +): asserts sessions is (PtyProcessInfo & Record)[] { + if (!Array.isArray(sessions) || sessions.length > MAX_SSH_PTY_PROCESS_LIST_ENTRIES) { + throw new Error('invalid_ssh_pty_process_list') + } + let aggregateBytes = 0 + let aggregateOwners = 0 + for (const session of sessions) { + if (typeof session !== 'object' || session === null) { + throw new Error('invalid_ssh_pty_process_list') + } + const idBytes = admittedSshRelayPtyIdBytes(session.id) + const cwdBytes = stringBytes(session.cwd, MAX_SSH_PTY_PROCESS_CWD_BYTES) + const titleBytes = stringBytes(session.title, MAX_SSH_PTY_PROCESS_TITLE_BYTES) + const worktreeBytes = optionalStringBytes( + session.worktreeId, + MAX_SSH_PTY_PROCESS_WORKTREE_ID_BYTES + ) + const terminalHandleBytes = optionalStringBytes( + session.terminalHandle, + MAX_SSH_PTY_PROCESS_TERMINAL_HANDLE_BYTES + ) + const wslDistroBytes = + session.wslDistro === null + ? 0 + : optionalStringBytes(session.wslDistro, MAX_SSH_PTY_PROCESS_WSL_DISTRO_BYTES) + if ( + idBytes === null || + cwdBytes === null || + titleBytes === null || + worktreeBytes === null || + terminalHandleBytes === null || + wslDistroBytes === null || + (session.incarnationId !== undefined && !isPtyIncarnationId(session.incarnationId)) + ) { + throw new Error('invalid_ssh_pty_process_list') + } + aggregateBytes += + idBytes + cwdBytes + titleBytes + worktreeBytes + terminalHandleBytes + wslDistroBytes + if ( + session.agentSessionOwners !== undefined && + (!Array.isArray(session.agentSessionOwners) || + session.agentSessionOwners.length > MAX_SSH_PTY_PROCESS_OWNERS_PER_ENTRY) + ) { + throw new Error('invalid_ssh_pty_process_list') + } + for (const owner of session.agentSessionOwners ?? []) { + const bytes = ownerBytes(owner) + if (bytes === null || owner.ptyId !== session.id) { + throw new Error('agent_session_ownership_unknown') + } + aggregateBytes += bytes + } + aggregateOwners += session.agentSessionOwners?.length ?? 0 + if ( + aggregateBytes > MAX_SSH_PTY_PROCESS_LIST_BYTES || + aggregateOwners > MAX_SSH_PTY_PROCESS_LIST_OWNERS + ) { + throw new Error('invalid_ssh_pty_process_list') + } + } +} export function mapSshPtyProcessList( - sessions: PtyProcessInfo[], + sessions: unknown, toAppPtyId: (id: string) => string ): PtyProcessInfo[] { + assertProcessList(sessions) return sessions.map((session) => { if (session.agentSessionOwners?.length && !isPtyIncarnationId(session.incarnationId)) { throw new Error('agent_session_ownership_unknown') } return { - ...session, id: toAppPtyId(session.id), - ...(session.agentSessionOwners + cwd: session.cwd, + title: session.title, + ...(session.incarnationId ? { incarnationId: session.incarnationId } : {}), + ...(session.worktreeId !== undefined ? { worktreeId: session.worktreeId } : {}), + ...(session.terminalHandle !== undefined ? { terminalHandle: session.terminalHandle } : {}), + ...(session.wslDistro !== undefined ? { wslDistro: session.wslDistro } : {}), + ...(session.agentSessionOwners !== undefined ? { agentSessionOwners: session.agentSessionOwners.map((owner) => { if (!isAgentSessionOwnerBinding(owner) || owner.ptyId !== session.id) { throw new Error('agent_session_ownership_unknown') } - return { ...owner, ptyId: toAppPtyId(owner.ptyId) } + return { + claim: { + digestVersion: owner.claim.digestVersion, + keyId: owner.claim.keyId, + identityDigest: owner.claim.identityDigest, + worktreeScopeDigest: owner.claim.worktreeScopeDigest, + agent: owner.claim.agent + }, + generation: owner.generation, + phase: owner.phase, + ptyId: toAppPtyId(owner.ptyId), + surface: { + worktreeId: owner.surface.worktreeId, + tabId: owner.surface.tabId, + leafId: owner.surface.leafId, + terminalHandle: owner.surface.terminalHandle + } + } }) } : {}) diff --git a/src/main/providers/ssh-pty-live-roster.test.ts b/src/main/providers/ssh-pty-live-roster.test.ts new file mode 100644 index 00000000000..aac1d5adcf4 --- /dev/null +++ b/src/main/providers/ssh-pty-live-roster.test.ts @@ -0,0 +1,134 @@ +import { describe, expect, it } from 'vitest' +import { + MAX_SSH_PTY_LIVE_ROSTER_ENTRIES, + MAX_SSH_PTY_LIVE_ROSTER_ID_BYTES, + SshPtyLiveRoster +} from './ssh-pty-live-roster' +import { MAX_SSH_APP_PTY_ID_BYTES } from './ssh-pty-wire-admission' + +describe('SshPtyLiveRoster', () => { + it('preserves ordinary listing and exit behavior', () => { + const roster = new SshPtyLiveRoster() + const listing = roster.beginListing() + + roster.reconcileListing(listing, ['pty-1']) + expect(roster.has('pty-1')).toBe(true) + + roster.recordExit('pty-1') + expect(roster.has('pty-1')).toBe(false) + }) + + it('gives a locally spawned PTY one stale-list grace', () => { + const roster = new SshPtyLiveRoster() + roster.recordSpawn('pty-new') + + roster.reconcileListing(roster.beginListing(), []) + expect(roster.has('pty-new')).toBe(true) + + roster.reconcileListing(roster.beginListing(), []) + expect(roster.has('pty-new')).toBe(false) + }) + + it('keeps a PTY observed while a listing is in flight', () => { + const roster = new SshPtyLiveRoster() + const listing = roster.beginListing() + + roster.recordNotification('pty-new') + roster.reconcileListing(listing, []) + + expect(roster.has('pty-new')).toBe(true) + }) + + it('ignores an older listing after a newer listing reconciles', () => { + const roster = new SshPtyLiveRoster() + roster.recordNotification('pty-live') + const older = roster.beginListing() + const newer = roster.beginListing() + + roster.reconcileListing(newer, ['pty-live']) + roster.reconcileListing(older, []) + + expect(roster.has('pty-live')).toBe(true) + }) + + it('does not resurrect an exit with an in-flight stale listing', () => { + const roster = new SshPtyLiveRoster() + roster.recordNotification('pty-exited') + const listing = roster.beginListing() + + roster.recordExit('pty-exited') + roster.reconcileListing(listing, ['pty-exited']) + + expect(roster.has('pty-exited')).toBe(false) + }) + + it('caps missed-exit unique-id churn and evicts the oldest live row', () => { + const roster = new SshPtyLiveRoster() + const total = MAX_SSH_PTY_LIVE_ROSTER_ENTRIES + 50 + + for (let index = 0; index < total; index += 1) { + roster.recordNotification(`pty-${index}`) + } + + expect(roster.has('pty-0')).toBe(false) + expect(roster.has(`pty-${total - MAX_SSH_PTY_LIVE_ROSTER_ENTRIES}`)).toBe(true) + expect(roster.has(`pty-${total - 1}`)).toBe(true) + }) + + it('does not let exit fences evict supported live entries', () => { + const roster = new SshPtyLiveRoster() + for (let index = 0; index < 50; index += 1) { + roster.recordNotification(`live-${index}`) + } + for (let index = 0; index < MAX_SSH_PTY_LIVE_ROSTER_ENTRIES * 2; index += 1) { + roster.recordExit(`gone-${index}`) + } + + for (let index = 0; index < 50; index += 1) { + expect(roster.has(`live-${index}`)).toBe(true) + } + }) + + it('rejects oversized ids and caps aggregate retained id bytes', () => { + const roster = new SshPtyLiveRoster() + const oversized = 'x'.repeat(MAX_SSH_APP_PTY_ID_BYTES + 1) + roster.recordNotification(oversized) + expect(roster.has(oversized)).toBe(false) + + const payload = 'x'.repeat(16 * 1024 - 32) + const total = Math.floor(MAX_SSH_PTY_LIVE_ROSTER_ID_BYTES / Buffer.byteLength(payload)) + 2 + for (let index = 0; index < total; index += 1) { + roster.recordNotification(`${index}:${payload}`) + } + + expect(roster.has(`0:${payload}`)).toBe(false) + expect(roster.has(`${total - 1}:${payload}`)).toBe(true) + }) + + it('ignores a listing response that completes after clear', () => { + const roster = new SshPtyLiveRoster() + const listing = roster.beginListing() + + roster.clear() + roster.reconcileListing(listing, ['pty-stale']) + + expect(roster.has('pty-stale')).toBe(false) + }) + + it('does not resurrect an exit after its tombstone is evicted', () => { + const roster = new SshPtyLiveRoster() + roster.recordSpawn('pty-reused') + const stale = roster.beginListing() + roster.recordExit('pty-reused') + for (let index = 0; index < MAX_SSH_PTY_LIVE_ROSTER_ENTRIES; index += 1) { + roster.recordExit(`gone-${index}`) + } + + roster.reconcileListing(stale, ['pty-reused']) + expect(roster.has('pty-reused')).toBe(false) + + const fresh = roster.beginListing() + roster.reconcileListing(fresh, ['pty-reused']) + expect(roster.has('pty-reused')).toBe(true) + }) +}) diff --git a/src/main/providers/ssh-pty-live-roster.ts b/src/main/providers/ssh-pty-live-roster.ts new file mode 100644 index 00000000000..cafbdb53e5e --- /dev/null +++ b/src/main/providers/ssh-pty-live-roster.ts @@ -0,0 +1,155 @@ +import { admittedSshAppPtyIdBytes } from './ssh-pty-wire-admission' + +export const MAX_SSH_PTY_LIVE_ROSTER_ENTRIES = 256 +export const MAX_SSH_PTY_LIVE_ROSTER_ID_BYTES = 2 * 1024 * 1024 + +export type SshPtyListingToken = { + sequence: number + observationRevision: number +} + +type RosterEntry = { + idBytes: number + live: boolean + missingListingGrace: boolean + observationRevision: number +} + +export class SshPtyLiveRoster { + private readonly entries = new Map() + private observationRevision = 0 + private listingSequence = 0 + private latestReconciledListing = 0 + private retainedIdBytes = 0 + private staleListingFenceRevision = 0 + + beginListing(): SshPtyListingToken { + return { + sequence: ++this.listingSequence, + observationRevision: this.observationRevision + } + } + + recordSpawn(id: string): void { + this.record(id, true) + } + + recordNotification(id: string): void { + this.record(id, false) + } + + recordExit(id: string): void { + const idBytes = admittedSshAppPtyIdBytes(id) + if (idBytes === null) { + return + } + this.setEntry(id, { + idBytes, + live: false, + missingListingGrace: false, + observationRevision: ++this.observationRevision + }) + } + + reconcileListing(token: SshPtyListingToken, listedIds: readonly string[]): void { + if (token.sequence <= this.latestReconciledListing) { + return + } + this.latestReconciledListing = token.sequence + if (token.observationRevision < this.staleListingFenceRevision) { + return + } + const listed = new Set() + for (const id of listedIds) { + if (listed.size < MAX_SSH_PTY_LIVE_ROSTER_ENTRIES && admittedSshAppPtyIdBytes(id) !== null) { + listed.add(id) + } + } + + for (const [id, entry] of this.entries) { + if (!entry.live || listed.has(id) || entry.observationRevision > token.observationRevision) { + continue + } + if (entry.missingListingGrace) { + entry.missingListingGrace = false + entry.observationRevision = ++this.observationRevision + continue + } + this.removeEntry(id) + } + + for (const id of listed) { + const current = this.entries.get(id) + // Why: an exit newer than this listing request must fence its stale response. + if (current && !current.live && current.observationRevision > token.observationRevision) { + continue + } + this.record(id, false) + } + } + + has(id: string): boolean { + return this.entries.get(id)?.live === true + } + + clear(): void { + this.entries.clear() + this.retainedIdBytes = 0 + // Why: responses from listings already in flight must not repopulate a disposed provider. + this.latestReconciledListing = this.listingSequence + } + + private record(id: string, missingListingGrace: boolean): void { + const idBytes = admittedSshAppPtyIdBytes(id) + if (idBytes === null) { + return + } + this.setEntry(id, { + idBytes, + live: true, + missingListingGrace, + observationRevision: ++this.observationRevision + }) + } + + private setEntry(id: string, entry: RosterEntry): void { + this.removeEntry(id) + this.entries.set(id, entry) + this.retainedIdBytes += entry.idBytes + this.capEntries() + } + + private capEntries(): void { + while ( + this.entries.size > MAX_SSH_PTY_LIVE_ROSTER_ENTRIES || + this.retainedIdBytes > MAX_SSH_PTY_LIVE_ROSTER_ID_BYTES + ) { + let inactive: string | undefined + for (const [id, entry] of this.entries) { + if (!entry.live) { + inactive = id + break + } + } + const oldest = this.entries.keys().next().value as string | undefined + const evictedId = inactive ?? oldest ?? '' + const evicted = this.entries.get(evictedId) + if (evicted && !evicted.live) { + this.staleListingFenceRevision = Math.max( + this.staleListingFenceRevision, + evicted.observationRevision + ) + } + this.removeEntry(evictedId) + } + } + + private removeEntry(id: string): void { + const entry = this.entries.get(id) + if (!entry) { + return + } + this.retainedIdBytes -= entry.idBytes + this.entries.delete(id) + } +} diff --git a/src/main/providers/ssh-pty-output-delivery-credit.test.ts b/src/main/providers/ssh-pty-output-delivery-credit.test.ts new file mode 100644 index 00000000000..e3df99902c9 --- /dev/null +++ b/src/main/providers/ssh-pty-output-delivery-credit.test.ts @@ -0,0 +1,281 @@ +import { describe, expect, it, vi } from 'vitest' +import { + MAX_SSH_PTY_CREDIT_SEGMENTS, + MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE, + MAX_SSH_PTY_CREDIT_STATE_ID_BYTES, + MAX_SSH_PTY_CREDIT_STATES, + SshPtyOutputDeliveryCredit +} from './ssh-pty-output-delivery-credit' +import { + MAX_SSH_PTY_DELIVERY_TOKEN_BYTES, + MAX_SSH_RELAY_PTY_ID_BYTES +} from './ssh-pty-wire-admission' + +function retainedStates(credit: SshPtyOutputDeliveryCredit): Map { + return Reflect.get(credit, 'states') as Map +} + +describe('SshPtyOutputDeliveryCredit', () => { + it('defers tokenized credit until the captured owner acknowledges renderer progress', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + let upstreamCredit: { charCount: number; acknowledge(chars: number): void } | undefined + + credit.ingest( + { + id: 'pty-1', + data: 'output', + rawLength: 10, + deliveryToken: 'delivery-1' + }, + (output) => { + upstreamCredit = output.upstreamCredit + } + ) + + expect(notifyAck).not.toHaveBeenCalled() + upstreamCredit?.acknowledge(4) + upstreamCredit?.acknowledge(20) + + expect(notifyAck).toHaveBeenNthCalledWith(1, { + id: 'pty-1', + charCount: 4, + deliveryToken: 'delivery-1' + }) + expect(notifyAck).toHaveBeenNthCalledWith(2, { + id: 'pty-1', + charCount: 6, + deliveryToken: 'delivery-1' + }) + }) + + it('acknowledges a later dropped segment without consuming earlier visible credit', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + let first: { acknowledge(chars: number): void } | undefined + let second: { acknowledge(chars: number): void } | undefined + credit.ingest({ id: 'pty-1', data: 'first', deliveryToken: 'delivery-1' }, (output) => { + first = output.upstreamCredit + }) + credit.ingest({ id: 'pty-1', data: 'second', deliveryToken: 'delivery-1' }, (output) => { + second = output.upstreamCredit + }) + + second?.acknowledge(6) + first?.acknowledge(5) + + expect(notifyAck.mock.calls).toEqual([ + [{ id: 'pty-1', charCount: 6, deliveryToken: 'delivery-1' }], + [{ id: 'pty-1', charCount: 5, deliveryToken: 'delivery-1' }] + ]) + }) + + it('walks token segments in order for legacy provider acknowledgement calls', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + credit.ingest({ id: 'pty-1', data: 'first', deliveryToken: 'token-1' }, () => {}) + credit.ingest({ id: 'pty-1', data: 'second', deliveryToken: 'token-1' }, () => {}) + + credit.acknowledgeLegacy('pty-1', 11) + + expect(notifyAck.mock.calls).toEqual([ + [{ id: 'pty-1', charCount: 5, deliveryToken: 'token-1' }], + [{ id: 'pty-1', charCount: 6, deliveryToken: 'token-1' }] + ]) + }) + + it('drops stale credit when attach rotates the delivery token', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + let stale: { acknowledge(chars: number): void } | undefined + let current: { acknowledge(chars: number): void } | undefined + credit.ingest({ id: 'pty-1', data: 'stale', deliveryToken: 'before-attach' }, (output) => { + stale = output.upstreamCredit + }) + credit.ingest({ id: 'pty-1', data: 'current', deliveryToken: 'after-attach' }, (output) => { + current = output.upstreamCredit + }) + + current?.acknowledge(7) + stale?.acknowledge(5) + credit.ingest({ id: 'pty-1', data: 'next', deliveryToken: 'after-attach' }, () => {}) + credit.acknowledgeLegacy('pty-1', 4) + + expect(notifyAck.mock.calls).toEqual([ + [{ id: 'pty-1', charCount: 7, deliveryToken: 'after-attach' }], + [{ id: 'pty-1', charCount: 4, deliveryToken: 'after-attach' }] + ]) + }) + + it('preserves legacy no-token acknowledgements', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + credit.ingest({ id: 'pty-1', data: 'legacy' }, () => {}) + + credit.acknowledgeLegacy('pty-1', 6) + + expect(notifyAck).toHaveBeenCalledWith({ id: 'pty-1', charCount: 6 }) + }) + + it('falls back to delivered length when raw length is invalid', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + let upstreamCredit: { acknowledge(chars: number): void } | undefined + credit.ingest( + { + id: 'pty-1', + data: 'output', + rawLength: Number.POSITIVE_INFINITY, + deliveryToken: 'delivery-2' + }, + (output) => { + upstreamCredit = output.upstreamCredit + } + ) + + upstreamCredit?.acknowledge(100) + + expect(notifyAck).toHaveBeenCalledWith({ + id: 'pty-1', + charCount: 6, + deliveryToken: 'delivery-2' + }) + }) + + it('makes captured credit inert after exit or disposal', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + let exited: { acknowledge(chars: number): void } | undefined + let disposed: { acknowledge(chars: number): void } | undefined + credit.ingest({ id: 'pty-1', data: 'exit', deliveryToken: 'delivery-3' }, (output) => { + exited = output.upstreamCredit + }) + credit.release('pty-1') + exited?.acknowledge(4) + credit.ingest({ id: 'pty-2', data: 'dispose', deliveryToken: 'delivery-4' }, (output) => { + disposed = output.upstreamCredit + }) + + credit.dispose() + disposed?.acknowledge(7) + + expect(notifyAck).not.toHaveBeenCalled() + }) + + it('returns delivery credit when listener ingestion throws', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + + expect(() => + credit.ingest({ id: 'pty-1', data: 'output', deliveryToken: 'delivery-5' }, () => { + throw new Error('listener failed') + }) + ).toThrow('listener failed') + expect(notifyAck).toHaveBeenCalledWith({ + id: 'pty-1', + charCount: 6, + deliveryToken: 'delivery-5' + }) + }) + + it('rejects oversized ids and delivery tokens before retention', () => { + const deliver = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(vi.fn()) + + credit.ingest( + { + id: 'x'.repeat(MAX_SSH_RELAY_PTY_ID_BYTES + 1), + data: 'output', + deliveryToken: 'token' + }, + deliver + ) + credit.ingest( + { + id: 'pty-1', + data: 'output', + deliveryToken: 'x'.repeat(MAX_SSH_PTY_DELIVERY_TOKEN_BYTES + 1) + }, + deliver + ) + + expect(deliver).not.toHaveBeenCalled() + expect(retainedStates(credit)).toHaveLength(0) + }) + + it('caps unique-id credit state churn and returns evicted credit', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + for (let index = 0; index < MAX_SSH_PTY_CREDIT_STATES + 50; index += 1) { + credit.ingest({ id: `pty-${index}`, data: 'x', deliveryToken: `token-${index}` }, () => {}) + } + + const states = retainedStates(credit) + expect(states.size).toBe(MAX_SSH_PTY_CREDIT_STATES) + expect(states.has('pty-0')).toBe(false) + expect(states.has(`pty-${MAX_SSH_PTY_CREDIT_STATES + 49}`)).toBe(true) + expect(notifyAck).toHaveBeenCalledWith({ + id: 'pty-0', + charCount: 1, + deliveryToken: 'token-0' + }) + }) + + it('caps aggregate credit-state id bytes', () => { + const credit = new SshPtyOutputDeliveryCredit(vi.fn()) + const suffix = 'x'.repeat(MAX_SSH_RELAY_PTY_ID_BYTES - 16) + for (let index = 0; index < MAX_SSH_PTY_CREDIT_STATES; index += 1) { + credit.ingest( + { id: `pty-${index}:${suffix}`, data: 'x', deliveryToken: `token-${index}` }, + () => {} + ) + } + + const retainedBytes = [...retainedStates(credit).keys()].reduce( + (total, id) => total + Buffer.byteLength(id), + 0 + ) + expect(retainedBytes).toBeLessThanOrEqual(MAX_SSH_PTY_CREDIT_STATE_ID_BYTES) + expect(retainedStates(credit).size).toBeLessThan(MAX_SSH_PTY_CREDIT_STATES) + }) + + it('caps retained segments per state and in aggregate', () => { + const credit = new SshPtyOutputDeliveryCredit(vi.fn()) + const perStateTotal = MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE + 50 + for (let index = 0; index < perStateTotal; index += 1) { + credit.ingest({ id: 'pty-one', data: 'x', deliveryToken: 'token-one' }, () => {}) + } + expect(retainedStates(credit).get('pty-one')?.segments).toHaveLength( + MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE + ) + + for (let state = 0; state < 8; state += 1) { + for (let segment = 0; segment < MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE; segment += 1) { + credit.ingest({ id: `pty-${state}`, data: 'x', deliveryToken: `token-${state}` }, () => {}) + } + } + + const retained = [...retainedStates(credit).values()].reduce( + (total, state) => total + state.segments.length, + 0 + ) + expect(retained).toBe(MAX_SSH_PTY_CREDIT_SEGMENTS) + }) + + it('returns unretained direct credit when state churn evicts its owner', () => { + const notifyAck = vi.fn() + const credit = new SshPtyOutputDeliveryCredit(notifyAck) + for (let index = 0; index <= MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE; index += 1) { + credit.ingest({ id: 'pty-direct', data: 'x', deliveryToken: 'token-direct' }, () => {}) + } + for (let index = 0; index < MAX_SSH_PTY_CREDIT_STATES; index += 1) { + credit.ingest({ id: `pty-${index}`, data: 'x', deliveryToken: `token-${index}` }, () => {}) + } + + expect(notifyAck).toHaveBeenCalledWith({ + id: 'pty-direct', + charCount: MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE + 1, + deliveryToken: 'token-direct' + }) + }) +}) diff --git a/src/main/providers/ssh-pty-output-delivery-credit.ts b/src/main/providers/ssh-pty-output-delivery-credit.ts new file mode 100644 index 00000000000..1c0c40ad007 --- /dev/null +++ b/src/main/providers/ssh-pty-output-delivery-credit.ts @@ -0,0 +1,280 @@ +import type { PtyDataUpstreamCredit } from './pty-provider-events' +import { + admittedSshRelayPtyIdBytes, + isAdmittedSshDeliveryToken, + isAdmittedSshRelayPtyId +} from './ssh-pty-wire-admission' + +export const MAX_SSH_PTY_CREDIT_STATES = 256 +export const MAX_SSH_PTY_CREDIT_STATE_ID_BYTES = 512 * 1024 +export const MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE = 4096 +export const MAX_SSH_PTY_CREDIT_SEGMENTS = 16_384 + +export type SshPtyOutputDelivery = { + relayId: string + data: string + rawLength?: number + transformed: boolean + seq?: number + upstreamCredit?: PtyDataUpstreamCredit +} + +type AckParams = { + id: string + charCount: number + deliveryToken?: string +} + +export class SshPtyOutputDeliveryCredit { + private readonly states = new Map() + private disposed = false + private retainedStateIdBytes = 0 + private retainedSegments = 0 + + constructor(private readonly notifyAck: (params: AckParams) => void) {} + + ingest(params: Record, deliver: (output: SshPtyOutputDelivery) => void): void { + const relayId = params.id + const data = params.data + if (!isAdmittedSshRelayPtyId(relayId) || typeof data !== 'string') { + return + } + if (params.deliveryToken !== undefined && !isAdmittedSshDeliveryToken(params.deliveryToken)) { + return + } + const deliveryToken = params.deliveryToken as string | undefined + const rawLength = + typeof params.rawLength === 'number' && + Number.isSafeInteger(params.rawLength) && + params.rawLength >= 0 + ? params.rawLength + : undefined + const seq = + typeof params.seq === 'number' && Number.isSafeInteger(params.seq) && params.seq >= 0 + ? params.seq + : undefined + const charCount = rawLength ?? data.length + const upstreamCredit = this.createCredit(relayId, charCount, deliveryToken) + try { + deliver({ + relayId, + data, + ...(rawLength === undefined ? {} : { rawLength }), + transformed: params.transformed === true, + ...(seq === undefined ? {} : { seq }), + ...(upstreamCredit ? { upstreamCredit } : {}) + }) + } catch (error) { + upstreamCredit?.acknowledge(charCount) + throw error + } + } + + acknowledgeLegacy(relayId: string, charCount: number): void { + if ( + !isAdmittedSshRelayPtyId(relayId) || + !Number.isFinite(charCount) || + charCount <= 0 || + this.disposed + ) { + return + } + const state = this.states.get(relayId) + if (!state) { + this.notifyAck({ id: relayId, charCount: Math.floor(charCount) }) + return + } + let remaining = Math.floor(charCount) + while (remaining > 0 && state.segments.length > 0) { + const segment = state.segments[0] + const acknowledged = Math.min(remaining, segment.remainingChars) + this.acknowledgeSegment(state, segment, acknowledged) + remaining -= acknowledged + } + if (remaining > 0 && !state.tokenized) { + this.notifyAck({ id: relayId, charCount: remaining }) + } + } + + release(relayId: string): void { + const state = this.states.get(relayId) + if (!state) { + return + } + this.removeState(state, false) + } + + dispose(): void { + this.disposed = true + for (const state of this.states.values()) { + state.active = false + this.clearSegments(state) + } + this.states.clear() + this.retainedStateIdBytes = 0 + } + + private createCredit( + relayId: string, + charCount: number, + deliveryToken: string | undefined + ): PtyDataUpstreamCredit | undefined { + if (this.disposed || charCount <= 0 || !deliveryToken) { + return undefined + } + let existing = this.states.get(relayId) + // Why: pty.attach rotates the relay token; stale renderer credit from the prior + // generation must not pin every later acknowledged segment behind it. + if (existing && existing.deliveryToken !== deliveryToken) { + this.removeState(existing, false) + existing = undefined + } + const state = + existing?.active === true + ? existing + : { + relayId, + idBytes: admittedSshRelayPtyIdBytes(relayId) ?? 0, + active: true, + tokenized: true, + deliveryToken, + segments: [], + directOutstandingChars: 0 + } + if (!existing) { + this.retainedStateIdBytes += state.idBytes + } else { + this.states.delete(relayId) + } + this.states.set(relayId, state) + this.capStates() + if ( + state.segments.length >= MAX_SSH_PTY_CREDIT_SEGMENTS_PER_STATE || + this.retainedSegments >= MAX_SSH_PTY_CREDIT_SEGMENTS + ) { + return this.createDirectCredit(state, charCount) + } + const segment: PtyCreditSegment = { + deliveryToken, + remainingChars: charCount + } + state.segments.push(segment) + this.retainedSegments++ + return { + charCount, + acknowledge: (requestedChars) => { + if (!Number.isFinite(requestedChars) || requestedChars <= 0) { + return + } + this.acknowledgeSegment(state, segment, Math.floor(requestedChars)) + } + } + } + + private acknowledgeSegment( + state: PtyCreditState, + segment: PtyCreditSegment, + requestedChars: number + ): void { + if (this.disposed || !state.active || requestedChars <= 0 || segment.remainingChars <= 0) { + return + } + const charCount = Math.min(requestedChars, segment.remainingChars) + segment.remainingChars -= charCount + this.notifyAck({ + id: state.relayId, + charCount, + ...(segment.deliveryToken ? { deliveryToken: segment.deliveryToken } : {}) + }) + while (state.segments[0]?.remainingChars === 0) { + state.segments.shift() + this.retainedSegments-- + } + } + + private createDirectCredit(state: PtyCreditState, charCount: number): PtyDataUpstreamCredit { + let remainingChars = charCount + state.directOutstandingChars += charCount + return { + charCount, + acknowledge: (requestedChars) => { + if ( + this.disposed || + !state.active || + state.deliveryToken.length === 0 || + !Number.isFinite(requestedChars) || + requestedChars <= 0 || + remainingChars <= 0 + ) { + return + } + const acknowledged = Math.min(Math.floor(requestedChars), remainingChars) + remainingChars -= acknowledged + state.directOutstandingChars -= acknowledged + this.notifyAck({ + id: state.relayId, + charCount: acknowledged, + deliveryToken: state.deliveryToken + }) + } + } + } + + private capStates(): void { + while ( + this.states.size > MAX_SSH_PTY_CREDIT_STATES || + this.retainedStateIdBytes > MAX_SSH_PTY_CREDIT_STATE_ID_BYTES + ) { + const oldest = this.states.values().next().value as PtyCreditState | undefined + if (!oldest) { + return + } + this.removeState(oldest, true) + } + } + + private removeState(state: PtyCreditState, returnOutstanding: boolean): void { + state.active = false + if (returnOutstanding) { + const outstandingChars = + state.directOutstandingChars + + state.segments.reduce((total, segment) => total + segment.remainingChars, 0) + if (outstandingChars > 0) { + try { + this.notifyAck({ + id: state.relayId, + charCount: outstandingChars, + deliveryToken: state.deliveryToken + }) + } catch { + // Best effort: eviction must still enforce the local memory bound. + } + } + } + state.directOutstandingChars = 0 + this.clearSegments(state) + if (this.states.delete(state.relayId)) { + this.retainedStateIdBytes -= state.idBytes + } + } + + private clearSegments(state: PtyCreditState): void { + this.retainedSegments -= state.segments.length + state.segments.length = 0 + } +} + +type PtyCreditSegment = { + deliveryToken?: string + remainingChars: number +} + +type PtyCreditState = { + relayId: string + idBytes: number + active: boolean + tokenized: boolean + deliveryToken: string + segments: PtyCreditSegment[] + directOutstandingChars: number +} diff --git a/src/main/providers/ssh-pty-provider-contract.ts b/src/main/providers/ssh-pty-provider-contract.ts index 9043133b03c..34d46a9148d 100644 --- a/src/main/providers/ssh-pty-provider-contract.ts +++ b/src/main/providers/ssh-pty-provider-contract.ts @@ -1,4 +1,5 @@ import type { PtyIncarnationId } from '../../shared/pty-incarnation' +import type { PtyDataEvent } from './pty-provider-events' export type RemoteCliBridgeEnv = { binDir: string @@ -8,13 +9,7 @@ export type RemoteCliBridgeEnv = { pathDelimiter?: ':' | ';' } -export type SshPtyDataCallback = (payload: { - id: string - data: string - sequenceChars?: number - transformed?: boolean - seq?: number -}) => void +export type SshPtyDataCallback = (payload: PtyDataEvent) => void export type SshPtyReplayCallback = (payload: { id: string; data: string }) => void export type SshPtyExitCallback = (payload: { id: string diff --git a/src/main/providers/ssh-pty-provider-exit-race.test.ts b/src/main/providers/ssh-pty-provider-exit-race.test.ts index 12616e4d133..bca158029eb 100644 --- a/src/main/providers/ssh-pty-provider-exit-race.test.ts +++ b/src/main/providers/ssh-pty-provider-exit-race.test.ts @@ -76,3 +76,34 @@ it('rejects an SSH reattach whose matching exit shares the attach reply batch', isReattach: true }) }) + +it('does not treat a malformed exit incarnation as a wildcard spawn race', async () => { + const mux = { + request: vi.fn(), + notify: vi.fn(), + onNotification: vi.fn(), + dispose: vi.fn(), + isDisposed: vi.fn().mockReturnValue(false) + } + const provider = new SshPtyProvider('conn-1', mux as never) + const exitListener = vi.fn() + provider.onExit(exitListener) + mux.request.mockImplementation(async (method: string) => { + if (method === 'pty.spawn') { + const notify = mux.onNotification.mock.calls[0]?.[0] + notify?.('pty.exit', { + id: 'pty-raced', + code: 0, + incarnationId: 'x'.repeat(129) + }) + return { id: 'pty-raced', incarnationId: 'incarnation-current' } + } + return undefined + }) + + await expect(provider.spawn({ cols: 80, rows: 24 })).resolves.toMatchObject({ + id: 'ssh:conn-1@@pty-raced', + incarnationId: 'incarnation-current' + }) + expect(exitListener).not.toHaveBeenCalled() +}) diff --git a/src/main/providers/ssh-pty-provider-notifications.test.ts b/src/main/providers/ssh-pty-provider-notifications.test.ts new file mode 100644 index 00000000000..3ccdd735d5a --- /dev/null +++ b/src/main/providers/ssh-pty-provider-notifications.test.ts @@ -0,0 +1,148 @@ +import { describe, expect, it, vi } from 'vitest' +import { SshPtyProviderNotifications } from './ssh-pty-provider-notifications' +import { MAX_SSH_RELAY_PTY_ID_BYTES } from './ssh-pty-wire-admission' + +type NotificationHandler = (method: string, params: Record) => void + +function createSubscription() { + const notify = vi.fn() + const onNotification = vi.fn() + const toAppPtyId = vi.fn((id: string) => `ssh:conn@@${id}`) + const recordNotification = vi.fn() + const recordExit = vi.fn() + + const notifications = new SshPtyProviderNotifications( + { onNotification, notify } as never, + toAppPtyId, + recordNotification, + recordExit + ) + + const handler = onNotification.mock.calls[0]?.[0] as NotificationHandler | undefined + if (!handler) { + throw new Error('notification handler was not registered') + } + + return { handler, toAppPtyId, recordNotification, recordExit, notifications } +} + +describe('SshPtyProviderNotifications', () => { + it('ignores non-PTY notifications without mapping params.id', () => { + const { handler, toAppPtyId } = createSubscription() + + expect(() => handler('workspace.changed', { snapshot: { revision: 1 } })).not.toThrow() + expect(() => + handler('fs.changed', { + events: [{ kind: 'update', absolutePath: '/tmp/repo/file.txt' }] + }) + ).not.toThrow() + expect(toAppPtyId).not.toHaveBeenCalled() + }) + + it('routes pty.data after validating the string id', () => { + const { handler, toAppPtyId, notifications } = createSubscription() + const onData = vi.fn() + notifications.onData(onData) + + handler('pty.data', { id: 'pty-1', data: 'hello', rawLength: 5, seq: 9 }) + + expect(toAppPtyId).toHaveBeenCalledWith('pty-1') + expect(onData).toHaveBeenCalledWith( + expect.objectContaining({ + id: 'ssh:conn@@pty-1', + data: 'hello', + sequenceChars: 5, + seq: 9 + }) + ) + }) + + it('records pty.exit with the validated relay id', () => { + const { handler, recordExit, notifications } = createSubscription() + const onExit = vi.fn() + notifications.onExit(onExit) + + handler('pty.exit', { id: 'pty-1', code: 0, incarnationId: 'incarnation-1' }) + + expect(recordExit).toHaveBeenCalledWith('pty-1', 'incarnation-1') + expect(onExit).toHaveBeenCalledWith({ + id: 'ssh:conn@@pty-1', + code: 0, + incarnationId: 'incarnation-1' + }) + }) + + // Why: every PTY method must reject a missing id at ingress. pty.data is also + // screened by the delivery-credit layer, so replay/exit are asserted too — + // otherwise removing the ingress guard leaves this suite green. + it.each([ + { method: 'pty.data', params: { data: 'orphan' } }, + { method: 'pty.replay', params: { data: 'orphan' } }, + { method: 'pty.exit', params: { code: 0 } } + ])('ignores $method with a missing id', ({ method, params }) => { + const { handler, toAppPtyId, recordNotification, recordExit, notifications } = + createSubscription() + const listener = vi.fn() + notifications.onData(listener) + notifications.onReplay(listener) + notifications.onExit(listener) + + expect(() => handler(method, params)).not.toThrow() + expect(toAppPtyId).not.toHaveBeenCalled() + expect(recordNotification).not.toHaveBeenCalled() + expect(recordExit).not.toHaveBeenCalled() + expect(listener).not.toHaveBeenCalled() + }) + + // Why: the relay id is attacker-influenced and is retained per live PTY, so an + // oversized id must be dropped at ingress rather than mapped and held. + it.each(['pty.data', 'pty.replay', 'pty.exit'])( + 'drops %s relay ids past the admission ceiling', + (method) => { + const { handler, toAppPtyId, recordNotification, recordExit } = createSubscription() + + handler(method, { + id: 'x'.repeat(MAX_SSH_RELAY_PTY_ID_BYTES + 1), + data: 'hello', + code: 0 + }) + + expect(toAppPtyId).not.toHaveBeenCalled() + expect(recordNotification).not.toHaveBeenCalled() + expect(recordExit).not.toHaveBeenCalled() + } + ) + + it('drops pty.exit payloads with a non-numeric code', () => { + const { handler, recordExit, notifications } = createSubscription() + const onExit = vi.fn() + notifications.onExit(onExit) + + handler('pty.exit', { id: 'pty-1', code: 'not-a-number' }) + + expect(recordExit).not.toHaveBeenCalled() + expect(onExit).not.toHaveBeenCalled() + }) + + it('drops pty.replay payloads with non-string data', () => { + const { handler, recordNotification, notifications } = createSubscription() + const onReplay = vi.fn() + notifications.onReplay(onReplay) + + handler('pty.replay', { id: 'pty-1', data: { not: 'a string' } }) + + expect(recordNotification).not.toHaveBeenCalled() + expect(onReplay).not.toHaveBeenCalled() + }) + + it('stops routing after dispose', () => { + const { handler, notifications } = createSubscription() + const onData = vi.fn() + notifications.onData(onData) + + notifications.dispose() + handler('pty.data', { id: 'pty-1', data: 'hello' }) + + expect(onData).not.toHaveBeenCalled() + }) +}) diff --git a/src/main/providers/ssh-pty-provider-notifications.ts b/src/main/providers/ssh-pty-provider-notifications.ts new file mode 100644 index 00000000000..7d6524f0818 --- /dev/null +++ b/src/main/providers/ssh-pty-provider-notifications.ts @@ -0,0 +1,124 @@ +import { isPtyIncarnationId } from '../../shared/pty-incarnation' +import type { SshChannelMultiplexer } from '../ssh/ssh-channel-multiplexer' +import type { + SshPtyDataCallback, + SshPtyExitCallback, + SshPtyReplayCallback +} from './ssh-pty-provider-contract' +import { SshPtyOutputDeliveryCredit } from './ssh-pty-output-delivery-credit' +import { isAdmittedSshRelayPtyId } from './ssh-pty-wire-admission' + +export class SshPtyProviderNotifications { + private readonly dataListeners = new Set() + private readonly replayListeners = new Set() + private readonly exitListeners = new Set() + private readonly outputDeliveryCredit: SshPtyOutputDeliveryCredit + private unsubscribe: (() => void) | null + + constructor( + mux: SshChannelMultiplexer, + private readonly toAppPtyId: (relayId: string) => string, + recordNotification: (relayId: string) => void, + recordExit: (relayId: string, incarnationId: unknown) => void + ) { + this.outputDeliveryCredit = new SshPtyOutputDeliveryCredit((params) => { + mux.notify('pty.ackData', params) + }) + this.unsubscribe = mux.onNotification((method, params) => { + const relayId = params.id + if (!isAdmittedSshRelayPtyId(relayId)) { + return + } + switch (method) { + case 'pty.data': + this.emitData(params, recordNotification) + break + case 'pty.replay': { + if (typeof params.data !== 'string') { + return + } + recordNotification(relayId) + for (const listener of this.replayListeners) { + listener({ + id: this.toAppPtyId(relayId), + data: params.data + }) + } + break + } + case 'pty.exit': { + const code = params.code + const incarnationId = params.incarnationId + if ( + typeof code !== 'number' || + !Number.isSafeInteger(code) || + (incarnationId !== undefined && !isPtyIncarnationId(incarnationId)) + ) { + return + } + this.outputDeliveryCredit.release(relayId) + recordExit(relayId, incarnationId) + for (const listener of this.exitListeners) { + listener({ + id: this.toAppPtyId(relayId), + code, + ...(isPtyIncarnationId(incarnationId) ? { incarnationId } : {}) + }) + } + break + } + } + }) + } + + private emitData( + params: Record, + recordNotification: (relayId: string) => void + ): void { + this.outputDeliveryCredit.ingest(params, (output) => { + recordNotification(output.relayId) + if (this.dataListeners.size === 0) { + output.upstreamCredit?.acknowledge(output.upstreamCredit.charCount) + return + } + for (const listener of this.dataListeners) { + listener({ + id: this.toAppPtyId(output.relayId), + data: output.data, + ...(output.rawLength === undefined ? {} : { sequenceChars: output.rawLength }), + ...(output.transformed ? { transformed: true } : {}), + ...(output.seq === undefined ? {} : { seq: output.seq }), + ...(output.upstreamCredit ? { upstreamCredit: output.upstreamCredit } : {}) + }) + } + }) + } + + acknowledgeLegacy(relayId: string, charCount: number): void { + this.outputDeliveryCredit.acknowledgeLegacy(relayId, charCount) + } + + onData(callback: SshPtyDataCallback): () => void { + this.dataListeners.add(callback) + return () => this.dataListeners.delete(callback) + } + + onReplay(callback: SshPtyReplayCallback): () => void { + this.replayListeners.add(callback) + return () => this.replayListeners.delete(callback) + } + + onExit(callback: SshPtyExitCallback): () => void { + this.exitListeners.add(callback) + return () => this.exitListeners.delete(callback) + } + + dispose(): void { + this.unsubscribe?.() + this.unsubscribe = null + this.dataListeners.clear() + this.replayListeners.clear() + this.exitListeners.clear() + this.outputDeliveryCredit.dispose() + } +} diff --git a/src/main/providers/ssh-pty-provider-process-events.test.ts b/src/main/providers/ssh-pty-provider-process-events.test.ts index 2d033aa40d3..4a2cf30c6b0 100644 --- a/src/main/providers/ssh-pty-provider-process-events.test.ts +++ b/src/main/providers/ssh-pty-provider-process-events.test.ts @@ -1,5 +1,14 @@ import { beforeEach, describe, expect, it, vi } from 'vitest' import { SshPtyProvider } from './ssh-pty-provider' +import { + MAX_SSH_PTY_PROCESS_CWD_BYTES, + MAX_SSH_PTY_PROCESS_LIST_BYTES, + MAX_SSH_PTY_PROCESS_LIST_ENTRIES, + MAX_SSH_PTY_PROCESS_LIST_OWNERS, + MAX_SSH_PTY_PROCESS_OWNERS_PER_ENTRY, + MAX_SSH_PTY_PROCESS_TITLE_BYTES +} from './ssh-agent-session-process-list' +import { MAX_SSH_RELAY_PTY_ID_BYTES } from './ssh-pty-wire-admission' type MockMultiplexer = { request: ReturnType @@ -89,6 +98,71 @@ describe('SshPtyProvider process listings and events', () => { ]) }) + it('drops unknown process and owner fields instead of retaining remote payloads', async () => { + const unknownPayload = 'x'.repeat(1024 * 1024) + mux.request.mockResolvedValue([ + { + id: 'pty-1', + incarnationId: 'incarnation-1', + cwd: '/home', + title: 'codex', + unknownPayload, + agentSessionOwners: [ + { + claim: { + digestVersion: 1, + keyId: 'key', + identityDigest: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + worktreeScopeDigest: 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb', + agent: 'codex', + unknownPayload + }, + generation: 'generation-1', + phase: 'live', + ptyId: 'pty-1', + surface: { + worktreeId: 'worktree', + tabId: 'tab', + leafId: '11111111-1111-4111-8111-111111111111', + terminalHandle: 'term_claimed', + unknownPayload + }, + unknownPayload + } + ] + } + ]) + + await expect(provider.listProcesses()).resolves.toEqual([ + { + id: scopedPty1, + incarnationId: 'incarnation-1', + cwd: '/home', + title: 'codex', + agentSessionOwners: [ + { + claim: { + digestVersion: 1, + keyId: 'key', + identityDigest: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + worktreeScopeDigest: 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb', + agent: 'codex' + }, + generation: 'generation-1', + phase: 'live', + ptyId: scopedPty1, + surface: { + worktreeId: 'worktree', + tabId: 'tab', + leafId: '11111111-1111-4111-8111-111111111111', + terminalHandle: 'term_claimed' + } + } + ] + } + ]) + }) + it('rejects recovered claimed owners without PTY incarnation proof', async () => { mux.request.mockResolvedValue([ { @@ -121,6 +195,126 @@ describe('SshPtyProvider process listings and events', () => { await expect(provider.listProcesses()).rejects.toThrow('agent_session_ownership_unknown') }) + it('rejects a non-array process list before mapping ids', async () => { + mux.request.mockResolvedValue({ id: 'pty-1' }) + + await expect(provider.listProcesses()).rejects.toThrow('invalid_ssh_pty_process_list') + }) + + it('rejects an oversized process list before mapping ids', async () => { + const id = { + toString: vi.fn(() => 'pty-never-map') + } + mux.request.mockResolvedValue( + Array.from({ length: MAX_SSH_PTY_PROCESS_LIST_ENTRIES + 1 }, () => ({ + id, + cwd: '/home', + title: 'shell' + })) + ) + + await expect(provider.listProcesses()).rejects.toThrow('invalid_ssh_pty_process_list') + expect(id.toString).not.toHaveBeenCalled() + }) + + it('rejects oversized owner arrays before mapping owners', async () => { + mux.request.mockResolvedValue([ + { + id: 'pty-1', + cwd: '/home', + title: 'shell', + agentSessionOwners: Array.from( + { length: MAX_SSH_PTY_PROCESS_OWNERS_PER_ENTRY + 1 }, + () => ({}) + ) + } + ]) + + await expect(provider.listProcesses()).rejects.toThrow('invalid_ssh_pty_process_list') + }) + + it('rejects owner amplification spread across otherwise valid sessions', async () => { + const owner = { + claim: { + digestVersion: 1, + keyId: 'key', + identityDigest: 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', + worktreeScopeDigest: 'bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb', + agent: 'codex' + }, + generation: 'generation-1', + phase: 'live', + ptyId: 'pty-1', + surface: { + worktreeId: 'worktree', + tabId: 'tab', + leafId: '11111111-1111-4111-8111-111111111111', + terminalHandle: 'term_claimed' + } + } + const ownersPerSession = MAX_SSH_PTY_PROCESS_OWNERS_PER_ENTRY + const sessionCount = Math.floor(MAX_SSH_PTY_PROCESS_LIST_OWNERS / ownersPerSession) + 1 + mux.request.mockResolvedValue( + Array.from({ length: sessionCount }, (_, sessionIndex) => { + const id = `pty-${sessionIndex}` + return { + id, + incarnationId: 'incarnation-1', + cwd: '/home', + title: 'shell', + agentSessionOwners: Array.from({ length: ownersPerSession }, (_, ownerIndex) => ({ + ...owner, + generation: `generation-${sessionIndex}-${ownerIndex}`, + ptyId: id + })) + } + }) + ) + + await expect(provider.listProcesses()).rejects.toThrow('invalid_ssh_pty_process_list') + }) + + it.each([ + { + field: 'id', + value: 'x'.repeat(MAX_SSH_RELAY_PTY_ID_BYTES + 1) + }, + { + field: 'cwd', + value: 'x'.repeat(MAX_SSH_PTY_PROCESS_CWD_BYTES + 1) + }, + { + field: 'title', + value: 'x'.repeat(MAX_SSH_PTY_PROCESS_TITLE_BYTES + 1) + } + ])('rejects an oversized $field before mapping', async ({ field, value }) => { + mux.request.mockResolvedValue([{ id: 'pty-1', cwd: '/home', title: 'shell', [field]: value }]) + + await expect(provider.listProcesses()).rejects.toThrow('invalid_ssh_pty_process_list') + }) + + it('rejects process-list strings above the aggregate byte budget', async () => { + const cwd = 'x'.repeat(MAX_SSH_PTY_PROCESS_CWD_BYTES) + const count = Math.floor(MAX_SSH_PTY_PROCESS_LIST_BYTES / Buffer.byteLength(cwd)) + 1 + mux.request.mockResolvedValue( + Array.from({ length: count }, (_, index) => ({ + id: `pty-${index}`, + cwd, + title: 'shell' + })) + ) + + await expect(provider.listProcesses()).rejects.toThrow('invalid_ssh_pty_process_list') + }) + + it('rejects an oversized spawn id before retention', async () => { + const oversized = 'x'.repeat(MAX_SSH_RELAY_PTY_ID_BYTES + 1) + mux.request.mockResolvedValue({ id: oversized }) + + await expect(provider.spawn({ cols: 80, rows: 24 })).rejects.toThrow('invalid_ssh_pty_id') + expect(provider.hasPty(`ssh:conn-1@@${oversized}`)).toBe(false) + }) + it('forwards data, replay, and incarnation-aware exit notifications', () => { const dataHandler = vi.fn() const replayHandler = vi.fn() @@ -133,7 +327,9 @@ describe('SshPtyProvider process listings and events', () => { notify('pty.data', { id: 'pty-1', data: 'output' }) notify('pty.data', { id: 'pty-1', data: '', rawLength: 9, seq: 9, transformed: true }) notify('pty.replay', { id: 'pty-1', data: 'buffered output' }) + expect(provider.hasPty(scopedPty1)).toBe(true) notify('pty.exit', { id: 'pty-1', code: 0, incarnationId: 'incarnation-1' }) + expect(provider.hasPty(scopedPty1)).toBe(false) expect(dataHandler).toHaveBeenNthCalledWith(1, { id: scopedPty1, data: 'output' }) expect(dataHandler).toHaveBeenNthCalledWith(2, { @@ -151,6 +347,51 @@ describe('SshPtyProvider process listings and events', () => { }) }) + it('drops malformed notification ids before routing or retention', () => { + const dataHandler = vi.fn() + const replayHandler = vi.fn() + const exitHandler = vi.fn() + provider.onData(dataHandler) + provider.onReplay(replayHandler) + provider.onExit(exitHandler) + const notify = mux.onNotification.mock.calls[0][0] + const oversized = 'x'.repeat(MAX_SSH_RELAY_PTY_ID_BYTES + 1) + + notify('pty.data', { id: 42, data: 'output', deliveryToken: 'token' }) + notify('pty.replay', { id: oversized, data: 'replay' }) + notify('pty.exit', { id: oversized, code: 0 }) + + expect(dataHandler).not.toHaveBeenCalled() + expect(replayHandler).not.toHaveBeenCalled() + expect(exitHandler).not.toHaveBeenCalled() + expect(provider.hasPty(`ssh:conn-1@@${oversized}`)).toBe(false) + expect(mux.notify).not.toHaveBeenCalled() + }) + + it('drops malformed notification payloads before liveness or exit-race mutation', () => { + const dataHandler = vi.fn() + const replayHandler = vi.fn() + const exitHandler = vi.fn() + provider.onData(dataHandler) + provider.onReplay(replayHandler) + provider.onExit(exitHandler) + const notify = mux.onNotification.mock.calls[0][0] + + notify('pty.data', { id: 'pty-bad-data', data: 42, deliveryToken: 'token' }) + notify('pty.replay', { id: 'pty-bad-replay', data: { nested: 'output' } }) + notify('pty.exit', { id: 'pty-bad-code', code: '0', incarnationId: 'incarnation-1' }) + notify('pty.exit', { id: 'pty-bad-incarnation', code: 0, incarnationId: 'x'.repeat(129) }) + + expect(dataHandler).not.toHaveBeenCalled() + expect(replayHandler).not.toHaveBeenCalled() + expect(exitHandler).not.toHaveBeenCalled() + expect(provider.hasPty('ssh:conn-1@@pty-bad-data')).toBe(false) + expect(provider.hasPty('ssh:conn-1@@pty-bad-replay')).toBe(false) + expect(provider.hasPty('ssh:conn-1@@pty-bad-code')).toBe(false) + expect(provider.hasPty('ssh:conn-1@@pty-bad-incarnation')).toBe(false) + expect(mux.notify).not.toHaveBeenCalled() + }) + it('supports listener removal, fanout, and connection namespaces', () => { const removed = vi.fn() const first = vi.fn() diff --git a/src/main/providers/ssh-pty-provider.test.ts b/src/main/providers/ssh-pty-provider.test.ts index 99bb651ed99..c83ba43908a 100644 --- a/src/main/providers/ssh-pty-provider.test.ts +++ b/src/main/providers/ssh-pty-provider.test.ts @@ -32,6 +32,16 @@ describe('SshPtyProvider', () => { provider = new SshPtyProvider('conn-1', mux as never) }) + function emitRelayNotification(method: string, params: Record): void { + const callback = mux.onNotification.mock.calls[0]?.[0] as + | ((method: string, params: Record) => void) + | undefined + if (!callback) { + throw new Error('Notification listener was not registered') + } + callback(method, params) + } + it('returns the connectionId', () => { expect(provider.getConnectionId()).toBe('conn-1') }) @@ -821,6 +831,47 @@ describe('SshPtyProvider', () => { expect(mux.notify).toHaveBeenCalledWith('pty.ackData', { id: 'pty-1', charCount: 1024 }) }) + it('defers tokenized relay output credit until downstream acknowledgement', () => { + const events: Parameters[0]>[0][] = [] + provider.onData((event) => { + events.push(event) + }) + + emitRelayNotification('pty.data', { + id: 'pty-1', + data: 'clean', + rawLength: 12, + deliveryToken: 'delivery-1' + }) + + expect(events).toHaveLength(1) + expect(events[0]).toMatchObject({ id: scopedPty1, data: 'clean', sequenceChars: 12 }) + expect(mux.notify).not.toHaveBeenCalled() + + events[0]?.upstreamCredit?.acknowledge(12) + + expect(mux.notify).toHaveBeenCalledOnce() + expect(mux.notify).toHaveBeenCalledWith('pty.ackData', { + id: 'pty-1', + charCount: 12, + deliveryToken: 'delivery-1' + }) + }) + + it('restores legacy acknowledgements after tokenized PTY exit', () => { + emitRelayNotification('pty.data', { + id: 'pty-1', + data: 'output', + deliveryToken: 'delivery-3' + }) + emitRelayNotification('pty.exit', { id: 'pty-1', code: 0 }) + mux.notify.mockClear() + + provider.acknowledgeDataEvent(scopedPty1, 6) + + expect(mux.notify).toHaveBeenCalledWith('pty.ackData', { id: 'pty-1', charCount: 6 }) + }) + it('hasChildProcesses sends request and returns result', async () => { mux.request.mockResolvedValue(true) const result = await provider.hasChildProcesses(scopedPty1) diff --git a/src/main/providers/ssh-pty-provider.ts b/src/main/providers/ssh-pty-provider.ts index 59f20bfdf0a..9173fd5b9ca 100644 --- a/src/main/providers/ssh-pty-provider.ts +++ b/src/main/providers/ssh-pty-provider.ts @@ -8,7 +8,6 @@ import type { SshPtyExitCallback, SshPtyReplayCallback } from './ssh-pty-provider-contract' -import { subscribeSshPtyNotifications } from './ssh-pty-notification-routing' import { validateClaimedSshSpawn } from './ssh-agent-session-claim-validation' import { assertSshAgentSessionCreateResult, @@ -23,6 +22,9 @@ import { import { buildSshPtySpawnRequest } from './ssh-pty-spawn-request' import { SshPtySpawnExitRaceTracker } from './ssh-pty-spawn-exit-race' import { SshAgentSessionCapabilities } from './ssh-agent-session-capabilities' +import { SshPtyProviderNotifications } from './ssh-pty-provider-notifications' +import { SshPtyLiveRoster } from './ssh-pty-live-roster' +import { isAdmittedSshRelayPtyId } from './ssh-pty-wire-admission' import type { PtyProcessInspection } from './pty-process-inspection' // Why: sequential relay teardown calls share one absolute budget; convert to the mux-relative timeout only at dispatch. @@ -34,15 +36,11 @@ function relayTimeoutOptions(deadlineMs: number | undefined): { timeoutMs: numbe export class SshPtyProvider implements IPtyProvider { private mux: SshChannelMultiplexer private connectionId: string - private dataListeners = new Set() - private replayListeners = new Set() - private exitListeners = new Set() - private livePtyIds = new Set() - // Why: stale notification callbacks must not outlive a disconnected provider. - private unsubscribeNotifications: (() => void) | null = null + private readonly livePtys = new SshPtyLiveRoster() readonly getAppliedSize: NonNullable private readonly agentSessionCapabilities: SshAgentSessionCapabilities private spawnExitRaces = new SshPtySpawnExitRaceTracker() + private readonly notifications: SshPtyProviderNotifications constructor( connectionId: string, @@ -53,28 +51,20 @@ export class SshPtyProvider implements IPtyProvider { this.mux = mux this.agentSessionCapabilities = new SshAgentSessionCapabilities(mux) this.getAppliedSize = createSshPtyAppliedSizeReader(mux, connectionId) - - this.unsubscribeNotifications = subscribeSshPtyNotifications({ + this.notifications = new SshPtyProviderNotifications( mux, - toAppPtyId: (id) => this.toAppPtyId(id), - dataListeners: this.dataListeners, - replayListeners: this.replayListeners, - exitListeners: this.exitListeners, - livePtyIds: this.livePtyIds, - recordExit: (relayPtyId, incarnationId) => - this.spawnExitRaces.recordExit(relayPtyId, incarnationId) - }) + (relayId) => this.toAppPtyId(relayId), + (relayId) => this.livePtys.recordNotification(this.toAppPtyId(relayId)), + (relayId, incarnationId) => { + this.spawnExitRaces.recordExit(relayId, incarnationId) + this.livePtys.recordExit(this.toAppPtyId(relayId)) + } + ) } dispose(): void { - if (this.unsubscribeNotifications) { - this.unsubscribeNotifications() - this.unsubscribeNotifications = null - } - this.dataListeners.clear() - this.replayListeners.clear() - this.exitListeners.clear() - this.livePtyIds.clear() + this.livePtys.clear() + this.notifications.dispose() } getConnectionId = (): string => this.connectionId @@ -108,7 +98,7 @@ export class SshPtyProvider implements IPtyProvider { options: opts, exitRaceTracker: this.spawnExitRaces }) - this.livePtyIds.add(result.id) + this.livePtys.recordSpawn(result.id) return result } @@ -138,6 +128,9 @@ export class SshPtyProvider implements IPtyProvider { assertSshAgentSessionCreateResult(result) } const spawnResult = result as PtySpawnResult + if (!isAdmittedSshRelayPtyId(spawnResult.id)) { + throw new Error('invalid_ssh_pty_id') + } if (this.spawnExitRaces.didMatchingExitArrive(operation, spawnResult)) { // Why: relay notification can share the response batch; no controller registration may follow. throw Object.assign(new Error('agent_session_exited_during_start'), { @@ -161,7 +154,7 @@ export class SshPtyProvider implements IPtyProvider { } } const id = this.toAppPtyId(spawnResult.id) - this.livePtyIds.add(id) + this.livePtys.recordSpawn(id) return { ...spawnResult, id, @@ -240,7 +233,7 @@ export class SshPtyProvider implements IPtyProvider { }, relayTimeoutOptions(opts.deadlineMs) ) - this.livePtyIds.delete(id) + this.livePtys.recordExit(id) } async sendSignal(id: string, signal: string): Promise { @@ -269,7 +262,8 @@ export class SshPtyProvider implements IPtyProvider { } acknowledgeDataEvent(id: string, charCount: number): void { - this.mux.notify('pty.ackData', { id: this.toRelayPtyId(id), charCount }) + const relayId = this.toRelayPtyId(id) + this.notifications.acknowledgeLegacy(relayId, charCount) } async hasChildProcesses(id: string): Promise { @@ -300,20 +294,22 @@ export class SshPtyProvider implements IPtyProvider { } async listProcesses(opts?: { deadlineMs?: number }): Promise { + const listing = this.livePtys.beginListing() const result = await this.mux.request( 'pty.listProcesses', undefined, relayTimeoutOptions(opts?.deadlineMs) ) - const processes = mapSshPtyProcessList(result as PtyProcessInfo[], (id) => this.toAppPtyId(id)) - for (const process of processes) { - this.livePtyIds.add(process.id) - } + const processes = mapSshPtyProcessList(result, (id) => this.toAppPtyId(id)) + this.livePtys.reconcileListing( + listing, + processes.map((process) => process.id) + ) return processes } hasPty(id: string): boolean { - return this.livePtyIds.has(id) + return this.livePtys.has(id) } async getDefaultShell(): Promise { @@ -327,17 +323,14 @@ export class SshPtyProvider implements IPtyProvider { } onData(callback: SshPtyDataCallback): () => void { - this.dataListeners.add(callback) - return () => this.dataListeners.delete(callback) + return this.notifications.onData(callback) } onReplay(callback: SshPtyReplayCallback): () => void { - this.replayListeners.add(callback) - return () => this.replayListeners.delete(callback) + return this.notifications.onReplay(callback) } onExit(callback: SshPtyExitCallback): () => void { - this.exitListeners.add(callback) - return () => this.exitListeners.delete(callback) + return this.notifications.onExit(callback) } } diff --git a/src/main/providers/ssh-pty-spawn-exit-race.test.ts b/src/main/providers/ssh-pty-spawn-exit-race.test.ts new file mode 100644 index 00000000000..ceace5c6aba --- /dev/null +++ b/src/main/providers/ssh-pty-spawn-exit-race.test.ts @@ -0,0 +1,106 @@ +import { describe, expect, it } from 'vitest' +import { + MAX_PENDING_SSH_PTY_SPAWN_EXIT_OPERATIONS, + MAX_SSH_PTY_SPAWN_EXIT_BYTES_PER_OPERATION, + MAX_SSH_PTY_SPAWN_EXIT_EVENTS_PER_OPERATION, + SshPtySpawnExitRaceTracker +} from './ssh-pty-spawn-exit-race' +import { MAX_SSH_RELAY_PTY_ID_BYTES } from './ssh-pty-wire-admission' + +describe('SshPtySpawnExitRaceTracker', () => { + it('matches the same PTY incarnation without fencing a replacement', () => { + const tracker = new SshPtySpawnExitRaceTracker() + const operation = tracker.begin() + + tracker.recordExit('pty-1', 'incarnation-old') + + expect( + tracker.didMatchingExitArrive(operation, { + id: 'pty-1', + incarnationId: 'incarnation-current' + }) + ).toBe(false) + expect( + tracker.didMatchingExitArrive(operation, { + id: 'pty-1', + incarnationId: 'incarnation-old' + }) + ).toBe(true) + }) + + it('does not turn malformed incarnation metadata into a wildcard exit', () => { + const tracker = new SshPtySpawnExitRaceTracker() + const operation = tracker.begin() + + tracker.recordExit('pty-1', 'x'.repeat(129)) + + expect( + tracker.didMatchingExitArrive(operation, { + id: 'pty-1', + incarnationId: 'incarnation-current' + }) + ).toBe(false) + }) + + it('supports all 50 relay sessions without overflowing', () => { + const tracker = new SshPtySpawnExitRaceTracker() + const operation = tracker.begin() + for (let index = 0; index < 50; index += 1) { + tracker.recordExit(`pty-${index}`, `incarnation-${index}`) + } + + expect( + tracker.didMatchingExitArrive(operation, { + id: 'pty-49', + incarnationId: 'incarnation-49' + }) + ).toBe(true) + expect(Reflect.get(operation, 'overflowed')).toBe(false) + }) + + it('fails closed and releases retained exits when the count cap is exceeded', () => { + const tracker = new SshPtySpawnExitRaceTracker() + const operation = tracker.begin() + for (let index = 0; index <= MAX_SSH_PTY_SPAWN_EXIT_EVENTS_PER_OPERATION; index += 1) { + tracker.recordExit(`pty-${index}`, `incarnation-${index}`) + } + + expect(Reflect.get(operation, 'overflowed')).toBe(true) + expect(Reflect.get(operation, 'exits')).toHaveLength(0) + expect( + tracker.didMatchingExitArrive(operation, { + id: 'unseen', + incarnationId: 'unseen-incarnation' + }) + ).toBe(true) + }) + + it('fails closed when retained exit identifiers exceed the byte cap', () => { + const tracker = new SshPtySpawnExitRaceTracker() + const operation = tracker.begin() + const suffix = 'x'.repeat(MAX_SSH_RELAY_PTY_ID_BYTES - 16) + const eventBytes = Buffer.byteLength(`pty-000:${suffix}`) + Buffer.byteLength('💥'.repeat(64)) + const eventsToOverflow = Math.floor(MAX_SSH_PTY_SPAWN_EXIT_BYTES_PER_OPERATION / eventBytes) + 1 + expect(eventsToOverflow).toBeLessThan(MAX_SSH_PTY_SPAWN_EXIT_EVENTS_PER_OPERATION) + + for (let index = 0; index < eventsToOverflow; index += 1) { + tracker.recordExit(`pty-${String(index).padStart(3, '0')}:${suffix}`, '💥'.repeat(64)) + } + + expect(Reflect.get(operation, 'overflowed')).toBe(true) + expect(Reflect.get(operation, 'retainedBytes')).toBe(0) + }) + + it('rejects before dispatch when concurrent operation tracking is saturated', () => { + const tracker = new SshPtySpawnExitRaceTracker() + const admitted = Array.from({ length: MAX_PENDING_SSH_PTY_SPAWN_EXIT_OPERATIONS }, () => + tracker.begin() + ) + + expect(() => tracker.begin()).toThrow('ssh_pty_spawn_exit_tracking_capacity') + + for (const operation of admitted) { + tracker.finish(operation) + } + }) +}) diff --git a/src/main/providers/ssh-pty-spawn-exit-race.ts b/src/main/providers/ssh-pty-spawn-exit-race.ts index 0658e3e23eb..f0f9d677496 100644 --- a/src/main/providers/ssh-pty-spawn-exit-race.ts +++ b/src/main/providers/ssh-pty-spawn-exit-race.ts @@ -1,24 +1,55 @@ import { isPtyIncarnationId, type PtyIncarnationId } from '../../shared/pty-incarnation' +import { admittedSshRelayPtyIdBytes } from './ssh-pty-wire-admission' + +export const MAX_SSH_PTY_SPAWN_EXIT_EVENTS_PER_OPERATION = 128 +export const MAX_SSH_PTY_SPAWN_EXIT_BYTES_PER_OPERATION = 512 * 1024 +export const MAX_PENDING_SSH_PTY_SPAWN_EXIT_OPERATIONS = 128 type PendingSshPtySpawn = { exits: { relayPtyId: string; incarnationId?: PtyIncarnationId }[] + retainedBytes: number + overflowed: boolean } export class SshPtySpawnExitRaceTracker { private pending = new Set() begin(): PendingSshPtySpawn { - const operation = { exits: [] } + if (this.pending.size >= MAX_PENDING_SSH_PTY_SPAWN_EXIT_OPERATIONS) { + throw new Error('ssh_pty_spawn_exit_tracking_capacity') + } + const operation: PendingSshPtySpawn = { + exits: [], + retainedBytes: 0, + overflowed: false + } this.pending.add(operation) return operation } recordExit(relayPtyId: string, incarnationId: unknown): void { + const idBytes = admittedSshRelayPtyIdBytes(relayPtyId) + if (idBytes === null || (incarnationId !== undefined && !isPtyIncarnationId(incarnationId))) { + return + } + const retainedBytes = + idBytes + (isPtyIncarnationId(incarnationId) ? Buffer.byteLength(incarnationId, 'utf8') : 0) for (const operation of this.pending) { + if ( + operation.overflowed || + operation.exits.length >= MAX_SSH_PTY_SPAWN_EXIT_EVENTS_PER_OPERATION || + operation.retainedBytes + retainedBytes > MAX_SSH_PTY_SPAWN_EXIT_BYTES_PER_OPERATION + ) { + operation.overflowed = true + operation.exits.length = 0 + operation.retainedBytes = 0 + continue + } operation.exits.push({ relayPtyId, ...(isPtyIncarnationId(incarnationId) ? { incarnationId } : {}) }) + operation.retainedBytes += retainedBytes } } @@ -26,16 +57,21 @@ export class SshPtySpawnExitRaceTracker { operation: PendingSshPtySpawn, result: { id: string; incarnationId?: PtyIncarnationId } ): boolean { - return operation.exits.some( - (exit) => - exit.relayPtyId === result.id && - (!exit.incarnationId || - !result.incarnationId || - exit.incarnationId === result.incarnationId) + return ( + operation.overflowed || + operation.exits.some( + (exit) => + exit.relayPtyId === result.id && + (!exit.incarnationId || + !result.incarnationId || + exit.incarnationId === result.incarnationId) + ) ) } finish(operation: PendingSshPtySpawn): void { this.pending.delete(operation) + operation.exits.length = 0 + operation.retainedBytes = 0 } } diff --git a/src/main/providers/ssh-pty-wire-admission.ts b/src/main/providers/ssh-pty-wire-admission.ts new file mode 100644 index 00000000000..af7ac167fc9 --- /dev/null +++ b/src/main/providers/ssh-pty-wire-admission.ts @@ -0,0 +1,29 @@ +import { hasUnsafeProviderSessionIdChars } from '../../shared/agent-session-resume' + +export const MAX_SSH_RELAY_PTY_ID_BYTES = 4 * 1024 +export const MAX_SSH_APP_PTY_ID_BYTES = 16 * 1024 +export const MAX_SSH_PTY_DELIVERY_TOKEN_BYTES = 1024 + +function admittedStringBytes(value: unknown, maxBytes: number): number | null { + if (typeof value !== 'string' || value.length === 0 || hasUnsafeProviderSessionIdChars(value)) { + return null + } + const bytes = Buffer.byteLength(value, 'utf8') + return bytes <= maxBytes ? bytes : null +} + +export function admittedSshRelayPtyIdBytes(value: unknown): number | null { + return admittedStringBytes(value, MAX_SSH_RELAY_PTY_ID_BYTES) +} + +export function admittedSshAppPtyIdBytes(value: unknown): number | null { + return admittedStringBytes(value, MAX_SSH_APP_PTY_ID_BYTES) +} + +export function isAdmittedSshRelayPtyId(value: unknown): value is string { + return admittedSshRelayPtyIdBytes(value) !== null +} + +export function isAdmittedSshDeliveryToken(value: unknown): value is string { + return admittedStringBytes(value, MAX_SSH_PTY_DELIVERY_TOKEN_BYTES) !== null +}