diff --git a/src/index.ts b/src/index.ts index 30745c3..6bcb68d 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,9 +1,6 @@ -import { exec } from "node:child_process"; +import { spawn } from "node:child_process"; import { mkdir, writeFile } from "node:fs/promises"; import { join } from "node:path"; -import { promisify } from "node:util"; - -const execAsync = promisify(exec); export interface WorkerPayload { taskType: "agent.run" | "command.run" | string; @@ -115,38 +112,143 @@ async function runCommand(payload: WorkerPayload, options: WorkerOptions): Promi if (!command) { return { ok: false, events: [], artifacts: [], error: "command.run requires input.command." }; } - const allowlist = options.commandAllowlist ?? commandAllowlistFromEnv(); - if (allowlist.length > 0 && !allowlist.some((allowed) => command === allowed || command.startsWith(`${allowed} `))) { + + let parsedCommand: ParsedCommand; + try { + parsedCommand = parseCommand(command); + } catch (error) { + const message = error instanceof Error ? error.message : String(error); return { ok: false, artifacts: [], - events: [{ type: "task.log", message: `Command is not allowed: ${command}`, payload: { stream: "stderr" } }], - error: "Command rejected by AGENTDISPATCH_COMMAND_ALLOWLIST." + events: [{ type: "task.log", message, payload: { stream: "stderr" } }], + error: message }; } - try { - const result = await execAsync(command, { timeout: Number(payload.input.timeoutSeconds ?? 900) * 1000 }); - const artifacts = await writeArtifacts(options.artifactDir, { command, stdout: result.stdout, stderr: result.stderr, exitCode: 0 }); + + const allowlist = options.commandAllowlist ?? commandAllowlistFromEnv(); + if (allowlist.length > 0 && !allowlist.includes(parsedCommand.executable)) { return { - ok: true, - output: result.stdout, - artifacts, - events: [ - { type: "task.heartbeat", message: "AgentDispatch worker heartbeat.", payload: { status: "running" } }, - { type: "task.log", message: result.stdout, payload: { stream: "stdout" } }, - { type: "task.log", message: result.stderr, payload: { stream: "stderr" } }, - { type: "task.result", payload: { exitCode: 0 } } - ] + ok: false, + artifacts: [], + events: [{ type: "task.log", message: `Command is not allowed: ${parsedCommand.executable}`, payload: { stream: "stderr" } }], + error: "Command rejected by AGENTDISPATCH_COMMAND_ALLOWLIST." }; - } catch (error) { - const message = error instanceof Error ? error.message : String(error); + } + + const result = await spawnCommand(parsedCommand, Number(payload.input.timeoutSeconds ?? 900) * 1000); + if (result.exitCode !== 0) { + const message = result.error ?? `Command exited with code ${result.exitCode}.`; return { ok: false, artifacts: [], - events: [{ type: "task.log", message, payload: { stream: "stderr" } }], + events: [{ type: "task.log", message: result.stderr || message, payload: { stream: "stderr", exitCode: result.exitCode } }], error: message }; } + + const artifacts = await writeArtifacts(options.artifactDir, { command, stdout: result.stdout, stderr: result.stderr, exitCode: 0 }); + return { + ok: true, + output: result.stdout, + artifacts, + events: [ + { type: "task.heartbeat", message: "AgentDispatch worker heartbeat.", payload: { status: "running" } }, + { type: "task.log", message: result.stdout, payload: { stream: "stdout" } }, + { type: "task.log", message: result.stderr, payload: { stream: "stderr" } }, + { type: "task.result", payload: { exitCode: 0 } } + ] + }; +} + +interface ParsedCommand { + executable: string; + args: string[]; +} + +interface SpawnCommandResult { + stdout: string; + stderr: string; + exitCode: number | null; + error?: string; +} + +function parseCommand(command: string): ParsedCommand { + const args: string[] = []; + let current = ""; + let quote: "'" | '"' | undefined; + let escaped = false; + + for (const char of command.trim()) { + if (escaped) { + current += char; + escaped = false; + continue; + } + if (char === "\\") { + escaped = true; + continue; + } + if (quote) { + if (char === quote) { + quote = undefined; + } else { + current += char; + } + continue; + } + if (char === "'" || char === '"') { + quote = char; + continue; + } + if (/\s/.test(char)) { + if (current) { + args.push(current); + current = ""; + } + continue; + } + current += char; + } + + if (escaped) current += "\\"; + if (quote) throw new Error("command.run input.command contains an unterminated quote."); + if (current) args.push(current); + const [executable, ...rest] = args; + if (!executable) throw new Error("command.run requires input.command."); + return { executable, args: rest }; +} + +function spawnCommand(command: ParsedCommand, timeoutMs: number): Promise { + return new Promise((resolve) => { + const child = spawn(command.executable, command.args, { shell: false }); + let stdout = ""; + let stderr = ""; + let settled = false; + const timeout = setTimeout(() => { + if (settled) return; + child.kill("SIGTERM"); + settled = true; + resolve({ stdout, stderr, exitCode: null, error: `Command timed out after ${timeoutMs}ms.` }); + }, timeoutMs); + + child.stdout.setEncoding("utf8"); + child.stderr.setEncoding("utf8"); + child.stdout.on("data", (chunk) => { stdout += chunk; }); + child.stderr.on("data", (chunk) => { stderr += chunk; }); + child.on("error", (error) => { + if (settled) return; + clearTimeout(timeout); + settled = true; + resolve({ stdout, stderr, exitCode: null, error: error.message }); + }); + child.on("close", (code) => { + if (settled) return; + clearTimeout(timeout); + settled = true; + resolve({ stdout, stderr, exitCode: code }); + }); + }); } async function writeArtifacts(artifactDir = process.env.AGENTDISPATCH_ARTIFACT_DIR ?? "/tmp/agentdispatch-artifacts", payload: Record): Promise { diff --git a/src/server.ts b/src/server.ts index 60cf7d2..b0db527 100644 --- a/src/server.ts +++ b/src/server.ts @@ -35,7 +35,7 @@ export async function handleRequest(request: IncomingMessage, response: ServerRe artifactDir: process.env.AGENTDISPATCH_ARTIFACT_DIR, commandAllowlist: process.env.AGENTDISPATCH_COMMAND_ALLOWLIST?.split(",").map((value) => value.trim()).filter(Boolean) }); - response.writeHead(result.ok ? 200 : 500, { "content-type": "application/json" }); + response.writeHead(200, { "content-type": "application/json" }); response.end(JSON.stringify(result)); } catch (error) { response.writeHead(400, { "content-type": "application/json" }); diff --git a/test/server.test.ts b/test/server.test.ts index 99cb51d..bd4e91d 100644 --- a/test/server.test.ts +++ b/test/server.test.ts @@ -33,6 +33,16 @@ describe("AgentCore worker HTTP server", () => { await expect(response.json()).resolves.toMatchObject({ ok: true, output: "Accepted instruction: work" }); }); + it("returns structured application failures with HTTP 200", async () => { + const response = await fetch(`${baseUrl}/invocations`, { + method: "POST", + headers: { "content-type": "application/json" }, + body: JSON.stringify({ taskType: "agent.run", input: { instruction: "work", framework: "missing-framework" } }) + }); + expect(response.status).toBe(200); + await expect(response.json()).resolves.toMatchObject({ ok: false, error: "Unsupported agent framework: missing-framework" }); + }); + it("rejects unsupported POST paths", async () => { const response = await fetch(`${baseUrl}/wrong`, { method: "POST", diff --git a/test/worker.test.ts b/test/worker.test.ts index 8c712cd..9ec593d 100644 --- a/test/worker.test.ts +++ b/test/worker.test.ts @@ -61,6 +61,24 @@ describe("worker contract", () => { expect(result.error).toContain("ALLOWLIST"); }); + it("does not execute shell metacharacters through allowed commands", async () => { + const result = await runAgentDispatchWorkerTask( + { taskType: "command.run", input: { command: "echo ok; uname -a", timeoutSeconds: 2 } }, + { commandAllowlist: ["echo"] } + ); + expect(result.ok).toBe(true); + expect(result.output).toBe("ok; uname -a\n"); + }); + + it("supports quoted command arguments without shell execution", async () => { + const result = await runAgentDispatchWorkerTask( + { taskType: "command.run", input: { command: "echo \"hello world\"", timeoutSeconds: 2 } }, + { commandAllowlist: ["echo"] } + ); + expect(result.ok).toBe(true); + expect(result.output).toBe("hello world\n"); + }); + it("runs allowed commands and writes artifacts", async () => { const artifactDir = await mkdtemp(join(tmpdir(), "agentdispatch-worker-")); const result = await runAgentDispatchWorkerTask(