import { EventEmitter } from "node:events"; import type { ConnectConfig } from "ssh2"; import { describe, expect, it, vi } from "vitest"; import type { ResolvedSshHost } from "./ssh-host-store"; import { AGENT_BOOTSTRAP_PROBE_COMMAND, verifyAgentInstallationId, } from "./ssh-agent-command"; import { SSH_REMOTE_PACKAGE_BOOTSTRAP_COMMAND } from "./ssh-remote-package-bootstrap"; import { Ssh2AuthenticatedConnection, SshConnectionPool, type AuthenticatedSshConnection, type SshConnectionFactory, type SshConnectionPoolTarget, type SshPoolConnectionIdentity, } from "./ssh-connection-pool"; const sshEd25519Algorithm = Buffer.from("ssh-ed25519"); const pinnedHostKey = Buffer.concat([ Buffer.from([0, 0, 0, sshEd25519Algorithm.byteLength]), sshEd25519Algorithm, Buffer.from([0, 0, 0, 32]), Buffer.alloc(32, 7), ]); const hostKey = { algorithm: "ssh-ed25519", publicKeyBase64: pinnedHostKey.toString("base64"), fingerprintSha256: `SHA256:${"A".repeat(43)}`, generation: 3, }; function createTarget( overrides: Partial = {}, identityOverrides: Partial< Pick > = {}, ): SshConnectionPoolTarget { return { host: { id: "00000000-0000-4000-8000-000000000201", name: "Build host", hostname: "build.example.com", port: 22, username: "builder", authentication: "password", password: "private password", hostKey, ...overrides, }, hostRevision: 7, hostKeyGeneration: 3, ...identityOverrides, }; } function createConnection(identity: SshPoolConnectionIdentity) { return { identity, isUsable: vi.fn(() => true), forwardLoopback: vi.fn(), onDisconnect: vi.fn(() => () => undefined), openTerminalShell: vi.fn(), openAgentAttach: vi.fn(), runAgentDoctor: vi.fn(), runAgentBootstrapProbe: vi.fn(), probeRemotePackageBootstrap: vi.fn(), createRemotePackageUploadStaging: vi.fn(), prepareRemotePackageBootstrap: vi.fn(), prepareUploadedRemotePackageBootstrap: vi.fn(), commitRemotePackageBootstrap: vi.fn(), cleanupRemotePackageBootstrap: vi.fn(), runAgentLifecycleAction: vi.fn(), runAgentRuntimeAction: vi.fn(), openStagedSftp: vi.fn(), listDirectories: vi.fn(), dispose: vi.fn(), } satisfies AuthenticatedSshConnection; } describe("SSH connection pool", () => { it("reuses one authenticated connection for matching identities and refcounts leases", async () => { vi.useFakeTimers(); try { const connections: ReturnType[] = []; const factory: SshConnectionFactory = vi.fn(async (_target, identity) => { const connection = createConnection(identity); connections.push(connection); return connection; }); const pool = new SshConnectionPool(factory, { idleTimeoutMs: 100, }); const [first, second] = await Promise.all([ pool.acquire(createTarget()), pool.acquire(createTarget()), ]); expect(factory).toHaveBeenCalledOnce(); expect(first.identity).toEqual(second.identity); first.release(); await vi.advanceTimersByTimeAsync(100); expect(connections[0]?.dispose).not.toHaveBeenCalled(); second.release(); await vi.advanceTimersByTimeAsync(99); expect(connections[0]?.dispose).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1); expect(connections[0]?.dispose).toHaveBeenCalledOnce(); } finally { vi.useRealTimers(); } }); it("isolates revision, host-key generation, credential, and host identities", async () => { const factory: SshConnectionFactory = vi.fn(async (_target, identity) => createConnection(identity), ); const pool = new SshConnectionPool(factory); const targets = [ createTarget(), createTarget({}, { hostRevision: 8 }), createTarget( { hostKey: { ...hostKey, generation: 4 } }, { hostKeyGeneration: 4 }, ), createTarget({ password: "rotated password" }), createTarget({ id: "00000000-0000-4000-8000-000000000202", }), ]; const leases = await Promise.all( targets.map((target) => pool.acquire(target)), ); expect(factory).toHaveBeenCalledTimes(5); expect( new Set( leases.map( (lease) => `${lease.identity.hostId}:${lease.identity.hostRevision}:` + `${lease.identity.hostKeyGeneration}:` + lease.identity.authenticationIdentity, ), ).size, ).toBe(5); for (const lease of leases) { lease.release(); } pool.dispose(); }); it("evicts the oldest idle entry for sequential distinct hosts beyond the pool limit", async () => { const connections: ReturnType[] = []; const factory: SshConnectionFactory = vi.fn(async (_target, identity) => { const connection = createConnection(identity); connections.push(connection); return connection; }); const pool = new SshConnectionPool(factory, { idleTimeoutMs: 60_000, }); for (let index = 0; index < 40; index += 1) { const hostId = `00000000-0000-4000-8000-${String(index).padStart(12, "0")}`; const lease = await pool.acquire(createTarget({ id: hostId })); lease.release(); } expect(factory).toHaveBeenCalledTimes(40); expect( connections .slice(0, 8) .every((connection) => connection.dispose.mock.calls.length === 1), ).toBe(true); expect( connections .slice(8) .every((connection) => connection.dispose.mock.calls.length === 0), ).toBe(true); pool.dispose(); }); it("rejects a distinct host when all pool entries have active leases", async () => { const connections: ReturnType[] = []; const factory: SshConnectionFactory = vi.fn(async (_target, identity) => { const connection = createConnection(identity); connections.push(connection); return connection; }); const pool = new SshConnectionPool(factory); const leases = await Promise.all( Array.from({ length: 32 }, (_, index) => pool.acquire( createTarget({ id: `00000000-0000-4000-8000-${String(index).padStart(12, "0")}`, }), ), ), ); await expect( pool.acquire( createTarget({ id: "00000000-0000-4000-8000-000000000032", }), ), ).rejects.toThrow("SSH 连接池已达到安全上限"); expect(factory).toHaveBeenCalledTimes(32); expect( connections.every( (connection) => connection.dispose.mock.calls.length === 0, ), ).toBe(true); for (const lease of leases) { lease.release(); } pool.dispose(); }); it("lets one canceled waiter leave without canceling a shared connection", async () => { let resolveConnection: ((connection: AuthenticatedSshConnection) => void) | undefined; const factory: SshConnectionFactory = vi.fn( (_target, identity) => new Promise((resolve) => { resolveConnection = resolve; expect(identity.authenticationIdentity).toMatch(/^[a-f0-9]{64}$/u); }), ); const pool = new SshConnectionPool(factory); const controller = new AbortController(); const canceled = pool.acquire(createTarget(), controller.signal); const surviving = pool.acquire(createTarget()); controller.abort(); await expect(canceled).rejects.toMatchObject({ name: "AbortError", }); const identity = (factory as ReturnType).mock .calls[0]?.[1] as SshPoolConnectionIdentity; resolveConnection?.(createConnection(identity)); const lease = await surviving; expect(lease.isUsable()).toBe(true); lease.release(); pool.dispose(); }); it("provides a dedicated bounded remote-package bootstrap lease", async () => { const connection = createConnection({ hostId: createTarget().host.id, hostRevision: 7, hostKeyGeneration: 3, authenticationIdentity: "a".repeat(64), }); connection.probeRemotePackageBootstrap.mockResolvedValue({ available: true, }); connection.createRemotePackageUploadStaging.mockResolvedValue({ created: true, operationId: "00000000-0000-4000-8000-000000000301", archivePath: "/home/builder/.goodbuddy/agent/staging/op-00000000-0000-4000-8000-000000000301/package.gbagent", bootstrapNodePath: "/home/builder/.goodbuddy/agent/staging/op-00000000-0000-4000-8000-000000000301/bootstrap/agent/node", }); connection.prepareUploadedRemotePackageBootstrap.mockResolvedValue({ prepared: false, unavailable: false, reason: "size-mismatch", }); const pool = new SshConnectionPool(async () => connection); const lease = await pool.acquireRemotePackageBootstrap(createTarget()); const candidate = { operationId: "00000000-0000-4000-8000-000000000301", urls: [ "https://goodbuddy.oss-cn-beijing.aliyuncs.com/agent-releases/v1.0.0/package.gbagent", ], size: 123, sha256: "a".repeat(64), }; await expect(lease.probe(candidate)).resolves.toEqual({ available: true, }); expect(connection.probeRemotePackageBootstrap).toHaveBeenCalledWith( candidate, undefined, ); await expect(lease.createUploadStaging(candidate)).resolves.toMatchObject({ created: true, }); expect(connection.createRemotePackageUploadStaging).toHaveBeenCalledWith( candidate, undefined, ); await expect(lease.prepareUploaded(candidate)).resolves.toMatchObject({ reason: "size-mismatch", }); expect( connection.prepareUploadedRemotePackageBootstrap, ).toHaveBeenCalledWith(candidate, undefined); expect(lease).not.toHaveProperty("exec"); expect(lease).not.toHaveProperty("sftp"); lease.release(); expect(() => lease.probe(candidate)).toThrow("租约已释放"); expect(() => lease.prepareUploaded(candidate)).toThrow("租约已释放"); pool.dispose(); }); it("disposes all host entries and prevents use after release", async () => { const connection = createConnection({ hostId: createTarget().host.id, hostRevision: 7, hostKeyGeneration: 3, authenticationIdentity: "a".repeat(64), }); const pool = new SshConnectionPool(async () => connection); const lease = await pool.acquire(createTarget()); const signal = new AbortController().signal; void lease.runAgentBootstrapProbe(signal); expect(connection.runAgentBootstrapProbe).toHaveBeenCalledWith(signal); void lease.listDirectories("/srv", signal); expect(connection.listDirectories).toHaveBeenCalledWith("/srv", signal); void lease.runAgentRuntimeAction( verifyAgentInstallationId("agent-v1"), { kind: "runtime-activate", runtimeId: "opencode", bundleDigest: `sha256:${"a".repeat(64)}`, architecture: "x64", }, signal, ); expect(connection.runAgentRuntimeAction).toHaveBeenCalledWith( "agent-v1", expect.objectContaining({ kind: "runtime-activate" }), signal, ); lease.release(); expect(() => lease.openAgentAttach(verifyAgentInstallationId("agent-v1")), ).toThrow("租约已释放"); expect(() => lease.runAgentBootstrapProbe()).toThrow("租约已释放"); expect(() => lease.listDirectories()).toThrow("租约已释放"); expect(lease).not.toHaveProperty("exec"); expect(lease).not.toHaveProperty("sftp"); pool.disposeHost(createTarget().host.id); expect(connection.dispose).toHaveBeenCalledOnce(); }); }); describe("authenticated ssh2 connection", () => { it("uses pinned algorithms, keepalive, no forwarding, and only fixed Agent exec", async () => { let connectConfig: ConnectConfig | undefined; const channel = Object.assign(new EventEmitter(), { stderr: new EventEmitter(), destroy: vi.fn(), end: vi.fn(), }); const client = Object.assign(new EventEmitter(), { connect: vi.fn(function connect( this: EventEmitter, config: ConnectConfig, ) { connectConfig = config; expect( (config.hostVerifier as ((key: Buffer) => boolean) | undefined)?.( pinnedHostKey, ), ).toBe(true); this.emit("ready"); }), end: vi.fn(), destroy: vi.fn(), exec: vi.fn( ( command: string, options: unknown, callback: (error: Error | undefined, stream: typeof channel) => void, ) => { if (command === SSH_REMOTE_PACKAGE_BOOTSTRAP_COMMAND) { channel.end.mockImplementationOnce(function end( this: EventEmitter, ) { queueMicrotask(() => { if ( (options as { allowHalfOpen?: boolean }).allowHalfOpen !== true ) { this.emit("close", null); return; } this.emit( "data", '{"type":"result","status":"available"}\n', ); this.emit("close", 0); }); }); } callback(undefined, channel); }, ), shell: vi.fn( ( _window: unknown, _options: unknown, callback: (error: Error | undefined, stream: typeof channel) => void, ) => callback(undefined, channel), ), sftp: vi.fn(), }); const target = createTarget(); const identity: SshPoolConnectionIdentity = { hostId: target.host.id, hostRevision: target.hostRevision, hostKeyGeneration: target.hostKeyGeneration, authenticationIdentity: "a".repeat(64), }; const connection = await Ssh2AuthenticatedConnection.connect( target, identity, new AbortController().signal, { createClient: () => client as never, supportedOpenSslCiphers: () => ["aes-256-gcm"], controlPlanePackageInstaller: 'module.exports = { source: "GoodBuddy" }\n', }, ); await connection.openAgentAttach(verifyAgentInstallationId("agent-v1")); await connection.openTerminalShell({ cols: 132, rows: 43 }); await expect( connection.probeRemotePackageBootstrap({ operationId: "00000000-0000-4000-8000-000000000301", urls: [ "https://goodbuddy.oss-cn-beijing.aliyuncs.com/agent-releases/v1.0.0/package.gbagent", ], size: 123, sha256: "a".repeat(64), }), ).resolves.toEqual({ available: true }); expect(connectConfig).toMatchObject({ host: "build.example.com", port: 22, username: "builder", password: "private password", agentForward: false, keepaliveInterval: 5_000, keepaliveCountMax: 2, algorithms: { cipher: ["aes256-gcm@openssh.com"], }, }); expect(client.exec).toHaveBeenCalledWith( 'exec "$HOME/.goodbuddy/agent/installations/' + 'agent-v1/goodbuddy-agent" attach-or-bootstrap ' + "--installation-id agent-v1", { env: {}, pty: false, x11: false, allowHalfOpen: false, }, expect.any(Function), ); expect(client.exec).toHaveBeenCalledWith( SSH_REMOTE_PACKAGE_BOOTSTRAP_COMMAND, { env: {}, pty: false, x11: false, allowHalfOpen: true, }, expect.any(Function), ); expect(client.shell).toHaveBeenCalledWith( { term: "xterm-256color", cols: 132, rows: 43, width: 0, height: 0, }, { env: {}, x11: false }, expect.any(Function), ); expect(connection).not.toHaveProperty("exec"); }); it("hard-fails a host-key mismatch and aborts pending authentication", async () => { let authenticationAttempted = false; const mismatchedClient = Object.assign(new EventEmitter(), { connect: vi.fn(function connect( this: EventEmitter, config: ConnectConfig, ) { const accepted = ( config.hostVerifier as ((key: Buffer) => boolean) | undefined )?.(Buffer.from("different-key")); if (accepted) { authenticationAttempted = true; } this.emit( "error", Object.assign(new Error("key rejected"), { level: "handshake", }), ); }), end: vi.fn(), destroy: vi.fn(), exec: vi.fn(), shell: vi.fn(), sftp: vi.fn(), }); const target = createTarget(); const identity = { hostId: target.host.id, hostRevision: target.hostRevision, hostKeyGeneration: target.hostKeyGeneration, authenticationIdentity: "a".repeat(64), }; expect(() => Ssh2AuthenticatedConnection.connect( target, { ...identity, hostRevision: identity.hostRevision + 1 }, new AbortController().signal, { createClient: () => mismatchedClient as never }, ), ).toThrow("认证目标不匹配"); await expect( Ssh2AuthenticatedConnection.connect( target, identity, new AbortController().signal, { createClient: () => mismatchedClient as never }, ), ).rejects.toThrow("主机密钥已变化"); expect(authenticationAttempted).toBe(false); expect(mismatchedClient.destroy).toHaveBeenCalledOnce(); const pendingClient = Object.assign(new EventEmitter(), { connect: vi.fn(), end: vi.fn(), destroy: vi.fn(), exec: vi.fn(), shell: vi.fn(), sftp: vi.fn(), }); const controller = new AbortController(); const pending = Ssh2AuthenticatedConnection.connect( target, identity, controller.signal, { createClient: () => pendingClient as never }, ); controller.abort(); await expect(pending).rejects.toMatchObject({ name: "AbortError", }); expect(pendingClient.destroy).toHaveBeenCalledOnce(); }); it("bounds fixed diagnostic output and destroys an overflowing channel", async () => { const channel = Object.assign(new EventEmitter(), { stderr: new EventEmitter(), destroy: vi.fn(), }); const client = Object.assign(new EventEmitter(), { connect: vi.fn(function connect(this: EventEmitter) { this.emit("ready"); }), end: vi.fn(), destroy: vi.fn(), exec: vi.fn( ( _command: string, _options: unknown, callback: (error: Error | undefined, stream: typeof channel) => void, ) => { callback(undefined, channel); setTimeout(() => { channel.emit("data", Buffer.alloc(64 * 1024 + 1)); }, 0); }, ), shell: vi.fn(), sftp: vi.fn(), }); const target = createTarget(); const connection = await Ssh2AuthenticatedConnection.connect( target, { hostId: target.host.id, hostRevision: target.hostRevision, hostKeyGeneration: target.hostKeyGeneration, authenticationIdentity: "a".repeat(64), }, new AbortController().signal, { createClient: () => client as never }, ); await expect( connection.runAgentDoctor(verifyAgentInstallationId("agent-v1")), ).rejects.toThrow("输出超过安全限制"); expect(client.exec.mock.calls[0]?.[0]).toContain( '/goodbuddy-agent" doctor', ); expect(channel.destroy).toHaveBeenCalledOnce(); }); it("bounds bootstrap probe output and destroys an overflowing channel", async () => { const channel = Object.assign(new EventEmitter(), { stderr: new EventEmitter(), destroy: vi.fn(), }); const client = Object.assign(new EventEmitter(), { connect: vi.fn(function connect(this: EventEmitter) { this.emit("ready"); }), end: vi.fn(), destroy: vi.fn(), exec: vi.fn( ( _command: string, _options: unknown, callback: (error: Error | undefined, stream: typeof channel) => void, ) => { callback(undefined, channel); setTimeout(() => { channel.emit("data", Buffer.alloc(64 * 1024 + 1)); }, 0); }, ), shell: vi.fn(), sftp: vi.fn(), }); const target = createTarget(); const connection = await Ssh2AuthenticatedConnection.connect( target, { hostId: target.host.id, hostRevision: target.hostRevision, hostKeyGeneration: target.hostKeyGeneration, authenticationIdentity: "a".repeat(64), }, new AbortController().signal, { createClient: () => client as never }, ); await expect(connection.runAgentBootstrapProbe()).rejects.toThrow( "输出超过安全限制", ); expect(client.exec.mock.calls[0]?.[0]).toBe(AGENT_BOOTSTRAP_PROBE_COMMAND); expect(channel.destroy).toHaveBeenCalledOnce(); }); it("runs the fixed bootstrap probe and collects split output", async () => { const channel = Object.assign(new EventEmitter(), { stderr: new EventEmitter(), destroy: vi.fn(), }); const output = [ "GOODBUDDY_AGENT_BOOTSTRAP_PROBE_V1", "home=/home/builder", "uid=1001", "os=Linux", "arch=aarch64", "shell=/bin/bash", "procfs=ready", "", ].join("\n"); const client = Object.assign(new EventEmitter(), { connect: vi.fn(function connect(this: EventEmitter) { this.emit("ready"); }), end: vi.fn(), destroy: vi.fn(), exec: vi.fn( ( _command: string, _options: unknown, callback: (error: Error | undefined, stream: typeof channel) => void, ) => { callback(undefined, channel); setTimeout(() => { channel.emit("data", output.slice(0, 17)); channel.emit("data", Buffer.from(output.slice(17))); channel.emit("close", 0); }, 0); }, ), shell: vi.fn(), sftp: vi.fn(), }); const target = createTarget(); const connection = await Ssh2AuthenticatedConnection.connect( target, { hostId: target.host.id, hostRevision: target.hostRevision, hostKeyGeneration: target.hostKeyGeneration, authenticationIdentity: "a".repeat(64), }, new AbortController().signal, { createClient: () => client as never }, ); await expect(connection.runAgentBootstrapProbe()).resolves.toMatchObject({ ready: true, architecture: "arm64", }); await expect(connection.runAgentBootstrapProbe()).resolves.toMatchObject({ ready: true, architecture: "arm64", }); expect(client.exec).toHaveBeenCalledWith( AGENT_BOOTSTRAP_PROBE_COMMAND, { env: {}, pty: false, x11: false, allowHalfOpen: false, }, expect.any(Function), ); expect(client.exec).toHaveBeenCalledOnce(); client.emit("close"); await expect(connection.runAgentBootstrapProbe()).rejects.toThrow( "SSH 连接已关闭", ); }); it("caches an incompatible bootstrap result on the connection", async () => { const channel = Object.assign(new EventEmitter(), { stderr: new EventEmitter(), destroy: vi.fn(), }); const output = [ "GOODBUDDY_AGENT_BOOTSTRAP_PROBE_V1", "home=/home/builder", "uid=1001", "os=Linux", "arch=riscv64", "shell=/bin/bash", "procfs=ready", "", ].join("\n"); const client = Object.assign(new EventEmitter(), { connect: vi.fn(function connect(this: EventEmitter) { this.emit("ready"); }), end: vi.fn(), destroy: vi.fn(), exec: vi.fn( ( _command: string, _options: unknown, callback: (error: Error | undefined, stream: typeof channel) => void, ) => { callback(undefined, channel); setTimeout(() => { channel.emit("data", output); channel.emit("close", 0); }, 0); }, ), shell: vi.fn(), sftp: vi.fn(), }); const target = createTarget(); const connection = await Ssh2AuthenticatedConnection.connect( target, { hostId: target.host.id, hostRevision: target.hostRevision, hostKeyGeneration: target.hostKeyGeneration, authenticationIdentity: "a".repeat(64), }, new AbortController().signal, { createClient: () => client as never }, ); await expect(connection.runAgentBootstrapProbe()).resolves.toEqual({ ready: false, reason: "unsupported-architecture", }); await expect(connection.runAgentBootstrapProbe()).resolves.toEqual({ ready: false, reason: "unsupported-architecture", }); expect(client.exec).toHaveBeenCalledOnce(); }); it("cancels a running bootstrap probe and destroys its channel", async () => { const channel = Object.assign(new EventEmitter(), { stderr: new EventEmitter(), destroy: vi.fn(), }); const client = Object.assign(new EventEmitter(), { connect: vi.fn(function connect(this: EventEmitter) { this.emit("ready"); }), end: vi.fn(), destroy: vi.fn(), exec: vi.fn( ( _command: string, _options: unknown, callback: (error: Error | undefined, stream: typeof channel) => void, ) => callback(undefined, channel), ), sftp: vi.fn(), }); const target = createTarget(); const connection = await Ssh2AuthenticatedConnection.connect( target, { hostId: target.host.id, hostRevision: target.hostRevision, hostKeyGeneration: target.hostKeyGeneration, authenticationIdentity: "a".repeat(64), }, new AbortController().signal, { createClient: () => client as never }, ); const controller = new AbortController(); const result = connection.runAgentBootstrapProbe(controller.signal); controller.abort(); await expect(result).rejects.toMatchObject({ name: "AbortError", }); expect(channel.destroy).toHaveBeenCalledOnce(); }); });