From 8ff66b342f014ab2dd864b18dbe0f1ccd20e12ae Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 07:04:51 -0400 Subject: [PATCH 1/6] feat: add secure code worker pairing --- .env.example | 1 + docs/adr/001-stateful-code-environments.md | 86 ++++ docs/remote-bridge/README.md | 39 +- packages/code/README.md | 40 +- packages/code/src/cli.ts | 151 +++++-- packages/code/src/identity.test.ts | 35 ++ packages/code/src/identity.ts | 66 +++ packages/code/src/index.ts | 3 + packages/code/src/pairing.test.ts | 36 ++ packages/code/src/pairing.ts | 73 +++ packages/code/src/protocol.ts | 14 + packages/code/src/storage.test.ts | 30 ++ packages/code/src/storage.ts | 57 +++ packages/code/src/worker.test.ts | 160 +++---- packages/code/src/worker.ts | 232 +++++----- service/rollup.config.js | 1 + service/src/api-server.ts | 2 +- service/src/bridge/index.ts | 16 + service/src/bridge/pairing.test.ts | 220 +++++++++ service/src/bridge/pairing.ts | 256 +++++++++++ service/src/bridge/router.test.ts | 125 ++++++ service/src/bridge/router.ts | 446 ++++++++++++------- service/src/config.test.ts | 7 + service/src/config.ts | 14 + service/src/sandbox-backend/remote-bridge.ts | 2 +- service/src/secure-startup.test.ts | 9 +- service/src/secure-startup.ts | 5 + service/src/service-api.ts | 2 +- service/tsconfig.json | 3 +- 29 files changed, 1729 insertions(+), 402 deletions(-) create mode 100644 docs/adr/001-stateful-code-environments.md create mode 100644 packages/code/src/identity.test.ts create mode 100644 packages/code/src/identity.ts create mode 100644 packages/code/src/pairing.test.ts create mode 100644 packages/code/src/pairing.ts create mode 100644 packages/code/src/storage.test.ts create mode 100644 packages/code/src/storage.ts create mode 100644 service/src/bridge/index.ts create mode 100644 service/src/bridge/pairing.test.ts create mode 100644 service/src/bridge/pairing.ts create mode 100644 service/src/bridge/router.test.ts diff --git a/.env.example b/.env.example index 3bd113e..98c9726 100644 --- a/.env.example +++ b/.env.example @@ -30,6 +30,7 @@ SANDBOX_OUTPUT_MAX_SIZE=65536 # CODEAPI_RUNTIME_SESSION_MODE=affinity # CODEAPI_BRIDGE_WORKER_ID=my-vm # CODEAPI_BRIDGE_TOKEN=replace-with-a-strong-random-secret +# CODEAPI_BRIDGE_AUTH_MODE=paired # Service Configuration PYTHON_CONCURRENCY=5 diff --git a/docs/adr/001-stateful-code-environments.md b/docs/adr/001-stateful-code-environments.md new file mode 100644 index 0000000..8163a95 --- /dev/null +++ b/docs/adr/001-stateful-code-environments.md @@ -0,0 +1,86 @@ +# ADR 001: Stateful code environments use an outbound Code API bridge + +- Status: Accepted for alpha +- Date: 2026-08-30 + +## Context + +LibreChat needs coding agents to reuse a workspace across conversation turns +while allowing the environment owner to choose the VM. Internet-facing +LibreChat instances cannot safely require inbound access to that VM, forward +end-user tokens to it, or treat an MCP connection as a sandbox boundary. + +The first alpha demonstrated a stable runtime-session ID, a single fenced +worker lease, and workspace persistence across turns. Its static shared worker +token was sufficient to prove execution flow but is not an acceptable hardened +enrollment mechanism. + +## Decision + +The product concept is a **stateful code environment**. Code API remains its +broker and policy boundary, and `remote-bridge` is a Code API sandbox backend. +The `@librechat/code` worker connects outbound from the chosen VM and forwards +assignments only to a loopback or private sandbox endpoint. + +Hardened workers enroll through a one-time pairing code: + +1. An administrator creates a code scoped to the configured worker ID. +2. The CLI generates an Ed25519 keypair locally and redeems the code with only + its public key. +3. Code API returns a five-minute credential bound to that public key. +4. Every worker request signs the method, path, body digest, timestamp, nonce, + and credential. +5. Code API rejects stale timestamps and replayed nonces and supports rotation + and immediate revocation. + +Static bearer authentication remains a non-hardened compatibility mode. + +## Ownership and state + +The alpha environment is deployment/operator owned and configured with one +worker ID. A future LibreChat control plane may persist deployment-, tenant-, +or user-owned environment records and issue the same pairing operation through +RBAC-protected APIs without changing the worker execution protocol. + +Workspace state belongs to the stable runtime session, not to a transient +assignment lease. For `remote-bridge`, that state currently survives turns on +the same worker and backing disk. It is not yet checkpointed or portable across +worker replacement; the UI and operator documentation must not imply otherwise. + +## Security invariants + +- The VM requires no inbound internet listener. +- Code API, not the worker, authenticates LibreChat users and normalizes work. +- A stolen short-lived credential is insufficient without the worker private + key; a stolen private key is insufficient after credential expiry or + revocation. +- Pairing codes and credentials are stored by digest where lookup permits. +- One configured worker has at most one active fenced assignment. +- Sandbox isolation and default-deny egress remain mandatory; pairing secures + the transport identity but does not make the host a sandbox. +- A compromised worker can lie about advertised capabilities. Capability + labels and policy digests are audit signals until enforcement is coupled to + an attested sandbox or trusted host policy. + +## Consequences + +- `@librechat/code` owns the provider-neutral protocol, identity handling, and + worker CLI; Code API owns enrollment, scheduling, and execution policy. +- LibreChat owns environment persistence, ownership, RBAC, and user experience. +- The Agents SDK keeps only its adapter until a second concrete consumer proves + which coding-tool abstractions are genuinely provider neutral. +- MCP may expose environment operations later, but it is not the worker + transport or isolation boundary. +- Multi-worker directories, checkpoint/restore, owner-scoped quotas, and + enforced network capability profiles remain follow-up decisions. + +## Alternatives rejected + +- **Inbound SSH/HTTP to the VM:** expands attack surface and complicates NAT and + firewall operation. +- **MCP as the worker protocol:** conflates tool discovery with leases, + cancellation, fencing, and sandbox policy. +- **Put the runtime in the Agents SDK:** couples provider-neutral execution to + one agent integration and makes non-agent consumers depend on agent internals. +- **Long-lived shared bearer token:** easy to bootstrap, but replayable and not + bound to a worker-held key. diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index f5632d5..c0b9661 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -12,7 +12,7 @@ LibreChat -> Code API -> Redis assignment Code API remains the public authentication, policy, manifest, timeout, and result-normalization boundary. The bridge worker has a separate operator -credential and never accepts end-user bearer tokens directly. +identity and never accepts end-user bearer tokens directly. ## Code API configuration @@ -23,7 +23,8 @@ CODEAPI_SANDBOX_BACKEND=remote-bridge CODEAPI_EXECUTION_PROFILE=stateful CODEAPI_RUNTIME_SESSION_MODE=affinity CODEAPI_BRIDGE_WORKER_ID=my-vm -CODEAPI_BRIDGE_TOKEN= +CODEAPI_BRIDGE_TOKEN= +CODEAPI_BRIDGE_AUTH_MODE=paired ``` Use `strict` instead of `affinity` if every request must include a runtime @@ -31,12 +32,20 @@ session hint. In hardened mode, startup requires the bridge token to be at least 32 bytes. `PTC_MODE=blocking` is rejected; replay mode is required because a remote execution cannot retain an open Code API process across tool callbacks. -Start the CLI beside a sandbox using the same worker ID and secret; see -[`@librechat/code`](../../packages/code/README.md). -Stateful deployments must also set `LIBRECHAT_CODE_STATEFUL_WORKSPACE=true` -and route the CLI's `{runtimeSessionId}` endpoint template to an isolated, -persistent local runner per session. A single sandbox endpoint is stateless and -is rejected for runtime-session assignments. +Create a single-use pairing code with the administrator secret: + +```bash +curl -fsS https://code.example.com/v1/bridge/pairings \ + -H "Authorization: Bearer $CODEAPI_BRIDGE_TOKEN" \ + -H 'Content-Type: application/json' \ + --data '{"workerId":"my-vm"}' +``` + +Redeem the returned code on the VM using +[`@librechat/code`](../../packages/code/README.md). The CLI generates its key +locally, proves possession on every request, and rotates its short-lived +credential before expiry. `CODEAPI_BRIDGE_AUTH_MODE=static` remains available +for non-hardened development compatibility only. ## LibreChat configuration @@ -63,6 +72,13 @@ execution. ## Lifecycle and fencing - Registration is ephemeral in Redis and must be refreshed by the worker. +- Pairing codes are stored hashed, expire after ten minutes, and are consumed + atomically on their first redemption attempt. +- Worker credentials expire after five minutes and are bound to an Ed25519 + public key. Exact-request signatures include the HTTP method, path, body + digest, timestamp, nonce, and credential. +- Accepted proof nonces cannot be replayed, credentials rotate before expiry, + and an administrator can revoke the active worker identity immediately. - Code API permits one active assignment per configured worker. - Each assignment has an absolute deadline, generation, and random lease token. - Settlements with the wrong worker, generation, token, or expired deadline are @@ -79,9 +95,10 @@ For internet-facing LibreChat deployments, use the hardened microVM/NsJail stack, default-deny sandbox egress, signed execution manifests, least-privilege host credentials, resource limits, and host/network monitoring. Bind the local sandbox endpoint to loopback or a private container network. Rotate a leaked -bridge token immediately; the initial protocol intentionally uses a static -operator secret and supports one configured worker per Code API deployment. +administrator token immediately. Pairing secures worker transport identity; it +cannot attest that a compromised VM truthfully reports or enforces its sandbox +capabilities. -The next control-plane layer can add short-lived pairing credentials and a +The next control-plane layer can add owner-scoped environment records and a multi-worker directory without changing the execution protocol or moving code tools into the Agents SDK. diff --git a/packages/code/README.md b/packages/code/README.md index a79a170..8195e03 100644 --- a/packages/code/README.md +++ b/packages/code/README.md @@ -9,7 +9,37 @@ untrusted internet traffic). It connects outbound to Code API, long-polls for assignments, forwards them to the local sandbox, and returns fenced results. The VM does not need an inbound public port. -## Run +## Pair + +Hardened deployments use a one-time code instead of copying a long-lived +worker secret onto the VM. After an administrator creates a code, run: + +```bash +librechat-code pair https://code.example.com/v1 '' \ + --worker-id my-vm +``` + +The CLI generates an Ed25519 key locally and writes its paired identity to +`~/.config/librechat/code/my-vm.json` with owner-only permissions. The private +key never leaves the VM. Worker requests carry an exact-request signature, +timestamp, and one-time nonce; the short-lived credential rotates +automatically. + +Then start the worker without a shared secret: + +```bash +LIBRECHAT_CODE_WORKER_ID=my-vm \ +LIBRECHAT_CODE_SANDBOX_ENDPOINT=http://127.0.0.1:2000/api/v2 \ +librechat-code run +``` + +Use `--identity ` while pairing and +`LIBRECHAT_CODE_IDENTITY_FILE=` while running to override the identity +file location. + +## Static compatibility mode + +Non-hardened development deployments may still run with a static token: ```bash npm install -g @librechat/code @@ -18,7 +48,7 @@ LIBRECHAT_CODE_URL=https://code.example.com/v1 \ LIBRECHAT_CODE_WORKER_TOKEN='' \ LIBRECHAT_CODE_WORKER_ID=my-vm \ LIBRECHAT_CODE_SANDBOX_ENDPOINT=http://127.0.0.1:2000/api/v2 \ -librechat-code +librechat-code run ``` Optional environment variables: @@ -38,6 +68,6 @@ A single built-in sandbox runner binds itself to one runtime session and must not be advertised as stateful. Use the default stateless capability until a session-routing supervisor is configured. -Use a unique worker ID and secret per Code API deployment, expose only the -sandbox loopback endpoint to the CLI, and enforce VM/container egress policy -independently of the bridge transport. +Static worker authentication is rejected when Code API hardened mode is +enabled. Expose only the sandbox loopback endpoint to the CLI, and enforce +VM/container egress policy independently of the bridge transport. diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index 411dda7..64811ce 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -1,11 +1,16 @@ #!/usr/bin/env node import { createHash } from 'node:crypto'; +import { homedir } from 'node:os'; +import { join } from 'node:path'; + +import { pairBridgeWorker } from './pairing.js'; +import { loadBridgeIdentity, saveBridgeIdentity } from './storage.js'; import { BridgeWorker } from './worker.js'; -function required(name: string): string { - const value = process.env[name]?.trim(); - if (!value) throw new Error(`${name} is required`); - return value; +function required(name: string, value = process.env[name]): string { + const normalized = value?.trim(); + if (!normalized) throw new Error(`${name} is required`); + return normalized; } function list(value: string | undefined): string[] { @@ -17,39 +22,117 @@ function list(value: string | undefined): string[] { ); } -const controller = new AbortController(); -process.once('SIGINT', () => controller.abort()); -process.once('SIGTERM', () => controller.abort()); - -const policy = process.env.LIBRECHAT_CODE_POLICY ?? 'default-deny'; -const statefulWorkspace = - process.env.LIBRECHAT_CODE_STATEFUL_WORKSPACE?.trim().toLowerCase() === 'true'; -const sandboxEndpoint = - process.env.LIBRECHAT_CODE_SANDBOX_ENDPOINT ?? - 'http://127.0.0.1:2000/api/v2'; -if (statefulWorkspace && !sandboxEndpoint.includes('{runtimeSessionId}')) { - throw new Error( - 'LIBRECHAT_CODE_STATEFUL_WORKSPACE requires LIBRECHAT_CODE_SANDBOX_ENDPOINT to contain {runtimeSessionId}', +function option(args: string[], name: string): string | undefined { + const index = args.indexOf(name); + if (index >= 0) return args[index + 1]; + return args.find((value) => value.startsWith(`${name}=`))?.slice(name.length + 1); +} + +function defaultIdentityPath(workerId: string): string { + const fileName = workerId.replace(/[^A-Za-z0-9._-]/g, '_'); + return join(homedir(), '.config', 'librechat', 'code', `${fileName}.json`); +} + +async function pair(args: string[]): Promise { + const codeApiUrl = required('instance URL', args[1]); + const code = required('one-time pairing code', args[2]); + const workerId = required( + '--worker-id or LIBRECHAT_CODE_WORKER_ID', + option(args, '--worker-id') ?? process.env.LIBRECHAT_CODE_WORKER_ID, + ); + const identityPath = + option(args, '--identity') ?? + process.env.LIBRECHAT_CODE_IDENTITY_FILE ?? + defaultIdentityPath(workerId); + const identity = await pairBridgeWorker({ codeApiUrl, workerId, code }); + await saveBridgeIdentity(identityPath, identity); + process.stdout.write( + `Paired worker ${workerId}. Identity saved to ${identityPath}\n`, ); } -const worker = new BridgeWorker({ - codeApiUrl: required('LIBRECHAT_CODE_URL'), - token: required('LIBRECHAT_CODE_WORKER_TOKEN'), - workerId: required('LIBRECHAT_CODE_WORKER_ID'), - sandboxEndpoint, - capabilities: { - statefulWorkspace, - sandboxProfile: process.env.LIBRECHAT_CODE_SANDBOX_PROFILE ?? 'nsjail', - runtimes: list(process.env.LIBRECHAT_CODE_RUNTIMES), - policyDigest: createHash('sha256').update(policy).digest('hex'), - }, - onError: (error) => { - const message = error instanceof Error ? error.message : 'unknown bridge error'; - process.stderr.write(`librechat-code: reconnecting after ${message}\n`); - }, -}); -worker.run(controller.signal).catch((error: Error) => { +async function run(): Promise { + const configuredWorkerId = process.env.LIBRECHAT_CODE_WORKER_ID?.trim(); + const configuredIdentityPath = process.env.LIBRECHAT_CODE_IDENTITY_FILE?.trim(); + const configuredToken = process.env.LIBRECHAT_CODE_WORKER_TOKEN?.trim(); + const identityPath = + configuredIdentityPath ?? + (configuredWorkerId && !configuredToken + ? defaultIdentityPath(configuredWorkerId) + : undefined); + const pairedIdentity = identityPath + ? await loadBridgeIdentity(identityPath) + : undefined; + const workerId = required( + 'LIBRECHAT_CODE_WORKER_ID', + configuredWorkerId ?? pairedIdentity?.workerId, + ); + if (pairedIdentity && pairedIdentity.workerId !== workerId) { + throw new Error( + `Identity belongs to ${pairedIdentity.workerId}, not configured worker ${workerId}`, + ); + } + const codeApiUrl = required( + 'LIBRECHAT_CODE_URL', + process.env.LIBRECHAT_CODE_URL ?? pairedIdentity?.codeApiUrl, + ); + const policy = process.env.LIBRECHAT_CODE_POLICY ?? 'default-deny'; + const workerIdentity = pairedIdentity + ? { + privateKey: pairedIdentity.privateKey, + credential: pairedIdentity.credential, + expiresAt: pairedIdentity.expiresAt, + } + : undefined; + const controller = new AbortController(); + process.once('SIGINT', () => controller.abort()); + process.once('SIGTERM', () => controller.abort()); + const worker = new BridgeWorker({ + codeApiUrl, + token: configuredToken, + identity: workerIdentity, + workerId, + sandboxEndpoint: + process.env.LIBRECHAT_CODE_SANDBOX_ENDPOINT ?? + 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: process.env.LIBRECHAT_CODE_SANDBOX_PROFILE ?? 'nsjail', + runtimes: list(process.env.LIBRECHAT_CODE_RUNTIMES), + policyDigest: createHash('sha256').update(policy).digest('hex'), + }, + onIdentityChange: + pairedIdentity && identityPath + ? async (identity) => { + await saveBridgeIdentity(identityPath, { + ...pairedIdentity, + credential: identity.credential, + expiresAt: identity.expiresAt, + }); + } + : undefined, + onError: (error) => { + const message = + error instanceof Error ? error.message : 'unknown bridge error'; + process.stderr.write(`librechat-code: reconnecting after ${message}\n`); + }, + }); + await worker.run(controller.signal); +} + +async function main(): Promise { + const args = process.argv.slice(2); + if (args[0] === 'pair') { + await pair(args); + return; + } + if (args[0] && args[0] !== 'run') { + throw new Error(`Unknown command: ${args[0]}`); + } + await run(); +} + +main().catch((error: Error) => { process.stderr.write(`librechat-code: ${error.message}\n`); process.exitCode = 1; }); diff --git a/packages/code/src/identity.test.ts b/packages/code/src/identity.test.ts new file mode 100644 index 0000000..4e92b91 --- /dev/null +++ b/packages/code/src/identity.test.ts @@ -0,0 +1,35 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { + createBridgeIdentity, + signBridgeRequest, + verifyBridgeRequest, +} from './identity.js'; + +test('worker identity proves possession for the exact HTTP request', () => { + const identity = createBridgeIdentity(); + const request = { + credential: 'short-lived-credential', + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce: 'single-use-request-nonce', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + + const signature = signBridgeRequest(identity.privateKey, request); + + assert.equal( + verifyBridgeRequest(identity.publicKey, request, signature), + true, + ); + assert.equal( + verifyBridgeRequest( + identity.publicKey, + { ...request, body: JSON.stringify({ protocolVersion: 1, waitMs: 0 }) }, + signature, + ), + false, + ); +}); diff --git a/packages/code/src/identity.ts b/packages/code/src/identity.ts new file mode 100644 index 0000000..ab11c86 --- /dev/null +++ b/packages/code/src/identity.ts @@ -0,0 +1,66 @@ +import { + createHash, + generateKeyPairSync, + sign, + verify, +} from 'node:crypto'; + +export interface BridgeIdentity { + publicKey: string; + privateKey: string; +} + +export interface BridgeRequestProofInput { + credential: string; + method: string; + path: string; + timestamp: string; + nonce: string; + body: string; +} + +export function createBridgeIdentity(): BridgeIdentity { + const { publicKey, privateKey } = generateKeyPairSync('ed25519', { + publicKeyEncoding: { type: 'spki', format: 'pem' }, + privateKeyEncoding: { type: 'pkcs8', format: 'pem' }, + }); + return { publicKey, privateKey }; +} + +function canonicalBridgeRequest(input: BridgeRequestProofInput): string { + const bodyDigest = createHash('sha256').update(input.body).digest('hex'); + return [ + input.method.toUpperCase(), + input.path, + input.timestamp, + input.nonce, + bodyDigest, + input.credential, + ].join('\n'); +} + +export function signBridgeRequest( + privateKey: string, + input: BridgeRequestProofInput, +): string { + return sign(null, Buffer.from(canonicalBridgeRequest(input)), privateKey).toString( + 'base64url', + ); +} + +export function verifyBridgeRequest( + publicKey: string, + input: BridgeRequestProofInput, + signature: string, +): boolean { + try { + return verify( + null, + Buffer.from(canonicalBridgeRequest(input)), + publicKey, + Buffer.from(signature, 'base64url'), + ); + } catch { + return false; + } +} diff --git a/packages/code/src/index.ts b/packages/code/src/index.ts index c5eeaaf..c65b9f1 100644 --- a/packages/code/src/index.ts +++ b/packages/code/src/index.ts @@ -1,2 +1,5 @@ export * from './protocol.js'; +export * from './identity.js'; +export * from './pairing.js'; +export * from './storage.js'; export * from './worker.js'; diff --git a/packages/code/src/pairing.test.ts b/packages/code/src/pairing.test.ts new file mode 100644 index 0000000..2b56439 --- /dev/null +++ b/packages/code/src/pairing.test.ts @@ -0,0 +1,36 @@ +import assert from 'node:assert/strict'; +import test from 'node:test'; + +import { pairBridgeWorker } from './pairing.js'; + +test('pairing binds a generated worker key to a single-use code', async () => { + const fetchImpl: typeof fetch = async (_input, init) => { + const body = JSON.parse(String(init?.body)) as { + workerId: string; + code: string; + publicKey: string; + }; + assert.equal(body.workerId, 'vm-1'); + assert.equal(body.code, 'one-time-code'); + assert.match(body.publicKey, /BEGIN PUBLIC KEY/); + return Response.json({ + protocolVersion: 1, + workerId: body.workerId, + credential: 'issued-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + }); + }; + + const paired = await pairBridgeWorker({ + codeApiUrl: 'https://code.example/v1/', + workerId: 'vm-1', + code: 'one-time-code', + fetchImpl, + }); + + assert.equal(paired.workerId, 'vm-1'); + assert.equal(paired.codeApiUrl, 'https://code.example/v1'); + assert.equal(paired.credential, 'issued-short-lived-credential-value'); + assert.match(paired.publicKey, /BEGIN PUBLIC KEY/); + assert.match(paired.privateKey, /BEGIN PRIVATE KEY/); +}); diff --git a/packages/code/src/pairing.ts b/packages/code/src/pairing.ts new file mode 100644 index 0000000..ed33fb1 --- /dev/null +++ b/packages/code/src/pairing.ts @@ -0,0 +1,73 @@ +import { createBridgeIdentity } from './identity.js'; +import { + BRIDGE_PROTOCOL_VERSION, + BridgeProtocolError, +} from './protocol.js'; + +import type { BridgeWorkerCredentialResponse } from './protocol.js'; + +export interface PairBridgeWorkerOptions { + codeApiUrl: string; + workerId: string; + code: string; + fetchImpl?: typeof fetch; +} + +export interface PairedBridgeWorkerIdentity + extends BridgeWorkerCredentialResponse { + codeApiUrl: string; + publicKey: string; + privateKey: string; +} + +function normalizedBaseUrl(value: string): string { + return value.replace(/\/+$/, ''); +} + +function errorMessage(value: object): string | undefined { + if ('error' in value && typeof value.error === 'string') return value.error; + return undefined; +} + +export async function pairBridgeWorker( + options: PairBridgeWorkerOptions, +): Promise { + const codeApiUrl = normalizedBaseUrl(options.codeApiUrl); + const identity = createBridgeIdentity(); + const response = await (options.fetchImpl ?? fetch)( + `${codeApiUrl}/bridge/pairings/redeem`, + { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: options.workerId, + code: options.code, + publicKey: identity.publicKey, + }), + }, + ); + const payload = (await response.json()) as object; + if (!response.ok) { + throw new BridgeProtocolError( + errorMessage(payload) ?? `Bridge pairing failed with HTTP ${response.status}`, + response.status, + ); + } + const credential = payload as BridgeWorkerCredentialResponse; + if ( + credential.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + credential.workerId !== options.workerId || + typeof credential.credential !== 'string' || + credential.credential.length < 32 || + !Number.isFinite(Date.parse(credential.expiresAt)) + ) { + throw new BridgeProtocolError('Code API returned an invalid worker credential'); + } + return { + ...credential, + codeApiUrl, + publicKey: identity.publicKey, + privateKey: identity.privateKey, + }; +} diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index a79d4ff..34358cb 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -24,6 +24,20 @@ export interface BridgeWorkerRegistrationResponse { leaseTtlMs: number; } +export interface BridgePairingRedemption { + protocolVersion: BridgeProtocolVersion; + workerId: string; + code: string; + publicKey: string; +} + +export interface BridgeWorkerCredentialResponse { + protocolVersion: BridgeProtocolVersion; + workerId: string; + credential: string; + expiresAt: string; +} + export interface BridgeSandboxRequest { body: TBody; headers: Record; diff --git a/packages/code/src/storage.test.ts b/packages/code/src/storage.test.ts new file mode 100644 index 0000000..8b5128d --- /dev/null +++ b/packages/code/src/storage.test.ts @@ -0,0 +1,30 @@ +import assert from 'node:assert/strict'; +import { mkdtemp, rm, stat } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import test from 'node:test'; + +import { loadBridgeIdentity, saveBridgeIdentity } from './storage.js'; + +test('paired identity is persisted atomically with owner-only permissions', async () => { + const directory = await mkdtemp(join(tmpdir(), 'librechat-code-')); + const path = join(directory, 'identity.json'); + const identity = { + protocolVersion: 1 as const, + workerId: 'vm-1', + codeApiUrl: 'https://code.example/v1', + credential: 'issued-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + publicKey: 'public-key', + privateKey: 'private-key', + }; + + try { + await saveBridgeIdentity(path, identity); + + assert.deepEqual(await loadBridgeIdentity(path), identity); + assert.equal((await stat(path)).mode & 0o777, 0o600); + } finally { + await rm(directory, { recursive: true, force: true }); + } +}); diff --git a/packages/code/src/storage.ts b/packages/code/src/storage.ts new file mode 100644 index 0000000..c44135e --- /dev/null +++ b/packages/code/src/storage.ts @@ -0,0 +1,57 @@ +import { randomBytes } from 'node:crypto'; +import { chmod, mkdir, open, readFile, rename, rm } from 'node:fs/promises'; +import { dirname } from 'node:path'; + +import { BRIDGE_PROTOCOL_VERSION, BridgeProtocolError } from './protocol.js'; + +import type { PairedBridgeWorkerIdentity } from './pairing.js'; + +function isRecord(value: unknown): value is Record { + return typeof value === 'object' && value !== null; +} + +function isPairedIdentity(value: unknown): value is PairedBridgeWorkerIdentity { + if (!isRecord(value)) return false; + return ( + value.protocolVersion === BRIDGE_PROTOCOL_VERSION && + typeof value.workerId === 'string' && + typeof value.codeApiUrl === 'string' && + typeof value.credential === 'string' && + typeof value.expiresAt === 'string' && + Number.isFinite(Date.parse(value.expiresAt)) && + typeof value.publicKey === 'string' && + typeof value.privateKey === 'string' + ); +} + +export async function saveBridgeIdentity( + path: string, + identity: PairedBridgeWorkerIdentity, +): Promise { + await mkdir(dirname(path), { recursive: true, mode: 0o700 }); + const temporaryPath = `${path}.${randomBytes(8).toString('hex')}.tmp`; + try { + const file = await open(temporaryPath, 'wx', 0o600); + try { + await file.writeFile(`${JSON.stringify(identity, null, 2)}\n`, 'utf8'); + await file.sync(); + } finally { + await file.close(); + } + await rename(temporaryPath, path); + await chmod(path, 0o600); + } catch (error) { + await rm(temporaryPath, { force: true }); + throw error; + } +} + +export async function loadBridgeIdentity( + path: string, +): Promise { + const identity = JSON.parse(await readFile(path, 'utf8')) as unknown; + if (!isPairedIdentity(identity)) { + throw new BridgeProtocolError(`Invalid bridge identity file: ${path}`); + } + return identity; +} diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 03061d9..be72a3d 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -1,6 +1,10 @@ import assert from 'node:assert/strict'; import test from 'node:test'; -import { BridgeWorker } from './worker.js'; +import { + createBridgeIdentity, + verifyBridgeRequest, +} from './identity.js'; +import { BridgeWorker, reconnectDelayMs } from './worker.js'; import type { BridgeAssignment } from './protocol.js'; @@ -27,9 +31,7 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' codeApiUrl: 'https://code.example/v1/', token: 'worker-secret', workerId: 'vm-1', - incarnationId: 'incarnation-00000001', - sandboxEndpoint: - 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2/', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2/', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -41,7 +43,6 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' protocolVersion: 1, assignmentId: 'assignment-1', workerId: 'vm-1', - incarnationId: 'incarnation-00000001', generation: 3, leaseToken: 'lease-token-that-is-long-enough-for-testing', expiresAt: new Date(Date.now() + 10_000).toISOString(), @@ -55,10 +56,7 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' await worker.executeAndSettle(assignment); assert.equal(requests.length, 2); - assert.equal( - requests[0].url, - 'http://127.0.0.1:2000/sessions/rt-user-1/api/v2/execute', - ); + assert.equal(requests[0].url, 'http://127.0.0.1:2000/api/v2/execute'); assert.equal( (requests[0].init?.headers as Record)[ 'X-Runtime-Session-Id' @@ -70,112 +68,116 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' protocolVersion: 1, generation: 3, leaseToken: 'lease-token-that-is-long-enough-for-testing', - incarnationId: 'incarnation-00000001', status: 'fulfilled', result: { session_id: 'run-1', files: [] }, }); }); -test('worker aborts sandbox execution at the absolute assignment deadline', async () => { - let settlement: Record | undefined; +test('paired worker proves possession on bridge requests', async () => { + const key = createBridgeIdentity(); + let bridgeRequest: { url: string; init?: RequestInit } | undefined; const fetchImpl: typeof fetch = async (input, init) => { - if (String(input).endsWith('/execute')) { - return await new Promise((_resolve, reject) => { - init?.signal?.addEventListener( - 'abort', - () => reject(new DOMException('aborted', 'AbortError')), - { once: true }, - ); - }); - } - settlement = JSON.parse(String(init?.body)) as Record; - return new Response( - JSON.stringify({ protocolVersion: 1, accepted: true }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, - ); + bridgeRequest = { url: String(input), init }; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', - token: 'worker-secret', workerId: 'vm-1', - incarnationId: 'incarnation-00000001', sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'issued-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), + }, capabilities: { - statefulWorkspace: false, + statefulWorkspace: true, sandboxProfile: 'nsjail', runtimes: ['bash'], }, fetchImpl, }); - await worker.executeAndSettle({ - protocolVersion: 1, - assignmentId: 'assignment-deadline', - workerId: 'vm-1', - incarnationId: 'incarnation-00000001', - generation: 1, - leaseToken: 'lease-token-that-is-long-enough-for-testing', - expiresAt: new Date(Date.now() + 30).toISOString(), - request: { body: { language: 'bash' }, headers: {} }, - }); + await worker.register(); - assert.equal(settlement?.status, 'rejected'); - assert.equal(settlement?.incarnationId, 'incarnation-00000001'); + assert.ok(bridgeRequest); + const headers = bridgeRequest.init?.headers as Record; + const body = String(bridgeRequest.init?.body); + assert.equal( + verifyBridgeRequest( + key.publicKey, + { + credential: 'issued-short-lived-credential-value', + method: 'POST', + path: '/v1/bridge/workers/register', + timestamp: headers['X-LibreChat-Code-Timestamp'], + nonce: headers['X-LibreChat-Code-Nonce'], + body, + }, + headers['X-LibreChat-Code-Signature'], + ), + true, + ); }); -test('worker refreshes its registration during a long assignment', async () => { - let registrations = 0; +test('paired worker rotates an expiring credential before registration', async () => { + const key = createBridgeIdentity(); + const requests: Array<{ url: string; init?: RequestInit }> = []; + let persistedCredential = ''; const fetchImpl: typeof fetch = async (input, init) => { const url = String(input); - if (url.endsWith('/workers/register')) { - registrations += 1; - return new Response( - JSON.stringify({ - protocolVersion: 1, - workerId: 'vm-1', - incarnationId: 'incarnation-00000001', - registeredAt: new Date().toISOString(), - leaseTtlMs: 50, - }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, - ); - } - if (url.endsWith('/execute')) { - await new Promise((resolve) => setTimeout(resolve, 90)); - return new Response(JSON.stringify({ session_id: 'run-1', files: [] }), { - status: 200, - headers: { 'Content-Type': 'application/json' }, + requests.push({ url, init }); + if (url.endsWith('/credentials/refresh')) { + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'rotated-short-lived-credential-value', + expiresAt: new Date(Date.now() + 300_000).toISOString(), }); } - return new Response( - JSON.stringify({ protocolVersion: 1, accepted: true, body: init?.body }), - { status: 200, headers: { 'Content-Type': 'application/json' } }, - ); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); }; const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', - token: 'worker-secret', workerId: 'vm-1', - incarnationId: 'incarnation-00000001', sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'original-short-lived-credential-value', + expiresAt: new Date(Date.now() + 30_000).toISOString(), + }, capabilities: { - statefulWorkspace: false, + statefulWorkspace: true, sandboxProfile: 'nsjail', runtimes: ['bash'], }, fetchImpl, + onIdentityChange: (identity) => { + persistedCredential = identity.credential; + }, }); + + await worker.refreshCredential(); await worker.register(); - await worker.executeAndSettle({ - protocolVersion: 1, - assignmentId: 'assignment-heartbeat', - workerId: 'vm-1', - incarnationId: 'incarnation-00000001', - generation: 1, - leaseToken: 'lease-token-that-is-long-enough-for-testing', - expiresAt: new Date(Date.now() + 1_000).toISOString(), - request: { body: { language: 'bash' }, headers: {} }, - }); - assert.ok(registrations >= 2); + assert.equal(persistedCredential, 'rotated-short-lived-credential-value'); + assert.equal( + (requests[1].init?.headers as Record).Authorization, + 'Bridge rotated-short-lived-credential-value', + ); +}); + +test('reconnect delay uses bounded exponential jitter', () => { + assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 0), 500); + assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 1), 1_000); + assert.equal(reconnectDelayMs(10, 1_000, 30_000, () => 1), 30_000); }); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 3108394..a4465b0 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -5,6 +5,7 @@ import { BridgeProtocolError, bridgeWorkerPath, } from './protocol.js'; +import { signBridgeRequest } from './identity.js'; import type { BridgeAssignment, @@ -12,27 +13,46 @@ import type { BridgeSettlement, BridgeSettlementResponse, BridgeWorkerCapabilities, + BridgeWorkerCredentialResponse, BridgeWorkerRegistrationResponse, } from './protocol.js'; export interface BridgeWorkerOptions { codeApiUrl: string; - token: string; + token?: string; + identity?: BridgeWorkerIdentity; workerId: string; sandboxEndpoint: string; capabilities: BridgeWorkerCapabilities; leaseWaitMs?: number; reconnectDelayMs?: number; + reconnectMaxDelayMs?: number; + reconnectRandom?: () => number; fetchImpl?: typeof fetch; onError?: (error: unknown) => void; - incarnationId?: string; + onIdentityChange?: (identity: BridgeWorkerIdentity) => void | Promise; +} + +export interface BridgeWorkerIdentity { + privateKey: string; + credential: string; + expiresAt: string; } const DEFAULT_LEASE_WAIT_MS = 25_000; const DEFAULT_RECONNECT_DELAY_MS = 1_000; -const DEFAULT_REGISTRATION_TTL_MS = 60_000; -const MIN_REGISTRATION_HEARTBEAT_MS = 25; -const RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; +const DEFAULT_RECONNECT_MAX_DELAY_MS = 30_000; +const CREDENTIAL_REFRESH_WINDOW_MS = 60_000; + +export function reconnectDelayMs( + attempt: number, + baseDelayMs = DEFAULT_RECONNECT_DELAY_MS, + maxDelayMs = DEFAULT_RECONNECT_MAX_DELAY_MS, + random: () => number = Math.random, +): number { + const cap = Math.min(maxDelayMs, baseDelayMs * 2 ** Math.max(0, attempt)); + return Math.floor(cap * (0.5 + Math.min(1, Math.max(0, random())) * 0.5)); +} function normalizedBaseUrl(value: string): string { return value.replace(/\/+$/, ''); @@ -47,37 +67,30 @@ export class BridgeWorker { private readonly fetchImpl: typeof fetch; private readonly codeApiUrl: string; private readonly sandboxEndpoint: string; - private readonly incarnationId: string; - private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; constructor(private readonly options: BridgeWorkerOptions) { + if (!options.token && !options.identity) { + throw new BridgeProtocolError( + 'Bridge worker requires a static token or paired identity', + ); + } this.fetchImpl = options.fetchImpl ?? fetch; this.codeApiUrl = normalizedBaseUrl(options.codeApiUrl); this.sandboxEndpoint = normalizedBaseUrl(options.sandboxEndpoint); - this.incarnationId = - options.incarnationId ?? randomBytes(18).toString('base64url'); } async register( signal?: AbortSignal, ): Promise { - const registration = await this.request( + return this.request( `${this.codeApiUrl}/bridge/workers/register`, { protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: this.options.workerId, - incarnationId: this.incarnationId, capabilities: this.options.capabilities, }, signal, ); - if (registration.incarnationId !== this.incarnationId) { - throw new BridgeProtocolError( - 'Code API registered a different worker incarnation', - ); - } - this.registrationTtlMs = registration.leaseTtlMs; - return registration; } async lease(signal?: AbortSignal): Promise { @@ -86,44 +99,74 @@ export class BridgeWorker { { protocolVersion: BRIDGE_PROTOCOL_VERSION, waitMs: this.options.leaseWaitMs ?? DEFAULT_LEASE_WAIT_MS, - incarnationId: this.incarnationId, }, signal, ); - if ( - response.assignment != null && - response.assignment.incarnationId !== this.incarnationId - ) { - throw new BridgeProtocolError( - 'Code API leased an assignment for a different worker incarnation', - ); - } return response.assignment; } async run(signal?: AbortSignal): Promise { + let reconnectAttempt = 0; while (!signal?.aborted) { try { + await this.refreshCredential(signal); await this.register(signal); const assignment = await this.lease(signal); + reconnectAttempt = 0; if (!assignment) continue; await this.executeAndSettle(assignment, signal); } catch (error) { if (signal?.aborted) return; if ( error instanceof BridgeProtocolError && - (error.status === 401 || error.status === 403 || error.status === 409) + (error.status === 401 || error.status === 403) ) { throw error; } this.options.onError?.(error); - const delay = - this.options.reconnectDelayMs ?? DEFAULT_RECONNECT_DELAY_MS; + const delay = reconnectDelayMs( + reconnectAttempt, + this.options.reconnectDelayMs, + this.options.reconnectMaxDelayMs, + this.options.reconnectRandom, + ); + reconnectAttempt += 1; await new Promise((resolve) => setTimeout(resolve, delay)); } } } + async refreshCredential(signal?: AbortSignal): Promise { + const identity = this.options.identity; + if (identity == null) return; + if ( + Date.parse(identity.expiresAt) - Date.now() > + CREDENTIAL_REFRESH_WINDOW_MS + ) { + return; + } + const credential = await this.request( + `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}` + + '/credentials/refresh', + { protocolVersion: BRIDGE_PROTOCOL_VERSION }, + signal, + ); + if ( + credential.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + credential.workerId !== this.options.workerId || + typeof credential.credential !== 'string' || + credential.credential.length < 32 || + !Number.isFinite(Date.parse(credential.expiresAt)) + ) { + throw new BridgeProtocolError( + 'Code API returned an invalid rotated worker credential', + ); + } + identity.credential = credential.credential; + identity.expiresAt = credential.expiresAt; + await this.options.onIdentityChange?.(identity); + } + async executeAndSettle( assignment: BridgeAssignment, signal?: AbortSignal, @@ -131,23 +174,6 @@ export class BridgeWorker { const executionController = new AbortController(); const abortExecution = (): void => executionController.abort(); signal?.addEventListener('abort', abortExecution, { once: true }); - const deadlineDelay = Math.max( - 0, - Date.parse(assignment.expiresAt) - Date.now(), - ); - const deadlineTimer = setTimeout( - () => executionController.abort(), - deadlineDelay, - ); - const heartbeatController = new AbortController(); - let heartbeatError: unknown; - const heartbeat = this.maintainRegistration( - heartbeatController.signal, - executionController, - ).catch((error) => { - heartbeatError = error; - executionController.abort(); - }); const cancellationController = new AbortController(); const cancellationWatcher = this.watchCancellation( assignment, @@ -162,20 +188,16 @@ export class BridgeWorker { ? { 'X-Runtime-Session-Id': assignment.runtimeSessionId } : {}), }; - const response = await this.fetchImpl( - `${this.sandboxEndpointFor(assignment)}/execute`, - { - method: 'POST', - headers: { - ...headers, - 'Content-Type': 'application/json', - }, - body: JSON.stringify(assignment.request.body), - signal: executionController.signal, + const response = await this.fetchImpl(`${this.sandboxEndpoint}/execute`, { + method: 'POST', + headers: { + ...headers, + 'Content-Type': 'application/json', }, - ); + body: JSON.stringify(assignment.request.body), + signal: executionController.signal, + }); const payload = (await response.json()) as object; - if (heartbeatError != null) throw heartbeatError; if (!response.ok) { throw new BridgeProtocolError( errorMessage(payload) ?? @@ -187,7 +209,6 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment.generation, leaseToken: assignment.leaseToken, - incarnationId: this.incarnationId, status: 'fulfilled', result: payload, }; @@ -196,16 +217,12 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment.generation, leaseToken: assignment.leaseToken, - incarnationId: this.incarnationId, status: 'rejected', error: error instanceof Error ? error.message : 'Sandbox execution failed', }; } - clearTimeout(deadlineTimer); - heartbeatController.abort(); - await heartbeat; cancellationController.abort(); await cancellationWatcher; signal?.removeEventListener('abort', abortExecution); @@ -216,54 +233,6 @@ export class BridgeWorker { ); } - private sandboxEndpointFor(assignment: BridgeAssignment): string { - if (assignment.runtimeSessionId == null) return this.sandboxEndpoint; - if ( - this.options.capabilities.statefulWorkspace !== true || - !this.sandboxEndpoint.includes(RUNTIME_SESSION_PLACEHOLDER) - ) { - throw new BridgeProtocolError( - 'Stateful assignments require a sandbox endpoint template containing {runtimeSessionId}', - ); - } - return this.sandboxEndpoint.replace( - RUNTIME_SESSION_PLACEHOLDER, - encodeURIComponent(assignment.runtimeSessionId), - ); - } - - private async maintainRegistration( - signal: AbortSignal, - executionController: AbortController, - ): Promise { - while (!signal.aborted && !executionController.signal.aborted) { - await this.delay( - Math.max( - MIN_REGISTRATION_HEARTBEAT_MS, - Math.floor(this.registrationTtlMs / 2), - ), - signal, - ); - if (signal.aborted || executionController.signal.aborted) return; - await this.register(signal); - } - } - - private async delay(ms: number, signal: AbortSignal): Promise { - if (signal.aborted) return; - await new Promise((resolve) => { - const onAbort = (): void => { - clearTimeout(timer); - resolve(); - }; - const timer = setTimeout(() => { - signal.removeEventListener('abort', onAbort); - resolve(); - }, ms); - signal.addEventListener('abort', onAbort, { once: true }); - }); - } - private assignmentUrl(assignment: BridgeAssignment, action: string): string { return ( `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}` + @@ -282,10 +251,7 @@ export class BridgeWorker { try { const response = await this.request<{ cancelled: boolean }>( this.assignmentUrl(assignment, 'cancellation'), - { - protocolVersion: BRIDGE_PROTOCOL_VERSION, - incarnationId: this.incarnationId, - }, + { protocolVersion: BRIDGE_PROTOCOL_VERSION }, signal, ); if (response.cancelled) { @@ -307,13 +273,14 @@ export class BridgeWorker { body: object, signal?: AbortSignal, ): Promise { + const requestBody = JSON.stringify(body); const response = await this.fetchImpl(url, { method: 'POST', headers: { - Authorization: `Bearer ${this.options.token}`, + ...this.authorizationHeaders(url, requestBody), 'Content-Type': 'application/json', }, - body: JSON.stringify(body), + body: requestBody, signal, }); const payload = (await response.json()) as object; @@ -326,4 +293,33 @@ export class BridgeWorker { } return payload as T; } + + private authorizationHeaders( + url: string, + body: string, + ): Record { + const identity = this.options.identity; + if (identity == null) { + return { Authorization: `Bearer ${this.options.token}` }; + } + const timestamp = new Date().toISOString(); + const nonce = randomBytes(18).toString('base64url'); + const proof = { + credential: identity.credential, + method: 'POST', + path: new URL(url).pathname, + timestamp, + nonce, + body, + }; + return { + Authorization: `Bridge ${identity.credential}`, + 'X-LibreChat-Code-Timestamp': timestamp, + 'X-LibreChat-Code-Nonce': nonce, + 'X-LibreChat-Code-Signature': signBridgeRequest( + identity.privateKey, + proof, + ), + }; + } } diff --git a/service/rollup.config.js b/service/rollup.config.js index 0e7a8d8..059c71d 100644 --- a/service/rollup.config.js +++ b/service/rollup.config.js @@ -42,6 +42,7 @@ export default { 'src/**/*.ts', '../shared/telemetry-core.ts', '../packages/code/src/protocol.ts', + '../packages/code/src/identity.ts', ], sourceMap: true, declaration: false, diff --git a/service/src/api-server.ts b/service/src/api-server.ts index 8578460..1e4634a 100644 --- a/service/src/api-server.ts +++ b/service/src/api-server.ts @@ -18,7 +18,7 @@ import { requestErrorLogger, requestNotFoundLogger } from './middleware/request- import { localAuth } from './auth/local'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; -import bridgeRouter from './bridge/router'; +import bridgeRouter from './bridge'; import { connection } from './queue'; import { metricsHandler } from './metrics'; import { httpMetricsMiddleware } from './middleware/httpMetrics'; diff --git a/service/src/bridge/index.ts b/service/src/bridge/index.ts new file mode 100644 index 0000000..dfba2a8 --- /dev/null +++ b/service/src/bridge/index.ts @@ -0,0 +1,16 @@ +import { connection } from '../queue'; +import { env } from '../config'; +import { RedisBridgePairingStore } from './pairing'; +import { createBridgeRouter } from './router'; +import { RedisBridgeStore } from './store'; + +export const bridgeStore = new RedisBridgeStore(connection); +export const bridgePairings = new RedisBridgePairingStore(connection); + +export default createBridgeRouter({ + store: bridgeStore, + pairings: bridgePairings, + authMode: env.BRIDGE_AUTH_MODE, + adminToken: env.BRIDGE_TOKEN, + configuredWorkerId: env.BRIDGE_WORKER_ID, +}); diff --git a/service/src/bridge/pairing.test.ts b/service/src/bridge/pairing.test.ts new file mode 100644 index 0000000..a57395f --- /dev/null +++ b/service/src/bridge/pairing.test.ts @@ -0,0 +1,220 @@ +import { afterEach, describe, expect, test } from 'bun:test'; +import RedisMock from 'ioredis-mock'; + +import type Redis from 'ioredis'; + +import { + createBridgeIdentity, + signBridgeRequest, +} from '../../../packages/code/src/identity'; +import { RedisBridgePairingStore } from './pairing'; + +const redis = new RedisMock() as unknown as Redis; +const pairings = new RedisBridgePairingStore(redis); + +afterEach(async () => { + await redis.flushall(); +}); + +describe('RedisBridgePairingStore', () => { + test('redeems a pairing code exactly once for the intended worker identity', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + + const credential = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + + expect(credential.workerId).toBe('vm-1'); + expect(credential.credential.length).toBeGreaterThanOrEqual(32); + await expect( + pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }), + ).rejects.toMatchObject({ code: 'PAIRING_INVALID' }); + }); + + test('authorizes a credential only with proof from its worker key', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce: 'request-nonce-1', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }), + ).resolves.toEqual({ workerId: 'vm-1' }); + }); + + test('rejects replay of an already accepted worker proof', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce: 'single-use-nonce', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + const request = { + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }; + + await pairings.authorize(request); + + await expect(pairings.authorize(request)).rejects.toMatchObject({ + code: 'PROOF_REPLAYED', + }); + }); + + test('rejects a correctly signed proof outside the clock window', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date(Date.now() - 5 * 60_000).toISOString(), + nonce: 'stale-request-nonce', + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }), + ).rejects.toMatchObject({ code: 'PROOF_INVALID' }); + }); + + test('revocation immediately invalidates the active worker credential', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const issued = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path: '/v1/bridge/workers/register', + timestamp: new Date().toISOString(), + nonce: 'post-revocation-request', + body: JSON.stringify({ protocolVersion: 1, workerId: 'vm-1' }), + }; + + await pairings.revoke('vm-1'); + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + }); + + test('rotation replaces rather than duplicates the active credential', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const original = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + + const rotated = await pairings.rotate('vm-1'); + + const proofFor = ( + credential: string, + nonce: string, + ): Parameters[0] => { + const proof = { + credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/lease', + timestamp: new Date().toISOString(), + nonce, + body: JSON.stringify({ protocolVersion: 1, waitMs: 25_000 }), + }; + return { + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }; + }; + + await expect( + pairings.authorize(proofFor(original.credential, 'old-credential')), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + await expect( + pairings.authorize(proofFor(rotated.credential, 'new-credential')), + ).resolves.toEqual({ workerId: 'vm-1' }); + }); + + test('repairing a worker invalidates its previously paired credential', async () => { + const firstIdentity = createBridgeIdentity(); + const firstPairing = await pairings.issue('vm-1'); + const first = await pairings.redeem({ + workerId: 'vm-1', + code: firstPairing.code, + publicKey: firstIdentity.publicKey, + }); + const nextIdentity = createBridgeIdentity(); + const nextPairing = await pairings.issue('vm-1'); + await pairings.redeem({ + workerId: 'vm-1', + code: nextPairing.code, + publicKey: nextIdentity.publicKey, + }); + const proof = { + credential: first.credential, + method: 'POST', + path: '/v1/bridge/workers/register', + timestamp: new Date().toISOString(), + nonce: 'superseded-pairing', + body: JSON.stringify({ protocolVersion: 1, workerId: 'vm-1' }), + }; + + await expect( + pairings.authorize({ + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(firstIdentity.privateKey, proof), + }), + ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + }); +}); diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts new file mode 100644 index 0000000..c417090 --- /dev/null +++ b/service/src/bridge/pairing.ts @@ -0,0 +1,256 @@ +import { + createHash, + createPublicKey, + randomBytes, +} from 'crypto'; + +import type Redis from 'ioredis'; + +import { verifyBridgeRequest } from '../../../packages/code/src/identity'; + +const PREFIX = 'codeapi:bridge:v1'; +const DEFAULT_PAIRING_TTL_SECONDS = 10 * 60; +const DEFAULT_CREDENTIAL_TTL_SECONDS = 5 * 60; +const PROOF_NONCE_TTL_SECONDS = 2 * 60; +const PROOF_CLOCK_SKEW_MS = 60_000; + +interface StoredPairing { + workerId: string; + expiresAt: string; +} + +interface StoredCredential { + workerId: string; + publicKey: string; + expiresAt: string; +} + +export interface BridgePairing { + workerId: string; + code: string; + expiresAt: string; +} + +export interface BridgeWorkerCredential { + workerId: string; + credential: string; + expiresAt: string; +} + +export class BridgePairingError extends Error { + constructor( + public readonly code: + | 'PAIRING_INVALID' + | 'PUBLIC_KEY_INVALID' + | 'CREDENTIAL_INVALID' + | 'PROOF_INVALID' + | 'PROOF_REPLAYED', + message: string, + ) { + super(message); + this.name = 'BridgePairingError'; + } +} + +function digest(value: string): string { + return createHash('sha256').update(value).digest('hex'); +} + +function pairingKey(code: string): string { + return `${PREFIX}:pairing:${digest(code)}`; +} + +function credentialKey(credential: string): string { + return credentialDigestKey(digest(credential)); +} + +function credentialDigestKey(credentialDigest: string): string { + return `${PREFIX}:credential:${credentialDigest}`; +} + +function workerIdentityKey(workerId: string): string { + return `${PREFIX}:identity:${workerId}`; +} + +function proofNonceKey(credential: string, nonce: string): string { + return `${PREFIX}:proof:${digest(credential)}:${digest(nonce)}`; +} + +function validEd25519PublicKey(publicKey: string): boolean { + try { + return createPublicKey(publicKey).asymmetricKeyType === 'ed25519'; + } catch { + return false; + } +} + +export class RedisBridgePairingStore { + constructor( + private readonly redis: Redis, + private readonly pairingTtlSeconds = DEFAULT_PAIRING_TTL_SECONDS, + private readonly credentialTtlSeconds = DEFAULT_CREDENTIAL_TTL_SECONDS, + ) {} + + async issue(workerId: string): Promise { + const code = randomBytes(24).toString('base64url'); + const expiresAt = new Date( + Date.now() + this.pairingTtlSeconds * 1000, + ).toISOString(); + const pairing: StoredPairing = { workerId, expiresAt }; + await this.redis.set( + pairingKey(code), + JSON.stringify(pairing), + 'EX', + this.pairingTtlSeconds, + ); + return { workerId, code, expiresAt }; + } + + async redeem(args: { + workerId: string; + code: string; + publicKey: string; + }): Promise { + const raw = await this.redis.getdel(pairingKey(args.code)); + if (raw == null) { + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing code is invalid or expired', + ); + } + const pairing = JSON.parse(raw) as StoredPairing; + if (pairing.workerId !== args.workerId) { + throw new BridgePairingError( + 'PAIRING_INVALID', + 'Pairing code does not authorize this worker', + ); + } + if (!validEd25519PublicKey(args.publicKey)) { + throw new BridgePairingError( + 'PUBLIC_KEY_INVALID', + 'Worker public key must be an Ed25519 key', + ); + } + + return await this.issueCredential(args.workerId, args.publicKey); + } + + async authorize(args: { + workerId: string; + credential: string; + method: string; + path: string; + timestamp: string; + nonce: string; + body: string; + signature: string; + }): Promise<{ workerId: string }> { + const proofTime = Date.parse(args.timestamp); + if ( + !Number.isFinite(proofTime) || + Math.abs(Date.now() - proofTime) > PROOF_CLOCK_SKEW_MS + ) { + throw new BridgePairingError( + 'PROOF_INVALID', + 'Worker request proof is outside the accepted clock window', + ); + } + const credentialDigest = digest(args.credential); + const [raw, activeDigest] = await this.redis.mget( + credentialDigestKey(credentialDigest), + workerIdentityKey(args.workerId), + ); + if (raw == null || activeDigest !== credentialDigest) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + const stored = JSON.parse(raw) as StoredCredential; + if (stored.workerId !== args.workerId) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential does not authorize this worker', + ); + } + if (!verifyBridgeRequest(stored.publicKey, args, args.signature)) { + throw new BridgePairingError( + 'PROOF_INVALID', + 'Worker request proof is invalid', + ); + } + const accepted = await this.redis.set( + proofNonceKey(args.credential, args.nonce), + '1', + 'EX', + PROOF_NONCE_TTL_SECONDS, + 'NX', + ); + if (accepted !== 'OK') { + throw new BridgePairingError( + 'PROOF_REPLAYED', + 'Worker request proof has already been used', + ); + } + return { workerId: stored.workerId }; + } + + async revoke(workerId: string): Promise { + const identityKey = workerIdentityKey(workerId); + const credentialDigest = await this.redis.get(identityKey); + if (credentialDigest == null) return; + await this.redis.del(identityKey, credentialDigestKey(credentialDigest)); + } + + async rotate(workerId: string): Promise { + const identityKey = workerIdentityKey(workerId); + const previousDigest = await this.redis.get(identityKey); + const previousRaw = + previousDigest == null + ? null + : await this.redis.get(credentialDigestKey(previousDigest)); + if (previousRaw == null || previousDigest == null) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + const previous = JSON.parse(previousRaw) as StoredCredential; + return await this.issueCredential( + workerId, + previous.publicKey, + previousDigest, + ); + } + + private async issueCredential( + workerId: string, + publicKey: string, + previousDigest?: string, + ): Promise { + const credential = randomBytes(32).toString('base64url'); + const credentialDigest = digest(credential); + const expiresAt = new Date( + Date.now() + this.credentialTtlSeconds * 1000, + ).toISOString(); + const stored: StoredCredential = { workerId, publicKey, expiresAt }; + const transaction = this.redis.multi(); + transaction.set( + credentialDigestKey(credentialDigest), + JSON.stringify(stored), + 'EX', + this.credentialTtlSeconds, + ); + transaction.set( + workerIdentityKey(workerId), + credentialDigest, + 'EX', + this.credentialTtlSeconds, + ); + if (previousDigest !== undefined) { + transaction.del(credentialDigestKey(previousDigest)); + } + await transaction.exec(); + return { workerId, credential, expiresAt }; + } +} diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts new file mode 100644 index 0000000..3f80c04 --- /dev/null +++ b/service/src/bridge/router.test.ts @@ -0,0 +1,125 @@ +import { createServer, type Server } from 'http'; + +import { afterEach, describe, expect, test } from 'bun:test'; +import express, { json } from 'express'; +import RedisMock from 'ioredis-mock'; + +import type Redis from 'ioredis'; + +import { + createBridgeIdentity, + signBridgeRequest, +} from '../../../packages/code/src/identity'; +import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; +import { RedisBridgePairingStore } from './pairing'; +import { createBridgeRouter } from './router'; +import { RedisBridgeStore } from './store'; + +const redis = new RedisMock() as unknown as Redis; +let server: Server | undefined; + +afterEach(async () => { + server?.close(); + server = undefined; + await redis.flushall(); +}); + +describe('paired bridge HTTP API', () => { + test('pairs a worker and accepts its proof-of-possession registration', async () => { + const app = express(); + app.use(json()); + app.use( + '/v1/bridge', + createBridgeRouter({ + store: new RedisBridgeStore(redis), + pairings: new RedisBridgePairingStore(redis), + authMode: 'paired', + adminToken: 'strong-administrator-bootstrap-token', + configuredWorkerId: 'vm-1', + }), + ); + server = createServer(app); + await new Promise((resolve) => server?.listen(0, '127.0.0.1', resolve)); + const address = server.address(); + if (address == null || typeof address === 'string') { + throw new Error('Expected TCP listener'); + } + const baseUrl = `http://127.0.0.1:${address.port}/v1/bridge`; + const pairingResponse = await fetch(`${baseUrl}/pairings`, { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ workerId: 'vm-1' }), + }); + const pairing = (await pairingResponse.json()) as { code: string }; + expect(pairingResponse.status).toBe(200); + + const identity = createBridgeIdentity(); + const redemptionResponse = await fetch(`${baseUrl}/pairings/redeem`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }), + }); + const issued = (await redemptionResponse.json()) as { + credential: string; + }; + expect(redemptionResponse.status).toBe(200); + + const path = '/v1/bridge/workers/register'; + const body = JSON.stringify({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: 'vm-1', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + }); + const proof = { + credential: issued.credential, + method: 'POST', + path, + timestamp: new Date().toISOString(), + nonce: 'http-registration-nonce', + body, + }; + const headers = { + Authorization: `Bridge ${issued.credential}`, + 'Content-Type': 'application/json', + 'X-LibreChat-Code-Timestamp': proof.timestamp, + 'X-LibreChat-Code-Nonce': proof.nonce, + 'X-LibreChat-Code-Signature': signBridgeRequest( + identity.privateKey, + proof, + ), + }; + const registrationUrl = `http://127.0.0.1:${address.port}${path}`; + const registrationResponse = await fetch(registrationUrl, { + method: 'POST', + headers, + body, + }); + + expect(registrationResponse.status).toBe(200); + await expect(registrationResponse.json()).resolves.toMatchObject({ + workerId: 'vm-1', + }); + + const replayResponse = await fetch(registrationUrl, { + method: 'POST', + headers, + body, + }); + expect(replayResponse.status).toBe(401); + await expect(replayResponse.json()).resolves.toMatchObject({ + code: 'PROOF_REPLAYED', + }); + }); +}); diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index c69e1e0..3b50aac 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -1,20 +1,27 @@ import { timingSafeEqual } from 'crypto'; import { Router } from 'express'; + import type { NextFunction, Request, Response } from 'express'; import type { BridgeWorkerRegistration } from '../../../packages/code/src/protocol'; import type { CodeBridgeSettlement } from './store'; import { BRIDGE_PROTOCOL_VERSION } from '../../../packages/code/src/protocol'; -import { connection } from '../queue'; -import { env } from '../config'; +import { BridgePairingError, RedisBridgePairingStore } from './pairing'; import { BridgeStoreError, RedisBridgeStore } from './store'; const WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; -const INCARNATION_ID_PATTERN = /^[A-Za-z0-9_-]{16,128}$/; const MAX_LEASE_WAIT_MS = 30_000; -export const bridgeStore = new RedisBridgeStore(connection); +export type BridgeAuthMode = 'static' | 'paired'; + +export interface BridgeRouterOptions { + store: RedisBridgeStore; + pairings: RedisBridgePairingStore; + authMode: BridgeAuthMode; + adminToken: string; + configuredWorkerId?: string; +} function sameToken(left: string, right: string): boolean { const leftBuffer = Buffer.from(left); @@ -25,31 +32,10 @@ function sameToken(left: string, right: string): boolean { ); } -function bridgeAuth(req: Request, res: Response, next: NextFunction): void { - if (!env.BRIDGE_TOKEN) { - res.status(503).json({ error: 'Code bridge is not configured' }); - return; - } - const token = - req - .header('Authorization') - ?.match(/^Bearer\s+(.+)$/i)?.[1] - ?.trim() ?? ''; - if (!token || !sameToken(token, env.BRIDGE_TOKEN)) { - res.status(401).json({ error: 'Invalid code bridge worker token' }); - return; - } - next(); -} - function validWorkerId(value: string): boolean { return WORKER_ID_PATTERN.test(value); } -function validIncarnationId(value: unknown): value is string { - return typeof value === 'string' && INCARNATION_ID_PATTERN.test(value); -} - function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } @@ -67,8 +53,7 @@ function isSettlement(value: unknown): value is CodeBridgeSettlement { !Number.isSafeInteger(value.generation) || value.generation < 1 || typeof value.leaseToken !== 'string' || - value.leaseToken.length < 32 || - !validIncarnationId(value.incarnationId) + value.leaseToken.length < 32 ) { return false; } @@ -82,143 +67,304 @@ function isSettlement(value: unknown): value is CodeBridgeSettlement { ); } -const router = Router(); -router.use(bridgeAuth); +export function createBridgeRouter(options: BridgeRouterOptions): Router { + const router = Router(); -router.post('/workers/register', async (req: Request, res: Response) => { - const registration = req.body as unknown; - if ( - !isRecord(registration) || - registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - typeof registration.workerId !== 'string' || - !validWorkerId(registration.workerId) || - !validIncarnationId(registration.incarnationId) || - !isRecord(registration.capabilities) || - typeof registration.capabilities.statefulWorkspace !== 'boolean' || - typeof registration.capabilities.sandboxProfile !== 'string' || - registration.capabilities.sandboxProfile.trim().length === 0 || - registration.capabilities.sandboxProfile.length > 128 || - !Array.isArray(registration.capabilities.runtimes) || - registration.capabilities.runtimes.length > 32 || - !registration.capabilities.runtimes.every( - (runtime) => - typeof runtime === 'string' && runtime.length > 0 && runtime.length <= 64, - ) || - (registration.capabilities.policyDigest !== undefined && - (typeof registration.capabilities.policyDigest !== 'string' || - !/^[a-f0-9]{64}$/.test(registration.capabilities.policyDigest))) - ) { - res.status(400).json({ error: 'Invalid bridge worker registration' }); - return; - } - if (env.BRIDGE_WORKER_ID && registration.workerId !== env.BRIDGE_WORKER_ID) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - try { - await bridgeStore.register( - registration as unknown as BridgeWorkerRegistration, - ); - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + const configuredWorker = (workerId: string): boolean => + options.configuredWorkerId == null || + options.configuredWorkerId === '' || + workerId === options.configuredWorkerId; + + const bearerToken = (req: Request): string => + req + .header('Authorization') + ?.match(/^Bearer\s+(.+)$/i)?.[1] + ?.trim() ?? ''; + + const adminAuth = ( + req: Request, + res: Response, + next: NextFunction, + ): void => { + if (!options.adminToken) { + res.status(503).json({ error: 'Code bridge is not configured' }); return; } - throw error; - } - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - workerId: registration.workerId, - incarnationId: registration.incarnationId, - registeredAt: new Date().toISOString(), - leaseTtlMs: 60_000, + const token = bearerToken(req); + if (!token || !sameToken(token, options.adminToken)) { + res.status(401).json({ error: 'Invalid code bridge administrator token' }); + return; + } + next(); + }; + + const staticWorkerAuth = ( + req: Request, + res: Response, + next: NextFunction, + ): void => { + const token = bearerToken(req); + if (!token || !sameToken(token, options.adminToken)) { + res.status(401).json({ error: 'Invalid code bridge worker token' }); + return; + } + next(); + }; + + const pairedWorkerAuth = ( + req: Request, + res: Response, + next: NextFunction, + ): void => { + const workerId = + req.params.workerId || + (isRecord(req.body) && typeof req.body.workerId === 'string' + ? req.body.workerId + : ''); + const credential = + req + .header('Authorization') + ?.match(/^Bridge\s+(.+)$/i)?.[1] + ?.trim() ?? ''; + const timestamp = req.header('X-LibreChat-Code-Timestamp') ?? ''; + const nonce = req.header('X-LibreChat-Code-Nonce') ?? ''; + const signature = req.header('X-LibreChat-Code-Signature') ?? ''; + if ( + !validWorkerId(workerId) || + !credential || + !timestamp || + !nonce || + !signature + ) { + res.status(401).json({ error: 'Invalid paired worker authorization' }); + return; + } + void options.pairings + .authorize({ + workerId, + credential, + method: req.method, + path: req.originalUrl.split('?')[0], + timestamp, + nonce, + body: JSON.stringify(req.body ?? {}), + signature, + }) + .then(() => next()) + .catch((error: unknown) => { + if (error instanceof BridgePairingError) { + res.status(401).json({ error: error.message, code: error.code }); + return; + } + next(error); + }); + }; + + const workerAuth = + options.authMode === 'paired' ? pairedWorkerAuth : staticWorkerAuth; + + router.post('/pairings', adminAuth, async (req: Request, res: Response) => { + if (options.authMode !== 'paired') { + res.status(409).json({ error: 'Paired worker authentication is disabled' }); + return; + } + const workerId = isRecord(req.body) ? req.body.workerId : undefined; + if ( + typeof workerId !== 'string' || + !validWorkerId(workerId) || + !configuredWorker(workerId) + ) { + res.status(400).json({ error: 'Invalid bridge worker ID' }); + return; + } + const pairing = await options.pairings.issue(workerId); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...pairing }); }); -}); -router.post('/workers/:workerId/lease', async (req: Request, res: Response) => { - const workerId = req.params.workerId; - const body = isRecord(req.body) ? req.body : {}; - const requestedWait = Number(body.waitMs ?? 25_000); - if ( - !validWorkerId(workerId) || - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) || - !Number.isFinite(requestedWait) || - requestedWait < 0 - ) { - res.status(400).json({ error: 'Invalid bridge lease request' }); - return; - } - if (env.BRIDGE_WORKER_ID && workerId !== env.BRIDGE_WORKER_ID) { - res.status(403).json({ - error: 'Worker is not authorized for this Code API deployment', - }); - return; - } - try { - const assignment = await bridgeStore.lease( - workerId, - body.incarnationId, - Math.min(requestedWait, MAX_LEASE_WAIT_MS), - ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); - } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + router.post('/pairings/redeem', async (req: Request, res: Response) => { + if (options.authMode !== 'paired') { + res.status(409).json({ error: 'Paired worker authentication is disabled' }); return; } - throw error; - } -}); - -router.post( - '/workers/:workerId/assignments/:assignmentId/settle', - async (req, res) => { - const settlement = req.body as unknown; - if (!isSettlement(settlement)) { - res.status(400).json({ error: 'Invalid bridge settlement' }); + const redemption = req.body as unknown; + if ( + !isRecord(redemption) || + redemption.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof redemption.workerId !== 'string' || + !validWorkerId(redemption.workerId) || + !configuredWorker(redemption.workerId) || + typeof redemption.code !== 'string' || + redemption.code.length < 16 || + typeof redemption.publicKey !== 'string' || + redemption.publicKey.length > 4096 + ) { + res.status(400).json({ error: 'Invalid bridge pairing redemption' }); return; } try { - await bridgeStore.settle( - req.params.workerId, - req.params.assignmentId, - settlement, - ); - res.json({ - protocolVersion: BRIDGE_PROTOCOL_VERSION, - accepted: true, + const credential = await options.pairings.redeem({ + workerId: redemption.workerId, + code: redemption.code, + publicKey: redemption.publicKey, }); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); } catch (error) { - if (error instanceof BridgeStoreError) { - sendStoreError(error, res); + if (error instanceof BridgePairingError) { + const status = error.code === 'PUBLIC_KEY_INVALID' ? 400 : 401; + res.status(status).json({ error: error.message, code: error.code }); return; } throw error; } - }, -); + }); -router.post( - '/workers/:workerId/assignments/:assignmentId/cancellation', - async (req, res) => { - const body = isRecord(req.body) ? req.body : {}; - if ( - body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || - !validIncarnationId(body.incarnationId) - ) { - res.status(400).json({ error: 'Invalid bridge cancellation request' }); - return; - } - const cancelled = await bridgeStore.cancelled( - req.params.workerId, - body.incarnationId, - req.params.assignmentId, - ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); - }, -); - -export default router; + router.post( + '/workers/:workerId/revoke', + adminAuth, + async (req: Request, res: Response) => { + if (!validWorkerId(req.params.workerId)) { + res.status(400).json({ error: 'Invalid bridge worker ID' }); + return; + } + await options.pairings.revoke(req.params.workerId); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, revoked: true }); + }, + ); + + router.post( + '/workers/:workerId/credentials/refresh', + workerAuth, + async (req: Request, res: Response) => { + try { + const credential = await options.pairings.rotate(req.params.workerId); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); + } catch (error) { + if (error instanceof BridgePairingError) { + res.status(401).json({ error: error.message, code: error.code }); + return; + } + throw error; + } + }, + ); + + router.post( + '/workers/register', + workerAuth, + async (req: Request, res: Response) => { + const registration = req.body as unknown; + if ( + !isRecord(registration) || + registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + typeof registration.workerId !== 'string' || + !validWorkerId(registration.workerId) || + !isRecord(registration.capabilities) || + registration.capabilities.statefulWorkspace !== true || + typeof registration.capabilities.sandboxProfile !== 'string' || + registration.capabilities.sandboxProfile.trim().length === 0 || + registration.capabilities.sandboxProfile.length > 128 || + !Array.isArray(registration.capabilities.runtimes) || + registration.capabilities.runtimes.length > 32 || + !registration.capabilities.runtimes.every( + (runtime) => + typeof runtime === 'string' && + runtime.length > 0 && + runtime.length <= 64, + ) || + (registration.capabilities.policyDigest !== undefined && + (typeof registration.capabilities.policyDigest !== 'string' || + !/^[a-f0-9]{64}$/.test(registration.capabilities.policyDigest))) + ) { + res.status(400).json({ error: 'Invalid bridge worker registration' }); + return; + } + if (!configuredWorker(registration.workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + await options.store.register( + registration as unknown as BridgeWorkerRegistration, + ); + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + workerId: registration.workerId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + }, + ); + + router.post( + '/workers/:workerId/lease', + workerAuth, + async (req: Request, res: Response) => { + const workerId = req.params.workerId; + const body = isRecord(req.body) ? req.body : {}; + const requestedWait = Number(body.waitMs ?? 25_000); + if ( + !validWorkerId(workerId) || + !Number.isFinite(requestedWait) || + requestedWait < 0 + ) { + res.status(400).json({ error: 'Invalid bridge lease request' }); + return; + } + if (!configuredWorker(workerId)) { + res.status(403).json({ + error: 'Worker is not authorized for this Code API deployment', + }); + return; + } + const assignment = await options.store.lease( + workerId, + Math.min(requestedWait, MAX_LEASE_WAIT_MS), + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); + }, + ); + + router.post( + '/workers/:workerId/assignments/:assignmentId/settle', + workerAuth, + async (req: Request, res: Response) => { + const settlement = req.body as unknown; + if (!isSettlement(settlement)) { + res.status(400).json({ error: 'Invalid bridge settlement' }); + return; + } + try { + await options.store.settle( + req.params.workerId, + req.params.assignmentId, + settlement, + ); + res.json({ + protocolVersion: BRIDGE_PROTOCOL_VERSION, + accepted: true, + }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } + }, + ); + + router.post( + '/workers/:workerId/assignments/:assignmentId/cancellation', + workerAuth, + async (req: Request, res: Response) => { + const cancelled = await options.store.cancelled( + req.params.workerId, + req.params.assignmentId, + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); + }, + ); + + return router; +} diff --git a/service/src/config.test.ts b/service/src/config.test.ts index 22878ea..87658f9 100644 --- a/service/src/config.test.ts +++ b/service/src/config.test.ts @@ -1,6 +1,7 @@ import { describe, expect, test } from 'bun:test'; import { parsePlanLimits, + resolveBridgeAuthMode, resolveRuntimeSessionMode, resolveSandboxBackend, } from './config'; @@ -9,6 +10,7 @@ describe('sandbox execution configuration', () => { test('defaults only unset backend and session mode values', () => { expect(resolveSandboxBackend(undefined)).toBe('http'); expect(resolveRuntimeSessionMode(undefined)).toBe('stateless'); + expect(resolveBridgeAuthMode(undefined)).toBe('static'); }); test('accepts every supported backend and session mode', () => { @@ -18,6 +20,8 @@ describe('sandbox execution configuration', () => { expect(resolveRuntimeSessionMode('stateless')).toBe('stateless'); expect(resolveRuntimeSessionMode('affinity')).toBe('affinity'); expect(resolveRuntimeSessionMode('strict')).toBe('strict'); + expect(resolveBridgeAuthMode('static')).toBe('static'); + expect(resolveBridgeAuthMode('paired')).toBe('paired'); }); test('rejects unknown values instead of silently changing execution semantics', () => { @@ -31,6 +35,9 @@ describe('sandbox execution configuration', () => { ); expect(() => resolveRuntimeSessionMode('')).toThrow('CODEAPI_RUNTIME_SESSION_MODE'); expect(() => resolveRuntimeSessionMode(' ')).toThrow('CODEAPI_RUNTIME_SESSION_MODE'); + expect(() => resolveBridgeAuthMode('token')).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE must be one of: static, paired', + ); }); }); diff --git a/service/src/config.ts b/service/src/config.ts index dbc6dfe..605c701 100644 --- a/service/src/config.ts +++ b/service/src/config.ts @@ -263,8 +263,20 @@ export function resolveRuntimeSessionMode( ); } +export function resolveBridgeAuthMode( + raw: string | undefined, +): 'static' | 'paired' { + return configuredChoice( + raw, + 'CODEAPI_BRIDGE_AUTH_MODE', + 'static', + ['static', 'paired'], + ); +} + const sandboxBackend = resolveSandboxBackend(process.env.CODEAPI_SANDBOX_BACKEND); const runtimeSessionMode = resolveRuntimeSessionMode(process.env.CODEAPI_RUNTIME_SESSION_MODE); +const bridgeAuthMode = resolveBridgeAuthMode(process.env.CODEAPI_BRIDGE_AUTH_MODE); export const env = { PORT: process.env.SERVICE_PORT ?? 3112, @@ -355,6 +367,8 @@ export const env = { SANDBOX_BACKEND: sandboxBackend, /** Outbound worker selected by the remote-bridge backend. */ BRIDGE_WORKER_ID: process.env.CODEAPI_BRIDGE_WORKER_ID ?? '', + /** Static compatibility auth or short-lived proof-of-possession credentials. */ + BRIDGE_AUTH_MODE: bridgeAuthMode, /** Enrollment and lease credential shared only with the configured worker. */ BRIDGE_TOKEN: process.env.CODEAPI_BRIDGE_TOKEN ?? '', /** diff --git a/service/src/sandbox-backend/remote-bridge.ts b/service/src/sandbox-backend/remote-bridge.ts index a18788a..a30ab74 100644 --- a/service/src/sandbox-backend/remote-bridge.ts +++ b/service/src/sandbox-backend/remote-bridge.ts @@ -7,7 +7,7 @@ import type { import type { RedisBridgeStore } from '../bridge/store'; import { env } from '../config'; -import { bridgeStore } from '../bridge/router'; +import { bridgeStore } from '../bridge'; import { BridgeStoreError } from '../bridge/store'; import { SandboxBackendError } from './types'; diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 113b4fc..d057129 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -15,6 +15,7 @@ const saved = { executionProfileSource: env.EXECUTION_PROFILE_SOURCE, sandboxBackend: env.SANDBOX_BACKEND, bridgeWorkerId: env.BRIDGE_WORKER_ID, + bridgeAuthMode: env.BRIDGE_AUTH_MODE, bridgeToken: env.BRIDGE_TOKEN, ptcMode: env.PTC_MODE, runtimeSessionMode: env.RUNTIME_SESSION_MODE, @@ -55,6 +56,7 @@ function restore(): void { env.EXECUTION_PROFILE_SOURCE = saved.executionProfileSource; env.SANDBOX_BACKEND = saved.sandboxBackend; env.BRIDGE_WORKER_ID = saved.bridgeWorkerId; + env.BRIDGE_AUTH_MODE = saved.bridgeAuthMode; env.BRIDGE_TOKEN = saved.bridgeToken; env.PTC_MODE = saved.ptcMode; env.RUNTIME_SESSION_MODE = saved.runtimeSessionMode; @@ -310,7 +312,7 @@ describe('sandbox backend policy', () => { expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); - test('remote bridge requires replay PTC and a strong token in hardened mode', () => { + test('hardened remote bridge requires replay PTC, paired auth, and a strong administrator token', () => { env.SANDBOX_BACKEND = 'remote-bridge'; env.RUNTIME_SESSION_MODE = 'affinity'; env.BRIDGE_WORKER_ID = 'engineering-vm'; @@ -325,6 +327,11 @@ describe('sandbox backend policy', () => { expect(() => validateSandboxBackendPolicy()).toThrow('at least 32 bytes'); env.BRIDGE_TOKEN = 'strong-remote-bridge-token-32-bytes'; + expect(() => validateSandboxBackendPolicy()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index dc9cb4c..e94899d 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -122,6 +122,11 @@ export function validateSandboxBackendPolicy(): void { requireValue('CODEAPI_BRIDGE_WORKER_ID', env.BRIDGE_WORKER_ID); if (env.HARDENED_SANDBOX_MODE) { requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); + if (env.BRIDGE_AUTH_MODE !== 'paired') { + throw new SecureStartupConfigError( + 'Hardened remote bridge deployments require CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + } } else { requireValue('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); } diff --git a/service/src/service-api.ts b/service/src/service-api.ts index 6a0a833..79db08d 100644 --- a/service/src/service-api.ts +++ b/service/src/service-api.ts @@ -5,7 +5,7 @@ import { requestErrorLogger, requestNotFoundLogger } from './middleware/request- import { executionProfileMiddleware } from './middleware/execution-profile'; import serviceRouter from './service/router'; import programmaticRouter from './service/programmatic-router'; -import bridgeRouter from './bridge/router'; +import bridgeRouter from './bridge'; import { connection } from './queue'; import { env } from './config'; import logger from './logger'; diff --git a/service/tsconfig.json b/service/tsconfig.json index dcddf82..13f3087 100644 --- a/service/tsconfig.json +++ b/service/tsconfig.json @@ -16,7 +16,8 @@ "include": [ "src/**/*.ts", "../shared/telemetry-core.ts", - "../packages/code/src/protocol.ts" + "../packages/code/src/protocol.ts", + "../packages/code/src/identity.ts" ], "exclude": [ "node_modules", From 38ce27cae099ba4d309e2f7bc756bba0fddf924c Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 20:34:10 -0400 Subject: [PATCH 2/6] fix: harden paired worker lifecycle --- packages/code/src/cli.ts | 17 +++-- packages/code/src/storage.test.ts | 13 +++- packages/code/src/storage.ts | 13 +++- packages/code/src/worker.test.ts | 80 ++++++++++++++++++++++++ packages/code/src/worker.ts | 29 +++++++-- service/src/bridge/pairing.test.ts | 49 +++++++++++++-- service/src/bridge/pairing.ts | 99 ++++++++++++++++++++++++++---- service/src/bridge/router.test.ts | 13 ++++ service/src/bridge/router.ts | 19 +++++- service/src/lifecycle.ts | 2 + service/src/secure-startup.test.ts | 13 ++++ service/src/secure-startup.ts | 15 +++++ 12 files changed, 327 insertions(+), 35 deletions(-) diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index 64811ce..b8fcf5b 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -1,10 +1,12 @@ #!/usr/bin/env node import { createHash } from 'node:crypto'; -import { homedir } from 'node:os'; -import { join } from 'node:path'; import { pairBridgeWorker } from './pairing.js'; -import { loadBridgeIdentity, saveBridgeIdentity } from './storage.js'; +import { + defaultBridgeIdentityPath, + loadBridgeIdentity, + saveBridgeIdentity, +} from './storage.js'; import { BridgeWorker } from './worker.js'; function required(name: string, value = process.env[name]): string { @@ -28,11 +30,6 @@ function option(args: string[], name: string): string | undefined { return args.find((value) => value.startsWith(`${name}=`))?.slice(name.length + 1); } -function defaultIdentityPath(workerId: string): string { - const fileName = workerId.replace(/[^A-Za-z0-9._-]/g, '_'); - return join(homedir(), '.config', 'librechat', 'code', `${fileName}.json`); -} - async function pair(args: string[]): Promise { const codeApiUrl = required('instance URL', args[1]); const code = required('one-time pairing code', args[2]); @@ -43,7 +40,7 @@ async function pair(args: string[]): Promise { const identityPath = option(args, '--identity') ?? process.env.LIBRECHAT_CODE_IDENTITY_FILE ?? - defaultIdentityPath(workerId); + defaultBridgeIdentityPath(workerId); const identity = await pairBridgeWorker({ codeApiUrl, workerId, code }); await saveBridgeIdentity(identityPath, identity); process.stdout.write( @@ -58,7 +55,7 @@ async function run(): Promise { const identityPath = configuredIdentityPath ?? (configuredWorkerId && !configuredToken - ? defaultIdentityPath(configuredWorkerId) + ? defaultBridgeIdentityPath(configuredWorkerId) : undefined); const pairedIdentity = identityPath ? await loadBridgeIdentity(identityPath) diff --git a/packages/code/src/storage.test.ts b/packages/code/src/storage.test.ts index 8b5128d..ddd782b 100644 --- a/packages/code/src/storage.test.ts +++ b/packages/code/src/storage.test.ts @@ -4,7 +4,18 @@ import { tmpdir } from 'node:os'; import { join } from 'node:path'; import test from 'node:test'; -import { loadBridgeIdentity, saveBridgeIdentity } from './storage.js'; +import { + defaultBridgeIdentityPath, + loadBridgeIdentity, + saveBridgeIdentity, +} from './storage.js'; + +test('default identity paths do not collide after worker ID sanitization', () => { + assert.notEqual( + defaultBridgeIdentityPath('vm:a'), + defaultBridgeIdentityPath('vm_a'), + ); +}); test('paired identity is persisted atomically with owner-only permissions', async () => { const directory = await mkdtemp(join(tmpdir(), 'librechat-code-')); diff --git a/packages/code/src/storage.ts b/packages/code/src/storage.ts index c44135e..a26c3eb 100644 --- a/packages/code/src/storage.ts +++ b/packages/code/src/storage.ts @@ -1,6 +1,7 @@ -import { randomBytes } from 'node:crypto'; +import { createHash, randomBytes } from 'node:crypto'; import { chmod, mkdir, open, readFile, rename, rm } from 'node:fs/promises'; -import { dirname } from 'node:path'; +import { homedir } from 'node:os'; +import { dirname, join } from 'node:path'; import { BRIDGE_PROTOCOL_VERSION, BridgeProtocolError } from './protocol.js'; @@ -24,6 +25,14 @@ function isPairedIdentity(value: unknown): value is PairedBridgeWorkerIdentity { ); } +export function defaultBridgeIdentityPath(workerId: string): string { + const readableName = workerId.replace(/[^A-Za-z0-9._-]/g, '_'); + const fileName = readableName === workerId + ? readableName + : `${readableName}-${createHash('sha256').update(workerId).digest('hex').slice(0, 16)}`; + return join(homedir(), '.config', 'librechat', 'code', `${fileName}.json`); +} + export async function saveBridgeIdentity( path: string, identity: PairedBridgeWorkerIdentity, diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index be72a3d..b363172 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -176,6 +176,86 @@ test('paired worker rotates an expiring credential before registration', async ( ); }); +test('paired worker refreshes before an assignment that outlives its credential', async () => { + const key = createBridgeIdentity(); + const requests: string[] = []; + const fetchImpl: typeof fetch = async (input) => { + const url = String(input); + requests.push(url); + if (url.endsWith('/credentials/refresh')) { + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'assignment-safe-rotated-credential-value', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }); + } + if (url.endsWith('/execute')) { + return Response.json({ session_id: 'run-long', files: [] }); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity: { + privateKey: key.privateKey, + credential: 'credential-too-short-for-assignment', + expiresAt: new Date(Date.now() + 90_000).toISOString(), + }, + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-long', + workerId: 'vm-1', + generation: 4, + leaseToken: 'assignment-long-lease-token-value', + expiresAt: new Date(Date.now() + 5 * 60_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.match(requests[0], /credentials\/refresh$/); + assert.equal(requests[1], 'http://127.0.0.1:2000/api/v2/execute'); +}); + +test('worker shutdown interrupts reconnect backoff', async () => { + const controller = new AbortController(); + let failed!: () => void; + const failure = new Promise((resolve) => { + failed = resolve; + }); + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: true, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async () => { + throw new Error('offline'); + }, + reconnectDelayMs: 30_000, + reconnectMaxDelayMs: 30_000, + onError: () => failed(), + }); + + const run = worker.run(controller.signal); + await failure; + controller.abort(); + await run; +}); + test('reconnect delay uses bounded exponential jitter', () => { assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 0), 500); assert.equal(reconnectDelayMs(0, 1_000, 30_000, () => 1), 1_000); diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index a4465b0..4ec69a7 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -63,6 +63,21 @@ function errorMessage(value: object): string | undefined { return undefined; } +async function abortableDelay(ms: number, signal?: AbortSignal): Promise { + if (signal?.aborted === true) return; + await new Promise((resolve) => { + const onAbort = (): void => { + clearTimeout(timer); + resolve(); + }; + const timer = setTimeout(() => { + signal?.removeEventListener('abort', onAbort); + resolve(); + }, ms); + signal?.addEventListener('abort', onAbort, { once: true }); + }); +} + export class BridgeWorker { private readonly fetchImpl: typeof fetch; private readonly codeApiUrl: string; @@ -131,17 +146,19 @@ export class BridgeWorker { this.options.reconnectRandom, ); reconnectAttempt += 1; - await new Promise((resolve) => setTimeout(resolve, delay)); + await abortableDelay(delay, signal); } } } - async refreshCredential(signal?: AbortSignal): Promise { + async refreshCredential( + signal?: AbortSignal, + validThroughMs = Date.now() + CREDENTIAL_REFRESH_WINDOW_MS, + ): Promise { const identity = this.options.identity; if (identity == null) return; if ( - Date.parse(identity.expiresAt) - Date.now() > - CREDENTIAL_REFRESH_WINDOW_MS + Date.parse(identity.expiresAt) > validThroughMs ) { return; } @@ -171,6 +188,10 @@ export class BridgeWorker { assignment: BridgeAssignment, signal?: AbortSignal, ): Promise { + await this.refreshCredential( + signal, + Date.parse(assignment.expiresAt) + CREDENTIAL_REFRESH_WINDOW_MS, + ); const executionController = new AbortController(); const abortExecution = (): void => executionController.abort(); signal?.addEventListener('abort', abortExecution, { once: true }); diff --git a/service/src/bridge/pairing.test.ts b/service/src/bridge/pairing.test.ts index a57395f..1366d47 100644 --- a/service/src/bridge/pairing.test.ts +++ b/service/src/bridge/pairing.test.ts @@ -61,7 +61,7 @@ describe('RedisBridgePairingStore', () => { workerId: 'vm-1', signature: signBridgeRequest(identity.privateKey, proof), }), - ).resolves.toEqual({ workerId: 'vm-1' }); + ).resolves.toMatchObject({ workerId: 'vm-1' }); }); test('rejects replay of an already accepted worker proof', async () => { @@ -147,7 +147,7 @@ describe('RedisBridgePairingStore', () => { ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); }); - test('rotation replaces rather than duplicates the active credential', async () => { + test('rotation keeps the prior same-identity credential usable for recovery', async () => { const identity = createBridgeIdentity(); const pairing = await pairings.issue('vm-1'); const original = await pairings.redeem({ @@ -179,10 +179,51 @@ describe('RedisBridgePairingStore', () => { await expect( pairings.authorize(proofFor(original.credential, 'old-credential')), - ).rejects.toMatchObject({ code: 'CREDENTIAL_INVALID' }); + ).resolves.toMatchObject({ workerId: 'vm-1' }); await expect( pairings.authorize(proofFor(rotated.credential, 'new-credential')), - ).resolves.toEqual({ workerId: 'vm-1' }); + ).resolves.toMatchObject({ workerId: 'vm-1' }); + }); + + test('recovers when a refresh response is lost after the server commits it', async () => { + const identity = createBridgeIdentity(); + const pairing = await pairings.issue('vm-1'); + const original = await pairings.redeem({ + workerId: 'vm-1', + code: pairing.code, + publicKey: identity.publicKey, + }); + const proofFor = ( + credential: string, + nonce: string, + ): Parameters[0] => { + const proof = { + credential, + method: 'POST', + path: '/v1/bridge/workers/vm-1/credentials/refresh', + timestamp: new Date().toISOString(), + nonce, + body: JSON.stringify({ protocolVersion: 1 }), + }; + return { + ...proof, + workerId: 'vm-1', + signature: signBridgeRequest(identity.privateKey, proof), + }; + }; + + await pairings.rotate('vm-1'); + const retryAuthorization = await pairings.authorize( + proofFor(original.credential, 'refresh-response-lost'), + ); + const recovered = await pairings.rotate( + 'vm-1', + retryAuthorization.credentialId, + ); + + await expect( + pairings.authorize(proofFor(recovered.credential, 'refresh-recovered')), + ).resolves.toMatchObject({ workerId: 'vm-1' }); }); test('repairing a worker invalidates its previously paired credential', async () => { diff --git a/service/src/bridge/pairing.ts b/service/src/bridge/pairing.ts index c417090..54aaf5a 100644 --- a/service/src/bridge/pairing.ts +++ b/service/src/bridge/pairing.ts @@ -10,9 +10,23 @@ import { verifyBridgeRequest } from '../../../packages/code/src/identity'; const PREFIX = 'codeapi:bridge:v1'; const DEFAULT_PAIRING_TTL_SECONDS = 10 * 60; -const DEFAULT_CREDENTIAL_TTL_SECONDS = 5 * 60; +const DEFAULT_CREDENTIAL_TTL_SECONDS = 15 * 60; const PROOF_NONCE_TTL_SECONDS = 2 * 60; const PROOF_CLOCK_SKEW_MS = 60_000; +const ROTATE_CREDENTIAL_SCRIPT = ` +local activeDigest = redis.call('GET', KEYS[1]) +local previous = redis.call('GET', KEYS[2]) +if not activeDigest or not previous then + return 0 +end +if activeDigest ~= ARGV[1] and redis.call('GET', KEYS[4]) ~= ARGV[5] then + return 0 +end +redis.call('SET', KEYS[3], ARGV[3], 'EX', ARGV[4]) +redis.call('SET', KEYS[1], ARGV[2], 'EX', ARGV[4]) +redis.call('SET', KEYS[4], ARGV[5], 'EX', ARGV[4]) +return 1 +`; interface StoredPairing { workerId: string; @@ -21,6 +35,7 @@ interface StoredPairing { interface StoredCredential { workerId: string; + identityId: string; publicKey: string; expiresAt: string; } @@ -72,6 +87,10 @@ function workerIdentityKey(workerId: string): string { return `${PREFIX}:identity:${workerId}`; } +function workerStableIdentityKey(workerId: string): string { + return `${PREFIX}:stable-identity:${workerId}`; +} + function proofNonceKey(credential: string, nonce: string): string { return `${PREFIX}:proof:${digest(credential)}:${digest(nonce)}`; } @@ -144,7 +163,12 @@ export class RedisBridgePairingStore { nonce: string; body: string; signature: string; - }): Promise<{ workerId: string }> { + }): Promise<{ + workerId: string; + credentialId: string; + activeCredentialId: string; + identityId: string; + }> { const proofTime = Date.parse(args.timestamp); if ( !Number.isFinite(proofTime) || @@ -160,13 +184,27 @@ export class RedisBridgePairingStore { credentialDigestKey(credentialDigest), workerIdentityKey(args.workerId), ); - if (raw == null || activeDigest !== credentialDigest) { + if (raw == null || activeDigest == null) { throw new BridgePairingError( 'CREDENTIAL_INVALID', 'Worker credential is invalid or expired', ); } const stored = JSON.parse(raw) as StoredCredential; + if (activeDigest !== credentialDigest) { + const activeRaw = await this.redis.get( + credentialDigestKey(activeDigest), + ); + const active = activeRaw == null + ? undefined + : JSON.parse(activeRaw) as StoredCredential; + if (active?.identityId !== stored.identityId) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + } if (stored.workerId !== args.workerId) { throw new BridgePairingError( 'CREDENTIAL_INVALID', @@ -192,19 +230,31 @@ export class RedisBridgePairingStore { 'Worker request proof has already been used', ); } - return { workerId: stored.workerId }; + return { + workerId: stored.workerId, + credentialId: credentialDigest, + activeCredentialId: activeDigest, + identityId: stored.identityId, + }; } async revoke(workerId: string): Promise { const identityKey = workerIdentityKey(workerId); const credentialDigest = await this.redis.get(identityKey); if (credentialDigest == null) return; - await this.redis.del(identityKey, credentialDigestKey(credentialDigest)); + await this.redis.del( + identityKey, + workerStableIdentityKey(workerId), + credentialDigestKey(credentialDigest), + ); } - async rotate(workerId: string): Promise { + async rotate( + workerId: string, + expectedCredentialId?: string, + ): Promise { const identityKey = workerIdentityKey(workerId); - const previousDigest = await this.redis.get(identityKey); + const previousDigest = expectedCredentialId ?? await this.redis.get(identityKey); const previousRaw = previousDigest == null ? null @@ -220,6 +270,7 @@ export class RedisBridgePairingStore { workerId, previous.publicKey, previousDigest, + previous.identityId, ); } @@ -227,13 +278,36 @@ export class RedisBridgePairingStore { workerId: string, publicKey: string, previousDigest?: string, + identityId = randomBytes(18).toString('base64url'), ): Promise { const credential = randomBytes(32).toString('base64url'); const credentialDigest = digest(credential); const expiresAt = new Date( Date.now() + this.credentialTtlSeconds * 1000, ).toISOString(); - const stored: StoredCredential = { workerId, publicKey, expiresAt }; + const stored: StoredCredential = { workerId, identityId, publicKey, expiresAt }; + if (previousDigest !== undefined) { + const rotated = await this.redis.eval( + ROTATE_CREDENTIAL_SCRIPT, + 4, + workerIdentityKey(workerId), + credentialDigestKey(previousDigest), + credentialDigestKey(credentialDigest), + workerStableIdentityKey(workerId), + previousDigest, + credentialDigest, + JSON.stringify(stored), + String(this.credentialTtlSeconds), + identityId, + ); + if (rotated !== 1) { + throw new BridgePairingError( + 'CREDENTIAL_INVALID', + 'Worker credential is invalid or expired', + ); + } + return { workerId, credential, expiresAt }; + } const transaction = this.redis.multi(); transaction.set( credentialDigestKey(credentialDigest), @@ -247,9 +321,12 @@ export class RedisBridgePairingStore { 'EX', this.credentialTtlSeconds, ); - if (previousDigest !== undefined) { - transaction.del(credentialDigestKey(previousDigest)); - } + transaction.set( + workerStableIdentityKey(workerId), + identityId, + 'EX', + this.credentialTtlSeconds, + ); await transaction.exec(); return { workerId, credential, expiresAt }; } diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 3f80c04..34f7906 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -112,6 +112,19 @@ describe('paired bridge HTTP API', () => { workerId: 'vm-1', }); + const crossDeploymentRevoke = await fetch( + `${baseUrl}/workers/another-deployments-worker/revoke`, + { + method: 'POST', + headers: { + Authorization: 'Bearer strong-administrator-bootstrap-token', + 'Content-Type': 'application/json', + }, + body: '{}', + }, + ); + expect(crossDeploymentRevoke.status).toBe(400); + const replayResponse = await fetch(registrationUrl, { method: 'POST', headers, diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index 3b50aac..b91037b 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -150,7 +150,10 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { body: JSON.stringify(req.body ?? {}), signature, }) - .then(() => next()) + .then((authorization) => { + res.locals.bridgeWorkerAuthorization = authorization; + next(); + }) .catch((error: unknown) => { if (error instanceof BridgePairingError) { res.status(401).json({ error: error.message, code: error.code }); @@ -222,7 +225,10 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { '/workers/:workerId/revoke', adminAuth, async (req: Request, res: Response) => { - if (!validWorkerId(req.params.workerId)) { + if ( + !validWorkerId(req.params.workerId) || + !configuredWorker(req.params.workerId) + ) { res.status(400).json({ error: 'Invalid bridge worker ID' }); return; } @@ -236,7 +242,14 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { workerAuth, async (req: Request, res: Response) => { try { - const credential = await options.pairings.rotate(req.params.workerId); + const credential = await options.pairings.rotate( + req.params.workerId, + ( + res.locals.bridgeWorkerAuthorization as + | { credentialId: string } + | undefined + )?.credentialId, + ); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, ...credential }); } catch (error) { if (error instanceof BridgePairingError) { diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 8fc0adc..068eadc 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -7,6 +7,7 @@ import { validateApiHardenedConfig, validateExecutionProfilePolicy, validateSandboxBackendPolicy, + validateApiBridgePolicy, validateWorkerHardenedConfig, } from './secure-startup'; import logger from './logger'; @@ -90,6 +91,7 @@ export async function startupApiOnly(): Promise { logger.info('Starting API service (no workers)...'); validateApiHardenedConfig(); validateExecutionProfilePolicy({ requireBackendMatch: false }); + validateApiBridgePolicy(); /* No validateSandboxBackendPolicy() here: an API-only pod authenticates and * enqueues jobs, it never constructs the Lambda backend or checkpoint store. * Validating that policy would force worker-only config (LAMBDA_MICROVM_* and diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index d057129..8f7d9e5 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -2,6 +2,7 @@ import { afterEach, describe, expect, test } from 'bun:test'; import { env } from './config'; import { validateApiHardenedConfig, + validateApiBridgePolicy, validateEgressGatewayHardenedConfig, validateExecutionProfilePolicy, validateSandboxBackendPolicy, @@ -335,6 +336,18 @@ describe('sandbox backend policy', () => { expect(() => validateSandboxBackendPolicy()).not.toThrow(); }); + test('API-only hardened bridge validation rejects static worker auth', () => { + env.SANDBOX_BACKEND = 'remote-bridge'; + env.HARDENED_SANDBOX_MODE = true; + env.BRIDGE_AUTH_MODE = 'static'; + expect(() => validateApiBridgePolicy()).toThrow( + 'CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + + env.BRIDGE_AUTH_MODE = 'paired'; + expect(() => validateApiBridgePolicy()).not.toThrow(); + }); + test('rejects blocking PTC on the lambda backend', () => { configureValidLambda(); env.PTC_MODE = 'blocking'; diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index e94899d..a28b113 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -226,6 +226,21 @@ export function validateSandboxBackendPolicy(): void { } } +/** API-only pods expose bridge pairing and worker routes even though they do + * not construct the sandbox backend, so enforce the bridge authentication + * invariant without requiring worker-only Lambda or checkpoint settings. */ +export function validateApiBridgePolicy(): void { + if ( + env.SANDBOX_BACKEND === 'remote-bridge' && + env.HARDENED_SANDBOX_MODE && + env.BRIDGE_AUTH_MODE !== 'paired' + ) { + throw new SecureStartupConfigError( + 'Hardened remote bridge deployments require CODEAPI_BRIDGE_AUTH_MODE=paired', + ); + } +} + export function validateEgressGatewayHardenedConfig(): void { if (!env.HARDENED_SANDBOX_MODE) return; rejectValue('CODEAPI_SYNTHETIC_ACCESS_TOKEN', process.env.CODEAPI_SYNTHETIC_ACCESS_TOKEN); From 6f5209fa502d7600799384c4323c06a28460f284 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 20:52:53 -0400 Subject: [PATCH 3/6] fix: require paired auth on hardened APIs --- service/src/secure-startup.test.ts | 2 +- service/src/secure-startup.ts | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 8f7d9e5..6101b1a 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -337,7 +337,7 @@ describe('sandbox backend policy', () => { }); test('API-only hardened bridge validation rejects static worker auth', () => { - env.SANDBOX_BACKEND = 'remote-bridge'; + env.SANDBOX_BACKEND = 'http'; env.HARDENED_SANDBOX_MODE = true; env.BRIDGE_AUTH_MODE = 'static'; expect(() => validateApiBridgePolicy()).toThrow( diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index a28b113..4d7709b 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -231,12 +231,11 @@ export function validateSandboxBackendPolicy(): void { * invariant without requiring worker-only Lambda or checkpoint settings. */ export function validateApiBridgePolicy(): void { if ( - env.SANDBOX_BACKEND === 'remote-bridge' && env.HARDENED_SANDBOX_MODE && env.BRIDGE_AUTH_MODE !== 'paired' ) { throw new SecureStartupConfigError( - 'Hardened remote bridge deployments require CODEAPI_BRIDGE_AUTH_MODE=paired', + 'Hardened API deployments require CODEAPI_BRIDGE_AUTH_MODE=paired because bridge routes are always exposed', ); } } From e35106a67a309dcf669d2701cbd1efd924003a83 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 21:00:57 -0400 Subject: [PATCH 4/6] fix: harden bridge pairing startup policy --- docs/adr/001-stateful-code-environments.md | 2 +- docs/remote-bridge/README.md | 2 +- service/src/secure-startup.test.ts | 4 ++++ service/src/secure-startup.ts | 7 +++---- 4 files changed, 9 insertions(+), 6 deletions(-) diff --git a/docs/adr/001-stateful-code-environments.md b/docs/adr/001-stateful-code-environments.md index 8163a95..8b9830d 100644 --- a/docs/adr/001-stateful-code-environments.md +++ b/docs/adr/001-stateful-code-environments.md @@ -27,7 +27,7 @@ Hardened workers enroll through a one-time pairing code: 1. An administrator creates a code scoped to the configured worker ID. 2. The CLI generates an Ed25519 keypair locally and redeems the code with only its public key. -3. Code API returns a five-minute credential bound to that public key. +3. Code API returns a fifteen-minute credential bound to that public key. 4. Every worker request signs the method, path, body digest, timestamp, nonce, and credential. 5. Code API rejects stale timestamps and replayed nonces and supports rotation diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index c0b9661..bb9da94 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -74,7 +74,7 @@ execution. - Registration is ephemeral in Redis and must be refreshed by the worker. - Pairing codes are stored hashed, expire after ten minutes, and are consumed atomically on their first redemption attempt. -- Worker credentials expire after five minutes and are bound to an Ed25519 +- Worker credentials expire after fifteen minutes and are bound to an Ed25519 public key. Exact-request signatures include the HTTP method, path, body digest, timestamp, nonce, and credential. - Accepted proof nonces cannot be replayed, credentials rotate before expiry, diff --git a/service/src/secure-startup.test.ts b/service/src/secure-startup.test.ts index 6101b1a..600f975 100644 --- a/service/src/secure-startup.test.ts +++ b/service/src/secure-startup.test.ts @@ -340,12 +340,16 @@ describe('sandbox backend policy', () => { env.SANDBOX_BACKEND = 'http'; env.HARDENED_SANDBOX_MODE = true; env.BRIDGE_AUTH_MODE = 'static'; + env.BRIDGE_TOKEN = 'strong-remote-bridge-token-32-bytes'; expect(() => validateApiBridgePolicy()).toThrow( 'CODEAPI_BRIDGE_AUTH_MODE=paired', ); env.BRIDGE_AUTH_MODE = 'paired'; expect(() => validateApiBridgePolicy()).not.toThrow(); + + env.BRIDGE_TOKEN = 'guessable'; + expect(() => validateApiBridgePolicy()).toThrow('at least 32 bytes'); }); test('rejects blocking PTC on the lambda backend', () => { diff --git a/service/src/secure-startup.ts b/service/src/secure-startup.ts index 4d7709b..b213dbf 100644 --- a/service/src/secure-startup.ts +++ b/service/src/secure-startup.ts @@ -230,10 +230,9 @@ export function validateSandboxBackendPolicy(): void { * not construct the sandbox backend, so enforce the bridge authentication * invariant without requiring worker-only Lambda or checkpoint settings. */ export function validateApiBridgePolicy(): void { - if ( - env.HARDENED_SANDBOX_MODE && - env.BRIDGE_AUTH_MODE !== 'paired' - ) { + if (!env.HARDENED_SANDBOX_MODE) return; + requireStrongSecret('CODEAPI_BRIDGE_TOKEN', env.BRIDGE_TOKEN); + if (env.BRIDGE_AUTH_MODE !== 'paired') { throw new SecureStartupConfigError( 'Hardened API deployments require CODEAPI_BRIDGE_AUTH_MODE=paired because bridge routes are always exposed', ); From f85a53789476dbf8b4b7da1d7bf8b6d10b38bb05 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 21:11:36 -0400 Subject: [PATCH 5/6] fix: preserve bridge fencing through pairing --- docs/remote-bridge/README.md | 5 + packages/code/src/cli.ts | 17 +++- packages/code/src/worker.test.ts | 151 +++++++++++++++++++++++++++++- packages/code/src/worker.ts | 120 +++++++++++++++++++++--- service/src/bridge/router.test.ts | 2 + service/src/bridge/router.ts | 56 +++++++++-- service/src/lifecycle.ts | 1 + 7 files changed, 322 insertions(+), 30 deletions(-) diff --git a/docs/remote-bridge/README.md b/docs/remote-bridge/README.md index bb9da94..e5bddc9 100644 --- a/docs/remote-bridge/README.md +++ b/docs/remote-bridge/README.md @@ -47,6 +47,11 @@ locally, proves possession on every request, and rotates its short-lived credential before expiry. `CODEAPI_BRIDGE_AUTH_MODE=static` remains available for non-hardened development compatibility only. +Stateful deployments must also set `LIBRECHAT_CODE_STATEFUL_WORKSPACE=true` +and route the CLI's `{runtimeSessionId}` endpoint template to an isolated, +persistent local runner per session. A single sandbox endpoint is stateless and +is rejected for runtime-session assignments. + ## LibreChat configuration Expose the Code API deployment as an environment under the Agents endpoint: diff --git a/packages/code/src/cli.ts b/packages/code/src/cli.ts index b8fcf5b..8cdbe6b 100644 --- a/packages/code/src/cli.ts +++ b/packages/code/src/cli.ts @@ -74,6 +74,17 @@ async function run(): Promise { process.env.LIBRECHAT_CODE_URL ?? pairedIdentity?.codeApiUrl, ); const policy = process.env.LIBRECHAT_CODE_POLICY ?? 'default-deny'; + const statefulWorkspace = + process.env.LIBRECHAT_CODE_STATEFUL_WORKSPACE?.trim().toLowerCase() === + 'true'; + const sandboxEndpoint = + process.env.LIBRECHAT_CODE_SANDBOX_ENDPOINT ?? + 'http://127.0.0.1:2000/api/v2'; + if (statefulWorkspace && !sandboxEndpoint.includes('{runtimeSessionId}')) { + throw new Error( + 'LIBRECHAT_CODE_STATEFUL_WORKSPACE requires LIBRECHAT_CODE_SANDBOX_ENDPOINT to contain {runtimeSessionId}', + ); + } const workerIdentity = pairedIdentity ? { privateKey: pairedIdentity.privateKey, @@ -89,11 +100,9 @@ async function run(): Promise { token: configuredToken, identity: workerIdentity, workerId, - sandboxEndpoint: - process.env.LIBRECHAT_CODE_SANDBOX_ENDPOINT ?? - 'http://127.0.0.1:2000/api/v2', + sandboxEndpoint, capabilities: { - statefulWorkspace: true, + statefulWorkspace, sandboxProfile: process.env.LIBRECHAT_CODE_SANDBOX_PROFILE ?? 'nsjail', runtimes: list(process.env.LIBRECHAT_CODE_RUNTIMES), policyDigest: createHash('sha256').update(policy).digest('hex'), diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index b363172..888da87 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -8,6 +8,8 @@ import { BridgeWorker, reconnectDelayMs } from './worker.js'; import type { BridgeAssignment } from './protocol.js'; +const incarnationId = 'incarnation-00000001'; + test('worker forwards a fenced assignment to the sandbox and settles the result', async () => { const requests: Array<{ url: string; init?: RequestInit }> = []; const fetchImpl: typeof fetch = async (input, init) => { @@ -31,7 +33,9 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' codeApiUrl: 'https://code.example/v1/', token: 'worker-secret', workerId: 'vm-1', - sandboxEndpoint: 'http://127.0.0.1:2000/api/v2/', + incarnationId, + sandboxEndpoint: + 'http://127.0.0.1:2000/sessions/{runtimeSessionId}/api/v2/', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -43,6 +47,7 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' protocolVersion: 1, assignmentId: 'assignment-1', workerId: 'vm-1', + incarnationId, generation: 3, leaseToken: 'lease-token-that-is-long-enough-for-testing', expiresAt: new Date(Date.now() + 10_000).toISOString(), @@ -56,7 +61,10 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' await worker.executeAndSettle(assignment); assert.equal(requests.length, 2); - assert.equal(requests[0].url, 'http://127.0.0.1:2000/api/v2/execute'); + assert.equal( + requests[0].url, + 'http://127.0.0.1:2000/sessions/rt-user-1/api/v2/execute', + ); assert.equal( (requests[0].init?.headers as Record)[ 'X-Runtime-Session-Id' @@ -68,11 +76,104 @@ test('worker forwards a fenced assignment to the sandbox and settles the result' protocolVersion: 1, generation: 3, leaseToken: 'lease-token-that-is-long-enough-for-testing', + incarnationId, status: 'fulfilled', result: { session_id: 'run-1', files: [] }, }); }); +test('worker aborts sandbox execution at the absolute assignment deadline', async () => { + let settlement: Record | undefined; + const fetchImpl: typeof fetch = async (input, init) => { + if (String(input).endsWith('/execute')) { + return await new Promise((_resolve, reject) => { + init?.signal?.addEventListener( + 'abort', + () => reject(new DOMException('aborted', 'AbortError')), + { once: true }, + ); + }); + } + settlement = JSON.parse(String(init?.body)) as Record; + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-deadline', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 30).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.equal(settlement?.status, 'rejected'); + assert.equal(settlement?.incarnationId, incarnationId); +}); + +test('worker refreshes its registration during a long assignment', async () => { + let registrations = 0; + const fetchImpl: typeof fetch = async (input) => { + const url = String(input); + if (url.endsWith('/workers/register')) { + registrations += 1; + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 50, + }); + } + if (url.endsWith('/execute')) { + await new Promise((resolve) => setTimeout(resolve, 90)); + return Response.json({ session_id: 'run-1', files: [] }); + } + return Response.json({ protocolVersion: 1, accepted: true }); + }; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl, + }); + await worker.register(); + await worker.executeAndSettle({ + protocolVersion: 1, + assignmentId: 'assignment-heartbeat', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 1_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + }); + + assert.ok(registrations >= 2); +}); + test('paired worker proves possession on bridge requests', async () => { const key = createBridgeIdentity(); let bridgeRequest: { url: string; init?: RequestInit } | undefined; @@ -81,6 +182,7 @@ test('paired worker proves possession on bridge requests', async () => { return Response.json({ protocolVersion: 1, workerId: 'vm-1', + incarnationId, registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }); @@ -88,6 +190,7 @@ test('paired worker proves possession on bridge requests', async () => { const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', identity: { privateKey: key.privateKey, @@ -142,6 +245,7 @@ test('paired worker rotates an expiring credential before registration', async ( return Response.json({ protocolVersion: 1, workerId: 'vm-1', + incarnationId, registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }); @@ -149,6 +253,7 @@ test('paired worker rotates an expiring credential before registration', async ( const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', identity: { privateKey: key.privateKey, @@ -176,6 +281,45 @@ test('paired worker rotates an expiring credential before registration', async ( ); }); +test('paired worker retries persistence before adopting a rotated credential', async () => { + const key = createBridgeIdentity(); + const identity = { + privateKey: key.privateKey, + credential: 'original-short-lived-credential-value', + expiresAt: new Date(Date.now() + 30_000).toISOString(), + }; + let persistenceAttempts = 0; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + identity, + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + fetchImpl: async () => + Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + credential: 'rotated-short-lived-credential-value', + expiresAt: new Date(Date.now() + 15 * 60_000).toISOString(), + }), + onIdentityChange: () => { + persistenceAttempts += 1; + if (persistenceAttempts === 1) throw new Error('disk unavailable'); + }, + }); + + await assert.rejects(worker.refreshCredential(), /disk unavailable/); + assert.equal(identity.credential, 'original-short-lived-credential-value'); + await worker.refreshCredential(); + assert.equal(identity.credential, 'rotated-short-lived-credential-value'); + assert.equal(persistenceAttempts, 2); +}); + test('paired worker refreshes before an assignment that outlives its credential', async () => { const key = createBridgeIdentity(); const requests: string[] = []; @@ -198,6 +342,7 @@ test('paired worker refreshes before an assignment that outlives its credential' const worker = new BridgeWorker({ codeApiUrl: 'https://code.example/v1', workerId: 'vm-1', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', identity: { privateKey: key.privateKey, @@ -216,6 +361,7 @@ test('paired worker refreshes before an assignment that outlives its credential' protocolVersion: 1, assignmentId: 'assignment-long', workerId: 'vm-1', + incarnationId, generation: 4, leaseToken: 'assignment-long-lease-token-value', expiresAt: new Date(Date.now() + 5 * 60_000).toISOString(), @@ -236,6 +382,7 @@ test('worker shutdown interrupts reconnect backoff', async () => { codeApiUrl: 'https://code.example/v1', token: 'worker-secret', workerId: 'vm-1', + incarnationId, sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', capabilities: { statefulWorkspace: true, diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index 4ec69a7..bfd1288 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -31,6 +31,7 @@ export interface BridgeWorkerOptions { fetchImpl?: typeof fetch; onError?: (error: unknown) => void; onIdentityChange?: (identity: BridgeWorkerIdentity) => void | Promise; + incarnationId?: string; } export interface BridgeWorkerIdentity { @@ -43,6 +44,9 @@ const DEFAULT_LEASE_WAIT_MS = 25_000; const DEFAULT_RECONNECT_DELAY_MS = 1_000; const DEFAULT_RECONNECT_MAX_DELAY_MS = 30_000; const CREDENTIAL_REFRESH_WINDOW_MS = 60_000; +const DEFAULT_REGISTRATION_TTL_MS = 60_000; +const MIN_REGISTRATION_HEARTBEAT_MS = 25; +const RUNTIME_SESSION_PLACEHOLDER = '{runtimeSessionId}'; export function reconnectDelayMs( attempt: number, @@ -82,6 +86,8 @@ export class BridgeWorker { private readonly fetchImpl: typeof fetch; private readonly codeApiUrl: string; private readonly sandboxEndpoint: string; + private readonly incarnationId: string; + private registrationTtlMs = DEFAULT_REGISTRATION_TTL_MS; constructor(private readonly options: BridgeWorkerOptions) { if (!options.token && !options.identity) { @@ -92,20 +98,30 @@ export class BridgeWorker { this.fetchImpl = options.fetchImpl ?? fetch; this.codeApiUrl = normalizedBaseUrl(options.codeApiUrl); this.sandboxEndpoint = normalizedBaseUrl(options.sandboxEndpoint); + this.incarnationId = + options.incarnationId ?? randomBytes(18).toString('base64url'); } async register( signal?: AbortSignal, ): Promise { - return this.request( + const registration = await this.request( `${this.codeApiUrl}/bridge/workers/register`, { protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: this.options.workerId, + incarnationId: this.incarnationId, capabilities: this.options.capabilities, }, signal, ); + if (registration.incarnationId !== this.incarnationId) { + throw new BridgeProtocolError( + 'Code API registered a different worker incarnation', + ); + } + this.registrationTtlMs = registration.leaseTtlMs; + return registration; } async lease(signal?: AbortSignal): Promise { @@ -114,9 +130,18 @@ export class BridgeWorker { { protocolVersion: BRIDGE_PROTOCOL_VERSION, waitMs: this.options.leaseWaitMs ?? DEFAULT_LEASE_WAIT_MS, + incarnationId: this.incarnationId, }, signal, ); + if ( + response.assignment != null && + response.assignment.incarnationId !== this.incarnationId + ) { + throw new BridgeProtocolError( + 'Code API leased an assignment for a different worker incarnation', + ); + } return response.assignment; } @@ -134,7 +159,7 @@ export class BridgeWorker { if (signal?.aborted) return; if ( error instanceof BridgeProtocolError && - (error.status === 401 || error.status === 403) + (error.status === 401 || error.status === 403 || error.status === 409) ) { throw error; } @@ -179,9 +204,14 @@ export class BridgeWorker { 'Code API returned an invalid rotated worker credential', ); } - identity.credential = credential.credential; - identity.expiresAt = credential.expiresAt; - await this.options.onIdentityChange?.(identity); + const rotatedIdentity: BridgeWorkerIdentity = { + ...identity, + credential: credential.credential, + expiresAt: credential.expiresAt, + }; + await this.options.onIdentityChange?.(rotatedIdentity); + identity.credential = rotatedIdentity.credential; + identity.expiresAt = rotatedIdentity.expiresAt; } async executeAndSettle( @@ -195,6 +225,23 @@ export class BridgeWorker { const executionController = new AbortController(); const abortExecution = (): void => executionController.abort(); signal?.addEventListener('abort', abortExecution, { once: true }); + const deadlineDelay = Math.max( + 0, + Date.parse(assignment.expiresAt) - Date.now(), + ); + const deadlineTimer = setTimeout( + () => executionController.abort(), + deadlineDelay, + ); + const heartbeatController = new AbortController(); + let heartbeatError: unknown; + const heartbeat = this.maintainRegistration( + heartbeatController.signal, + executionController, + ).catch((error) => { + heartbeatError = error; + executionController.abort(); + }); const cancellationController = new AbortController(); const cancellationWatcher = this.watchCancellation( assignment, @@ -209,16 +256,20 @@ export class BridgeWorker { ? { 'X-Runtime-Session-Id': assignment.runtimeSessionId } : {}), }; - const response = await this.fetchImpl(`${this.sandboxEndpoint}/execute`, { - method: 'POST', - headers: { - ...headers, - 'Content-Type': 'application/json', + const response = await this.fetchImpl( + `${this.sandboxEndpointFor(assignment)}/execute`, + { + method: 'POST', + headers: { + ...headers, + 'Content-Type': 'application/json', + }, + body: JSON.stringify(assignment.request.body), + signal: executionController.signal, }, - body: JSON.stringify(assignment.request.body), - signal: executionController.signal, - }); + ); const payload = (await response.json()) as object; + if (heartbeatError != null) throw heartbeatError; if (!response.ok) { throw new BridgeProtocolError( errorMessage(payload) ?? @@ -230,6 +281,7 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment.generation, leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, status: 'fulfilled', result: payload, }; @@ -238,12 +290,16 @@ export class BridgeWorker { protocolVersion: BRIDGE_PROTOCOL_VERSION, generation: assignment.generation, leaseToken: assignment.leaseToken, + incarnationId: this.incarnationId, status: 'rejected', error: error instanceof Error ? error.message : 'Sandbox execution failed', }; } + clearTimeout(deadlineTimer); + heartbeatController.abort(); + await heartbeat; cancellationController.abort(); await cancellationWatcher; signal?.removeEventListener('abort', abortExecution); @@ -254,6 +310,39 @@ export class BridgeWorker { ); } + private sandboxEndpointFor(assignment: BridgeAssignment): string { + if (assignment.runtimeSessionId == null) return this.sandboxEndpoint; + if ( + this.options.capabilities.statefulWorkspace !== true || + !this.sandboxEndpoint.includes(RUNTIME_SESSION_PLACEHOLDER) + ) { + throw new BridgeProtocolError( + 'Stateful assignments require a sandbox endpoint template containing {runtimeSessionId}', + ); + } + return this.sandboxEndpoint.replace( + RUNTIME_SESSION_PLACEHOLDER, + encodeURIComponent(assignment.runtimeSessionId), + ); + } + + private async maintainRegistration( + signal: AbortSignal, + executionController: AbortController, + ): Promise { + while (!signal.aborted && !executionController.signal.aborted) { + await abortableDelay( + Math.max( + MIN_REGISTRATION_HEARTBEAT_MS, + Math.floor(this.registrationTtlMs / 2), + ), + signal, + ); + if (signal.aborted || executionController.signal.aborted) return; + await this.register(signal); + } + } + private assignmentUrl(assignment: BridgeAssignment, action: string): string { return ( `${this.codeApiUrl}${bridgeWorkerPath(this.options.workerId)}` + @@ -272,7 +361,10 @@ export class BridgeWorker { try { const response = await this.request<{ cancelled: boolean }>( this.assignmentUrl(assignment, 'cancellation'), - { protocolVersion: BRIDGE_PROTOCOL_VERSION }, + { + protocolVersion: BRIDGE_PROTOCOL_VERSION, + incarnationId: this.incarnationId, + }, signal, ); if (response.cancelled) { diff --git a/service/src/bridge/router.test.ts b/service/src/bridge/router.test.ts index 34f7906..077dcd1 100644 --- a/service/src/bridge/router.test.ts +++ b/service/src/bridge/router.test.ts @@ -76,6 +76,7 @@ describe('paired bridge HTTP API', () => { const body = JSON.stringify({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: 'vm-1', + incarnationId: 'incarnation-00000001', capabilities: { statefulWorkspace: true, sandboxProfile: 'nsjail', @@ -110,6 +111,7 @@ describe('paired bridge HTTP API', () => { expect(registrationResponse.status).toBe(200); await expect(registrationResponse.json()).resolves.toMatchObject({ workerId: 'vm-1', + incarnationId: 'incarnation-00000001', }); const crossDeploymentRevoke = await fetch( diff --git a/service/src/bridge/router.ts b/service/src/bridge/router.ts index b91037b..8ce58bf 100644 --- a/service/src/bridge/router.ts +++ b/service/src/bridge/router.ts @@ -11,6 +11,7 @@ import { BridgePairingError, RedisBridgePairingStore } from './pairing'; import { BridgeStoreError, RedisBridgeStore } from './store'; const WORKER_ID_PATTERN = /^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$/; +const INCARNATION_ID_PATTERN = /^[A-Za-z0-9_-]{16,128}$/; const MAX_LEASE_WAIT_MS = 30_000; export type BridgeAuthMode = 'static' | 'paired'; @@ -36,6 +37,10 @@ function validWorkerId(value: string): boolean { return WORKER_ID_PATTERN.test(value); } +function validIncarnationId(value: unknown): value is string { + return typeof value === 'string' && INCARNATION_ID_PATTERN.test(value); +} + function isRecord(value: unknown): value is Record { return typeof value === 'object' && value !== null; } @@ -53,7 +58,8 @@ function isSettlement(value: unknown): value is CodeBridgeSettlement { !Number.isSafeInteger(value.generation) || value.generation < 1 || typeof value.leaseToken !== 'string' || - value.leaseToken.length < 32 + value.leaseToken.length < 32 || + !validIncarnationId(value.incarnationId) ) { return false; } @@ -271,8 +277,9 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { registration.protocolVersion !== BRIDGE_PROTOCOL_VERSION || typeof registration.workerId !== 'string' || !validWorkerId(registration.workerId) || + !validIncarnationId(registration.incarnationId) || !isRecord(registration.capabilities) || - registration.capabilities.statefulWorkspace !== true || + typeof registration.capabilities.statefulWorkspace !== 'boolean' || typeof registration.capabilities.sandboxProfile !== 'string' || registration.capabilities.sandboxProfile.trim().length === 0 || registration.capabilities.sandboxProfile.length > 128 || @@ -297,12 +304,21 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }); return; } - await options.store.register( - registration as unknown as BridgeWorkerRegistration, - ); + try { + await options.store.register( + registration as unknown as BridgeWorkerRegistration, + ); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, workerId: registration.workerId, + incarnationId: registration.incarnationId, registeredAt: new Date().toISOString(), leaseTtlMs: 60_000, }); @@ -318,6 +334,8 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { const requestedWait = Number(body.waitMs ?? 25_000); if ( !validWorkerId(workerId) || + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) || !Number.isFinite(requestedWait) || requestedWait < 0 ) { @@ -330,11 +348,20 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { }); return; } - const assignment = await options.store.lease( - workerId, - Math.min(requestedWait, MAX_LEASE_WAIT_MS), - ); - res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); + try { + const assignment = await options.store.lease( + workerId, + body.incarnationId, + Math.min(requestedWait, MAX_LEASE_WAIT_MS), + ); + res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, assignment }); + } catch (error) { + if (error instanceof BridgeStoreError) { + sendStoreError(error, res); + return; + } + throw error; + } }, ); @@ -371,8 +398,17 @@ export function createBridgeRouter(options: BridgeRouterOptions): Router { '/workers/:workerId/assignments/:assignmentId/cancellation', workerAuth, async (req: Request, res: Response) => { + const body = isRecord(req.body) ? req.body : {}; + if ( + body.protocolVersion !== BRIDGE_PROTOCOL_VERSION || + !validIncarnationId(body.incarnationId) + ) { + res.status(400).json({ error: 'Invalid bridge cancellation request' }); + return; + } const cancelled = await options.store.cancelled( req.params.workerId, + body.incarnationId, req.params.assignmentId, ); res.json({ protocolVersion: BRIDGE_PROTOCOL_VERSION, cancelled }); diff --git a/service/src/lifecycle.ts b/service/src/lifecycle.ts index 068eadc..596193f 100644 --- a/service/src/lifecycle.ts +++ b/service/src/lifecycle.ts @@ -153,6 +153,7 @@ async function gracefulStartup(): Promise { validateWorkerHardenedConfig(); validateExecutionProfilePolicy(); validateSandboxBackendPolicy(); + validateApiBridgePolicy(); await validateLifecycleAuthConfig(); configureProfileMetrics(); From 86640e9644524bc1a6f72f42665c78d55c610b70 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sun, 30 Aug 2026 21:23:56 -0400 Subject: [PATCH 6/6] fix: distinguish assignment settlement conflicts --- packages/code/src/protocol.ts | 1 + packages/code/src/worker.test.ts | 66 ++++++++++++++++++++++++++++++++ packages/code/src/worker.ts | 8 +++- 3 files changed, 74 insertions(+), 1 deletion(-) diff --git a/packages/code/src/protocol.ts b/packages/code/src/protocol.ts index 34358cb..e518fe7 100644 --- a/packages/code/src/protocol.ts +++ b/packages/code/src/protocol.ts @@ -95,6 +95,7 @@ export class BridgeProtocolError extends Error { constructor( message: string, public readonly status?: number, + public readonly code?: string, ) { super(message); this.name = 'BridgeProtocolError'; diff --git a/packages/code/src/worker.test.ts b/packages/code/src/worker.test.ts index 888da87..75eb78c 100644 --- a/packages/code/src/worker.test.ts +++ b/packages/code/src/worker.test.ts @@ -126,6 +126,72 @@ test('worker aborts sandbox execution at the absolute assignment deadline', asyn assert.equal(settlement?.incarnationId, incarnationId); }); +test('worker continues after an expired assignment settlement conflict', async () => { + const controller = new AbortController(); + let registrations = 0; + let leases = 0; + const worker = new BridgeWorker({ + codeApiUrl: 'https://code.example/v1', + token: 'worker-secret', + workerId: 'vm-1', + incarnationId, + sandboxEndpoint: 'http://127.0.0.1:2000/api/v2', + capabilities: { + statefulWorkspace: false, + sandboxProfile: 'nsjail', + runtimes: ['bash'], + }, + reconnectDelayMs: 0, + reconnectMaxDelayMs: 0, + fetchImpl: async (input) => { + const url = String(input); + if (url.endsWith('/workers/register')) { + registrations += 1; + if (registrations === 2) controller.abort(); + return Response.json({ + protocolVersion: 1, + workerId: 'vm-1', + incarnationId, + registeredAt: new Date().toISOString(), + leaseTtlMs: 60_000, + }); + } + if (url.endsWith('/lease')) { + leases += 1; + return Response.json({ + protocolVersion: 1, + assignment: leases === 1 + ? { + protocolVersion: 1, + assignmentId: 'assignment-expired', + workerId: 'vm-1', + incarnationId, + generation: 1, + leaseToken: 'lease-token-that-is-long-enough-for-testing', + expiresAt: new Date(Date.now() + 10_000).toISOString(), + request: { body: { language: 'bash' }, headers: {} }, + } + : undefined, + }); + } + if (url.endsWith('/execute')) { + return Response.json({ session_id: 'run-1', files: [] }); + } + if (url.endsWith('/settle')) { + return Response.json( + { error: 'Bridge assignment has expired', code: 'ASSIGNMENT_EXPIRED' }, + { status: 409 }, + ); + } + return Response.json({ cancelled: false }); + }, + }); + + await worker.run(controller.signal); + + assert.equal(registrations, 2); +}); + test('worker refreshes its registration during a long assignment', async () => { let registrations = 0; const fetchImpl: typeof fetch = async (input) => { diff --git a/packages/code/src/worker.ts b/packages/code/src/worker.ts index bfd1288..5e528d6 100644 --- a/packages/code/src/worker.ts +++ b/packages/code/src/worker.ts @@ -159,7 +159,10 @@ export class BridgeWorker { if (signal?.aborted) return; if ( error instanceof BridgeProtocolError && - (error.status === 401 || error.status === 403 || error.status === 409) + (error.status === 401 || + error.status === 403 || + error.code === 'WORKER_FENCED' || + error.code === 'WORKER_QUARANTINED') ) { throw error; } @@ -402,6 +405,9 @@ export class BridgeWorker { errorMessage(payload) ?? `Bridge request failed with HTTP ${response.status}`, response.status, + 'code' in payload && typeof payload.code === 'string' + ? payload.code + : undefined, ); } return payload as T;