mirror of
https://github.com/KeygraphHQ/shannon.git
synced 2026-09-20 00:42:24 +02:00
fix(cli): make scan shutdown verifiable
- preselect and persist workflow identity before worker launch - cancel first, then verify bounded Temporal termination - reconcile Docker workers with Temporal open workflows - fail closed on stale images and unavailable lifecycle state - mark cancellation only after confirmed shutdown
This commit is contained in:
@@ -1,5 +1,5 @@
|
||||
/**
|
||||
* Thin Temporal client for reading one scan's state.
|
||||
* Thin Temporal client for reading scan state and controlling scan workflow lifecycle.
|
||||
*
|
||||
* A running scan is queried live (getProgress) and read via pendingActivities for
|
||||
* the in-flight agents; a closed scan is read once from its result. Everything goes
|
||||
@@ -13,10 +13,20 @@ import { ACTIVITY_TO_PROGRESS, type PipelineState } from './scan/pipeline.js';
|
||||
|
||||
const ADDRESS = '127.0.0.1:7233';
|
||||
const NAMESPACE = 'default';
|
||||
const LIFECYCLE_RPC_DEADLINE_MS = 3_000;
|
||||
const OPEN_SCAN_WORKFLOW_QUERY =
|
||||
"WorkflowType = 'pentestPipelineWorkflow' AND (ExecutionStatus = 'Running' OR ExecutionStatus = 'Paused')";
|
||||
|
||||
// WorkflowExecutionStatusName values that mean the scan has closed. RUNNING (and the unused
|
||||
// CONTINUED_AS_NEW) are the only non-terminal states.
|
||||
const TERMINAL_STATUSES: ReadonlySet<string> = new Set(['COMPLETED', 'FAILED', 'CANCELLED', 'TERMINATED', 'TIMED_OUT']);
|
||||
// WorkflowExecutionStatusName values that positively prove this execution has closed.
|
||||
// PAUSED is open; UNSPECIFIED and UNKNOWN are not safe closure evidence.
|
||||
const TERMINAL_STATUSES: ReadonlySet<string> = new Set([
|
||||
'COMPLETED',
|
||||
'FAILED',
|
||||
'CANCELLED',
|
||||
'TERMINATED',
|
||||
'CONTINUED_AS_NEW',
|
||||
'TIMED_OUT',
|
||||
]);
|
||||
|
||||
export interface RunningAgent {
|
||||
readonly agent: string;
|
||||
@@ -66,11 +76,28 @@ export type TerminalOutcome =
|
||||
| { readonly kind: 'success'; readonly state: PipelineState }
|
||||
| { readonly kind: 'failed'; readonly message: string };
|
||||
|
||||
/**
|
||||
* The authoritative Temporal state used by lifecycle commands. Transport failures deliberately
|
||||
* remain errors instead of being represented as a closed workflow: callers must not report a
|
||||
* scan stopped unless Temporal has positively confirmed it.
|
||||
*/
|
||||
export type WorkflowLifecycleState =
|
||||
| { readonly kind: 'open'; readonly status: 'RUNNING' | 'PAUSED' }
|
||||
| { readonly kind: 'terminal'; readonly status: string }
|
||||
| { readonly kind: 'unknown'; readonly status: string }
|
||||
| { readonly kind: 'not-found' };
|
||||
|
||||
/** A scan workflow returned by Temporal's eventually consistent open-workflow visibility query. */
|
||||
export interface RunningScanWorkflow {
|
||||
readonly workflowId: string;
|
||||
readonly taskQueue: string;
|
||||
}
|
||||
|
||||
let clientPromise: Promise<Client> | null = null;
|
||||
|
||||
function getClient(): Promise<Client> {
|
||||
if (!clientPromise) {
|
||||
const pending = Connection.connect({ address: ADDRESS }).then(
|
||||
const pending = Connection.connect({ address: ADDRESS, connectTimeout: LIFECYCLE_RPC_DEADLINE_MS }).then(
|
||||
(connection) => new Client({ connection, namespace: NAMESPACE }),
|
||||
);
|
||||
// A rejected connect must not be cached forever: clear the memo so the next call rebuilds
|
||||
@@ -95,6 +122,95 @@ function resetClient(only?: Promise<Client>): void {
|
||||
previous?.then((client) => client.connection.close()).catch(() => {});
|
||||
}
|
||||
|
||||
/** Close the current channel and establish another before a termination retry. */
|
||||
export async function refreshWorkflowLifecycleConnection(): Promise<void> {
|
||||
const previous = clientPromise;
|
||||
if (previous !== null) {
|
||||
if (clientPromise === previous) clientPromise = null;
|
||||
try {
|
||||
const client = await previous;
|
||||
await client.connection.close();
|
||||
} catch {
|
||||
// A failed prior connection is already detached. The new connection below is authoritative.
|
||||
}
|
||||
}
|
||||
await getClient();
|
||||
}
|
||||
|
||||
/**
|
||||
* Run a bounded lifecycle RPC and discard the connection when Temporal did not positively say
|
||||
* that the workflow is absent. A fresh connection is important after a gRPC timeout or transport
|
||||
* failure: reusing a wedged channel can turn a recoverable stop into an indefinitely ambiguous one.
|
||||
*/
|
||||
async function runLifecycleRpc<T>(operation: (client: Client) => Promise<T>): Promise<T> {
|
||||
const pending = getClient();
|
||||
try {
|
||||
const client = await pending;
|
||||
return await client.withDeadline(Date.now() + LIFECYCLE_RPC_DEADLINE_MS, () => operation(client));
|
||||
} catch (err) {
|
||||
if (!(err instanceof WorkflowNotFoundError)) resetClient(pending);
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/** Describe a workflow for lifecycle control without reading its progress or pending activities. */
|
||||
export async function describeWorkflowLifecycle(workflowId: string): Promise<WorkflowLifecycleState> {
|
||||
try {
|
||||
const desc = await runLifecycleRpc((client) => client.workflow.getHandle(workflowId).describe());
|
||||
if (desc.status.name === 'RUNNING' || desc.status.name === 'PAUSED') {
|
||||
return { kind: 'open', status: desc.status.name };
|
||||
}
|
||||
if (TERMINAL_STATUSES.has(desc.status.name)) return { kind: 'terminal', status: desc.status.name };
|
||||
return { kind: 'unknown', status: desc.status.name };
|
||||
} catch (err) {
|
||||
if (err instanceof WorkflowNotFoundError) return { kind: 'not-found' };
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/** Request cooperative cancellation. This confirms request acceptance, not workflow closure. */
|
||||
export async function requestWorkflowCancellation(workflowId: string): Promise<'requested' | 'not-found'> {
|
||||
try {
|
||||
await runLifecycleRpc((client) => client.workflow.getHandle(workflowId).cancel());
|
||||
return 'requested';
|
||||
} catch (err) {
|
||||
if (err instanceof WorkflowNotFoundError) return 'not-found';
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/** Request forced termination. This confirms request acceptance, not workflow closure. */
|
||||
export async function requestWorkflowTermination(
|
||||
workflowId: string,
|
||||
reason: string,
|
||||
): Promise<'requested' | 'not-found'> {
|
||||
try {
|
||||
await runLifecycleRpc((client) => client.workflow.getHandle(workflowId).terminate(reason));
|
||||
return 'requested';
|
||||
} catch (err) {
|
||||
if (err instanceof WorkflowNotFoundError) return 'not-found';
|
||||
throw err;
|
||||
}
|
||||
}
|
||||
|
||||
/** List currently open Shannon scan workflows through Temporal visibility. */
|
||||
export async function listRunningScanWorkflows(): Promise<readonly RunningScanWorkflow[]> {
|
||||
return runLifecycleRpc(async (client) => {
|
||||
const workflows: RunningScanWorkflow[] = [];
|
||||
for await (const execution of client.workflow.list({ query: OPEN_SCAN_WORKFLOW_QUERY })) {
|
||||
// Visibility is eventually consistent. Keep only the open scan rows returned by this page;
|
||||
// each discovered workflow is described directly before `stop` accepts its closure.
|
||||
if (
|
||||
(execution.status.name === 'RUNNING' || execution.status.name === 'PAUSED') &&
|
||||
execution.type === 'pentestPipelineWorkflow'
|
||||
) {
|
||||
workflows.push({ workflowId: execution.workflowId, taskQueue: execution.taskQueue });
|
||||
}
|
||||
}
|
||||
return workflows;
|
||||
});
|
||||
}
|
||||
|
||||
/** Describe a scan: status, timing, and the agents currently running (from pendingActivities). Null if not found. */
|
||||
export async function describeScan(workflowId: string): Promise<ScanDescription | null> {
|
||||
const client = await getClient();
|
||||
|
||||
Reference in New Issue
Block a user