From 4a61217056377b4438d52ffd88e8087334e6028b Mon Sep 17 00:00:00 2001 From: moloch-- <875022+moloch--@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:44:55 -0700 Subject: [PATCH 1/3] feat: add cancellable operations, tunnel recovery, and typed beacon tasks --- README.md | 45 +++ e2e/README.md | 9 +- e2e/groups/06-session-shell.ts | 207 ++++++++++ e2e/groups/07-beacon-core.ts | 34 +- src/beaconTasks.ts | 247 ++++++++++++ src/client.ts | 607 ++++++++++++++++++----------- src/index.ts | 10 +- src/internal/portForward.ts | 4 +- src/internal/timeout.ts | 143 ++++++- src/internal/tunnelManager.ts | 201 +++++++--- src/tests/beacon-task-api.test.ts | 128 ++++++ src/tests/beacon-tasks.test.ts | 307 +++++++++++++++ src/tests/cancellation-api.test.ts | 156 ++++++++ src/tests/client-connect.test.ts | 57 +++ src/tests/managed-shell.test.ts | 70 ++++ src/tests/timeout.test.ts | 247 +++++++++++- src/tests/tunnel-recovery.test.ts | 253 ++++++++++++ 17 files changed, 2415 insertions(+), 310 deletions(-) create mode 100644 e2e/groups/06-session-shell.ts create mode 100644 src/beaconTasks.ts create mode 100644 src/tests/beacon-task-api.test.ts create mode 100644 src/tests/beacon-tasks.test.ts create mode 100644 src/tests/cancellation-api.test.ts create mode 100644 src/tests/tunnel-recovery.test.ts diff --git a/README.md b/README.md index 6e1feac..c64d042 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,45 @@ same path on Linux/amd64, Windows/amd64, and macOS/arm64. See Operator connections use direct mTLS. Configurations containing a WireGuard operator profile are rejected before credentials or channels are created. WireGuard listener and implant C2 APIs remain available. +### Cancellation and deadlines + +Existing numeric timeout arguments remain supported. Convenience RPCs, `connect`, +and shell startup also accept `{ timeoutSeconds, signal }` in the same argument +position. Omitting `timeoutSeconds` preserves the method's default; zero disables +the local deadline while still allowing cancellation. `OperationAbortedError` +and `OperationTimeoutError` distinguish caller cancellation from deadline expiry. +Shell startup shares one deadline across tunnel creation, the server's bind +acknowledgement, and shell creation. Cancellation stops waiting and requests +transport cancellation; it does not undo a remote operation already accepted. + +```typescript +const controller = new AbortController() +const sessions = await client.getSessions({ timeoutSeconds: 15, signal: controller.signal }) +``` + +### Typed beacon tasks + +Existing beacon commands keep their queued-response return values. Attach a +typed handle using the task ID and the generated response codec: + +```typescript +import { sliverpb } from "sliver-script" + +const queued = await client.pingBeacon(beaconId, 42) +const task = client.beaconTask(queued.Response!.TaskID, sliverpb.Ping, { beaconId }) +const result = await task.wait({ timeoutSeconds: 120, signal: controller.signal }) +console.log(result.Nonce) +const metadata = await task.status() +``` + +`wait()` uses events and polling, checks task identity and target errors, and +supports repeated waits or attachment to a persisted ID. A handle created through +`SliverClient.beaconTask` resolves the current channel after reconnect. Canceling +a wait leaves the remote task queued; `task.cancel()` explicitly requests server +cancellation. `InteractiveBeacon.beaconTask` binds the expected beacon ID, and +the existing `lsTask()` handle also exposes `status()` and `cancel()`. Decoded task +responses retain the existing 64 KiB limit. + ### Stateful forwarding The client provides stateful TCP port-forward, reverse-port-forward, and SOCKS5 @@ -68,6 +107,12 @@ idempotent `close()`. Use `listPortForwards()` / `stopPortForward(id)` and inventory, including safe cleanup of legacy listeners whose target metadata is not trusted. Sliver's generic forwarding protocol currently supports TCP; the application traffic carried over it can be HTTP, RDP, or any other TCP protocol. +`tunnelStreamState$` reports `idle`, `connecting`, `connected`, `failed`, and +`stopped`, with a stream generation number. After a tunnel transport failure, +existing tunnels fail and a new tunnel starts a fresh stream lazily. Old data is +never replayed. `isConnected` continues to describe the control connection. +Tunnel identities are retained to reject stale handles, bounded to 65,536 IDs +per connection; reaching that limit requires disconnecting and reconnecting. These tunnels use full-close TCP semantics: a local FIN retires both directions, so protocols that send a delayed response only after the client half-closes are not supported. A SOCKS connection's `open` event means its Sliver lifecycle diff --git a/e2e/README.md b/e2e/README.md index 0cd0468..a55e0ae 100644 --- a/e2e/README.md +++ b/e2e/README.md @@ -19,6 +19,7 @@ keeping each area of client behavior independently identifiable in diagnostics. | `04-session-core` | Session inventory, direct and interactive ping, rename/restore, and environment read/set/unset | | `05-session-filesystem` | Working-directory changes, directory listing and globbing, upload/download ranges, recursive archives and grep, copy/move/timestamps/removal, and Unix permission/ownership calls | | `06-session-process-network` | Interfaces, live TCP/UDP socket inventory, mounts, basic/full process inventory, synchronous/background execution, child tracking and termination, Linux memfiles, and Windows token/privilege/service reads plus isolated HKCU registry CRUD | +| `06-session-shell` | Native managed shell output and ordered writes, Unix PTY sizing and resize, Windows piped-shell resize rejection, natural exit, disconnect teardown, owned-child cleanup, and client reconnect | | `07-beacon-core` | Beacon metadata and rename, direct and interactive queued pings, derived task-result events, 10-second reconfiguration, task cancellation, fetch, and history/count reconciliation | | `08-beacon-filesystem` | Queued working-directory, mkdir, cd, upload/download, list, and recursive removal calls with completed-task verification | | `09-beacon-process-network` | Queued process, interface, socket, and environment calls plus synchronous/background execution, child tracking, termination, and task-history verification | @@ -97,9 +98,11 @@ so diagnostic logs do not inherit unrelated CI secrets. Generated binary bytes are cleared after being written, and failure cleanup removes profiles, builds, listeners, sessions, beacons, processes, and the temporary root. -The suite does not run interactive shell tunnels, screenshots, process dumps, -service mutation, token mutation, registry-hive export, injection, or arbitrary -payload execution. Forwarding coverage uses deterministic, bounded loopback TCP +The suite does not run screenshots, process dumps, service mutation, token +mutation, registry-hive export, injection, or arbitrary payload execution. +Interactive shell coverage uses bounded native fixture commands and verifies +cleanup of only the shell children it created. Forwarding coverage uses +deterministic, bounded loopback TCP fixtures only. In particular, hosted CI does not install an RDP client, use RDP credentials, or connect to port 3389; RDP interoperability is validated in the separate Proxmox/GOAD range. diff --git a/e2e/groups/06-session-shell.ts b/e2e/groups/06-session-shell.ts new file mode 100644 index 0000000..8097f34 --- /dev/null +++ b/e2e/groups/06-session-shell.ts @@ -0,0 +1,207 @@ +import assert from "node:assert/strict"; +import { access, readFile } from "node:fs/promises"; +import path from "node:path"; +import { setTimeout as delay } from "node:timers/promises"; + +import type { ShellSessionHandle } from "../../lib"; +import type { E2ESuiteContext, LiveImplant } from "../context"; + +export const name = "06-session-shell"; +const operationSeconds = 30; +const outputLimitBytes = 128 * 1024; +const waitMilliseconds = 20_000; + +interface OutputObservation { + text: string; + finished: boolean; + error?: unknown; + done: Promise; +} + +interface OwnedShell { + handle: ShellSessionHandle; + output: OutputObservation; + executable: string; +} + +export async function run(context: E2ESuiteContext): Promise { + const implant = context.session; + assert.ok(implant?.session?.ID, "live session fixture"); + const sessionId = implant.session.ID; + const sentinel = path.join(implant.root, "outside-test-sentinel.txt"); + assert.equal(await readFile(sentinel, "utf8"), "must-survive\n", "shell fixture sentinel"); + const windows = context.environment.expectedOS === "windows"; + const systemRoot = process.env.SystemRoot ?? process.env.SYSTEMROOT; + if (windows) assert.ok(systemRoot && path.win32.isAbsolute(systemRoot), "absolute Windows system root"); + const shellPath = windows ? path.win32.join(systemRoot!, "System32", "cmd.exe") : "/bin/sh"; + await access(shellPath); // Do not let Sliver silently select its fallback shell. + const prefix = `SSE2E_${process.pid}_`; + const markerCommand = (suffix: string): string => windows + ? `echo %SSE2E_PREFIX%${suffix}\r\n` + : `printf '%s%s\\n' '${prefix}' '${suffix}'\n`; + const owned: OwnedShell[] = []; + let failure: unknown; + try { + const first = await startOwnedShell(context, implant, shellPath, !windows); + owned.push(first); + await write(first.handle, windows + ? `@echo off\r\nset "SSE2E_PREFIX=${prefix}"\r\n` + : "PS1=''; PS2=''; /bin/stty -echo\n"); + await write(first.handle, markerCommand("READY")); + await waitForOutput(first.output, `${prefix}READY\n`); + + // Queue separate writes together to exercise the handle's ordering guarantee. + await bounded(Promise.all(["FIRST", "SECOND", "THIRD"].map((part) => + write(first.handle, markerCommand(part)), + )), "ordered shell writes"); + await waitForOutput(first.output, `${prefix}THIRD\n`); + const positions = ["FIRST", "SECOND", "THIRD"].map((part) => first.output.text.indexOf(`${prefix}${part}\n`)); + assert.ok(positions[0]! >= 0 && positions[0]! < positions[1]! && positions[1]! < positions[2]!, "shell output follows write order"); + + if (windows) { + await assert.rejects(bounded(first.handle.resize(41, 101), "piped shell resize"), /requires a PTY/u, "piped Windows shell cannot resize"); + } else { + await write(first.handle, `printf '%s%s' '${prefix}' 'INITIAL:'; /bin/stty size\n`); + await waitForOutput(first.output, `${prefix}INITIAL:24 80\n`); + await bounded(first.handle.resize(41, 101), "native PTY resize"); + await write(first.handle, `printf '%s%s' '${prefix}' 'RESIZED:'; /bin/stty size\n`); + await waitForOutput(first.output, `${prefix}RESIZED:41 101\n`); + } + + await write(first.handle, windows ? "exit\r\n" : "exit\n"); + await bounded(first.output.done, "natural shell EOF"); + assert.equal(first.output.error, undefined, "natural shell exit completes the output iterator"); + await waitForShellExit(context, implant, first); + await assert.rejects(write(first.handle, "already closed"), /closed/u, "exited shell rejects writes"); + await bounded(first.handle.close(), "exited shell close"); + await bounded(first.handle.close(), "idempotent shell close"); + + const second = await startOwnedShell(context, implant, shellPath, !windows); + owned.push(second); + await write(second.handle, windows + ? `@echo off\r\nset "SSE2E_PREFIX=${prefix}"\r\n` + : "PS1=''; PS2=''; /bin/stty -echo\n"); + await write(second.handle, markerCommand("DISCONNECT_READY")); + await waitForOutput(second.output, `${prefix}DISCONNECT_READY\n`); + await bounded(context.client.disconnect(), "client disconnect with an open shell"); + await bounded(second.output.done, "disconnected shell output cleanup"); + assert.ok(second.output.error instanceof Error, "disconnect fails the open output iterator"); + assert.match(second.output.error.message, /Tunnel transport disconnected/u); + await assert.rejects(write(second.handle, "already disconnected"), /closed/u); + assert.equal(context.client.isConnected, false, "client disconnected after shell teardown"); + + await context.client.connect(); + const ping = await context.client.pingSession(sessionId, 0x53e110, operationSeconds); + assert.equal(ping.Response?.Err ?? "", "", "session remains usable after shell reconnect"); + assert.equal(ping.Nonce, 0x53e110); + await waitForShellExit(context, implant, second); + } catch (error) { + failure = error; + } + + const cleanupErrors: unknown[] = []; + try { + if (!context.client.isConnected) await context.client.connect(); + for (const shell of owned) { + try { + await bounded(shell.handle.close(), "owned shell cleanup"); + const current = await ownedProcess(context, implant, shell.handle.pid); + if (current) { + assert.equal(current.Executable, shell.executable, "cleanup shell executable identity"); + const terminated = await context.client.terminateSessionProcess(sessionId, shell.handle.pid, false, operationSeconds); + assert.equal(terminated.Response?.Err ?? "", "", "owned shell fallback termination"); + assert.equal(terminated.Pid, shell.handle.pid, "only the returned shell PID is terminated"); + } + await waitForShellExit(context, implant, shell); + await bounded(shell.output.done, "owned shell output cleanup"); + } catch (error) { + cleanupErrors.push(error); + } + } + assert.equal(await readFile(sentinel, "utf8"), "must-survive\n", "shell fixture sentinel preserved"); + } catch (error) { + cleanupErrors.push(error); + } + if (failure && cleanupErrors.length) throw new AggregateError([failure, ...cleanupErrors], "shell E2E and cleanup failed"); + if (failure) throw failure; + if (cleanupErrors.length) throw new AggregateError(cleanupErrors, "shell E2E cleanup failed"); +} + +async function startOwnedShell(context: E2ESuiteContext, implant: LiveImplant, shellPath: string, pty: boolean): Promise { + const handle = await context.client.startShellSession(implant.session!.ID, { + path: shellPath, pty, rows: 24, cols: 80, outputBufferBytes: outputLimitBytes, + }, operationSeconds); + const output = observeOutput(handle); + try { + assert.ok(handle.pid > 1 && handle.pid !== implant.pid, "distinct owned shell child PID"); + const process = await ownedProcess(context, implant, handle.pid); + assert.ok(process, "new shell is in the native process inventory"); + return { handle, output, executable: process.Executable }; + } catch (error) { + await bounded(handle.close(), "failed shell startup cleanup"); + throw error; + } +} + +function observeOutput(handle: ShellSessionHandle): OutputObservation { + const observed: OutputObservation = { text: "", finished: false, done: Promise.resolve() }; + observed.done = (async () => { + let bytes = 0; + try { + for await (const chunk of handle.output) { + bytes += chunk.byteLength; + assert.ok(bytes <= outputLimitBytes, "bounded shell fixture output"); + observed.text += Buffer.from(chunk).toString("utf8").replace(/\r/gu, ""); + } + } catch (error) { + observed.error = error; + } finally { + observed.finished = true; + } + })(); + return observed; +} + +async function waitForOutput(output: OutputObservation, marker: string): Promise { + const deadline = Date.now() + waitMilliseconds; + while (Date.now() < deadline) { + if (output.text.includes(marker)) return; + if (output.error) throw output.error; + assert.equal(output.finished, false, `shell closed before ${marker.trim()}`); + await delay(25); + } + throw new Error(`Timed out waiting for shell marker ${marker.trim()}; output=${JSON.stringify(output.text.slice(-2_000))}`); +} + +async function ownedProcess(context: E2ESuiteContext, implant: LiveImplant, pid: number) { + assert.ok(pid > 1 && pid !== implant.pid, "safe shell child PID"); + const listed = await context.client.psSession(implant.session!.ID, true, 5); + assert.equal(listed.Response?.Err ?? "", "", "shell process inventory"); + const current = listed.Processes.find((process) => process.Pid === pid); + if (current) assert.equal(current.Ppid, implant.pid, "shell process belongs to this fixture implant"); + return current; +} + +async function waitForShellExit(context: E2ESuiteContext, implant: LiveImplant, shell: OwnedShell): Promise { + const deadline = Date.now() + waitMilliseconds; + while (Date.now() < deadline) { + if (!await ownedProcess(context, implant, shell.handle.pid)) return; + await delay(100); + } + throw new Error(`Owned shell PID ${shell.handle.pid} did not exit`); +} + +function write(handle: ShellSessionHandle, data: string): Promise { + return bounded(handle.write(data), "shell input"); +} + +async function bounded(operation: Promise, label: string): Promise { + let timer: ReturnType | undefined; + try { + return await Promise.race([operation, new Promise((_resolve, reject) => { + timer = setTimeout(() => reject(new Error(`Timed out during ${label}`)), waitMilliseconds); + })]); + } finally { + if (timer) clearTimeout(timer); + } +} diff --git a/e2e/groups/07-beacon-core.ts b/e2e/groups/07-beacon-core.ts index 4493aed..82f4d47 100644 --- a/e2e/groups/07-beacon-core.ts +++ b/e2e/groups/07-beacon-core.ts @@ -79,14 +79,22 @@ export async function run(context: E2ESuiteContext): Promise { } const interactiveNonce = nonce + 7; - const interactivePing = await context.runBeaconTask( - beacon.ID, - () => context.client.interactBeacon(beacon.ID).ping(interactiveNonce, rpcTimeoutSeconds), - sliverScript.sliverpb.Ping.decode, - "interactive beacon ping", - ); - assert.equal(interactivePing.response.Nonce, interactiveNonce, "interactive beacon ping nonce"); - assert.equal(interactivePing.response.Response?.Err ?? "", "", "interactive beacon ping response error"); + const interactivePing = await context.client.interactBeacon(beacon.ID).ping(interactiveNonce, rpcTimeoutSeconds); + const interactiveMetadata = requireQueuedResponse(interactivePing.Response, beacon.ID, "interactive beacon ping"); + const interactiveHandle = context.client.beaconTask(interactiveMetadata.TaskID, sliverScript.sliverpb.Ping, { + beaconId: beacon.ID, + }); + const interactiveResult = await interactiveHandle.wait({ timeoutSeconds: 120 }); + assert.equal(interactiveResult.Nonce, interactiveNonce, "interactive beacon ping nonce"); + assert.equal(interactiveResult.Response?.Err ?? "", "", "interactive beacon ping response error"); + const interactiveStatus = await interactiveHandle.status(rpcTimeoutSeconds); + assert.equal(interactiveStatus.ID, interactiveMetadata.TaskID, "interactive handle task ID"); + assert.equal(interactiveStatus.BeaconID, beacon.ID, "interactive handle beacon ID"); + assert.equal(interactiveStatus.State, "completed", "interactive handle completed state"); + const reattachedResult = await context.client.beaconTask( + interactiveMetadata.TaskID, sliverScript.sliverpb.Ping, { beaconId: beacon.ID }, + ).wait(rpcTimeoutSeconds); + assert.equal(reattachedResult.Nonce, interactiveNonce, "reattached completed beacon task result"); const reconfigured = await context.runBeaconTask( beacon.ID, @@ -113,7 +121,13 @@ export async function run(context: E2ESuiteContext): Promise { // task remains deterministically pending on the loopback fixture. const canceledQueue = await context.client.pingBeacon(beacon.ID, nonce + 1, rpcTimeoutSeconds); const canceledMetadata = requireQueuedResponse(canceledQueue.Response, beacon.ID, "canceled beacon ping"); - const canceled = await context.client.cancelBeaconTask(canceledMetadata.TaskID, rpcTimeoutSeconds); + const canceledHandle = context.client.beaconTask(canceledMetadata.TaskID, sliverScript.sliverpb.Ping, { + beaconId: beacon.ID, + }); + const localWaitAbort = new AbortController(); + localWaitAbort.abort(); + await assert.rejects(canceledHandle.wait({ signal: localWaitAbort.signal }), sliverScript.OperationAbortedError); + const canceled = await canceledHandle.cancel(rpcTimeoutSeconds); assert.equal(canceled.ID, canceledMetadata.TaskID, "canceled beacon task ID"); assert.equal(canceled.BeaconID, beacon.ID, "canceled beacon ID"); assert.equal(canceled.State, "canceled", "canceled beacon task state"); @@ -128,7 +142,7 @@ export async function run(context: E2ESuiteContext): Promise { const history: clientpb.BeaconTasks = await context.client.getBeaconTasks(beacon.ID, rpcTimeoutSeconds); assertTaskSummary(history, pingTaskId, "completed"); - assertTaskSummary(history, interactivePing.task.ID, "completed"); + assertTaskSummary(history, interactiveMetadata.TaskID, "completed"); assertTaskSummary(history, reconfigured.task.ID, "completed"); assertTaskSummary(history, canceled.ID, "canceled"); diff --git a/src/beaconTasks.ts b/src/beaconTasks.ts new file mode 100644 index 0000000..dcc340c --- /dev/null +++ b/src/beaconTasks.ts @@ -0,0 +1,247 @@ +import { Subscription, type Observable } from "rxjs"; + +import { + OperationTimeoutError, + operationOptions, + validateTimeoutSeconds, + withTimeoutSignal, + type OperationTimeout, +} from "./internal/timeout"; +import { BEACON_TASK_MAX_PAYLOAD_BYTES, BEACON_TASK_RPC_MESSAGE_BYTES } from "./messageBudget"; +import { BeaconTask, type Event } from "./pb/clientpb/client"; +import type { SliverRPCClient } from "./pb/rpcpb/services"; + +/** Generated protobuf message codecs implement this interface. */ +export interface BeaconTaskDecoder { + decode(bytes: Buffer): T; +} + +/** A queued task can be awaited repeatedly or reattached using its persisted ID. */ +export interface BeaconTaskHandle { + readonly id: string; + /** Canceling this local wait does not cancel the server's task. */ + wait(options?: OperationTimeout): Promise; + status(options?: OperationTimeout): Promise; + /** Explicitly ask the server to cancel this task. */ + cancel(options?: OperationTimeout): Promise; +} + +export class BeaconTaskStateError extends Error { + constructor(readonly taskId: string, readonly state: "canceled" | "failed") { + super(`Beacon task ${taskId} ended in state ${state}`); + this.name = "BeaconTaskStateError"; + } +} + +export class BeaconTaskResponseError extends Error { + constructor(readonly taskId: string, message: string) { + super(message); + this.name = "BeaconTaskResponseError"; + } +} + +export class BeaconTaskIdentityError extends Error { + constructor(readonly taskId: string) { + super(`Beacon task response did not match task ${taskId}`); + this.name = "BeaconTaskIdentityError"; + } +} + +export class BeaconTaskTimeoutError extends Error { + constructor(readonly taskId: string, readonly timeoutSeconds: number) { + super(`Timeout waiting for beacon task result: ${taskId}`); + this.name = "BeaconTaskTimeoutError"; + } +} + +/** @internal Dependencies are supplied by SliverClient to retain channel ownership. */ +export interface BeaconTaskHandleDependencies { + taskId: string; + beaconId?: string; + decoder: BeaconTaskDecoder; + /** Resolve on every poll so a handle can survive a client reconnection. */ + getRpc: () => Pick | undefined; + cancel: (taskId: string, options: OperationTimeout) => Promise; + events: Observable | (() => Observable); + defaultTimeoutSeconds?: number; + pollIntervalMilliseconds?: number; +} + +/** @internal Construct a typed handle without exposing a method-name dispatcher. */ +export function createBeaconTaskHandle(dependencies: BeaconTaskHandleDependencies): BeaconTaskHandle { + const { taskId, beaconId, decoder, getRpc } = dependencies; + validateIdentifier(taskId, "Beacon task id"); + if (beaconId !== undefined) validateIdentifier(beaconId, "Beacon id"); + if (!decoder || typeof decoder.decode !== "function") throw new TypeError("Beacon task decoder is required"); + const defaultTimeout = validateTimeoutSeconds(dependencies.defaultTimeoutSeconds ?? 30); + const pollInterval = dependencies.pollIntervalMilliseconds ?? 1_000; + if (!Number.isSafeInteger(pollInterval) || pollInterval < 1 || pollInterval > 2_147_483_647) { + throw new RangeError("Beacon task poll interval must be a positive timer-safe integer"); + } + + const optionsFor = (options?: OperationTimeout) => operationOptions( + options === undefined ? defaultTimeout : options, + defaultTimeout, + ); + + const validateIdentity = (task: BeaconTask): BeaconTask => { + if (task.ID !== taskId || (beaconId !== undefined && task.BeaconID !== beaconId)) { + throw new BeaconTaskIdentityError(taskId); + } + return task; + }; + + const fetch = async (signal: AbortSignal): Promise => { + const rpc = getRpc(); + if (!rpc) throw new Error("SliverClient is not connected"); + const task = await rpc.getBeaconTaskContent({ ID: taskId }, { signal }); + return validateIdentity(task); + }; + + const decodeCompleted = (task: BeaconTask): T => { + if (task.Response.length === 0) { + throw new BeaconTaskResponseError(taskId, `Beacon task ${taskId} returned an empty response`); + } + if (task.Response.length > BEACON_TASK_MAX_PAYLOAD_BYTES) { + throw new BeaconTaskResponseError(taskId, `Beacon task ${taskId} response exceeds the bounded task limit`); + } + // Protobuf decoders may return slices of their input. Each wait owns a copy: + // neither clearing a failed decode nor a consumer mutating its result can + // damage a concurrent wait or the RPC's shared response bytes. + const owned = Buffer.from(task.Response); + let decoded: T; + try { + decoded = decoder.decode(owned); + } catch { + owned.fill(0); + throw new BeaconTaskResponseError(taskId, `Beacon task ${taskId} returned an invalid response`); + } + if (hasTargetError(decoded)) { + owned.fill(0); + throw new BeaconTaskResponseError(taskId, `Beacon task ${taskId} was rejected by the target`); + } + return decoded; + }; + + const wait = async (options?: OperationTimeout): Promise => { + try { + return await withTimeoutSignal(optionsFor(options), (signal) => new Promise((resolve, reject) => { + let settled = false; + let refreshing = false; + let refreshPending = false; + let pollTimer: ReturnType | undefined; + const subscription = new Subscription(); + const finish = (settle: () => void): void => { + if (settled) return; + settled = true; + if (pollTimer !== undefined) clearInterval(pollTimer); + signal.removeEventListener("abort", onAbort); + subscription.unsubscribe(); + settle(); + }; + const onAbort = (): void => finish(() => reject(signal.reason)); + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) { + onAbort(); + return; + } + + const refresh = async (): Promise => { + if (settled || refreshing) return; + refreshing = true; + try { + while (refreshPending && !settled) { + refreshPending = false; + // A disconnected client may reconnect before this wait's deadline. + const rpc = getRpc(); + if (!rpc) continue; + let task: BeaconTask; + try { + task = await rpc.getBeaconTaskContent({ ID: taskId }, { signal }); + } catch (error) { + if (settled) return; + if (isTransientRpcError(error)) continue; + throw error; + } + if (settled) return; + validateIdentity(task); + const state = task.State.toLowerCase(); + if (state === "canceled" || state === "failed") { + throw new BeaconTaskStateError(taskId, state); + } + if (state === "completed" || (state === "" && task.Response.length > 0)) { + const decoded = decodeCompleted(task); + finish(() => resolve(decoded)); + } + } + } catch (error) { + finish(() => reject(error)); + } finally { + refreshing = false; + if (refreshPending && !settled) void refresh(); + } + }; + const requestRefresh = (): void => { + if (settled) return; + refreshPending = true; + if (!refreshing) void refresh(); + }; + + try { + const events = typeof dependencies.events === "function" ? dependencies.events() : dependencies.events; + subscription.add(events.subscribe({ + next: (event) => { + if (event.EventType !== "beacon-taskresult" || event.Data.length > BEACON_TASK_RPC_MESSAGE_BYTES) return; + try { + const task = BeaconTask.decode(event.Data); + if (task.ID === taskId && (beaconId === undefined || task.BeaconID === beaconId)) requestRefresh(); + } catch { + // Unrelated malformed events must not fail a valid pending task. + } + }, + // Polling remains authoritative if the optional event stream fails. + error: () => undefined, + })); + } catch { + // Events optimize latency; a reconnecting stream is not required. + } + if (settled) return; + pollTimer = setInterval(requestRefresh, pollInterval); + requestRefresh(); + })); + } catch (error) { + if (error instanceof OperationTimeoutError) throw new BeaconTaskTimeoutError(taskId, error.timeoutSeconds); + throw error; + } + }; + + return Object.freeze({ + id: taskId, + wait, + status: (options?: OperationTimeout) => withTimeoutSignal(optionsFor(options), fetch), + cancel: async (options?: OperationTimeout) => { + const normalized = optionsFor(options); + return withTimeoutSignal(normalized, async (signal) => validateIdentity( + await dependencies.cancel(taskId, { ...normalized, signal }), + )); + }, + }); +} + +function validateIdentifier(value: string, label: string): void { + if (typeof value !== "string" || value.length === 0 || value.length > 256 || value.trim() !== value) { + throw new TypeError(`${label} must be a non-empty identifier of at most 256 characters`); + } +} + +function hasTargetError(value: unknown): boolean { + if (value === null || typeof value !== "object" || !("Response" in value)) return false; + const response = value.Response; + return response !== null && typeof response === "object" && "Err" in response && Boolean(response.Err); +} + +function isTransientRpcError(error: unknown): boolean { + if (!error || typeof error !== "object" || !("code" in error)) return false; + // CANCELLED / UNAVAILABLE may occur when an old channel closes on disconnect. + return error.code === 1 || error.code === 14; +} diff --git a/src/client.ts b/src/client.ts index 413ed60..e09334d 100644 --- a/src/client.ts +++ b/src/client.ts @@ -24,8 +24,10 @@ import { requestStopReversePortForward, } from "./internal/reversePortForward"; import { SocksTunnelLeases, startSocks5Proxy, type ManagedSocks5Proxy } from "./internal/socks5Proxy"; -import { timeoutSecondsToNanoseconds, validateTimeoutSeconds, withTimeoutSignal } from "./internal/timeout"; -import { TunnelManager } from "./internal/tunnelManager"; +import { operationOptions, operationTimeoutSeconds, OperationAbortedError, OperationTimeoutError, timeoutSecondsToNanoseconds, withTimeoutSignal } from "./internal/timeout"; +import { TunnelManager, type TunnelTransportState } from "./internal/tunnelManager"; +import type { OperationTimeout } from "./internal/timeout"; +import { createBeaconTaskHandle, type BeaconTaskDecoder, type BeaconTaskHandle } from "./beaconTasks"; import { BEACON_TASK_MAX_PAYLOAD_BYTES, RPC_MESSAGE_BUDGETS, @@ -340,7 +342,7 @@ class BaseCommands { protected readonly responseArtifactRpc: SliverRPCClient = artifactRpc, ) {} - protected request(timeoutSeconds: number): CommonRequest { + protected request(timeoutSeconds: OperationTimeout): CommonRequest { return { Async: false, Timeout: timeoutSecondsToNanoseconds(timeoutSeconds), @@ -349,29 +351,29 @@ class BaseCommands { }; } - protected async unary(timeoutSeconds: number, fn: (signal: AbortSignal) => Promise): Promise { + protected async unary(timeoutSeconds: OperationTimeout, fn: (signal: AbortSignal) => Promise): Promise { return withTimeoutSignal(timeoutSeconds, fn); } - ping(nonce: number, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + ping(nonce: number, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.ping({ Nonce: nonce, Request: this.request(timeoutSeconds) }, { signal }), ); } - ps(fullInfo = false, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + ps(fullInfo = false, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.ps({ FullInfo: fullInfo, Request: this.request(timeoutSeconds) }, { signal }), ); } - ls(path = ".", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + ls(path = ".", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.ls({ Path: path, Request: this.request(timeoutSeconds) }, { signal }), ); } - download(path: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + download(path: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return this.unary(timeoutSeconds, async (signal) => { const download = await this.artifactRpc.download( { Path: path, Request: this.request(timeoutSeconds) }, @@ -395,7 +397,7 @@ class BaseCommands { }); } - upload(path: string, data: Buffer, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + upload(path: string, data: Buffer, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { if (!Buffer.isBuffer(data)) throw new Error("Upload data must be bytes"); assertBoundedArtifact(data, "Upload"); const ownedData = Buffer.from(data); @@ -417,51 +419,51 @@ class BaseCommands { }).finally(() => ownedData.fill(0)); } - terminate(pid: number, force = false, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + terminate(pid: number, force = false, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.terminate({ Pid: pid, Force: force, Request: this.request(timeoutSeconds) }, { signal }), ); } - ifconfig(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + ifconfig(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.ifconfig({ Request: this.request(timeoutSeconds) }, { signal }), ); } - netstat(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + netstat(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.netstat({ Request: this.request(timeoutSeconds) }, { signal }), ); } - cd(path: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + cd(path: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.cd({ Path: path, Request: this.request(timeoutSeconds) }, { signal })); } - pwd(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + pwd(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.pwd({ Request: this.request(timeoutSeconds) }, { signal })); } - rm(path: string, recursive = false, force = false, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + rm(path: string, recursive = false, force = false, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.rm({ Path: path, Recursive: recursive, Force: force, Request: this.request(timeoutSeconds) }, { signal }), ); } - mkdir(path: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + mkdir(path: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.mkdir({ Path: path, Request: this.request(timeoutSeconds) }, { signal }), ); } - processDump(pid: number, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + processDump(pid: number, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.artifactRpc.processDump({ Pid: pid, Request: this.request(timeoutSeconds) }, { signal }), ); } - runAs(userName: string, processName: string, args: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + runAs(userName: string, processName: string, args: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.runAs( { Username: userName, ProcessName: processName, Args: args, Request: this.request(timeoutSeconds) }, @@ -470,17 +472,17 @@ class BaseCommands { ); } - impersonate(userName: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + impersonate(userName: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.impersonate({ Username: userName, Request: this.request(timeoutSeconds) }, { signal }), ); } - revToSelf(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + revToSelf(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.revToSelf({ Request: this.request(timeoutSeconds) }, { signal })); } - getSystem(hostingProcess: string, config: ImplantConfig, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + getSystem(hostingProcess: string, config: ImplantConfig, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.getSystem( { HostingProcess: hostingProcess, Config: config, Request: this.request(timeoutSeconds) }, @@ -500,7 +502,7 @@ class BaseCommands { shellcode: Buffer, encoder = "", rwxPages = false, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return this.unary(timeoutSeconds, (signal) => this.artifactRpc.task( @@ -522,7 +524,7 @@ class BaseCommands { lport: number, encoder = "", iterations = 0, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return this.unary(timeoutSeconds, (signal) => this.rpc.msf( @@ -539,7 +541,7 @@ class BaseCommands { lport: number, encoder = "", iterations = 0, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return this.unary(timeoutSeconds, (signal) => this.rpc.msfRemote( @@ -549,7 +551,7 @@ class BaseCommands { ); } - executeAssembly(assembly: Buffer, args: string[] = [], process = "", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + executeAssembly(assembly: Buffer, args: string[] = [], process = "", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.artifactRpc.executeAssembly( { Assembly: assembly, Arguments: args, Process: process, Request: this.request(timeoutSeconds) }, @@ -558,13 +560,13 @@ class BaseCommands { ); } - migrate(pid: number, config: ImplantConfig, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + migrate(pid: number, config: ImplantConfig, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.rpc.migrate({ Pid: pid, Config: config, Request: this.request(timeoutSeconds) }, { signal }), ); } - execute(exe: string, args: string[] = [], output = true, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + execute(exe: string, args: string[] = [], output = true, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { const request = this.request(timeoutSeconds); const rpc = !request.Async && output ? this.responseArtifactRpc : this.rpc; return this.unary(timeoutSeconds, (signal) => @@ -577,7 +579,7 @@ class BaseCommands { processName: string, args: string[] = [], entryPoint: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return this.unary(timeoutSeconds, (signal) => this.artifactRpc.sideload( @@ -592,7 +594,7 @@ class BaseCommands { entrypoint: string, processName: string, args: string[] = [], - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return this.unary(timeoutSeconds, (signal) => this.artifactRpc.spawnDll( @@ -602,7 +604,7 @@ class BaseCommands { ); } - screenshot(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + screenshot(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.unary(timeoutSeconds, (signal) => this.artifactRpc.screenshot({ Request: this.request(timeoutSeconds) }, { signal }), ); @@ -651,7 +653,7 @@ export class InteractiveBeacon extends BaseCommands { ?? (usingSeparateArtifactClient ? taskResultsOrBeaconId : taskContentRpcOrTaskResultsOrBeaconId) as string; } - protected request(timeoutSeconds: number): CommonRequest { + protected request(timeoutSeconds: OperationTimeout): CommonRequest { return { Async: true, Timeout: timeoutSecondsToNanoseconds(timeoutSeconds), @@ -660,7 +662,23 @@ export class InteractiveBeacon extends BaseCommands { }; } - async lsTask(path = ".", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + /** Attach to a queued task without queueing another remote operation. */ + beaconTask( + taskId: string, + decoder: BeaconTaskDecoder, + options: { readonly timeoutSeconds?: number } = {}, + ): BeaconTaskHandle { + return createBeaconTaskHandle({ + taskId, beaconId: this.beaconId, decoder, + getRpc: () => this.taskContentRpc, + events: this.taskResult$, + defaultTimeoutSeconds: options.timeoutSeconds, + cancel: (id, options) => withTimeoutSignal(options, (signal) => + this.artifactRpc.cancelBeaconTask({ ID: id }, { signal })), + }); + } + + async lsTask(path = ".", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { const lsTask = await super.ls(path, timeoutSeconds); if (lsTask.Response?.Err) { throw new Error(lsTask.Response.Err); @@ -670,6 +688,7 @@ export class InteractiveBeacon extends BaseCommands { throw new Error("Missing beacon task id"); } return { + ...this.beaconTask(taskId, Ls, { timeoutSeconds: operationTimeoutSeconds(timeoutSeconds) }), id: taskId, wait: async (waitTimeoutSeconds = timeoutSeconds) => { const taskContent = await waitForBeaconTaskContent( @@ -683,12 +702,21 @@ export class InteractiveBeacon extends BaseCommands { }; } - async ls(path = ".", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { - const task = await this.lsTask(path, timeoutSeconds); - return task.wait(timeoutSeconds); + async ls(path = ".", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { + return withTimeoutSignal(timeoutSeconds, async (signal) => { + const operation = { timeoutSeconds: operationTimeoutSeconds(timeoutSeconds), signal }; + const task = await this.lsTask(path, operation); + return task.wait(operation); + }); + } + + override async download(path: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { + return withTimeoutSignal(timeoutSeconds, (signal) => this.downloadWithOptions(path, { + timeoutSeconds: operationTimeoutSeconds(timeoutSeconds), signal, + })); } - override async download(path: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + private async downloadWithOptions(path: string, timeoutSeconds: OperationTimeout): Promise { const queued = await this.unary(timeoutSeconds, (signal) => this.artifactRpc.download( { @@ -775,7 +803,7 @@ export class InteractiveSession extends BaseCommands { ?? (usingSeparateArtifactClient ? tunnelsOrSessionId : responseArtifactRpcOrTunnelsOrSessionId) as string; } - protected request(timeoutSeconds: number): CommonRequest { + protected request(timeoutSeconds: OperationTimeout): CommonRequest { return { Async: false, Timeout: timeoutSecondsToNanoseconds(timeoutSeconds), @@ -784,17 +812,25 @@ export class InteractiveSession extends BaseCommands { }; } - async shell(path: string, pty = true, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { - const tunnel = await this.unary(timeoutSeconds, (signal) => - this.rpc.createTunnel({ SessionID: this.sessionId }, { signal }), - ); + async shell(path: string, pty = true, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { + return withTimeoutSignal(timeoutSeconds, (signal) => this.openShell(path, pty, { + timeoutSeconds: operationTimeoutSeconds(timeoutSeconds), signal, + }), (handle) => handle.close()); + } + + private async openShell(path: string, pty: boolean, timeoutSeconds: OperationTimeout): Promise { + // Keep observing late creation results after the outer deadline settles so + // their owned tunnel IDs can still be retired and closed. + const signal = operationOptions(timeoutSeconds).signal!; + const tunnel = await this.rpc.createTunnel({ SessionID: this.sessionId }, { signal }); if (!tunnel?.TunnelID) { throw new Error("Failed to create tunnel"); } + if (tunnel.SessionID && tunnel.SessionID !== this.sessionId) throw new Error("Mismatched shell session"); const tunnelId = tunnel.TunnelID; // Subscribe first so we don't miss early data. - const stdout$ = this.tunnels.subscribe(tunnelId).pipe( + const stdout$ = this.tunnels.subscribe(tunnelId, { exclusive: true }).pipe( filter((msg) => msg.TunnelID === tunnelId), filter((msg) => msg.Data.length > 0), map((msg) => msg.Data), @@ -812,11 +848,7 @@ export class InteractiveSession extends BaseCommands { try { // Bind tunnel to the tunnel stream. - await this.tunnels.send({ - TunnelID: tunnelId, - SessionID: this.sessionId, - Data: Buffer.alloc(0), - }); + await this.tunnels.bind(tunnelId, this.sessionId, operationOptions(timeoutSeconds).signal); // Ask the implant to open a shell on the bound tunnel. const shell = await this.unary(timeoutSeconds, (signal) => @@ -834,11 +866,12 @@ export class InteractiveSession extends BaseCommands { ), ); if (shell.Response?.Err) throw new Error("Shell was rejected by the target"); - } catch { + } catch (error) { this.tunnels.cancelTunnel(tunnelId); await this.unary(DEFAULT_TIMEOUT_SECONDS, (signal) => this.rpc.closeTunnel({ TunnelID: tunnelId, SessionID: this.sessionId }, { signal }), ).catch(() => undefined); + if (error instanceof OperationAbortedError || error instanceof OperationTimeoutError) throw error; throw new Error("Unable to start shell session"); } @@ -905,6 +938,11 @@ export class SliverClient { private eventsAbort: AbortController | null = null; private tunnels: TunnelManager | null = null; + private tunnelStateSubscription?: { unsubscribe(): void }; + private readonly tunnelStreamStateSubject = new BehaviorSubject({ + status: "stopped", generation: 0, + }); + readonly tunnelStreamState$ = this.tunnelStreamStateSubject.asObservable(); private lifecycleTail: Promise = Promise.resolve(); private readonly managedPortForwards = new Set(); private readonly managedSocks5Proxies = new Set(); @@ -970,7 +1008,7 @@ export class SliverClient { return this.clientFor("task-content"); } - private sessionRequest(sessionId: string, timeoutSeconds: number): CommonRequest { + private sessionRequest(sessionId: string, timeoutSeconds: OperationTimeout): CommonRequest { return { Async: false, Timeout: timeoutSecondsToNanoseconds(timeoutSeconds), @@ -979,7 +1017,7 @@ export class SliverClient { }; } - private beaconRequest(beaconId: string, timeoutSeconds: number): CommonRequest { + private beaconRequest(beaconId: string, timeoutSeconds: OperationTimeout): CommonRequest { return { Async: true, Timeout: timeoutSecondsToNanoseconds(timeoutSeconds), @@ -988,12 +1026,12 @@ export class SliverClient { }; } - private m4SessionRequest(sessionId: string, timeoutSeconds: number): CommonRequest { + private m4SessionRequest(sessionId: string, timeoutSeconds: OperationTimeout): CommonRequest { assertNonEmptyString(sessionId, "Session id"); return this.sessionRequest(sessionId, timeoutSeconds); } - private m4BeaconRequest(beaconId: string, timeoutSeconds: number): CommonRequest { + private m4BeaconRequest(beaconId: string, timeoutSeconds: OperationTimeout): CommonRequest { assertNonEmptyString(beaconId, "Beacon id"); return this.beaconRequest(beaconId, timeoutSeconds); } @@ -1008,15 +1046,25 @@ export class SliverClient { return result; } - connect(): Promise { - return this.enqueueLifecycle(() => this.connectUnlocked()); + connect(options: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { + let owner: AbortController | undefined; + return withTimeoutSignal(options, (signal) => this.enqueueLifecycle(() => this.connectUnlocked({ + timeoutSeconds: operationTimeoutSeconds(options), signal, + }, (connection) => { owner = connection; })), () => this.enqueueLifecycle(async () => { + if (owner && this.eventsAbort === owner) await this.disconnectUnlocked(); + })); } - private async connectUnlocked(): Promise { + private async connectUnlocked(options: OperationTimeout, onCreated: (owner: AbortController) => void): Promise { + operationOptions(options).signal?.throwIfAborted(); if (this.rpcClients.control) return this; this.eventsAbort = new AbortController(); + onCreated(this.eventsAbort); this.tunnels = new TunnelManager(); + this.tunnelStateSubscription = this.tunnels.transportState$.subscribe((state) => { + this.tunnelStreamStateSubject.next(state); + }); try { const credentials = createSliverRpcCredentials(this.config); @@ -1032,11 +1080,12 @@ export class SliverClient { } // Ensure auth and connectivity are working before we start streams. - await this.getVersion(); + await this.getVersion(options); this.tunnels.start(this.clientFor("tunnel-stream")); this.startEventsStream(); - await this.reconcileManagedReversePortForwards(); + await this.reconcileManagedReversePortForwards(operationOptions(options)); + operationOptions(options).signal?.throwIfAborted(); return this; } catch (error) { await this.disconnectUnlocked(); @@ -1062,6 +1111,8 @@ export class SliverClient { await this.tunnels?.stop(); this.tunnels = null; + this.tunnelStateSubscription?.unsubscribe(); + this.tunnelStateSubscription = undefined; for (const domain of RPC_MESSAGE_DOMAINS) { delete this.rpcClients[domain]; @@ -1124,35 +1175,35 @@ export class SliverClient { // --- Convenience APIs (typed, promise-based) --- - getVersion(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getVersion(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getVersion(this.empty, { signal })); } - getOperators(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getOperators(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getOperators(this.empty, { signal })); } - getSessions(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getSessions(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.inventoryRpc.getSessions(this.empty, { signal })); } - getBeacons(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getBeacons(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.inventoryRpc.getBeacons(this.empty, { signal })); } - renameSession(sessionId: string, name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + renameSession(sessionId: string, name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.rename({ SessionID: sessionId, BeaconID: "", Name: name }, { signal }), ); } - renameBeacon(beaconId: string, name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + renameBeacon(beaconId: string, name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.rename({ SessionID: "", BeaconID: beaconId, Name: name }, { signal }), ); } - pingSession(sessionId: string, nonce: number, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + pingSession(sessionId: string, nonce: number, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.ping( { Nonce: nonce, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1161,7 +1212,7 @@ export class SliverClient { ); } - pingBeacon(beaconId: string, nonce: number, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + pingBeacon(beaconId: string, nonce: number, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.ping( { Nonce: nonce, Request: this.beaconRequest(beaconId, timeoutSeconds) }, @@ -1170,7 +1221,7 @@ export class SliverClient { ); } - getEnvSession(sessionId: string, name = "", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getEnvSession(sessionId: string, name = "", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getEnv( { Name: name, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1179,7 +1230,7 @@ export class SliverClient { ); } - getEnvBeacon(beaconId: string, name = "", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getEnvBeacon(beaconId: string, name = "", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getEnv( { Name: name, Request: this.beaconRequest(beaconId, timeoutSeconds) }, @@ -1192,7 +1243,7 @@ export class SliverClient { sessionId: string, key: string, value: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.setEnv( @@ -1209,7 +1260,7 @@ export class SliverClient { beaconId: string, key: string, value: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.setEnv( @@ -1222,7 +1273,7 @@ export class SliverClient { ); } - unsetEnvSession(sessionId: string, name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + unsetEnvSession(sessionId: string, name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.unsetEnv( { Name: name, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1231,7 +1282,7 @@ export class SliverClient { ); } - unsetEnvBeacon(beaconId: string, name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + unsetEnvBeacon(beaconId: string, name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.unsetEnv( { Name: name, Request: this.beaconRequest(beaconId, timeoutSeconds) }, @@ -1242,7 +1293,8 @@ export class SliverClient { // --- Explicit M2 session workbench APIs --- - currentTokenOwnerSession(sessionId: string, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + currentTokenOwnerSession(sessionId: string, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.currentTokenOwner( { Request: this.m4SessionRequest(sessionId, timeoutSeconds) }, @@ -1251,7 +1303,8 @@ export class SliverClient { ); } - currentTokenOwnerBeacon(beaconId: string, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + currentTokenOwnerBeacon(beaconId: string, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.currentTokenOwner( { Request: this.m4BeaconRequest(beaconId, timeoutSeconds) }, @@ -1260,16 +1313,16 @@ export class SliverClient { ); } - listEnvSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + listEnvSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return this.getEnvSession(sessionId, "", timeoutSeconds); } - revealEnvSession(sessionId: string, exactName: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + revealEnvSession(sessionId: string, exactName: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { assertNonEmptyString(exactName, "Environment variable name"); return this.getEnvSession(sessionId, exactName, timeoutSeconds); } - ifconfigSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + ifconfigSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.ifconfig( { Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1281,7 +1334,7 @@ export class SliverClient { netstatSession( sessionId: string, options: SessionNetstatOptions, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.netstat( @@ -1298,7 +1351,7 @@ export class SliverClient { ); } - pwdSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + pwdSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.pwd( { Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1307,7 +1360,7 @@ export class SliverClient { ); } - cdSession(sessionId: string, path: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + cdSession(sessionId: string, path: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.cd( { Path: path, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1316,7 +1369,7 @@ export class SliverClient { ); } - lsSession(sessionId: string, path = ".", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + lsSession(sessionId: string, path = ".", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.ls( { Path: path, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1329,7 +1382,7 @@ export class SliverClient { sessionId: string, path: string, options: SessionDownloadFileOptions = {}, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { const maxBytes = boundedArtifactByteCount(options.maxBytes); const start = boundedNonNegativeInteger(options.start ?? 0, "Download start"); @@ -1373,7 +1426,7 @@ export class SliverClient { path: string, data: Buffer, options: SessionUploadOptions = {}, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { assertBoundedArtifact(data, "Upload"); const ownedData = Buffer.from(data); @@ -1410,7 +1463,7 @@ export class SliverClient { path: string, searchPattern: string, options: SessionGrepOptions = {}, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { const linesBefore = boundedNonNegativeInt32(options.linesBefore ?? 0, "Grep lines before"); const linesAfter = boundedNonNegativeInt32(options.linesAfter ?? 0, "Grep lines after"); @@ -1429,7 +1482,7 @@ export class SliverClient { ); } - cpSession(sessionId: string, source: string, destination: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + cpSession(sessionId: string, source: string, destination: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.cp( { Src: source, Dst: destination, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1438,7 +1491,7 @@ export class SliverClient { ); } - mvSession(sessionId: string, source: string, destination: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + mvSession(sessionId: string, source: string, destination: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.mv( { Src: source, Dst: destination, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1447,7 +1500,7 @@ export class SliverClient { ); } - mkdirSession(sessionId: string, path: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + mkdirSession(sessionId: string, path: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.mkdir( { Path: path, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1461,7 +1514,7 @@ export class SliverClient { path: string, recursive = false, force = false, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.rm( @@ -1471,7 +1524,7 @@ export class SliverClient { ); } - mountsSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + mountsSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.mount( { Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1480,7 +1533,7 @@ export class SliverClient { ); } - memfilesListSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + memfilesListSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.memfilesList( { Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1489,7 +1542,7 @@ export class SliverClient { ); } - memfilesAddSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + memfilesAddSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.memfilesAdd( { Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1498,7 +1551,7 @@ export class SliverClient { ); } - memfilesRmSession(sessionId: string, fd: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + memfilesRmSession(sessionId: string, fd: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.memfilesRm( { Fd: fd, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1512,7 +1565,7 @@ export class SliverClient { path: string, fileMode: string, recursive = false, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.chmod( @@ -1528,7 +1581,7 @@ export class SliverClient { uid: string, gid: string, recursive = false, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.chown( @@ -1543,7 +1596,7 @@ export class SliverClient { path: string, accessTime: string, modificationTime: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.chtimes( @@ -1553,7 +1606,7 @@ export class SliverClient { ); } - psSession(sessionId: string, fullInfo = false, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + psSession(sessionId: string, fullInfo = false, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.ps( { FullInfo: fullInfo, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1566,7 +1619,7 @@ export class SliverClient { sessionId: string, pid: number, force = false, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.terminate( @@ -1580,9 +1633,10 @@ export class SliverClient { sessionId: string, pid: number, dumpTimeoutSeconds = 60, - timeoutSeconds = dumpTimeoutSeconds + DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = dumpTimeoutSeconds + DEFAULT_TIMEOUT_SECONDS, ) { const validatedDumpTimeout = boundedPositiveInt32(dumpTimeoutSeconds, "Process dump timeout"); + timeoutSeconds = operationOptions(timeoutSeconds, validatedDumpTimeout + DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, async (signal) => { const response = await this.workbenchArtifactRpc.processDump( { @@ -1603,7 +1657,7 @@ export class SliverClient { }); } - screenshotSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + screenshotSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, async (signal) => { const response = await this.workbenchArtifactRpc.screenshot( { Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1620,7 +1674,7 @@ export class SliverClient { }); } - servicesSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + servicesSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.services( { Hostname: "", Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -1629,7 +1683,7 @@ export class SliverClient { ); } - serviceDetailSession(sessionId: string, serviceName: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + serviceDetailSession(sessionId: string, serviceName: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.serviceDetail( { @@ -1641,7 +1695,7 @@ export class SliverClient { ); } - startServiceSession(sessionId: string, serviceName: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + startServiceSession(sessionId: string, serviceName: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.startServiceByName( { @@ -1653,7 +1707,7 @@ export class SliverClient { ); } - stopServiceSession(sessionId: string, serviceName: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + stopServiceSession(sessionId: string, serviceName: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.stopService( { @@ -1670,7 +1724,7 @@ export class SliverClient { hive: string, path: string, key: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.registryRead( @@ -1684,7 +1738,7 @@ export class SliverClient { sessionId: string, hive: string, path: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.registryListSubKeys( @@ -1698,7 +1752,7 @@ export class SliverClient { sessionId: string, hive: string, path: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.registryListValues( @@ -1713,7 +1767,7 @@ export class SliverClient { rootHive: string, requestedHive: string, maxBytes = WORKBENCH_ARTIFACT_MAX_PAYLOAD_BYTES, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { const boundedMaxBytes = boundedArtifactByteCount(maxBytes); return withTimeoutSignal(timeoutSeconds, async (signal) => { @@ -1738,7 +1792,7 @@ export class SliverClient { path: string, key: string, value: SessionRegistryWriteValue, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { const fields = registryWriteFields(value); const dispose = () => fields.ByteValue.fill(0); @@ -1766,7 +1820,7 @@ export class SliverClient { hive: string, path: string, key: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.registryCreateKey( @@ -1781,7 +1835,7 @@ export class SliverClient { hive: string, path: string, key: string, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.registryDeleteKey( @@ -1793,7 +1847,7 @@ export class SliverClient { // --- Explicit M4 execution and privilege APIs --- - private executeTarget(request: CommonRequest, options: ExecuteOptions, timeoutSeconds: number) { + private executeTarget(request: CommonRequest, options: ExecuteOptions, timeoutSeconds: OperationTimeout) { assertNonEmptyString(options.path, "Executable path"); const args = m4StringArray(options.args); const background = options.background ?? false; @@ -1853,20 +1907,23 @@ export class SliverClient { executeSession( sessionId: string, options: ExecuteOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.executeTarget(this.m4SessionRequest(sessionId, timeoutSeconds), options, timeoutSeconds); } executeBeacon( beaconId: string, options: ExecuteOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.executeTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), options, timeoutSeconds); } - executeChildrenSession(sessionId: string, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + executeChildrenSession(sessionId: string, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.executeChildren( { Request: this.m4SessionRequest(sessionId, timeoutSeconds) }, @@ -1875,7 +1932,8 @@ export class SliverClient { ); } - executeChildrenBeacon(beaconId: string, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + executeChildrenBeacon(beaconId: string, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.executeChildren( { Request: this.m4BeaconRequest(beaconId, timeoutSeconds) }, @@ -1888,7 +1946,7 @@ export class SliverClient { request: CommonRequest, assembly: Buffer, options: ExecuteAssemblyOptions, - timeoutSeconds: number, + timeoutSeconds: OperationTimeout, ) { const isDll = options.isDll ?? false; if (isDll && (!options.className?.trim() || !options.method?.trim())) { @@ -1931,8 +1989,9 @@ export class SliverClient { sessionId: string, assembly: Buffer, options: ExecuteAssemblyOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.executeAssemblyTarget( this.m4SessionRequest(sessionId, timeoutSeconds), assembly, @@ -1945,8 +2004,9 @@ export class SliverClient { beaconId: string, assembly: Buffer, options: ExecuteAssemblyOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.executeAssemblyTarget( this.m4BeaconRequest(beaconId, timeoutSeconds), assembly, @@ -1959,7 +2019,7 @@ export class SliverClient { request: CommonRequest, shellcode: Buffer, options: ExecuteShellcodeOptions, - timeoutSeconds: number, + timeoutSeconds: OperationTimeout, ) { return withTimeoutSignal(timeoutSeconds, async (signal) => { const ownedShellcode = m4ArtifactCopy(shellcode, "Shellcode"); @@ -1984,8 +2044,9 @@ export class SliverClient { sessionId: string, shellcode: Buffer, options: ExecuteShellcodeOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.executeShellcodeTarget( this.m4SessionRequest(sessionId, timeoutSeconds), shellcode, @@ -1998,8 +2059,9 @@ export class SliverClient { beaconId: string, shellcode: Buffer, options: ExecuteShellcodeOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.executeShellcodeTarget( this.m4BeaconRequest(beaconId, timeoutSeconds), shellcode, @@ -2012,7 +2074,7 @@ export class SliverClient { request: CommonRequest, data: Buffer, options: SideloadOptions, - timeoutSeconds: number, + timeoutSeconds: OperationTimeout, ) { return withTimeoutSignal(timeoutSeconds, async (signal) => { const ownedData = m4ArtifactCopy(data, "Sideload library"); @@ -2042,8 +2104,9 @@ export class SliverClient { sessionId: string, data: Buffer, options: SideloadOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.sideloadTarget(this.m4SessionRequest(sessionId, timeoutSeconds), data, options, timeoutSeconds); } @@ -2051,8 +2114,9 @@ export class SliverClient { beaconId: string, data: Buffer, options: SideloadOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.sideloadTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), data, options, timeoutSeconds); } @@ -2060,7 +2124,7 @@ export class SliverClient { request: CommonRequest, data: Buffer, options: SpawnDllOptions, - timeoutSeconds: number, + timeoutSeconds: OperationTimeout, ) { return withTimeoutSignal(timeoutSeconds, async (signal) => { const ownedData = m4ArtifactCopy(data, "Reflective DLL"); @@ -2088,8 +2152,9 @@ export class SliverClient { sessionId: string, data: Buffer, options: SpawnDllOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.spawnDllTarget(this.m4SessionRequest(sessionId, timeoutSeconds), data, options, timeoutSeconds); } @@ -2097,16 +2162,18 @@ export class SliverClient { beaconId: string, data: Buffer, options: SpawnDllOptions = {}, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.spawnDllTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), data, options, timeoutSeconds); } - getShellcodeEncoderMap(timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + getShellcodeEncoderMap(timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.shellcodeEncoderMap({}, { signal })); } - private migrateTarget(request: CommonRequest, options: MigrateOptions, timeoutSeconds: number) { + private migrateTarget(request: CommonRequest, options: MigrateOptions, timeoutSeconds: OperationTimeout) { const pid = boundedUint32(options.pid ?? 0, "Migration process id"); const processName = options.processName?.trim() ?? ""; if (pid === 0 && processName === "") { @@ -2135,20 +2202,22 @@ export class SliverClient { migrateSession( sessionId: string, options: MigrateOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.migrateTarget(this.m4SessionRequest(sessionId, timeoutSeconds), options, timeoutSeconds); } migrateBeacon( beaconId: string, options: MigrateOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.migrateTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), options, timeoutSeconds); } - private msfTarget(request: CommonRequest, options: MsfOptions, timeoutSeconds: number) { + private msfTarget(request: CommonRequest, options: MsfOptions, timeoutSeconds: OperationTimeout) { assertNonEmptyString(options.lhost, "Metasploit listen host"); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.msf( @@ -2165,15 +2234,17 @@ export class SliverClient { ); } - msfSession(sessionId: string, options: MsfOptions, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + msfSession(sessionId: string, options: MsfOptions, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.msfTarget(this.m4SessionRequest(sessionId, timeoutSeconds), options, timeoutSeconds); } - msfBeacon(beaconId: string, options: MsfOptions, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + msfBeacon(beaconId: string, options: MsfOptions, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.msfTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), options, timeoutSeconds); } - private msfRemoteTarget(request: CommonRequest, options: MsfRemoteOptions, timeoutSeconds: number) { + private msfRemoteTarget(request: CommonRequest, options: MsfRemoteOptions, timeoutSeconds: OperationTimeout) { assertNonEmptyString(options.lhost, "Metasploit listen host"); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.msfRemote( @@ -2194,24 +2265,27 @@ export class SliverClient { msfRemoteSession( sessionId: string, options: MsfRemoteOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.msfRemoteTarget(this.m4SessionRequest(sessionId, timeoutSeconds), options, timeoutSeconds); } msfRemoteBeacon( beaconId: string, options: MsfRemoteOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.msfRemoteTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), options, timeoutSeconds); } runSshSession( sessionId: string, options: SshCommandOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); assertNonEmptyString(options.username, "SSH username"); assertNonEmptyString(options.hostname, "SSH hostname"); if (options.kerberosRealm && !options.kerberosKeytab) { @@ -2249,8 +2323,9 @@ export class SliverClient { startRemoteServiceSession( sessionId: string, options: StartRemoteServiceOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); assertBoundedNonEmptyString( options.hostname, "Remote service hostname", @@ -2295,8 +2370,9 @@ export class SliverClient { removeRemoteServiceSession( sessionId: string, options: RemoveRemoteServiceOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); assertBoundedNonEmptyString( options.hostname, "Remote service hostname", @@ -2322,7 +2398,7 @@ export class SliverClient { ); } - private runAsTarget(request: CommonRequest, options: RunAsOptions, timeoutSeconds: number) { + private runAsTarget(request: CommonRequest, options: RunAsOptions, timeoutSeconds: OperationTimeout) { assertNonEmptyString(options.username, "Run-as username"); assertNonEmptyString(options.processName, "Run-as process path"); return withTimeoutSignal(timeoutSeconds, (signal) => @@ -2345,20 +2421,22 @@ export class SliverClient { runAsSession( sessionId: string, options: RunAsOptions, - timeoutSeconds = M4_IDENTITY_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_IDENTITY_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_IDENTITY_TIMEOUT_SECONDS); return this.runAsTarget(this.m4SessionRequest(sessionId, timeoutSeconds), options, timeoutSeconds); } runAsBeacon( beaconId: string, options: RunAsOptions, - timeoutSeconds = M4_IDENTITY_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_IDENTITY_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_IDENTITY_TIMEOUT_SECONDS); return this.runAsTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), options, timeoutSeconds); } - private makeTokenTarget(request: CommonRequest, options: MakeTokenOptions, timeoutSeconds: number) { + private makeTokenTarget(request: CommonRequest, options: MakeTokenOptions, timeoutSeconds: OperationTimeout) { assertNonEmptyString(options.username, "Token username"); assertNonEmptyString(options.password, "Token password"); const logonType = options.logonType ?? 9; @@ -2380,20 +2458,22 @@ export class SliverClient { makeTokenSession( sessionId: string, options: MakeTokenOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.makeTokenTarget(this.m4SessionRequest(sessionId, timeoutSeconds), options, timeoutSeconds); } makeTokenBeacon( beaconId: string, options: MakeTokenOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return this.makeTokenTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), options, timeoutSeconds); } - private impersonateTarget(request: CommonRequest, username: string, timeoutSeconds: number) { + private impersonateTarget(request: CommonRequest, username: string, timeoutSeconds: OperationTimeout) { assertNonEmptyString(username, "Impersonation username"); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.impersonate({ Username: username, Request: request }, { signal }), @@ -2403,26 +2483,30 @@ export class SliverClient { impersonateSession( sessionId: string, username: string, - timeoutSeconds = M4_IDENTITY_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_IDENTITY_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_IDENTITY_TIMEOUT_SECONDS); return this.impersonateTarget(this.m4SessionRequest(sessionId, timeoutSeconds), username, timeoutSeconds); } impersonateBeacon( beaconId: string, username: string, - timeoutSeconds = M4_IDENTITY_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_IDENTITY_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_IDENTITY_TIMEOUT_SECONDS); return this.impersonateTarget(this.m4BeaconRequest(beaconId, timeoutSeconds), username, timeoutSeconds); } - revToSelfSession(sessionId: string, timeoutSeconds = M4_IDENTITY_TIMEOUT_SECONDS) { + revToSelfSession(sessionId: string, timeoutSeconds: OperationTimeout = M4_IDENTITY_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_IDENTITY_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.revToSelf({ Request: this.m4SessionRequest(sessionId, timeoutSeconds) }, { signal }), ); } - revToSelfBeacon(beaconId: string, timeoutSeconds = M4_IDENTITY_TIMEOUT_SECONDS) { + revToSelfBeacon(beaconId: string, timeoutSeconds: OperationTimeout = M4_IDENTITY_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_IDENTITY_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.revToSelf({ Request: this.m4BeaconRequest(beaconId, timeoutSeconds) }, { signal }), ); @@ -2431,8 +2515,9 @@ export class SliverClient { getSystemSession( sessionId: string, options: GetSystemOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); const config = m4ImplantConfig(options.config); return withTimeoutSignal(timeoutSeconds, async (signal) => { try { @@ -2451,13 +2536,15 @@ export class SliverClient { }).finally(config.dispose); } - getPrivsSession(sessionId: string, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + getPrivsSession(sessionId: string, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getPrivs({ Request: this.m4SessionRequest(sessionId, timeoutSeconds) }, { signal }), ); } - getPrivsBeacon(beaconId: string, timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS) { + getPrivsBeacon(beaconId: string, timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getPrivs({ Request: this.m4BeaconRequest(beaconId, timeoutSeconds) }, { signal }), ); @@ -2466,8 +2553,9 @@ export class SliverClient { backdoorSession( sessionId: string, options: BackdoorOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); assertNonEmptyString(options.filePath, "Backdoor remote file path"); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.backdoor( @@ -2485,8 +2573,9 @@ export class SliverClient { hijackDllSession( sessionId: string, options: HijackDllOptions, - timeoutSeconds = M4_DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = M4_DEFAULT_TIMEOUT_SECONDS, ) { + timeoutSeconds = operationOptions(timeoutSeconds, M4_DEFAULT_TIMEOUT_SECONDS); assertNonEmptyString(options.referenceDllPath, "Reference DLL path"); assertNonEmptyString(options.targetLocation, "DLL target location"); if (options.targetDll && options.profileName) { @@ -2518,7 +2607,7 @@ export class SliverClient { }); } - killSession(sessionId: string, force = false, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + killSession(sessionId: string, force = false, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.kill( { Force: force, Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -2527,7 +2616,7 @@ export class SliverClient { ); } - killBeacon(beaconId: string, force = false, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + killBeacon(beaconId: string, force = false, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.kill( { Force: force, Request: this.beaconRequest(beaconId, timeoutSeconds) }, @@ -2539,7 +2628,7 @@ export class SliverClient { reconfigureBeacon( beaconId: string, options: BeaconReconfigureOptions, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.reconfigure( @@ -2559,7 +2648,7 @@ export class SliverClient { beaconId: string, c2s: string[], delayNanoseconds = "0", - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.openSession( @@ -2573,7 +2662,7 @@ export class SliverClient { ); } - closeSession(sessionId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + closeSession(sessionId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.closeSession( { Request: this.sessionRequest(sessionId, timeoutSeconds) }, @@ -2582,39 +2671,58 @@ export class SliverClient { ); } - getBeaconTasks(beaconId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getBeaconTasks(beaconId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.inventoryRpc.getBeaconTasks({ ID: beaconId }, { signal }), ); } - fetchBeaconTask(taskId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + /** + * Attach a typed handle to any queued beacon operation. The codec is a + * generated protobuf message (for example sliverpb.Ping), not an RPC name. + * Handles resolve the current channel on every fetch and survive reconnects. + */ + beaconTask( + taskId: string, + decoder: BeaconTaskDecoder, + options: { readonly beaconId?: string; readonly timeoutSeconds?: number } = {}, + ): BeaconTaskHandle { + return createBeaconTaskHandle({ + taskId, decoder, beaconId: options.beaconId, + defaultTimeoutSeconds: options.timeoutSeconds, + getRpc: () => this.rpcClients["task-content"], + events: this.taskResult$, + cancel: (id, operation) => this.cancelBeaconTask(id, operation), + }); + } + + fetchBeaconTask(taskId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.taskContentRpc.getBeaconTaskContent({ ID: taskId }, { signal }), ); } - cancelBeaconTask(taskId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + cancelBeaconTask(taskId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.cancelBeaconTask({ ID: taskId }, { signal }), ); } - getJobs(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getJobs(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getJobs(this.empty, { signal })); } - killJob(jobId: number, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + killJob(jobId: number, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.killJob({ ID: jobId }, { signal })); } - restartJobs(jobIds: number[], timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + restartJobs(jobIds: number[], timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.rpc.restartJobs({ JobIDs: jobIds }, { signal }); }); } - startMTLSListener(host: string, port: number, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + startMTLSListener(host: string, port: number, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.startMTLSListener({ Host: host, Port: port }, { signal })); } @@ -2624,7 +2732,7 @@ export class SliverClient { tunIP: string, nPort: number, keyPort: number, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.startWGListener({ Host: host, Port: port, TunIP: tunIP, NPort: nPort, KeyPort: keyPort }, { signal }), @@ -2637,7 +2745,7 @@ export class SliverClient { host: string, port: number, enforceOTP = false, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.startDNSListener( @@ -2653,12 +2761,12 @@ export class SliverClient { port: number, website = "", enforceOTP = true, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return this.startHTTPListenerWithOptions({ domain, host, port, website, enforceOTP }, timeoutSeconds); } - startHTTPListenerWithOptions(options: HTTPListenerOptions, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + startHTTPListenerWithOptions(options: HTTPListenerOptions, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { const request: HTTPListenerReq = { Domain: options.domain ?? "", Host: options.host, @@ -2686,7 +2794,7 @@ export class SliverClient { cert?: Buffer, key?: Buffer, enforceOTP = true, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, ) { return this.startHTTPSListenerWithOptions( { domain, host, port, website, acme, cert, key, enforceOTP }, @@ -2694,7 +2802,7 @@ export class SliverClient { ); } - startHTTPSListenerWithOptions(options: HTTPSListenerOptions, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + startHTTPSListenerWithOptions(options: HTTPSListenerOptions, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { const request: HTTPListenerReq = { Domain: options.domain ?? "", Host: options.host, @@ -2713,126 +2821,126 @@ export class SliverClient { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.startHTTPSListener(request, { signal })); } - startTCPStagerListener(host: string, port: number, data: Buffer, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + startTCPStagerListener(host: string, port: number, data: Buffer, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return this.startTCPStagerListenerWithOptions({ Protocol: 0, Host: host, Port: port, Data: data, ProfileName: "" }, timeoutSeconds); } - startTCPStagerListenerWithOptions(request: StagerListenerReq, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + startTCPStagerListenerWithOptions(request: StagerListenerReq, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.startTCPStagerListener(request, { signal })); } - getCompiler(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + getCompiler(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.inventoryRpc.getCompiler(this.empty, { signal })); } - generateUniqueIP(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + generateUniqueIP(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.generateUniqueIP(this.empty, { signal })); } - generateImplant(config: ImplantConfig, name = "", timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + generateImplant(config: ImplantConfig, name = "", timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.generate({ Config: config, Name: name }, { signal })); } - async generate(config: ImplantConfig, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + async generate(config: ImplantConfig, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { const res = await this.generateImplant(config, "", timeoutSeconds); return res.File; } - generateSpoofMetadata(req: GenerateSpoofMetadataReq, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + generateSpoofMetadata(req: GenerateSpoofMetadataReq, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.artifactRpc.generateSpoofMetadata(req, { signal }); }); } - async regenerate(implantName: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + async regenerate(implantName: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { const res = await this.regenerateImplant(implantName, timeoutSeconds); return res.File; } - regenerateImplant(implantName: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + regenerateImplant(implantName: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.regenerate({ ImplantName: implantName }, { signal }), ); } - generateStage(request: GenerateStageReq, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + generateStage(request: GenerateStageReq, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.generateStage(request, { signal })); } - stageImplantBuild(buildNames: string[], timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + stageImplantBuild(buildNames: string[], timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.rpc.stageImplantBuild({ Build: buildNames }, { signal }); }); } - implantBuilds(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + implantBuilds(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.inventoryRpc.implantBuilds(this.empty, { signal })); } - deleteImplantBuild(name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + deleteImplantBuild(name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.rpc.deleteImplantBuild({ Name: name }, { signal }); }); } - canaries(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + canaries(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.canaries(this.empty, { signal })); } - implantProfiles(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + implantProfiles(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.inventoryRpc.implantProfiles(this.empty, { signal })); } - saveImplantProfile(profile: ImplantProfile, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + saveImplantProfile(profile: ImplantProfile, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.saveImplantProfile(profile, { signal })); } - deleteImplantProfile(name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + deleteImplantProfile(name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.rpc.deleteImplantProfile({ Name: name }, { signal }); }); } - lootAll(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + lootAll(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, async (signal) => { const res = await this.inventoryRpc.lootAll(this.empty, { signal }); return res.Loot; }); } - lootAdd(loot: Loot, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + lootAdd(loot: Loot, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.lootAdd(loot, { signal })); } - lootUpdate(loot: Loot, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + lootUpdate(loot: Loot, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.lootUpdate(loot, { signal })); } - lootRemove(lootId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + lootRemove(lootId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.rpc.lootRm({ ID: lootId }, { signal }); }); } - lootContent(lootId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + lootContent(lootId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.lootContent({ ID: lootId }, { signal })); } - credentialsAll(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + credentialsAll(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { const response = await this.inventoryRpc.creds(this.empty, { signal }); return response.Credentials; }); } - credentialById(credentialId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + credentialById(credentialId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { assertBoundedNonEmptyString(credentialId, "Credential id", CREDENTIAL_ID_MAX_CHARACTERS); return withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.getCredByID({ ID: credentialId }, { signal }), ); } - credentialAdd(credential: Credential, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + credentialAdd(credential: Credential, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { assertBoundedString(credential.Username, "Credential username", CREDENTIAL_METADATA_MAX_CHARACTERS); assertBoundedString(credential.Collection, "Credential collection", CREDENTIAL_METADATA_MAX_CHARACTERS); assertCredentialSecret(credential.Plaintext, "Credential plaintext"); @@ -2858,14 +2966,14 @@ export class SliverClient { }); } - credentialRemove(credentialId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + credentialRemove(credentialId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { assertBoundedNonEmptyString(credentialId, "Credential id", CREDENTIAL_ID_MAX_CHARACTERS); return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.rpc.credsRm({ Credentials: [{ ID: credentialId }] }, { signal }); }); } - credentialSniffHashType(hash: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + credentialSniffHashType(hash: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { assertCredentialSecret(hash, "Credential hash"); if (!hash) throw new Error("Credential hash must not be empty"); return withTimeoutSignal(timeoutSeconds, async (signal) => { @@ -2874,36 +2982,36 @@ export class SliverClient { }); } - websites(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + websites(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, async (signal) => { const res = await this.inventoryRpc.websites(this.empty, { signal }); return res.Websites; }); } - website(name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + website(name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.website({ Name: name }, { signal })); } - websiteRemove(name: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + websiteRemove(name: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { return withTimeoutSignal(timeoutSeconds, async (signal) => { await this.rpc.websiteRemove({ Name: name }, { signal }); }); } - websiteAddContent(name: string, contents: Record, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + websiteAddContent(name: string, contents: Record, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.websiteAddContent({ Name: name, Contents: contents }, { signal }), ); } - websiteUpdateContent(name: string, contents: Record, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + websiteUpdateContent(name: string, contents: Record, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.websiteUpdateContent({ Name: name, Contents: contents }, { signal }), ); } - websiteRemoveContent(name: string, paths: string[], timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + websiteRemoveContent(name: string, paths: string[], timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { return withTimeoutSignal(timeoutSeconds, (signal) => this.artifactRpc.websiteRemoveContent({ Name: name, Paths: paths }, { signal }), ); @@ -2911,22 +3019,22 @@ export class SliverClient { // --- High-level helpers (ergonomic wrappers) --- - async operators(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + async operators(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { const res = await this.getOperators(timeoutSeconds); return res.Operators; } - async sessions(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + async sessions(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { const res = await this.getSessions(timeoutSeconds); return res.Sessions; } - async beacons(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + async beacons(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { const res = await this.getBeacons(timeoutSeconds); return res.Beacons; } - async jobs(timeoutSeconds = DEFAULT_TIMEOUT_SECONDS) { + async jobs(timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS) { const res = await this.getJobs(timeoutSeconds); return res.Active; } @@ -2947,7 +3055,7 @@ export class SliverClient { sessionId: normalizedSessionId, options, operation, - request: (timeoutSeconds) => this.sessionRequest(normalizedSessionId, timeoutSeconds), + request: (timeoutSeconds: OperationTimeout) => this.sessionRequest(normalizedSessionId, timeoutSeconds), onClosed: (closed) => this.managedPortForwards.delete(closed), }); this.managedPortForwards.add(forward); @@ -3024,7 +3132,7 @@ export class SliverClient { sessionId: normalizedSessionId, options, operation, - request: (timeoutSeconds) => this.sessionRequest(normalizedSessionId, timeoutSeconds), + request: (timeoutSeconds: OperationTimeout) => this.sessionRequest(normalizedSessionId, timeoutSeconds), authorizeCreatedListenerCleanup: (candidate) => !this.managedReversePortForwards.has(reversePortForwardKey(candidate)), callbacks: { @@ -3054,7 +3162,7 @@ export class SliverClient { const listeners = await requestReversePortForwards( this.rpc, normalizedSessionId, - (timeoutSeconds) => this.sessionRequest(normalizedSessionId, timeoutSeconds), + (timeoutSeconds: OperationTimeout) => this.sessionRequest(normalizedSessionId, timeoutSeconds), operation, ); return Object.freeze(listeners); @@ -3074,7 +3182,7 @@ export class SliverClient { this.rpc, normalizedSessionId, listenerId, - (timeoutSeconds) => this.sessionRequest(normalizedSessionId, timeoutSeconds), + (timeoutSeconds: OperationTimeout) => this.sessionRequest(normalizedSessionId, timeoutSeconds), operation, ); this.managedReversePortForwards.get(reversePortForwardKey({ @@ -3093,7 +3201,7 @@ export class SliverClient { const listeners = await requestReversePortForwards( this.rpc, forward.sessionId, - (timeoutSeconds) => this.sessionRequest(forward.sessionId, timeoutSeconds), + (timeoutSeconds: OperationTimeout) => this.sessionRequest(forward.sessionId, timeoutSeconds), operation, ); const listener = listeners.find((candidate) => candidate.id === forward.id); @@ -3117,12 +3225,12 @@ export class SliverClient { this.rpc, forward.sessionId, forward.id, - (timeoutSeconds) => this.sessionRequest(forward.sessionId, timeoutSeconds), + (timeoutSeconds: OperationTimeout) => this.sessionRequest(forward.sessionId, timeoutSeconds), ); }); } - private async reconcileManagedReversePortForwards(): Promise { + private async reconcileManagedReversePortForwards(operation?: ForwardingOperationOptions): Promise { const bySession = new Map(); for (const forward of this.managedReversePortForwards.values()) { const forwards = bySession.get(forward.sessionId) ?? []; @@ -3134,7 +3242,8 @@ export class SliverClient { const listeners = await requestReversePortForwards( this.rpc, sessionId, - (timeoutSeconds) => this.sessionRequest(sessionId, timeoutSeconds), + (timeoutSeconds: OperationTimeout) => this.sessionRequest(sessionId, timeoutSeconds), + operation, ); const byId = new Map(listeners.map((listener) => [listener.id, listener])); for (const forward of forwards) forward.reconcile(byId.get(forward.id)); @@ -3152,7 +3261,17 @@ export class SliverClient { async startShellSession( sessionId: string, options: ShellSessionOptions, - timeoutSeconds = DEFAULT_TIMEOUT_SECONDS, + timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS, + ): Promise { + return withTimeoutSignal(timeoutSeconds, (signal) => this.startShellSessionOperation(sessionId, options, { + timeoutSeconds: operationTimeoutSeconds(timeoutSeconds), signal, + }), (handle) => handle.close()); + } + + private async startShellSessionOperation( + sessionId: string, + options: ShellSessionOptions, + timeoutSeconds: OperationTimeout, ): Promise { assertNonEmptyString(sessionId, "Session id"); if (!options || typeof options !== "object") throw new Error("Shell options are required"); @@ -3180,7 +3299,7 @@ export class SliverClient { const bestEffortRemoteClose = async (): Promise => { if (!tunnelManagerOwned || !tunnelId) return; - await withTimeoutSignal(timeoutSeconds, (signal) => + await withTimeoutSignal(5, (signal) => this.rpc.closeTunnel({ TunnelID: tunnelId, SessionID: sessionId }, { signal }), ).catch(() => undefined); }; @@ -3238,11 +3357,14 @@ export class SliverClient { }; try { - const created = await withTimeoutSignal(timeoutSeconds, (signal) => - this.rpc.createTunnel({ SessionID: sessionId }, { signal }), + // The public startup deadline settles promptly while this continuation + // retains ownership of a late creation result for cleanup. + const created = await this.rpc.createTunnel( + { SessionID: sessionId }, { signal: operationOptions(timeoutSeconds).signal }, ); tunnelId = created.TunnelID.trim(); if (!tunnelId) throw new Error("Missing tunnel id"); + if (created.SessionID && created.SessionID !== sessionId) throw new Error("Mismatched shell session"); const output = tunnels.openOutput(tunnelId, { maxBufferedBytes: outputBufferBytes, @@ -3257,13 +3379,9 @@ export class SliverClient { }); tunnelManagerOwned = true; - // The zero-data bind must be pulled by TunnelData before Shell can race - // an early prompt back to this process. - await tunnels.send({ - TunnelID: tunnelId, - SessionID: sessionId, - Data: Buffer.alloc(0), - }); + // Await the server's ownership acknowledgement under the shared startup + // deadline before opening the shell on another HTTP/2 stream. + await tunnels.bind(tunnelId, sessionId, operationOptions(timeoutSeconds).signal); const shell = await withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.shell( @@ -3360,11 +3478,12 @@ export class SliverClient { }, close: () => closeManagedTunnel(true), }; - } catch { + } catch (error) { // The Shell RPC can lose its response after the implant has already // spawned the child. Once a tunnel exists, mirror the canonical client // and attempt a bounded graceful exit before revoking the transport. await closeManagedTunnel(tunnelManagerOwned && Boolean(tunnelId)); + if (error instanceof OperationAbortedError || error instanceof OperationTimeoutError) throw error; throw new Error("Unable to start shell session"); } } @@ -3381,7 +3500,7 @@ export class SliverClient { return new InteractiveBeacon(this.rpc, this.artifactRpc, this.taskContentRpc, this.taskResult$, beaconId); } - async rmBeacon(beaconId: string, timeoutSeconds = DEFAULT_TIMEOUT_SECONDS): Promise { + async rmBeacon(beaconId: string, timeoutSeconds: OperationTimeout = DEFAULT_TIMEOUT_SECONDS): Promise { await withTimeoutSignal(timeoutSeconds, (signal) => this.rpc.rmBeacon({ ID: beaconId }, { signal })); } } @@ -3888,9 +4007,23 @@ async function waitForBeaconTaskContent( taskResult$: Observable, taskContentRpc: SliverRPCClient, taskId: string, - timeoutSeconds: number, + timeoutSeconds: OperationTimeout, +): Promise { + return withTimeoutSignal(timeoutSeconds, (signal) => waitForBeaconTaskContentWithSignal( + taskResult$, taskContentRpc, taskId, operationTimeoutSeconds(timeoutSeconds), signal, + )).catch((error: unknown) => { + if (error instanceof OperationTimeoutError) throw new Error(`Timeout waiting for beacon task result: ${taskId}`); + throw error; + }); +} + +async function waitForBeaconTaskContentWithSignal( + taskResult$: Observable, + taskContentRpc: SliverRPCClient, + taskId: string, + validatedTimeoutSeconds: number, + signal: AbortSignal, ): Promise { - const validatedTimeoutSeconds = validateTimeoutSeconds(timeoutSeconds); return new Promise((resolve, reject) => { let settled = false; let unsubscribePending = false; @@ -3908,6 +4041,7 @@ async function waitForBeaconTaskContent( const finish = (settle: () => void) => { if (settled) return; settled = true; + signal.removeEventListener("abort", onAbort); if (deadlineTimer !== undefined) clearTimeout(deadlineTimer); if (pollTimer !== undefined) clearInterval(pollTimer); controller.abort(); @@ -3919,6 +4053,13 @@ async function waitForBeaconTaskContent( settle(); }; + const onAbort = () => finish(() => reject(signal.reason)); + signal.addEventListener("abort", onAbort, { once: true }); + if (signal.aborted) { + onAbort(); + return; + } + const refresh = async () => { if (settled || refreshing) return; refreshing = true; diff --git a/src/index.ts b/src/index.ts index 0c108fb..6481ba5 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,7 +1,15 @@ export * from './client'; export * from './config'; export * from './forwarding'; -export { timeoutSecondsToNanoseconds, withTimeoutSignal } from "./internal/timeout"; +export { + timeoutSecondsToNanoseconds, withTimeoutSignal, OperationAbortedError, OperationTimeoutError, + type OperationOptions, type OperationTimeout, +} from "./internal/timeout"; +export { + type BeaconTaskHandle, type BeaconTaskDecoder, + BeaconTaskStateError, BeaconTaskResponseError, BeaconTaskIdentityError, BeaconTaskTimeoutError, +} from "./beaconTasks"; +export type { TunnelTransportState } from "./internal/tunnelManager"; export * from './messageBudget'; export * as commonpb from "./pb/commonpb/common"; export * as clientpb from "./pb/clientpb/client"; diff --git a/src/internal/portForward.ts b/src/internal/portForward.ts index 7bf488d..db86c29 100644 --- a/src/internal/portForward.ts +++ b/src/internal/portForward.ts @@ -32,7 +32,7 @@ const PORT_FORWARD_PROTOCOL_TCP = 1; type PortForwardRpc = Pick; export interface PortForwardTunnelTransport { - onTransportFailure(listener: () => void): () => void; + onTransportFailure(listener: () => void, options?: { readonly nextGeneration?: boolean }): () => void; openOutput( tunnelId: string, options: { @@ -160,7 +160,7 @@ export class ManagedPortForward implements PortForward { this.publishState("listening"); this.removeTransportFailureListener = this.tunnels.onTransportFailure(() => { void this.closeWithReason("transport-disconnected", true); - }); + }, { nextGeneration: true }); if (this.closePromise) { await this.closePromise; throw new Error("Tunnel transport is disconnected"); diff --git a/src/internal/timeout.ts b/src/internal/timeout.ts index 442f6cd..5f8b362 100644 --- a/src/internal/timeout.ts +++ b/src/internal/timeout.ts @@ -8,12 +8,57 @@ const MAX_TIMER_TIMEOUT_SECONDS = Math.floor(MAX_TIMER_DELAY_MILLISECONDS / 1_00 const MAX_PROTOBUF_TIMEOUT_SECONDS = Number((MAX_PROTOBUF_INT64 + 1n) / NANOSECONDS_PER_SECOND); const MAX_TIMEOUT_SECONDS = Math.min(MAX_TIMER_TIMEOUT_SECONDS, MAX_PROTOBUF_TIMEOUT_SECONDS); +/** A local operation deadline and optional caller cancellation signal. */ +export interface OperationOptions { + /** Defaults to 30 seconds. Zero disables the local deadline. */ + timeoutSeconds?: number; + /** Cancel this operation without disconnecting unrelated operations. */ + signal?: AbortSignal; +} + +/** Numeric timeouts retain their existing whole-second behavior. */ +export type OperationTimeout = number | OperationOptions; + +/** Caller cancellation; a dispatched remote operation may still complete. */ +export class OperationAbortedError extends Error { + constructor(reason?: unknown) { + super("Operation was aborted", { cause: reason }); + this.name = "OperationAbortedError"; + } +} + +/** Local deadline expiry; a dispatched remote operation may still complete. */ +export class OperationTimeoutError extends Error { + readonly timeoutSeconds: number; + + constructor(timeoutSeconds: number) { + super(`Operation timed out after ${timeoutSeconds} seconds`); + this.name = "OperationTimeoutError"; + this.timeoutSeconds = timeoutSeconds; + } +} + +/** Resolve and validate the whole-second timeout shared by local and wire deadlines. */ +export function operationTimeoutSeconds(timeout: OperationTimeout, defaultSeconds = 30): number { + if (typeof timeout === "number") return validateTimeoutSeconds(timeout); + if (!timeout || typeof timeout !== "object" || Array.isArray(timeout)) { + throw new TypeError("Operation timeout must be a number or options object"); + } + return validateTimeoutSeconds(timeout.timeoutSeconds === undefined ? defaultSeconds : timeout.timeoutSeconds); +} + +/** Preserve a method's default deadline when the caller supplies only a signal. */ +export function operationOptions(timeout: OperationTimeout, defaultSeconds = 30): OperationOptions { + const timeoutSeconds = operationTimeoutSeconds(timeout, defaultSeconds); + return typeof timeout === "number" ? { timeoutSeconds } : { timeoutSeconds, signal: timeout.signal }; +} + /** * Convert the public timeout unit to the nanoseconds expected by * commonpb.Request.Timeout without passing an imprecise number to protobuf. */ -export function timeoutSecondsToNanoseconds(timeoutSeconds: number): string { - const validatedTimeoutSeconds = validateTimeoutSeconds(timeoutSeconds); +export function timeoutSecondsToNanoseconds(timeout: OperationTimeout): string { + const validatedTimeoutSeconds = operationTimeoutSeconds(timeout); // Match Sliver's canonical Go clients: leave zero unset, and make the // server-side deadline expire just before the same whole-second local @@ -28,23 +73,95 @@ export function timeoutSecondsToNanoseconds(timeoutSeconds: number): string { return timeoutNanoseconds.toString(); } +/** + * Settle on caller cancellation or deadline expiry even if the callback ignores + * its signal. Late callback failures are observed. Supply onLateResult to release + * an owned resource or buffer returned after cancellation; its cleanup failures + * are also observed without replacing the operation's cancellation error. The + * default behavior cannot safely dispose an arbitrary T. + */ export async function withTimeoutSignal( - timeoutSeconds: number, + timeout: OperationTimeout, fn: (signal: AbortSignal) => Promise, + onLateResult?: (value: T) => void | Promise, ): Promise { - const validatedTimeoutSeconds = validateTimeoutSeconds(timeoutSeconds); - if (validatedTimeoutSeconds === 0) { - return fn(new AbortController().signal); + const timeoutSeconds = operationTimeoutSeconds(timeout); + const callerSignal = typeof timeout === "number" ? undefined : timeout.signal; + if (callerSignal !== undefined && ( + !callerSignal + || typeof callerSignal.aborted !== "boolean" + || typeof callerSignal.addEventListener !== "function" + || typeof callerSignal.removeEventListener !== "function" + )) { + throw new TypeError("Operation signal must be an AbortSignal"); } + if (callerSignal?.aborted) throw new OperationAbortedError(callerSignal.reason); const controller = new AbortController(); - const timeoutMs = validatedTimeoutSeconds * 1_000; - const timer = setTimeout(() => controller.abort(), timeoutMs); - try { - return await fn(controller.signal); - } finally { - clearTimeout(timer); - } + return new Promise((resolve, reject) => { + let settled = false; + let aborted = false; + let timer: ReturnType | undefined; + const cleanup = () => { + if (timer !== undefined) clearTimeout(timer); + callerSignal?.removeEventListener("abort", onCallerAbort); + }; + const abort = (error: OperationAbortedError | OperationTimeoutError) => { + if (settled) return; + settled = true; + aborted = true; + cleanup(); + // Claim the result before notifying callback abort handlers, which can + // synchronously reject with transport-specific cancellation errors. + reject(error); + controller.abort(error); + }; + const onCallerAbort = () => abort(new OperationAbortedError(callerSignal?.reason)); + + callerSignal?.addEventListener("abort", onCallerAbort, { once: true }); + // Also handle signals that became aborted while the listener was attached. + if (callerSignal?.aborted) { + onCallerAbort(); + return; + } + if (timeoutSeconds > 0) { + timer = setTimeout(() => abort(new OperationTimeoutError(timeoutSeconds)), timeoutSeconds * 1_000); + } + + try { + Promise.resolve(fn(controller.signal)).then( + (value) => { + if (settled) { + if (aborted && onLateResult) { + try { + // The caller has already received cancellation. Cleanup owns + // any late resource, including observing its own rejection. + Promise.resolve(onLateResult(value)).catch(() => undefined); + } catch { + // A synchronous cleanup failure must not become an unhandled + // rejection or expose callback details after cancellation. + } + } + return; + } + settled = true; + cleanup(); + resolve(value); + }, + (error: unknown) => { + if (settled) return; + settled = true; + cleanup(); + reject(error); + }, + ); + } catch (error) { + if (settled) return; + settled = true; + cleanup(); + reject(error); + } + }); } export function validateTimeoutSeconds(timeoutSeconds: number): number { diff --git a/src/internal/tunnelManager.ts b/src/internal/tunnelManager.ts index ffbd17d..24dcabf 100644 --- a/src/internal/tunnelManager.ts +++ b/src/internal/tunnelManager.ts @@ -1,4 +1,4 @@ -import { Subject, type Observable } from "rxjs"; +import { BehaviorSubject, Subject, type Observable } from "rxjs"; import { TUNNEL_STREAM_MAX_PAYLOAD_BYTES } from "../messageBudget"; import type { SliverRPCClient, DeepPartial } from "../pb/rpcpb/services"; @@ -11,6 +11,8 @@ export const TUNNEL_MANAGER_MAX_BUFFERED_OUTPUT_FRAMES = 1_024; export const TUNNEL_MANAGER_MAX_QUEUED_FRAMES_PER_TUNNEL = 8; export const TUNNEL_MANAGER_MAX_QUEUED_BYTES_PER_TUNNEL = 256 * 1024; export const TUNNEL_MANAGER_MAX_TOTAL_QUEUED_BYTES = 4 * 1024 * 1024; +/** Bounds retired identities retained to reject stale handles; reconnect resets this budget. */ +export const TUNNEL_MANAGER_MAX_LIFETIME_TUNNELS = 65_536; const SAFE_TUNNEL_CLOSED_ERROR = "Tunnel is closed"; const SAFE_TUNNEL_OVERFLOW_ERROR = "Tunnel output exceeded its bounded buffer"; @@ -28,8 +30,22 @@ export interface TunnelManagerStats { readonly queuedBytes: number; } +export interface TunnelTransportState { + readonly status: "stopped" | "idle" | "connecting" | "connected" | "failed"; + readonly generation: number; +} + +interface TunnelTransportGeneration { + readonly id: number; + readonly outgoing: FairTunnelSendQueue; + readonly abort: AbortController; + running: Promise | null; + failed: boolean; +} + interface TunnelState { readonly subject: Subject; + readonly generation: TunnelTransportGeneration; output?: AsyncQueue; onClosed?: () => void; onFailure?: (reason: "cancelled" | "overflow" | "transport") => void; @@ -256,56 +272,99 @@ class FairTunnelSendQueue implements AsyncIterable> { } export class TunnelManager { - private readonly outgoing = new FairTunnelSendQueue(); private readonly byTunnelId = new Map(); - private readonly abort = new AbortController(); + private readonly issuedTunnelIds = new Set(); private readonly transportFailureListeners = new Set<() => void>(); + private readonly transportStateSubject = new BehaviorSubject( + Object.freeze({ status: "stopped", generation: 0 }), + ); + readonly transportState$ = this.transportStateSubject.asObservable(); private rpc: SliverRPCClient | null = null; - private running: Promise | null = null; - private transportFailed = false; + private generation: TunnelTransportGeneration | null = null; + private stopped = false; + + get transportState(): TunnelTransportState { + return this.transportStateSubject.value; + } start(rpc: SliverRPCClient): void { + if (this.stopped) throw new Error("Tunnel manager is stopped"); if (this.rpc) return; this.rpc = rpc; + this.publishTransportState("idle", 0); } - private ensureRunning(): void { - if (this.running) return; + private ensureRunning(): TunnelTransportGeneration { const rpc = this.rpc; - if (!rpc) throw new Error("Tunnel manager is not connected"); + if (this.stopped || !rpc) throw new Error("Tunnel manager is not connected"); + if (this.generation && !this.generation.failed) return this.generation; + + // Recovery is driven only by a newly-created tunnel. Failed tunnels and + // their queued frames are never replayed onto the replacement stream. + const generation: TunnelTransportGeneration = { + id: (this.generation?.id ?? 0) + 1, + outgoing: new FairTunnelSendQueue(), + abort: new AbortController(), + running: null, + failed: false, + }; + this.generation = generation; + this.publishTransportState("connecting", generation.id); + if (this.stopped) throw new Error("Tunnel manager is not connected"); - this.running = (async () => { + generation.running = (async () => { try { - const stream = rpc.tunnelData(this.outgoing, { signal: this.abort.signal }); - for await (const msg of stream) this.receive(msg); - if (!this.abort.signal.aborted) throw new Error(SAFE_TUNNEL_TRANSPORT_ERROR); + const stream = rpc.tunnelData(generation.outgoing, { signal: generation.abort.signal }); + for await (const msg of stream) this.receive(msg, generation); + if (!generation.abort.signal.aborted) throw new Error(SAFE_TUNNEL_TRANSPORT_ERROR); } catch { - if (!this.abort.signal.aborted) { - this.transportFailed = true; - this.outgoing.fail(); - for (const tunnelId of [...this.byTunnelId.keys()]) { - this.failTunnel(tunnelId, "transport"); - } - for (const listener of [...this.transportFailureListeners]) invokeClosed(listener); - } + // A transport implementation may throw before returning its iterable. + // Finish registering the initiating output before publishing failure. + await Promise.resolve(); + this.failGeneration(generation); } })(); + return generation; } - /** Registers a listener for an unexpected terminal TunnelData failure. */ - onTransportFailure(listener: () => void): () => void { - if (this.transportFailed) { + /** + * Registers for the current stream's terminal failure. A new listener owner + * may opt into the next generation when the preceding stream already failed. + * Listeners are retired with their generation and never called again. + */ + onTransportFailure(listener: () => void, options?: { readonly nextGeneration?: boolean }): () => void { + if (this.stopped || (this.generation?.failed && !options?.nextGeneration)) { invokeClosed(listener); return () => undefined; } - this.transportFailureListeners.add(listener); - return () => this.transportFailureListeners.delete(listener); + const registration = () => listener(); + this.transportFailureListeners.add(registration); + return () => this.transportFailureListeners.delete(registration); + } + + private publishTransportState(status: TunnelTransportState["status"], generation: number): void { + this.transportStateSubject.next(Object.freeze({ status, generation })); + } + + private failGeneration(generation: TunnelTransportGeneration): void { + if (this.stopped || this.generation !== generation || generation.failed) return; + generation.failed = true; + const failedTunnels = [...this.byTunnelId.entries()].filter(([, state]) => state.generation === generation); + const listeners = [...this.transportFailureListeners]; + this.transportFailureListeners.clear(); + // Revoke every old identity before callbacks can create a new generation. + for (const [id] of failedTunnels) this.byTunnelId.delete(id); + generation.outgoing.fail(); + generation.abort.abort(); + this.publishTransportState("failed", generation.id); + for (const [, state] of failedTunnels) this.failTunnelState(state, "transport"); + for (const listener of listeners) invokeClosed(listener); } /** Legacy RxJS compatibility. New main-process callers use openOutput(). */ - subscribe(tunnelId: string): Observable { - return this.ensureTunnel(tunnelId).subject.asObservable(); + subscribe(tunnelId: string, options?: { readonly exclusive?: boolean }): Observable { + return this.ensureTunnel(tunnelId, options?.exclusive ?? false).subject.asObservable(); } openOutput(tunnelId: string, options: TunnelOutputOptions): AsyncIterable { @@ -328,6 +387,7 @@ export class TunnelManager { overflowError: () => new Error(SAFE_TUNNEL_OVERFLOW_ERROR), dispose: clearBytes, onConsumerCancel: () => { + if (this.byTunnelId.get(tunnelId.trim()) !== state) return; this.cancelTunnel(tunnelId); invokeFailure(options.onFailure, "cancelled"); }, @@ -336,7 +396,9 @@ export class TunnelManager { } send(msg: DeepPartial): Promise { - return this.outgoing.enqueue(msg); + const state = this.byTunnelId.get(msg.TunnelID?.trim() ?? ""); + if (!state || this.stopped) return Promise.reject(new Error(SAFE_TUNNEL_CLOSED_ERROR)); + return state.generation.outgoing.enqueue(msg); } /** @@ -382,10 +444,11 @@ export class TunnelManager { } cancelTunnel(tunnelId: string): void { - const state = this.byTunnelId.get(tunnelId); + const normalized = tunnelId.trim(); + const state = this.byTunnelId.get(normalized); if (!state) return; - this.byTunnelId.delete(tunnelId); - this.outgoing.unregister(tunnelId); + this.byTunnelId.delete(normalized); + state.generation.outgoing.unregister(normalized); if (state.bind) settleTunnelBind(state.bind, "reject", new Error(SAFE_TUNNEL_CLOSED_ERROR)); state.output?.cancel(); state.subject.complete(); @@ -394,41 +457,64 @@ export class TunnelManager { stats(): TunnelManagerStats { return { activeTunnels: this.byTunnelId.size, - ...this.outgoing.stats(), + ...(this.generation?.outgoing.stats() ?? { queuedFrames: 0, queuedBytes: 0 }), }; } async stop(): Promise { - this.abort.abort(); - this.outgoing.close(); - if (this.running) await this.running; - this.running = null; + if (this.stopped) { + if (this.generation?.running) await this.generation.running; + return; + } + this.stopped = true; this.rpc = null; - for (const tunnelId of [...this.byTunnelId.keys()]) this.failTunnel(tunnelId, "transport"); + const generation = this.generation; this.transportFailureListeners.clear(); + generation?.abort.abort(); + generation?.outgoing.close(); + this.publishTransportState("stopped", generation?.id ?? 0); + for (const tunnelId of [...this.byTunnelId.keys()]) this.failTunnel(tunnelId, "transport"); + if (generation?.running) await generation.running; } private ensureTunnel(tunnelId: string, exclusive = false): TunnelState { const normalized = tunnelId.trim(); if (!normalized) throw new Error("Tunnel id is required"); - // Establish TunnelData only when a caller actually registers a tunnel; - // regular control and artifact RPCs do not need a long-lived stream. - this.ensureRunning(); - let state = this.byTunnelId.get(normalized); if (state) { if (exclusive) throw new Error("Tunnel output is already registered"); return state; } + if (this.stopped || !this.rpc) throw new Error("Tunnel manager is not connected"); + if (this.issuedTunnelIds.has(normalized)) throw new Error("Tunnel identity has already been retired"); + if (this.issuedTunnelIds.size >= TUNNEL_MANAGER_MAX_LIFETIME_TUNNELS) { + throw new Error("Tunnel identity capacity is exhausted; reconnect the client"); + } if (this.byTunnelId.size >= TUNNEL_MANAGER_MAX_ACTIVE_TUNNELS) { throw new Error("Tunnel capacity is exhausted"); } - state = { subject: new Subject() }; + // Establish TunnelData only for a new identity; stale handles cannot + // restart a stream or gain authority over a replacement tunnel. + const generation = this.ensureRunning(); + // State observers run synchronously and may have registered a tunnel while + // handling the new generation's connecting notification. + state = this.byTunnelId.get(normalized); + if (state) { + if (exclusive) throw new Error("Tunnel output is already registered"); + return state; + } + if (this.issuedTunnelIds.has(normalized)) throw new Error("Tunnel identity has already been retired"); + if (this.issuedTunnelIds.size >= TUNNEL_MANAGER_MAX_LIFETIME_TUNNELS) { + throw new Error("Tunnel identity capacity is exhausted; reconnect the client"); + } + if (this.byTunnelId.size >= TUNNEL_MANAGER_MAX_ACTIVE_TUNNELS) throw new Error("Tunnel capacity is exhausted"); + state = { subject: new Subject(), generation }; this.byTunnelId.set(normalized, state); try { - this.outgoing.register(normalized); + generation.outgoing.register(normalized); + this.issuedTunnelIds.add(normalized); } catch (error) { this.byTunnelId.delete(normalized); throw error; @@ -436,10 +522,17 @@ export class TunnelManager { return state; } - private receive(msg: TunnelData): void { + private receive(msg: TunnelData, generation: TunnelTransportGeneration): void { + if (this.stopped || this.generation !== generation || generation.failed) { + msg.Data.fill(0); + return; + } const tunnelId = msg.TunnelID; const state = this.byTunnelId.get(tunnelId); - if (!state) return; + if (!state || state.generation !== generation) { + msg.Data.fill(0); + return; + } // Once bind() establishes an owner, TunnelID alone is no longer a valid // routing capability. Ignore foreign-session frames before overflow, @@ -457,6 +550,15 @@ export class TunnelManager { return; } + if (this.transportState.status === "connecting") { + this.publishTransportState("connected", generation.id); + // Observers may synchronously disconnect while handling readiness. + if (this.stopped || this.byTunnelId.get(tunnelId) !== state) { + msg.Data.fill(0); + return; + } + } + if (state.bind && msg.SessionID === state.bind.sessionId && msg.Data.length === 0 && !msg.Closed) { settleTunnelBind(state.bind, "resolve"); } @@ -472,7 +574,7 @@ export class TunnelManager { if (msg.Closed) { this.byTunnelId.delete(tunnelId); - this.outgoing.unregister(tunnelId); + generation.outgoing.unregister(tunnelId); if (state.bind) settleTunnelBind(state.bind, "reject", new Error(SAFE_TUNNEL_CLOSED_ERROR)); state.output?.close(); state.subject.complete(); @@ -485,7 +587,12 @@ export class TunnelManager { if (!state) return; this.byTunnelId.delete(tunnelId); const error = new Error(reason === "overflow" ? SAFE_TUNNEL_OVERFLOW_ERROR : SAFE_TUNNEL_TRANSPORT_ERROR); - this.outgoing.unregister(tunnelId, error); + state.generation.outgoing.unregister(tunnelId, error); + this.failTunnelState(state, reason); + } + + private failTunnelState(state: TunnelState, reason: "overflow" | "transport"): void { + const error = new Error(reason === "overflow" ? SAFE_TUNNEL_OVERFLOW_ERROR : SAFE_TUNNEL_TRANSPORT_ERROR); if (state.bind) settleTunnelBind(state.bind, "reject", error); state.output?.fail(error); state.subject.error(error); diff --git a/src/tests/beacon-task-api.test.ts b/src/tests/beacon-task-api.test.ts new file mode 100644 index 0000000..4ec7ffe --- /dev/null +++ b/src/tests/beacon-task-api.test.ts @@ -0,0 +1,128 @@ +import { + BeaconTaskIdentityError, + OperationAbortedError, + SliverClient, + clientpb, + sliverpb, + type SliverClientConfig, +} from "../index"; + +const taskId = "task-api"; +const beaconId = "beacon-api"; + +function metadata(overrides: Partial = {}): clientpb.BeaconTask { + return clientpb.BeaconTask.create({ + ID: taskId, + BeaconID: beaconId, + State: "completed", + Response: Buffer.from(sliverpb.Ping.encode(sliverpb.Ping.create({ Nonce: 91 })).finish()), + ...overrides, + }); +} + +function clientWithRpc(clients: Record) { + const config: SliverClientConfig = { + operator: "beacon-task-test", token: "fixture-token", lhost: "localhost", lport: 31337, + ca_certificate: "", certificate: "", private_key: "", + }; + const client = new SliverClient(config); + const internals = client as unknown as { rpcClients: Record }; + Object.assign(internals.rpcClients, clients); + return { client, clients: internals.rpcClients }; +} + +afterEach(() => jest.useRealTimers()); + +test("SliverClient task handles use replacement task-content channels after reconnect", async () => { + jest.useFakeTimers(); + const initial = { + getBeaconTaskContent: jest.fn(async () => metadata({ State: "pending", Response: Buffer.alloc(0) })), + }; + const replacement = { getBeaconTaskContent: jest.fn(async () => metadata()) }; + const { client, clients } = clientWithRpc({ "task-content": initial }); + const handle = client.beaconTask(taskId, sliverpb.Ping, { beaconId }); + const waiting = handle.wait({ timeoutSeconds: 5 }); + await jest.advanceTimersByTimeAsync(0); + delete clients["task-content"]; + await jest.advanceTimersByTimeAsync(1_000); + clients["task-content"] = replacement; + await jest.advanceTimersByTimeAsync(1_000); + + await expect(waiting).resolves.toMatchObject({ Nonce: 91 }); + expect(initial.getBeaconTaskContent).toHaveBeenCalledTimes(1); + expect(replacement.getBeaconTaskContent).toHaveBeenCalledTimes(1); + await expect(handle.status(3)).resolves.toMatchObject({ State: "completed" }); + await expect(client.beaconTask(taskId, sliverpb.Ping).wait(3)).resolves.toMatchObject({ Nonce: 91 }); + expect(jest.getTimerCount()).toBe(0); +}); + +test("SliverClient handle cancellation uses only the artifact channel", async () => { + const controlCancel = jest.fn(); + const contentCancel = jest.fn(); + const artifactCancel = jest.fn(async () => metadata({ State: "canceled", Response: Buffer.alloc(0) })); + const { client } = clientWithRpc({ + control: { cancelBeaconTask: controlCancel }, + "task-content": { cancelBeaconTask: contentCancel }, + artifact: { cancelBeaconTask: artifactCancel }, + }); + await expect(client.beaconTask(taskId, sliverpb.Ping, { beaconId }).cancel(4)) + .resolves.toMatchObject({ ID: taskId, State: "canceled" }); + expect(artifactCancel).toHaveBeenCalledWith({ ID: taskId }, { signal: expect.any(AbortSignal) }); + expect(controlCancel).not.toHaveBeenCalled(); + expect(contentCancel).not.toHaveBeenCalled(); +}); + +test("InteractiveBeacon binds handles to its own beacon identity", async () => { + const getBeaconTaskContent = jest.fn(async () => metadata({ BeaconID: "different-beacon" })); + const { client } = clientWithRpc({ control: {}, artifact: {}, "task-content": { getBeaconTaskContent } }); + const handle = client.interactBeacon(beaconId).beaconTask(taskId, sliverpb.Ping); + await expect(handle.wait(2)).rejects.toBeInstanceOf(BeaconTaskIdentityError); + await expect(handle.status(2)).rejects.toBeInstanceOf(BeaconTaskIdentityError); +}); + +test("legacy numeric lsTask calls retain wait and gain status and cancellation", async () => { + const ls = jest.fn(async () => sliverpb.Ls.create({ + Response: { Async: true, TaskID: taskId, BeaconID: beaconId, Err: "" }, + })); + const response = metadata({ + Response: Buffer.from(sliverpb.Ls.encode(sliverpb.Ls.create({ Path: "/fixture", Exists: true })).finish()), + }); + const getBeaconTaskContent = jest.fn(async () => response); + const cancelBeaconTask = jest.fn(async () => metadata({ State: "canceled", Response: Buffer.alloc(0) })); + const { client } = clientWithRpc({ + control: { ls }, artifact: { cancelBeaconTask }, "task-content": { getBeaconTaskContent }, + }); + const handle = await client.interactBeacon(beaconId).lsTask("/fixture", 7); + expect(handle.id).toBe(taskId); + expect(ls).toHaveBeenCalledWith({ + Path: "/fixture", + Request: { Async: true, Timeout: "6999999999", BeaconID: beaconId, SessionID: "" }, + }, { signal: expect.any(AbortSignal) }); + await expect(handle.wait(7)).resolves.toMatchObject({ Path: "/fixture", Exists: true }); + await expect(handle.status(7)).resolves.toMatchObject({ ID: taskId, State: "completed" }); + await expect(handle.cancel(7)).resolves.toMatchObject({ ID: taskId, State: "canceled" }); + expect(cancelBeaconTask).toHaveBeenCalledTimes(1); +}); + +test("canceling a public handle's local wait leaves another wait and remote task intact", async () => { + jest.useFakeTimers(); + let complete = false; + const getBeaconTaskContent = jest.fn(async () => complete + ? metadata() + : metadata({ State: "pending", Response: Buffer.alloc(0) })); + const cancelBeaconTask = jest.fn(); + const { client } = clientWithRpc({ "task-content": { getBeaconTaskContent }, artifact: { cancelBeaconTask } }); + const handle = client.beaconTask(taskId, sliverpb.Ping, { beaconId }); + const controller = new AbortController(); + const abandoned = handle.wait({ signal: controller.signal }); + const rejected = expect(abandoned).rejects.toBeInstanceOf(OperationAbortedError); + const retained = handle.wait(5); + await jest.advanceTimersByTimeAsync(0); + controller.abort(); + await rejected; + complete = true; + await jest.advanceTimersByTimeAsync(1_000); + await expect(retained).resolves.toMatchObject({ Nonce: 91 }); + expect(cancelBeaconTask).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); +}); diff --git a/src/tests/beacon-tasks.test.ts b/src/tests/beacon-tasks.test.ts new file mode 100644 index 0000000..996d369 --- /dev/null +++ b/src/tests/beacon-tasks.test.ts @@ -0,0 +1,307 @@ +import { Subject, of, throwError } from "rxjs"; + +import { + BeaconTaskIdentityError, + BeaconTaskResponseError, + BeaconTaskStateError, + BeaconTaskTimeoutError, + createBeaconTaskHandle, + type BeaconTaskHandleDependencies, +} from "../beaconTasks"; +import { OperationAbortedError } from "../internal/timeout"; +import { BEACON_TASK_MAX_PAYLOAD_BYTES } from "../messageBudget"; +import { BeaconTask, Event } from "../pb/clientpb/client"; +import { Ping } from "../pb/sliverpb/sliver"; + +const taskId = "task-42"; +const beaconId = "beacon-42"; + +function task(overrides: Partial = {}): BeaconTask { + return BeaconTask.create({ + ID: taskId, + BeaconID: beaconId, + State: "completed", + Response: Buffer.from(Ping.encode(Ping.create({ Nonce: 42, Response: { Err: "" } })).finish()), + ...overrides, + }); +} + +function setup(overrides: Partial> = {}) { + const events = new Subject(); + const getBeaconTaskContent = jest.fn(async (_request: unknown, _options?: { signal?: AbortSignal }) => task()); + const cancel = jest.fn(async (_id: string, _options: unknown) => task({ State: "canceled" })); + const rpc = { getBeaconTaskContent }; + const handle = createBeaconTaskHandle({ + taskId, beaconId, decoder: Ping, getRpc: () => rpc, events, cancel, + ...overrides, + }); + return { events, rpc, cancel, handle }; +} + +async function flush(): Promise { + for (let i = 0; i < 5; i += 1) await Promise.resolve(); +} + +afterEach(() => { + jest.useRealTimers(); +}); + +test("fetches a task that completed before subscribing and can wait again", async () => { + const { handle, rpc, events } = setup(); + const completed = task(); + rpc.getBeaconTaskContent.mockResolvedValue(completed); + + await expect(handle.wait()).resolves.toMatchObject({ Nonce: 42 }); + await expect(handle.wait()).resolves.toMatchObject({ Nonce: 42 }); + expect(handle.id).toBe(taskId); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(2); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledWith({ ID: taskId }, { signal: expect.any(AbortSignal) }); + expect(events.observers).toHaveLength(0); + expect(completed.Response.some((byte) => byte !== 0)).toBe(true); +}); + +test("reattaches by a persisted task id without a known beacon id", async () => { + const { handle } = setup({ beaconId: undefined }); + await expect(handle.wait()).resolves.toMatchObject({ Nonce: 42 }); +}); + +test("polling recovers a completed task when its result event was lost", async () => { + jest.useFakeTimers(); + const { handle, rpc, events } = setup(); + rpc.getBeaconTaskContent.mockResolvedValueOnce(task({ State: "pending", Response: Buffer.alloc(0) })); + const waiting = handle.wait(5); + await flush(); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toMatchObject({ Nonce: 42 }); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(2); + expect(events.observers).toHaveLength(0); + expect(jest.getTimerCount()).toBe(0); +}); + +test("task result events trigger an early refresh without waiting for polling", async () => { + const { handle, rpc, events } = setup(); + rpc.getBeaconTaskContent.mockResolvedValueOnce(task({ State: "pending", Response: Buffer.alloc(0) })); + const waiting = handle.wait(5); + await flush(); + events.next(Event.create({ EventType: "beacon-taskresult", Data: Buffer.from(BeaconTask.encode(task()).finish()) })); + await expect(waiting).resolves.toMatchObject({ Nonce: 42 }); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(2); +}); + +test("a synchronous event at subscription time does not lose or duplicate completion", async () => { + const event = Event.create({ EventType: "beacon-taskresult", Data: Buffer.from(BeaconTask.encode(task()).finish()) }); + const { handle, rpc } = setup({ events: of(event) }); + await expect(handle.wait()).resolves.toMatchObject({ Nonce: 42 }); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(1); +}); + +test.each(["canceled", "failed"] as const)("rejects terminal %s tasks with a typed state error", async (state) => { + const { handle, rpc } = setup(); + rpc.getBeaconTaskContent.mockResolvedValue(task({ State: state, Response: Buffer.alloc(0) })); + await expect(handle.wait()).rejects.toMatchObject({ + name: "BeaconTaskStateError", taskId, state, message: `Beacon task ${taskId} ended in state ${state}`, + }); + await expect(handle.wait()).rejects.toBeInstanceOf(BeaconTaskStateError); +}); + +test.each([ + { ID: "other-task" }, + { BeaconID: "other-beacon" }, + { BeaconID: "" }, +])("rejects mismatched task metadata %j", async (fields) => { + const { handle, rpc } = setup(); + rpc.getBeaconTaskContent.mockResolvedValue(task(fields)); + await expect(handle.wait()).rejects.toBeInstanceOf(BeaconTaskIdentityError); + await expect(handle.status()).rejects.toBeInstanceOf(BeaconTaskIdentityError); +}); + +test("caller abort cleans up a stalled wait without canceling the remote task", async () => { + jest.useFakeTimers(); + const { handle, rpc, events, cancel } = setup(); + rpc.getBeaconTaskContent.mockImplementation(() => new Promise(() => undefined)); + const controller = new AbortController(); + const waiting = handle.wait({ timeoutSeconds: 0, signal: controller.signal }); + const rejected = expect(waiting).rejects.toBeInstanceOf(OperationAbortedError); + await flush(); + const rpcSignal = rpc.getBeaconTaskContent.mock.calls[0]?.[1]?.signal; + controller.abort("cancel local waiting"); + await rejected; + expect(rpcSignal?.aborted).toBe(true); + expect(events.observers).toHaveLength(0); + expect(jest.getTimerCount()).toBe(0); + expect(cancel).not.toHaveBeenCalled(); +}); + +test("pre-aborted waits perform no fetch or event subscription", async () => { + const { handle, rpc, events } = setup(); + const controller = new AbortController(); + controller.abort(); + await expect(handle.wait({ signal: controller.signal })).rejects.toBeInstanceOf(OperationAbortedError); + expect(rpc.getBeaconTaskContent).not.toHaveBeenCalled(); + expect(events.observers).toHaveLength(0); +}); + +test("an aborted wait can be resumed on the same handle", async () => { + const { handle, rpc } = setup(); + rpc.getBeaconTaskContent.mockResolvedValueOnce(task({ State: "pending", Response: Buffer.alloc(0) })); + const controller = new AbortController(); + const waiting = handle.wait({ signal: controller.signal }); + const rejected = expect(waiting).rejects.toBeInstanceOf(OperationAbortedError); + await flush(); + controller.abort(); + await rejected; + await expect(handle.wait()).resolves.toMatchObject({ Nonce: 42 }); +}); + +test("a zero default timeout is preserved when only a caller signal is supplied", async () => { + jest.useFakeTimers(); + const { handle, rpc } = setup({ defaultTimeoutSeconds: 0 }); + rpc.getBeaconTaskContent.mockResolvedValue(task({ State: "pending", Response: Buffer.alloc(0) })); + const controller = new AbortController(); + let settled = false; + const waiting = handle.wait({ signal: controller.signal }); + void waiting.then(() => { settled = true; }, () => { settled = true; }); + await jest.advanceTimersByTimeAsync(31_000); + expect(settled).toBe(false); + rpc.getBeaconTaskContent.mockResolvedValue(task()); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toMatchObject({ Nonce: 42 }); + expect(jest.getTimerCount()).toBe(0); +}); + +test("time limits produce a typed task timeout and clean up polling", async () => { + jest.useFakeTimers(); + const { handle, rpc, events } = setup({ defaultTimeoutSeconds: 2 }); + rpc.getBeaconTaskContent.mockResolvedValue(task({ State: "pending", Response: Buffer.alloc(0) })); + const waiting = handle.wait({}); + const rejected = expect(waiting).rejects.toMatchObject({ + name: "BeaconTaskTimeoutError", taskId, timeoutSeconds: 2, + message: `Timeout waiting for beacon task result: ${taskId}`, + }); + await jest.advanceTimersByTimeAsync(2_000); + await rejected; + await expect(waiting).rejects.toBeInstanceOf(BeaconTaskTimeoutError); + expect(events.observers).toHaveLength(0); + expect(jest.getTimerCount()).toBe(0); +}); + +test("malformed and unrelated events do not fail or refresh another task", async () => { + jest.useFakeTimers(); + const { handle, rpc, events } = setup(); + rpc.getBeaconTaskContent.mockResolvedValueOnce(task({ State: "pending", Response: Buffer.alloc(0) })); + const waiting = handle.wait(); + await flush(); + events.next(Event.create({ EventType: "beacon-taskresult", Data: Buffer.from([0xff]) })); + events.next(Event.create({ EventType: "beacon-taskresult", Data: Buffer.from(BeaconTask.encode(task({ ID: "other" })).finish()) })); + await flush(); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(1); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toMatchObject({ Nonce: 42 }); +}); + +test("polling works when the event observable has failed", async () => { + const { handle } = setup({ events: throwError(() => new Error("stream unavailable")) }); + await expect(handle.wait()).resolves.toMatchObject({ Nonce: 42 }); +}); + +test("decoded target errors are typed and do not reveal target error text", async () => { + const { handle, rpc } = setup(); + const targetError = "sensitive target diagnostic"; + rpc.getBeaconTaskContent.mockResolvedValue(task({ + Response: Buffer.from(Ping.encode(Ping.create({ Response: { Err: targetError } })).finish()), + })); + await expect(handle.wait()).rejects.toMatchObject({ + name: "BeaconTaskResponseError", message: `Beacon task ${taskId} was rejected by the target`, + }); +}); + +test.each([ + Buffer.alloc(0), + Buffer.from([0xff]), + Buffer.alloc(BEACON_TASK_MAX_PAYLOAD_BYTES + 1), +])("rejects empty, malformed, or oversized terminal task bytes", async (bytes) => { + const { handle, rpc } = setup(); + rpc.getBeaconTaskContent.mockResolvedValue(task({ Response: bytes })); + await expect(handle.wait()).rejects.toBeInstanceOf(BeaconTaskResponseError); +}); + +test("each concurrent decoder receives independent response bytes", async () => { + const shared = task({ Response: Buffer.from([7, 8, 9]) }); + const rpc = { getBeaconTaskContent: jest.fn(async () => shared) }; + const handle = createBeaconTaskHandle({ + taskId, beaconId, getRpc: () => rpc, events: of(), + decoder: { decode: (bytes: Buffer) => bytes.subarray(0) }, + cancel: async () => shared, + }); + const [first, second] = await Promise.all([handle.wait(), handle.wait()]); + first.fill(0); + expect(second).toEqual(Buffer.from([7, 8, 9])); + expect(shared.Response).toEqual(Buffer.from([7, 8, 9])); +}); + +test("resolves a new RPC client after a temporary disconnection", async () => { + jest.useFakeTimers(); + const first = { getBeaconTaskContent: jest.fn(async () => task({ State: "pending", Response: Buffer.alloc(0) })) }; + const second = { getBeaconTaskContent: jest.fn(async () => task()) }; + let current: typeof first | undefined = first; + const { handle } = setup({ getRpc: () => current }); + const waiting = handle.wait(10); + await flush(); + current = undefined; + await jest.advanceTimersByTimeAsync(1_000); + current = second; + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toMatchObject({ Nonce: 42 }); + expect(first.getBeaconTaskContent).toHaveBeenCalledTimes(1); + expect(second.getBeaconTaskContent).toHaveBeenCalledTimes(1); +}); + +test.each([1, 14])("retries transient RPC status %s while waiting", async (code) => { + jest.useFakeTimers(); + const { handle, rpc } = setup(); + rpc.getBeaconTaskContent.mockRejectedValueOnce(Object.assign(new Error("transport closed"), { code })); + const waiting = handle.wait(); + await jest.advanceTimersByTimeAsync(1_000); + await expect(waiting).resolves.toMatchObject({ Nonce: 42 }); +}); + +test("permanent RPC errors fail immediately without leaking a polling timer", async () => { + jest.useFakeTimers(); + const { handle, rpc, events } = setup(); + const denied = Object.assign(new Error("permission denied"), { code: 7 }); + rpc.getBeaconTaskContent.mockRejectedValue(denied); + await expect(handle.wait()).rejects.toBe(denied); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(1); + expect(events.observers).toHaveLength(0); + expect(jest.getTimerCount()).toBe(0); +}); + +test("status fetches metadata and explicit cancel calls only the server cancellation callback", async () => { + const { handle, rpc, cancel } = setup(); + await expect(handle.status(4)).resolves.toMatchObject({ ID: taskId, State: "completed" }); + await expect(handle.cancel({ timeoutSeconds: 7 })).resolves.toMatchObject({ ID: taskId, State: "canceled" }); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(1); + expect(cancel).toHaveBeenCalledWith(taskId, { timeoutSeconds: 7, signal: expect.any(AbortSignal) }); +}); + +test("pre-aborted cancellation does not dispatch a remote cancellation", async () => { + const { handle, cancel } = setup(); + const controller = new AbortController(); + controller.abort(); + await expect(handle.cancel({ signal: controller.signal })).rejects.toBeInstanceOf(OperationAbortedError); + expect(cancel).not.toHaveBeenCalled(); +}); + +test("status fails clearly while disconnected and cancel verifies returned identity", async () => { + const { handle } = setup({ + getRpc: () => undefined, + cancel: async () => task({ ID: "wrong-id" }), + }); + await expect(handle.status()).rejects.toThrow("SliverClient is not connected"); + await expect(handle.cancel()).rejects.toBeInstanceOf(BeaconTaskIdentityError); +}); + +test.each(["", " ", " task", "task ", "x".repeat(257)])("rejects invalid task id %j before using RPCs", (id) => { + expect(() => setup({ taskId: id })).toThrow(TypeError); +}); diff --git a/src/tests/cancellation-api.test.ts b/src/tests/cancellation-api.test.ts new file mode 100644 index 0000000..3a2ea95 --- /dev/null +++ b/src/tests/cancellation-api.test.ts @@ -0,0 +1,156 @@ +import { EMPTY } from "rxjs"; + +import { SliverClient, type ShellSessionOptions } from "../client"; +import type { SliverClientConfig } from "../config"; +import { OperationAbortedError, OperationTimeoutError } from "../internal/timeout"; +import type { TunnelManager } from "../internal/tunnelManager"; +import type { Request } from "../pb/commonpb/common"; + +const shellOptions: ShellSessionOptions = { path: "/bin/sh", pty: false, rows: 24, cols: 80 }; + +function clientWithRpc(clients: Record, tunnels?: object): SliverClient { + const config: SliverClientConfig = { + operator: "cancellation-test", token: "token", lhost: "localhost", lport: 31337, + ca_certificate: "", certificate: "", private_key: "", + }; + const client = new SliverClient(config); + const internals = client as unknown as { rpcClients: Record; tunnels: TunnelManager }; + Object.assign(internals.rpcClients, clients); + if (tunnels) internals.tunnels = tunnels as TunnelManager; + return client; +} + +afterEach(() => jest.useRealTimers()); + +test.each([ + ["method default", 60, "59999999999"], + ["derived dump default", 90, "89999999999"], +] as const)("signal-only options preserve the %s deadline", async (kind, seconds, wireTimeout) => { + jest.useFakeTimers(); + let sentRequest: Request | undefined; + const pending = jest.fn((request: { Request?: Request }) => { + sentRequest = request.Request; + return new Promise(() => undefined); + }); + const client = clientWithRpc({ + control: { currentTokenOwner: pending }, + "workbench-artifact": { processDump: pending }, + }); + const options = { signal: new AbortController().signal }; + const operation = kind === "method default" + ? client.currentTokenOwnerSession("session", options) + : client.processDumpSession("session", 42, 60, options); + let settled = false; + const result = operation.catch((error: unknown) => { settled = true; return error; }); + expect(sentRequest?.Timeout).toBe(wireTimeout); + await jest.advanceTimersByTimeAsync((seconds - 1) * 1_000); + expect(settled).toBe(false); + await jest.advanceTimersByTimeAsync(1_000); + expect(await result).toBeInstanceOf(OperationTimeoutError); + expect(jest.getTimerCount()).toBe(0); +}); + +test("pre-aborted public operations dispatch no RPC", async () => { + const ping = jest.fn(); + const createTunnel = jest.fn(); + const client = clientWithRpc({ control: { ping, createTunnel } }); + const controller = new AbortController(); + controller.abort(); + const options = { signal: controller.signal }; + await expect(client.pingSession("session", 1, options)).rejects.toBeInstanceOf(OperationAbortedError); + await expect(client.startShellSession("session", shellOptions, options)).rejects.toBeInstanceOf(OperationAbortedError); + expect(ping).not.toHaveBeenCalled(); + expect(createTunnel).not.toHaveBeenCalled(); +}); + +test("canceling one public request leaves another request and its channel usable", async () => { + const signals: AbortSignal[] = []; + let resolveRetained!: (response: { Nonce: number }) => void; + const ping = jest.fn((request: { Nonce: number }, options: { signal: AbortSignal }) => { + signals.push(options.signal); + if (request.Nonce === 2) return new Promise<{ Nonce: number }>((resolve) => { resolveRetained = resolve; }); + return new Promise(() => undefined); + }); + const client = clientWithRpc({ control: { ping } }); + const controller = new AbortController(); + const abandoned = client.pingSession("session", 1, { signal: controller.signal }); + const rejected = expect(abandoned).rejects.toBeInstanceOf(OperationAbortedError); + const retained = client.pingSession("session", 2, 0); + controller.abort(); + await rejected; + expect(signals[0]?.aborted).toBe(true); + expect(signals[1]?.aborted).toBe(false); + expect(client.isConnected).toBe(true); + resolveRetained({ Nonce: 2 }); + await expect(retained).resolves.toEqual({ Nonce: 2 }); +}); + +test.each(["managed", "legacy"] as const)("closes a late CreateTunnel result after %s startup was canceled", async (kind) => { + jest.useFakeTimers(); + let resolveCreate!: (response: { TunnelID: string }) => void; + const createTunnel = jest.fn(() => new Promise<{ TunnelID: string }>((resolve) => { resolveCreate = resolve; })); + const closeTunnel = jest.fn(async () => ({})); + const shell = jest.fn(); + const tunnels = { + openOutput: jest.fn(() => EMPTY), + subscribe: jest.fn(() => EMPTY), + bind: jest.fn(async (_id: string, _session: string, signal?: AbortSignal) => signal?.throwIfAborted()), + send: jest.fn(async () => undefined), + cancelTunnel: jest.fn(), + }; + const client = clientWithRpc({ control: { createTunnel, closeTunnel, shell }, artifact: {} }, tunnels); + const controller = new AbortController(); + const options = { timeoutSeconds: 0, signal: controller.signal }; + const opening = kind === "managed" + ? client.startShellSession("session", shellOptions, options) + : client.interactSession("session").shell("/bin/sh", false, options); + const rejected = expect(opening).rejects.toBeInstanceOf(OperationAbortedError); + controller.abort(); + await rejected; + expect(closeTunnel).not.toHaveBeenCalled(); + + resolveCreate({ TunnelID: `late-${kind}` }); + await jest.advanceTimersByTimeAsync(2_100); + expect(shell).not.toHaveBeenCalled(); + expect(tunnels.cancelTunnel).toHaveBeenCalledWith(`late-${kind}`); + expect(closeTunnel).toHaveBeenCalledTimes(1); + expect(closeTunnel).toHaveBeenCalledWith( + { TunnelID: `late-${kind}`, SessionID: "session" }, { signal: expect.any(AbortSignal) }, + ); + expect(jest.getTimerCount()).toBe(0); +}); + +test.each(["managed", "legacy"] as const)("disposes a %s shell whose startup result races caller cancellation", async (kind) => { + jest.useFakeTimers({ doNotFake: ["queueMicrotask"] }); + const controller = new AbortController(); + const tunnelId = `handoff-${kind}`; + const closeTunnel = jest.fn(async () => ({})); + const shell = jest.fn(() => new Promise((resolve) => { + resolve({ TunnelID: tunnelId, Pid: 42, Response: { Err: "" } }); + // The RPC fulfills first; abort while its continuations are handing the + // successful shell handle back to the outer startup operation. + queueMicrotask(() => queueMicrotask(() => controller.abort())); + })); + const tunnels = { + openOutput: jest.fn(() => EMPTY), subscribe: jest.fn(() => EMPTY), + bind: jest.fn(async () => undefined), send: jest.fn(async () => undefined), + cancelTunnel: jest.fn(), + }; + const client = clientWithRpc({ + control: { createTunnel: jest.fn(async () => ({ TunnelID: tunnelId })), closeTunnel, shell }, + artifact: {}, + }, tunnels); + const options = { timeoutSeconds: 0, signal: controller.signal }; + const opening = kind === "managed" + ? client.startShellSession("session", shellOptions, options) + : client.interactSession("session").shell("/bin/sh", false, options); + await expect(opening).rejects.toBeInstanceOf(OperationAbortedError); + await jest.advanceTimersByTimeAsync(2_100); + expect(shell).toHaveBeenCalledTimes(1); + expect(tunnels.cancelTunnel).toHaveBeenCalledWith(tunnelId); + expect(closeTunnel).toHaveBeenCalledTimes(1); + expect(closeTunnel).toHaveBeenCalledWith( + { TunnelID: tunnelId, SessionID: "session" }, { signal: expect.any(AbortSignal) }, + ); + expect(jest.getTimerCount()).toBe(0); +}); diff --git a/src/tests/client-connect.test.ts b/src/tests/client-connect.test.ts index 5002ef4..1832892 100644 --- a/src/tests/client-connect.test.ts +++ b/src/tests/client-connect.test.ts @@ -13,6 +13,7 @@ jest.mock("../internal/credentials", () => ({ import { SliverClient } from "../client"; import type { SliverClientConfig } from "../config"; +import { OperationAbortedError, OperationTimeoutError } from "../internal/timeout"; function emptyStream(): AsyncIterable { return { @@ -130,3 +131,59 @@ test("SliverClient serializes concurrent connect and disconnect operations", asy expect(createSliverRpcCredentials).toHaveBeenCalledTimes(1); expect(client.isConnected).toBe(false); }); + +test("pre-aborted connect never creates credentials or channels", async () => { + const controller = new AbortController(); + controller.abort(); + await expect(new SliverClient(dummyConfig()).connect({ signal: controller.signal })) + .rejects.toBeInstanceOf(OperationAbortedError); + expect(createSliverRpcCredentials).not.toHaveBeenCalled(); + expect(createChannel).not.toHaveBeenCalled(); +}); + +test("connect deadline includes reconciliation and closes old channels before a queued reconnect", async () => { + jest.useFakeTimers(); + const channels: Array<{ close: jest.Mock }> = []; + const ordering: string[] = []; + createChannel.mockImplementation(() => { + const index = channels.length; + ordering.push(`create-${index}`); + const channel = { close: jest.fn(() => ordering.push(`close-${index}`)) }; + channels.push(channel); + return channel; + }); + const getRportFwdListeners = jest.fn() + .mockImplementationOnce((_request: unknown, options: { signal: AbortSignal }) => new Promise((_resolve, reject) => { + options.signal.addEventListener("abort", () => reject(options.signal.reason), { once: true }); + })) + .mockResolvedValue({ Listeners: [] }); + createClient.mockImplementation(() => ({ + getVersion: jest.fn(async () => ({})), + getRportFwdListeners, + events: jest.fn(() => emptyStream()), + tunnelData: jest.fn(() => emptyStream()), + })); + const client = new SliverClient(dummyConfig()); + const forward = { id: 1, sessionId: "session", markDetached: jest.fn(), reconcile: jest.fn() }; + (client as unknown as { managedReversePortForwards: Map }) + .managedReversePortForwards.set("fixture", forward); + try { + const first = client.connect(1); + const firstRejected = expect(first).rejects.toBeInstanceOf(OperationTimeoutError); + await jest.advanceTimersByTimeAsync(0); + expect(getRportFwdListeners).toHaveBeenCalledTimes(1); + const reconnected = client.connect(0); + await jest.advanceTimersByTimeAsync(1_000); + await firstRejected; + await expect(reconnected).resolves.toBe(client); + expect(client.isConnected).toBe(true); + expect(channels).toHaveLength(12); + expect(channels.slice(0, 6).every((channel) => channel.close.mock.calls.length === 1)).toBe(true); + expect(channels.slice(6).every((channel) => channel.close.mock.calls.length === 0)).toBe(true); + expect(ordering.indexOf("close-5")).toBeLessThan(ordering.indexOf("create-6")); + expect(getRportFwdListeners).toHaveBeenCalledTimes(2); + } finally { + await client.disconnect(); + jest.useRealTimers(); + } +}); diff --git a/src/tests/managed-shell.test.ts b/src/tests/managed-shell.test.ts index 1a6fe98..c9aaedd 100644 --- a/src/tests/managed-shell.test.ts +++ b/src/tests/managed-shell.test.ts @@ -15,6 +15,7 @@ import { TUNNEL_MANAGER_MAX_QUEUED_FRAMES_PER_TUNNEL, TunnelManager, } from "../internal/tunnelManager"; +import { OperationAbortedError, OperationTimeoutError } from "../internal/timeout"; const activeHarnesses = new Set(); afterEach(async () => { @@ -22,6 +23,65 @@ afterEach(async () => { }); describe("bounded managed session shell", () => { + test("uses one deadline across creation and a stalled bind acknowledgement", async () => { + jest.useFakeTimers(); + try { + const harness = new TunnelHarness(); + const shell = jest.fn(); + const closeTunnel = jest.fn(async () => ({})); + const client = clientWithShellRpc(harness.manager, { + createTunnel: jest.fn(() => new Promise((resolve) => setTimeout(() => resolve({ + TunnelID: "deadline-shell", SessionID: "deadline-session", + }), 700))), + shell, closeTunnel, + }); + harness.acknowledgeBinds = false; + const started = client.startShellSession("deadline-session", { + path: "/bin/sh", pty: true, rows: 24, cols: 80, + }, 1); + const rejected = expect(started).rejects.toBeInstanceOf(OperationTimeoutError); + await jest.advanceTimersByTimeAsync(700); + expect(harness.outbound).toHaveLength(1); + expect(shell).not.toHaveBeenCalled(); + await jest.advanceTimersByTimeAsync(300); + await rejected; + await jest.advanceTimersByTimeAsync(1_100); + expect(closeTunnel).toHaveBeenCalledTimes(1); + expect(harness.manager.stats().activeTunnels).toBe(0); + } finally { + jest.useRealTimers(); + } + }); + + test.each(["managed", "legacy"])("cancels %s startup while the transport does not consume the bind", async (kind) => { + jest.useFakeTimers(); + try { + const harness = new TunnelHarness(false); + const shell = jest.fn(); + const closeTunnel = jest.fn(async () => ({})); + const client = clientWithShellRpc(harness.manager, { + createTunnel: jest.fn(async () => ({ TunnelID: `cancel-${kind}`, SessionID: "cancel-session" })), + shell, closeTunnel, + }); + const controller = new AbortController(); + const operation = { timeoutSeconds: 0, signal: controller.signal }; + const started = kind === "managed" + ? client.startShellSession("cancel-session", { path: "/bin/sh", pty: true, rows: 24, cols: 80 }, operation) + : client.interactSession("cancel-session").shell("/bin/sh", true, operation); + const rejected = expect(started).rejects.toBeInstanceOf(OperationAbortedError); + await jest.advanceTimersByTimeAsync(1); + expect(harness.manager.stats().queuedFrames).toBe(1); + controller.abort(); + await rejected; + expect(shell).not.toHaveBeenCalled(); + await jest.advanceTimersByTimeAsync(2_100); + expect(closeTunnel).toHaveBeenCalledTimes(1); + expect(harness.manager.stats().activeTunnels).toBe(0); + } finally { + jest.useRealTimers(); + } + }); + test("registers early output before the exact bind and shell sequence", async () => { const order: string[] = []; const harness = new TunnelHarness(true, (message) => { @@ -248,6 +308,9 @@ describe("bounded managed session shell", () => { }; const tunnels = { openOutput: jest.fn(() => output), + bind: jest.fn(async (tunnelId: string, sessionId: string) => send({ + TunnelID: tunnelId, SessionID: sessionId, Data: Buffer.alloc(0), + })), send, cancelTunnel: jest.fn(), }; @@ -431,6 +494,7 @@ describe("legacy interactive session shell compatibility", () => { }); const tunnels = { subscribe: jest.fn(() => EMPTY), + bind: jest.fn(async () => undefined), send, cancelTunnel: jest.fn(), }; @@ -897,6 +961,7 @@ describe("bounded fair tunnel manager", () => { }); class TunnelHarness { + acknowledgeBinds = false; readonly manager = new TunnelManager(); readonly incoming = new PushAsyncIterable(); readonly outbound: Array> = []; @@ -954,6 +1019,9 @@ class TunnelHarness { : result.value; this.outbound.push(message); this.onOutbound?.(message); + if (this.acknowledgeBinds && message.Data?.length === 0 && message.SessionID) { + this.incoming.push(tunnelMessage(message.TunnelID ?? "", message.SessionID, Buffer.alloc(0))); + } if (this.echoRemoteClose && Buffer.from(message.Data ?? []).toString("utf8") === "logout\n") { this.incoming.push(tunnelMessage( message.TunnelID ?? "", message.SessionID ?? "", Buffer.alloc(0), true, @@ -1022,6 +1090,8 @@ function clientWithShellRpc( manager: TunnelManager, control: Record, ): SliverClient { + const harness = [...activeHarnesses].find((candidate) => candidate.manager === manager); + if (harness) harness.acknowledgeBinds = true; const client = new SliverClient(baseConfig()); const internals = client as unknown as { rpcClients: Record; tunnels: TunnelManager }; internals.rpcClients.control = control; diff --git a/src/tests/timeout.test.ts b/src/tests/timeout.test.ts index 27f3630..eb41569 100644 --- a/src/tests/timeout.test.ts +++ b/src/tests/timeout.test.ts @@ -1,7 +1,14 @@ import { SliverClient } from "../client"; import type { SliverClientConfig } from "../config"; import { Version } from "../pb/clientpb/client"; -import { timeoutSecondsToNanoseconds, withTimeoutSignal } from "../internal/timeout"; +import { + OperationAbortedError, + OperationTimeoutError, + operationOptions, + operationTimeoutSeconds, + timeoutSecondsToNanoseconds, + withTimeoutSignal, +} from "../internal/timeout"; const MAX_TIMEOUT_SECONDS = 2_147_483; const invalidTimeouts = [ @@ -47,6 +54,244 @@ describe("Sliver operation timeout encoding", () => { await expect(withTimeoutSignal(MAX_TIMEOUT_SECONDS, callback)).resolves.toBe(false); expect(callback).toHaveBeenCalledTimes(2); }); + + test("uses the same validated default for options and protobuf encoding", () => { + const controller = new AbortController(); + expect(operationTimeoutSeconds({})).toBe(30); + expect(operationTimeoutSeconds({ signal: controller.signal }, 60)).toBe(60); + expect(operationTimeoutSeconds({ timeoutSeconds: 0 }, 60)).toBe(0); + expect(timeoutSecondsToNanoseconds({})).toBe("29999999999"); + expect(timeoutSecondsToNanoseconds({ timeoutSeconds: 1 })).toBe("999999999"); + expect(operationOptions({ signal: controller.signal }, 60)).toEqual({ + timeoutSeconds: 60, + signal: controller.signal, + }); + expect(operationOptions(5, 60)).toEqual({ timeoutSeconds: 5 }); + expect(() => operationTimeoutSeconds({}, -1)).toThrow(RangeError); + }); + + test.each(invalidTimeouts)("rejects invalid options timeout %s before dispatch", async (timeoutSeconds) => { + const callback = jest.fn(async () => "called"); + await expect(withTimeoutSignal({ timeoutSeconds }, callback)).rejects.toBeInstanceOf(RangeError); + expect(() => timeoutSecondsToNanoseconds({ timeoutSeconds })).toThrow(RangeError); + expect(callback).not.toHaveBeenCalled(); + }); +}); + +describe("operation cancellation and deadlines", () => { + beforeEach(() => jest.useFakeTimers()); + afterEach(() => { + jest.useRealTimers(); + jest.restoreAllMocks(); + }); + + test("rejects pre-aborted callers without dispatching or registering resources", async () => { + const controller = new AbortController(); + const reason = new Error("caller cancelled"); + controller.abort(reason); + const add = jest.spyOn(controller.signal, "addEventListener"); + const callback = jest.fn(async () => "called"); + + await expect(withTimeoutSignal({ signal: controller.signal }, callback)).rejects.toMatchObject({ + name: "OperationAbortedError", + cause: reason, + }); + expect(callback).not.toHaveBeenCalled(); + expect(add).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + }); + + test("aborts a stalled callback with no deadline and observes its late rejection", async () => { + const controller = new AbortController(); + const remove = jest.spyOn(controller.signal, "removeEventListener"); + let derived!: AbortSignal; + let rejectLate!: (reason: Error) => void; + const operation = withTimeoutSignal({ timeoutSeconds: 0, signal: controller.signal }, (signal) => { + derived = signal; + return new Promise((_resolve, reject) => { rejectLate = reject; }); + }); + const result = operation.catch((error: unknown) => error); + expect(jest.getTimerCount()).toBe(0); + + const reason = new Error("stop this one"); + controller.abort(reason); + const failure = await result; + expect(failure).toBeInstanceOf(OperationAbortedError); + expect(failure).toMatchObject({ cause: reason }); + expect(derived.aborted).toBe(true); + expect(derived.reason).toBe(failure); + expect(remove).toHaveBeenCalledTimes(1); + + rejectLate(new Error("late transport failure")); + await jest.advanceTimersByTimeAsync(0); + expect(jest.getTimerCount()).toBe(0); + }); + + test("classifies a numeric deadline even when transport rejects on abort", async () => { + let derived!: AbortSignal; + const operation = withTimeoutSignal(2, (signal) => { + derived = signal; + return new Promise((_resolve, reject) => { + signal.addEventListener("abort", () => reject(new Error("transport cancelled")), { once: true }); + }); + }); + const result = operation.catch((error: unknown) => error); + + await jest.advanceTimersByTimeAsync(1_999); + expect(derived.aborted).toBe(false); + await jest.advanceTimersByTimeAsync(1); + const failure = await result; + expect(failure).toBeInstanceOf(OperationTimeoutError); + expect(failure).toMatchObject({ timeoutSeconds: 2 }); + expect(derived.reason).toBe(failure); + expect(jest.getTimerCount()).toBe(0); + }); + + test("expires signal-only options after the default 30 seconds without aborting the caller", async () => { + const controller = new AbortController(); + const remove = jest.spyOn(controller.signal, "removeEventListener"); + const operation = withTimeoutSignal({ signal: controller.signal }, () => new Promise(() => {})); + const result = operation.catch((error: unknown) => error); + + await jest.advanceTimersByTimeAsync(30_000); + expect(await result).toMatchObject({ name: "OperationTimeoutError", timeoutSeconds: 30 }); + expect(controller.signal.aborted).toBe(false); + expect(remove).toHaveBeenCalledTimes(1); + expect(jest.getTimerCount()).toBe(0); + }); + + test.each(["success", "rejection", "throw"] as const)("removes the caller listener and timer after callback %s", async (outcome) => { + const controller = new AbortController(); + const add = jest.spyOn(controller.signal, "addEventListener"); + const remove = jest.spyOn(controller.signal, "removeEventListener"); + const failure = new Error("callback failure"); + let derived!: AbortSignal; + const operation = withTimeoutSignal({ timeoutSeconds: 60, signal: controller.signal }, (signal) => { + derived = signal; + if (outcome === "throw") throw failure; + if (outcome === "rejection") return Promise.reject(failure); + return Promise.resolve("value"); + }); + if (outcome === "success") await expect(operation).resolves.toBe("value"); + else await expect(operation).rejects.toBe(failure); + + expect(remove).toHaveBeenCalledWith("abort", add.mock.calls[0]![1]); + expect(jest.getTimerCount()).toBe(0); + controller.abort(); + await jest.advanceTimersByTimeAsync(60_000); + expect(derived.aborted).toBe(false); + }); + + test("keeps sibling operations independent when one deadline expires", async () => { + const controller = new AbortController(); + let firstSignal!: AbortSignal; + let secondSignal!: AbortSignal; + let resolveSecond!: (value: string) => void; + const first = withTimeoutSignal({ timeoutSeconds: 1, signal: controller.signal }, (signal) => { + firstSignal = signal; + return new Promise(() => {}); + }); + const firstResult = first.catch((error: unknown) => error); + const second = withTimeoutSignal({ timeoutSeconds: 5, signal: controller.signal }, (signal) => { + secondSignal = signal; + return new Promise((resolve) => { resolveSecond = resolve; }); + }); + + await jest.advanceTimersByTimeAsync(1_000); + expect(await firstResult).toBeInstanceOf(OperationTimeoutError); + expect(firstSignal.aborted).toBe(true); + expect(secondSignal.aborted).toBe(false); + expect(controller.signal.aborted).toBe(false); + resolveSecond("still running"); + await expect(second).resolves.toBe("still running"); + expect(jest.getTimerCount()).toBe(0); + }); + + test("caller cancellation wins even when the callback synchronously returns success afterward", async () => { + const controller = new AbortController(); + const operation = withTimeoutSignal({ signal: controller.signal }, async () => { + controller.abort("cancel during dispatch"); + return "late success"; + }); + await expect(operation).rejects.toBeInstanceOf(OperationAbortedError); + expect(jest.getTimerCount()).toBe(0); + }); + + test("caller cancellation clears its pending deadline", async () => { + const controller = new AbortController(); + const operation = withTimeoutSignal({ timeoutSeconds: 5, signal: controller.signal }, () => new Promise(() => {})); + const result = operation.catch((error: unknown) => error); + expect(jest.getTimerCount()).toBe(1); + controller.abort(); + expect(await result).toBeInstanceOf(OperationAbortedError); + expect(jest.getTimerCount()).toBe(0); + await jest.advanceTimersByTimeAsync(5_000); + expect(await result).toBeInstanceOf(OperationAbortedError); + }); + + test("rejects malformed options and signals before dispatch", async () => { + const callback = jest.fn(async () => "called"); + for (const options of [null, [], "5", { signal: {} }, { timeoutSeconds: null }]) { + await expect(withTimeoutSignal(options as never, callback)).rejects.toBeInstanceOf(Error); + } + expect(callback).not.toHaveBeenCalled(); + expect(jest.getTimerCount()).toBe(0); + }); + + test("does not dispose a resource delivered successfully", async () => { + const resource = { close: jest.fn(async () => undefined) }; + const dispose = jest.fn(async (value: typeof resource) => value.close()); + await expect(withTimeoutSignal(5, async () => resource, dispose)).resolves.toBe(resource); + await jest.advanceTimersByTimeAsync(5_000); + expect(dispose).not.toHaveBeenCalled(); + expect(resource.close).not.toHaveBeenCalled(); + }); + + test.each(["caller", "deadline"] as const)("disposes a late resource exactly once after %s cancellation", async (cause) => { + const controller = new AbortController(); + const resource = { close: jest.fn(async () => undefined) }; + const dispose = jest.fn(async (value: typeof resource) => value.close()); + let resolveLate!: (value: typeof resource) => void; + const operation = withTimeoutSignal({ timeoutSeconds: 1, signal: controller.signal }, () => + new Promise((resolve) => { resolveLate = resolve; }), dispose, + ); + const result = operation.catch((error: unknown) => error); + if (cause === "caller") controller.abort(); + else await jest.advanceTimersByTimeAsync(1_000); + expect(await result).toBeInstanceOf(cause === "caller" ? OperationAbortedError : OperationTimeoutError); + expect(dispose).not.toHaveBeenCalled(); + + resolveLate(resource); + await jest.advanceTimersByTimeAsync(0); + expect(dispose).toHaveBeenCalledTimes(1); + expect(dispose).toHaveBeenCalledWith(resource); + expect(resource.close).toHaveBeenCalledTimes(1); + controller.abort(); + await jest.advanceTimersByTimeAsync(1_000); + expect(dispose).toHaveBeenCalledTimes(1); + expect(jest.getTimerCount()).toBe(0); + }); + + test.each(["throw", "reject"] as const)("observes a late disposer %s without changing cancellation", async (outcome) => { + const controller = new AbortController(); + let resolveLate!: (value: object) => void; + const dispose = jest.fn(() => { + if (outcome === "throw") throw new Error("private cleanup detail"); + return Promise.reject(new Error("private cleanup detail")); + }); + const operation = withTimeoutSignal({ signal: controller.signal }, () => + new Promise((resolve) => { resolveLate = resolve; }), dispose, + ); + const result = operation.catch((error: unknown) => error); + controller.abort("cancel"); + const cancellation = await result; + expect(cancellation).toBeInstanceOf(OperationAbortedError); + resolveLate({}); + await jest.advanceTimersByTimeAsync(0); + expect(dispose).toHaveBeenCalledTimes(1); + expect(await result).toBe(cancellation); + expect(jest.getTimerCount()).toBe(0); + }); }); function clientWithControlRpc(controlRpc: object): SliverClient { diff --git a/src/tests/tunnel-recovery.test.ts b/src/tests/tunnel-recovery.test.ts new file mode 100644 index 0000000..0c4786e --- /dev/null +++ b/src/tests/tunnel-recovery.test.ts @@ -0,0 +1,253 @@ +import type { SliverRPCClient, DeepPartial } from "../pb/rpcpb/services"; +import { TunnelData } from "../pb/sliverpb/sliver"; +import { + TUNNEL_MANAGER_MAX_LIFETIME_TUNNELS, + TunnelManager, + type TunnelTransportState, +} from "../internal/tunnelManager"; + +const managers = new Set(); +afterEach(async () => { + await Promise.all([...managers].map((manager) => manager.stop())); + managers.clear(); +}); + +test("fails old output and queued writes, then lazily binds a fresh tunnel without replay", async () => { + const { manager, streams } = harness(); + const states: TunnelTransportState[] = []; + manager.transportState$.subscribe((state) => states.push(state)); + const oldFailure = jest.fn(); + const oldOutput = manager.openOutput("old", { maxBufferedBytes: 64, onFailure: oldFailure }); + const oldRead = oldOutput[Symbol.asyncIterator]().next().catch((error: Error) => error.message); + const oldWrite = manager.send(frame("old", "queued-secret")).catch((error: Error) => error.message); + + streams[0]!.incoming.fail(); + await flush(); + expect(await oldRead).toBe("Tunnel transport disconnected"); + expect(await oldWrite).toBe("Tunnel transport disconnected"); + expect(oldFailure).toHaveBeenCalledWith("transport"); + expect(manager.stats()).toEqual({ activeTunnels: 0, queuedFrames: 0, queuedBytes: 0 }); + expect(manager.transportState).toEqual({ status: "failed", generation: 1 }); + await flush(); + expect(streams).toHaveLength(1); + + const output = manager.openOutput("new", { maxBufferedBytes: 64 }); + const binding = manager.bind("new", "session"); + expect(streams).toHaveLength(2); + const first = await streams[1]!.outgoing.next(); + expect(first.value).toMatchObject({ TunnelID: "new", SessionID: "session", Data: Buffer.alloc(0) }); + const successor = streams[1]!.outgoing.next().catch(() => ({ done: true })); + streams[1]!.incoming.push(frame("new")); + await binding; + streams[1]!.incoming.push(frame("new", "fresh")); + await expect(output[Symbol.asyncIterator]().next()).resolves.toEqual({ value: Uint8Array.from(Buffer.from("fresh")), done: false }); + expect(states.map(({ status, generation }) => `${generation}:${status}`)).toEqual([ + "0:idle", "1:connecting", "1:failed", "2:connecting", "2:connected", + ]); + expect(states.every(Object.isFrozen)).toBe(true); + await manager.stop(); + await successor; +}); + +test("retires old listeners while allowing new owners to await the next generation", async () => { + const { manager, streams } = harness(); + const oldListener = jest.fn(); + manager.onTransportFailure(oldListener); + manager.openOutput("first", { maxBufferedBytes: 64 }); + streams[0]!.incoming.fail(); + await flush(); + expect(oldListener).toHaveBeenCalledTimes(1); + const lateListener = jest.fn(); + manager.onTransportFailure(lateListener); + expect(lateListener).toHaveBeenCalledTimes(1); + + const newListener = jest.fn(); + manager.onTransportFailure(newListener, { nextGeneration: true }); + expect(newListener).not.toHaveBeenCalled(); + expect(streams).toHaveLength(1); + manager.openOutput("second", { maxBufferedBytes: 64 }); + streams[1]!.incoming.fail(); + await flush(); + expect(newListener).toHaveBeenCalledTimes(1); + expect(oldListener).toHaveBeenCalledTimes(1); + expect(lateListener).toHaveBeenCalledTimes(1); + manager.openOutput("third", { maxBufferedBytes: 64 }); + expect(manager.transportState).toEqual({ status: "connecting", generation: 3 }); + expect(streams).toHaveLength(3); +}); + +test("stale identities and old consumer cleanup cannot restart or affect a new generation", async () => { + const { manager, streams } = harness(); + const oldFailure = jest.fn(); + const old = manager.openOutput("old", { maxBufferedBytes: 64, onFailure: oldFailure }); + const iterator = old[Symbol.asyncIterator](); + streams[0]!.incoming.fail(); + await flush(); + expect(() => manager.openOutput("old", { maxBufferedBytes: 64 })).toThrow("retired"); + await expect(manager.bind("old", "session")).rejects.toThrow("retired"); + await expect(manager.send(frame("old", "late"))).rejects.toThrow("Tunnel is closed"); + expect(streams).toHaveLength(1); + + const current = manager.openOutput("new", { maxBufferedBytes: 64 }); + manager.cancelTunnel("old"); + await iterator.return?.(); + expect(oldFailure).toHaveBeenCalledTimes(1); + const read = current[Symbol.asyncIterator]().next(); + // The failed stream retains its own iterator and cannot route into its replacement. + streams[0]!.incoming.push(frame("new", "stale", true)); + streams[1]!.incoming.push(frame("new", "current")); + await expect(read).resolves.toEqual({ value: Uint8Array.from(Buffer.from("current")), done: false }); + expect(manager.stats().activeTunnels).toBe(1); +}); + +test("reentrant failure callbacks can open a replacement without old cleanup retiring it", async () => { + const { manager, streams } = harness(); + let replacement: AsyncIterable | undefined; + manager.openOutput("old", { + maxBufferedBytes: 64, + onFailure: () => { replacement = manager.openOutput("replacement", { maxBufferedBytes: 64 }); }, + }); + manager.onTransportFailure(() => manager.cancelTunnel("old")); + streams[0]!.incoming.fail(); + await flush(); + expect(streams).toHaveLength(2); + expect(manager.transportState).toEqual({ status: "connecting", generation: 2 }); + streams[1]!.incoming.push(frame("replacement", "ready")); + await expect(replacement![Symbol.asyncIterator]().next()).resolves.toMatchObject({ done: false }); +}); + +test("cancelled binds do not recover automatically and stop permanently prevents restart", async () => { + const { manager, streams } = harness(); + manager.openOutput("cancel", { maxBufferedBytes: 64 }); + const controller = new AbortController(); + const binding = manager.bind("cancel", "session", controller.signal); + controller.abort(); + await expect(binding).rejects.toThrow("cancelled"); + manager.cancelTunnel("cancel"); + await flush(); + expect(streams).toHaveLength(1); + await manager.stop(); + expect(manager.transportState).toEqual({ status: "stopped", generation: 1 }); + expect(() => manager.openOutput("after-stop", { maxBufferedBytes: 64 })).toThrow("not connected"); + expect(() => manager.start({} as SliverRPCClient)).toThrow("stopped"); + await expect(manager.send(frame("cancel"))).rejects.toThrow("closed"); + expect(streams).toHaveLength(1); +}); + +test("synchronous transport startup failure reaches the initiating output and can recover", async () => { + const manager = new TunnelManager(); + managers.add(manager); + const incoming = new Incoming(); + const tunnelData = jest.fn() + .mockImplementationOnce(() => { throw new Error("private transport detail"); }) + .mockImplementation((_outgoing, options: { signal: AbortSignal }) => { + options.signal.addEventListener("abort", () => incoming.close(), { once: true }); + return incoming; + }); + manager.start({ tunnelData } as unknown as SliverRPCClient); + const first = manager.openOutput("first", { maxBufferedBytes: 64 }); + await expect(first[Symbol.asyncIterator]().next()).rejects.toThrow("Tunnel transport disconnected"); + const second = manager.openOutput("second", { maxBufferedBytes: 64 }); + incoming.push(frame("second", "ok")); + await expect(second[Symbol.asyncIterator]().next()).resolves.toMatchObject({ done: false }); + expect(tunnelData).toHaveBeenCalledTimes(2); +}); + +test("bounds retired identity storage and requests reconnect when its lifetime budget is exhausted", () => { + const { manager, streams } = harness(); + for (let index = 0; index < TUNNEL_MANAGER_MAX_LIFETIME_TUNNELS; index += 1) { + const id = `retired-${index}`; + manager.subscribe(id); + manager.cancelTunnel(id); + } + expect(manager.stats()).toEqual({ activeTunnels: 0, queuedFrames: 0, queuedBytes: 0 }); + expect(() => manager.openOutput("over-budget", { maxBufferedBytes: 64 })).toThrow("reconnect the client"); + expect(streams).toHaveLength(1); +}); + +test("exclusive legacy subscriptions reject collisions without changing the original owner", async () => { + const { manager, streams } = harness(); + const original: string[] = []; + const originalClosed = jest.fn(); + const originalFailed = jest.fn(); + manager.subscribe("owned", { exclusive: true }).subscribe({ + next: (message) => original.push(message.Data.toString()), + complete: originalClosed, + error: originalFailed, + }); + expect(() => manager.subscribe(" owned ", { exclusive: true })).toThrow("already registered"); + const observer: string[] = []; + manager.subscribe("owned").subscribe({ + next: (message) => observer.push(message.Data.toString()), + error: () => undefined, + }); + const independent: string[] = []; + manager.subscribe("new", { exclusive: true }).subscribe({ + next: (message) => independent.push(message.Data.toString()), + error: () => undefined, + }); + streams[0]!.incoming.push(frame("owned", "original-still-live")); + streams[0]!.incoming.push(frame("new", "new-owner")); + await flush(); + expect(original).toEqual(["original-still-live"]); + expect(observer).toEqual(original); + expect(independent).toEqual(["new-owner"]); + expect(originalClosed).not.toHaveBeenCalled(); + expect(originalFailed).not.toHaveBeenCalled(); + expect(manager.stats().activeTunnels).toBe(2); + expect(streams).toHaveLength(1); +}); + +function harness() { + const manager = new TunnelManager(); + managers.add(manager); + const streams: Array<{ incoming: Incoming; outgoing: AsyncIterator> }> = []; + manager.start({ + tunnelData: (outgoing: AsyncIterable>, options: { signal: AbortSignal }) => { + const incoming = new Incoming(); + options.signal.addEventListener("abort", () => incoming.close(), { once: true }); + streams.push({ incoming, outgoing: outgoing[Symbol.asyncIterator]() }); + return incoming; + }, + } as unknown as SliverRPCClient); + return { manager, streams }; +} + +class Incoming implements AsyncIterable { + private queued: TunnelData[] = []; + private waiter?: { resolve: (result: IteratorResult) => void; reject: (error: Error) => void }; + private stopped = false; + push(value: TunnelData): void { + if (this.stopped) return; + const waiter = this.waiter; + this.waiter = undefined; + if (waiter) waiter.resolve({ value, done: false }); + else this.queued.push(value); + } + fail(): void { + this.stopped = true; + this.waiter?.reject(new Error("private transport detail")); + this.waiter = undefined; + } + close(): void { + this.stopped = true; + this.waiter?.resolve({ value: undefined as never, done: true }); + this.waiter = undefined; + } + [Symbol.asyncIterator](): AsyncIterator { + return { next: () => { + const value = this.queued.shift(); + if (value) return Promise.resolve({ value, done: false }); + if (this.stopped) return Promise.resolve({ value: undefined as never, done: true }); + return new Promise((resolve, reject) => { this.waiter = { resolve, reject }; }); + } }; + } +} + +function frame(id: string, data = "", closed = false): TunnelData { + return TunnelData.create({ TunnelID: id, SessionID: "session", Data: Buffer.from(data), Closed: closed }); +} + +async function flush(): Promise { + for (let index = 0; index < 8; index += 1) await Promise.resolve(); +} From fe7f958d938087ee086e3754321c5d01ab422303 Mon Sep 17 00:00:00 2001 From: moloch-- <875022+moloch--@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:44:55 -0700 Subject: [PATCH 2/3] fix(deps): remove vulnerable development dependency paths --- package-lock.json | 66 ++++++++++++++++------------------------------- package.json | 3 +++ 2 files changed, 25 insertions(+), 44 deletions(-) diff --git a/package-lock.json b/package-lock.json index a7893e0..41c5afe 100644 --- a/package-lock.json +++ b/package-lock.json @@ -2603,14 +2603,11 @@ } }, "node_modules/argparse": { - "version": "1.0.10", - "resolved": "https://registry.npmjs.org/argparse/-/argparse-1.0.10.tgz", - "integrity": "sha512-o5Roy6tNG4SL/FOkCAN6RzjiakZS25RLYFrcMttJqbdd8BWrnA+fGz57iN5Pb06pvBGvl5gQ0B48dJlslXvoTg==", + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/argparse/-/argparse-2.0.1.tgz", + "integrity": "sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q==", "dev": true, - "license": "MIT", - "dependencies": { - "sprintf-js": "~1.0.2" - } + "license": "Python-2.0" }, "node_modules/babel-jest": { "version": "30.5.1", @@ -3148,20 +3145,6 @@ "node": ">=8" } }, - "node_modules/esprima": { - "version": "4.0.1", - "resolved": "https://registry.npmjs.org/esprima/-/esprima-4.0.1.tgz", - "integrity": "sha512-eGuFFw7Upda+g4p+QHvnW0RyTX/SVeJBDM/gCtMARO0cLuT2HcEKnTPvhjV6aGeqrCB/sbNop0Kszm0jsaWU4A==", - "dev": true, - "license": "BSD-2-Clause", - "bin": { - "esparse": "bin/esparse.js", - "esvalidate": "bin/esvalidate.js" - }, - "engines": { - "node": ">=4" - } - }, "node_modules/execa": { "version": "5.1.1", "resolved": "https://registry.npmjs.org/execa/-/execa-5.1.1.tgz", @@ -3361,13 +3344,13 @@ "license": "ISC" }, "node_modules/handlebars": { - "version": "4.7.9", - "resolved": "https://registry.npmjs.org/handlebars/-/handlebars-4.7.9.tgz", - "integrity": "sha512-4E71E0rpOaQuJR2A3xDZ+GM1HyWYv1clR58tC8emQNeQe3RH7MAzSbat+V0wG78LQBo6m6bzSG/L4pBuCsgnUQ==", + "version": "4.7.10", + "resolved": "https://registry.npmjs.org/handlebars/-/handlebars-4.7.10.tgz", + "integrity": "sha512-P5VJMVM7qgBn6vjXMw8WG9uVI+ncf2pi72j4de4yz5ZULLj2RGqLYaKOYGsgyrViQ0tePOVlN1tDCCXXtFqXKg==", "dev": true, "license": "MIT", "dependencies": { - "minimist": "^1.2.5", + "minimist": "^1.2.8", "neo-async": "^2.6.2", "source-map": "^0.6.1", "wordwrap": "^1.0.0" @@ -4190,14 +4173,23 @@ "license": "MIT" }, "node_modules/js-yaml": { - "version": "3.15.2", - "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-3.15.2.tgz", - "integrity": "sha512-6EuL879VkRA+1Cz578mKMiKvjPNEuk6+r1JaFzoSWejZmtf7xWbIyw1e3KkxlkzTIt9Taw6JBhEppG7utc1P+w==", + "version": "4.3.2", + "resolved": "https://registry.npmjs.org/js-yaml/-/js-yaml-4.3.2.tgz", + "integrity": "sha512-SFNOvSJ+Dgf/9An904Yx+CgSlIPCkIpao4qo51lpee25TIRejdH3rhR4EZMGoNx3/TP3O+wzWuiTFl4sqbltzA==", "dev": true, + "funding": [ + { + "type": "github", + "url": "https://github.com/sponsors/puzrin" + }, + { + "type": "github", + "url": "https://github.com/sponsors/nodeca" + } + ], "license": "MIT", "dependencies": { - "argparse": "^1.0.7", - "esprima": "^4.0.0" + "argparse": "^2.0.1" }, "bin": { "js-yaml": "bin/js-yaml.js" @@ -4386,13 +4378,6 @@ "markdown-it": "bin/markdown-it.mjs" } }, - "node_modules/markdown-it/node_modules/argparse": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/argparse/-/argparse-2.0.1.tgz", - "integrity": "sha512-8+9WqebbFzpX9OR+Wa6O29asIogeRMzcGtAINdpMHHyAg10f05aSFVBbcEqGf/PXw1EjAZ+q2/bEBg3DvurK3Q==", - "dev": true, - "license": "Python-2.0" - }, "node_modules/mdurl": { "version": "2.1.0", "resolved": "https://registry.npmjs.org/mdurl/-/mdurl-2.1.0.tgz", @@ -4924,13 +4909,6 @@ "node": ">=0.10.0" } }, - "node_modules/sprintf-js": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/sprintf-js/-/sprintf-js-1.0.3.tgz", - "integrity": "sha512-D9cPgkvLlV3t3IzL0D0YLvGA9Ahk4PcvVwUbN0dSGr1aP0Nrt4AEnTUbuGvquEC0mA64Gqt1fzirlRs5ibXx8g==", - "dev": true, - "license": "BSD-3-Clause" - }, "node_modules/stack-utils": { "version": "2.0.6", "resolved": "https://registry.npmjs.org/stack-utils/-/stack-utils-2.0.6.tgz", diff --git a/package.json b/package.json index dab815c..3d0f2cc 100644 --- a/package.json +++ b/package.json @@ -69,6 +69,9 @@ "rxjs": "^7.8.2" }, "overrides": { + "@istanbuljs/load-nyc-config": { + "js-yaml": "4.3.2" + }, "dprint-node": "1.0.8", "ts-poet": "6.12.0" } From ec0b5f946da9e3495b98ac2c824eb71376f74a7b Mon Sep 17 00:00:00 2001 From: moloch-- <875022+moloch--@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:46:28 -0700 Subject: [PATCH 3/3] fix: verify beacon task ownership before cancellation --- src/beaconTasks.ts | 11 ++-- src/internal/timeout.ts | 2 +- src/tests/beacon-task-api.test.ts | 9 +++- src/tests/beacon-tasks.test.ts | 87 ++++++++++++++++++++++++++++--- 4 files changed, 95 insertions(+), 14 deletions(-) diff --git a/src/beaconTasks.ts b/src/beaconTasks.ts index dcc340c..9f28ebb 100644 --- a/src/beaconTasks.ts +++ b/src/beaconTasks.ts @@ -221,9 +221,14 @@ export function createBeaconTaskHandle(dependencies: BeaconTaskHandleDependen status: (options?: OperationTimeout) => withTimeoutSignal(optionsFor(options), fetch), cancel: async (options?: OperationTimeout) => { const normalized = optionsFor(options); - return withTimeoutSignal(normalized, async (signal) => validateIdentity( - await dependencies.cancel(taskId, { ...normalized, signal }), - )); + return withTimeoutSignal(normalized, async (signal) => { + // Verify the beacon binding before a remote mutation, under the same + // deadline that covers cancellation. Checking only its response would + // detect a mismatched beacon after that beacon's task was canceled. + if (beaconId !== undefined) await fetch(signal); + signal.throwIfAborted(); + return validateIdentity(await dependencies.cancel(taskId, { ...normalized, signal })); + }); }, }); } diff --git a/src/internal/timeout.ts b/src/internal/timeout.ts index 5f8b362..c80746a 100644 --- a/src/internal/timeout.ts +++ b/src/internal/timeout.ts @@ -10,7 +10,7 @@ const MAX_TIMEOUT_SECONDS = Math.min(MAX_TIMER_TIMEOUT_SECONDS, MAX_PROTOBUF_TIM /** A local operation deadline and optional caller cancellation signal. */ export interface OperationOptions { - /** Defaults to 30 seconds. Zero disables the local deadline. */ + /** Defaults to the called method's timeout. Zero disables the local deadline. */ timeoutSeconds?: number; /** Cancel this operation without disconnecting unrelated operations. */ signal?: AbortSignal; diff --git a/src/tests/beacon-task-api.test.ts b/src/tests/beacon-task-api.test.ts index 4ec7ffe..69fbca0 100644 --- a/src/tests/beacon-task-api.test.ts +++ b/src/tests/beacon-task-api.test.ts @@ -60,24 +60,29 @@ test("SliverClient handle cancellation uses only the artifact channel", async () const controlCancel = jest.fn(); const contentCancel = jest.fn(); const artifactCancel = jest.fn(async () => metadata({ State: "canceled", Response: Buffer.alloc(0) })); + const getBeaconTaskContent = jest.fn(async () => metadata()); const { client } = clientWithRpc({ control: { cancelBeaconTask: controlCancel }, - "task-content": { cancelBeaconTask: contentCancel }, + "task-content": { cancelBeaconTask: contentCancel, getBeaconTaskContent }, artifact: { cancelBeaconTask: artifactCancel }, }); await expect(client.beaconTask(taskId, sliverpb.Ping, { beaconId }).cancel(4)) .resolves.toMatchObject({ ID: taskId, State: "canceled" }); expect(artifactCancel).toHaveBeenCalledWith({ ID: taskId }, { signal: expect.any(AbortSignal) }); + expect(getBeaconTaskContent).toHaveBeenCalledTimes(1); expect(controlCancel).not.toHaveBeenCalled(); expect(contentCancel).not.toHaveBeenCalled(); }); test("InteractiveBeacon binds handles to its own beacon identity", async () => { const getBeaconTaskContent = jest.fn(async () => metadata({ BeaconID: "different-beacon" })); - const { client } = clientWithRpc({ control: {}, artifact: {}, "task-content": { getBeaconTaskContent } }); + const cancelBeaconTask = jest.fn(); + const { client } = clientWithRpc({ control: {}, artifact: { cancelBeaconTask }, "task-content": { getBeaconTaskContent } }); const handle = client.interactBeacon(beaconId).beaconTask(taskId, sliverpb.Ping); await expect(handle.wait(2)).rejects.toBeInstanceOf(BeaconTaskIdentityError); await expect(handle.status(2)).rejects.toBeInstanceOf(BeaconTaskIdentityError); + await expect(handle.cancel(2)).rejects.toBeInstanceOf(BeaconTaskIdentityError); + expect(cancelBeaconTask).not.toHaveBeenCalled(); }); test("legacy numeric lsTask calls retain wait and gain status and cancellation", async () => { diff --git a/src/tests/beacon-tasks.test.ts b/src/tests/beacon-tasks.test.ts index 996d369..e0efa4f 100644 --- a/src/tests/beacon-tasks.test.ts +++ b/src/tests/beacon-tasks.test.ts @@ -8,7 +8,7 @@ import { createBeaconTaskHandle, type BeaconTaskHandleDependencies, } from "../beaconTasks"; -import { OperationAbortedError } from "../internal/timeout"; +import { OperationAbortedError, OperationTimeoutError } from "../internal/timeout"; import { BEACON_TASK_MAX_PAYLOAD_BYTES } from "../messageBudget"; import { BeaconTask, Event } from "../pb/clientpb/client"; import { Ping } from "../pb/sliverpb/sliver"; @@ -277,11 +277,11 @@ test("permanent RPC errors fail immediately without leaking a polling timer", as expect(jest.getTimerCount()).toBe(0); }); -test("status fetches metadata and explicit cancel calls only the server cancellation callback", async () => { +test("status and explicit cancellation verify the bound task metadata", async () => { const { handle, rpc, cancel } = setup(); await expect(handle.status(4)).resolves.toMatchObject({ ID: taskId, State: "completed" }); await expect(handle.cancel({ timeoutSeconds: 7 })).resolves.toMatchObject({ ID: taskId, State: "canceled" }); - expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(1); + expect(rpc.getBeaconTaskContent).toHaveBeenCalledTimes(2); expect(cancel).toHaveBeenCalledWith(taskId, { timeoutSeconds: 7, signal: expect.any(AbortSignal) }); }); @@ -293,15 +293,86 @@ test("pre-aborted cancellation does not dispatch a remote cancellation", async ( expect(cancel).not.toHaveBeenCalled(); }); -test("status fails clearly while disconnected and cancel verifies returned identity", async () => { - const { handle } = setup({ - getRpc: () => undefined, - cancel: async () => task({ ID: "wrong-id" }), - }); +test("disconnected status and bound cancellation fail without a remote mutation", async () => { + const { handle, cancel } = setup({ getRpc: () => undefined }); await expect(handle.status()).rejects.toThrow("SliverClient is not connected"); + await expect(handle.cancel()).rejects.toThrow("SliverClient is not connected"); + expect(cancel).not.toHaveBeenCalled(); +}); + +test("cancel still verifies the returned mutation response", async () => { + const { handle } = setup({ cancel: async () => task({ ID: "wrong-id" }) }); await expect(handle.cancel()).rejects.toBeInstanceOf(BeaconTaskIdentityError); }); +test.each([{ ID: "other-task" }, { BeaconID: "other-beacon" }])( + "mismatched task metadata %j prevents remote cancellation", + async (fields) => { + const { handle, rpc, cancel } = setup(); + rpc.getBeaconTaskContent.mockResolvedValue(task(fields)); + await expect(handle.cancel()).rejects.toBeInstanceOf(BeaconTaskIdentityError); + expect(cancel).not.toHaveBeenCalled(); + }, +); + +test("cancellation identity lookup uses the current RPC client", async () => { + const old = { getBeaconTaskContent: jest.fn(async () => task({ BeaconID: "stale-beacon" })) }; + const replacement = { getBeaconTaskContent: jest.fn(async () => task()) }; + let current = old; + const { handle, cancel } = setup({ getRpc: () => current }); + current = replacement; + await expect(handle.cancel()).resolves.toMatchObject({ State: "canceled" }); + expect(old.getBeaconTaskContent).not.toHaveBeenCalled(); + expect(replacement.getBeaconTaskContent).toHaveBeenCalledTimes(1); + expect(cancel).toHaveBeenCalledTimes(1); +}); + +test("identity lookup and cancellation share one whole-operation deadline", async () => { + jest.useFakeTimers(); + let lookupSignal: AbortSignal | undefined; + let mutationSignal: AbortSignal | undefined; + const cancel = jest.fn((_id: string, options: unknown) => { + mutationSignal = (options as { signal: AbortSignal }).signal; + return new Promise(() => undefined); + }); + const { handle, rpc } = setup({ cancel }); + rpc.getBeaconTaskContent.mockImplementation((_request, options) => { + lookupSignal = options?.signal; + return new Promise((resolve) => setTimeout(() => resolve(task()), 700)); + }); + const canceling = handle.cancel(1); + const rejected = expect(canceling).rejects.toBeInstanceOf(OperationTimeoutError); + await jest.advanceTimersByTimeAsync(700); + expect(cancel).toHaveBeenCalledTimes(1); + expect(mutationSignal).toBe(lookupSignal); + expect(mutationSignal?.aborted).toBe(false); + await jest.advanceTimersByTimeAsync(300); + await rejected; + expect(mutationSignal?.aborted).toBe(true); + expect(jest.getTimerCount()).toBe(0); +}); + +test("a late identity lookup cannot dispatch cancellation after the caller aborts", async () => { + let resolveLookup!: (value: BeaconTask) => void; + const { handle, rpc, cancel } = setup(); + rpc.getBeaconTaskContent.mockImplementation(() => new Promise((resolve) => { resolveLookup = resolve; })); + const controller = new AbortController(); + const canceling = handle.cancel({ signal: controller.signal }); + const rejected = expect(canceling).rejects.toBeInstanceOf(OperationAbortedError); + controller.abort(); + await rejected; + resolveLookup(task()); + await flush(); + expect(cancel).not.toHaveBeenCalled(); +}); + +test("an unbound handle can cancel without fetching task content", async () => { + const { handle, rpc, cancel } = setup({ beaconId: undefined }); + await expect(handle.cancel()).resolves.toMatchObject({ State: "canceled" }); + expect(rpc.getBeaconTaskContent).not.toHaveBeenCalled(); + expect(cancel).toHaveBeenCalledTimes(1); +}); + test.each(["", " ", " task", "task ", "x".repeat(257)])("rejects invalid task id %j before using RPCs", (id) => { expect(() => setup({ taskId: id })).toThrow(TypeError); });