From 52db0b2e7f3382da94d0286433375cdc9e0b51ee Mon Sep 17 00:00:00 2001 From: Ben Taylor Date: Sun, 5 Jan 2025 21:29:35 +1100 Subject: [PATCH] Add Timeout to Sandbox (#307) * WIP: add sandboxing features * WIP: sandboxing types * Fix up support for timeouts via workers * Bump version and include docs --- packages/runtime/lib/types.ts | 10 +- packages/wasi/lib/worker/wasi-host.ts | 3 + sandbox/README.md | 20 ++- sandbox/cli/bin/main.ts | 28 ++-- sandbox/cli/deno.json | 2 +- sandbox/cli/deno.lock | 8 +- sandbox/cli/lib/host.ts | 128 +++++++++++++++++ sandbox/cli/lib/runtime.ts | 184 +++++++++++++++++++------ sandbox/cli/lib/types.ts | 10 +- sandbox/cli/lib/worker.ts | 191 ++++++++++++++++++++++++++ sandbox/cli/tests/sandbox.test.ts | 5 + sandbox/pyproject.toml | 2 +- sandbox/runno/__init__.py | 3 +- sandbox/runno/main.py | 26 +++- sandbox/runno/types.py | 28 +++- sandbox/tests/test_code.py | 6 + 16 files changed, 585 insertions(+), 69 deletions(-) create mode 100644 sandbox/cli/lib/host.ts create mode 100644 sandbox/cli/lib/worker.ts diff --git a/packages/runtime/lib/types.ts b/packages/runtime/lib/types.ts index ff7d6fec..04d268af 100644 --- a/packages/runtime/lib/types.ts +++ b/packages/runtime/lib/types.ts @@ -40,7 +40,15 @@ export type TerminatedResult = { resultType: "terminated"; }; -export type RunResult = CompleteResult | CrashResult | TerminatedResult; +export type TimeoutResult = { + resultType: "timeout"; +}; + +export type RunResult = + | CompleteResult + | CrashResult + | TerminatedResult + | TimeoutResult; export type RuntimeMethods = { showControls: () => void; diff --git a/packages/wasi/lib/worker/wasi-host.ts b/packages/wasi/lib/worker/wasi-host.ts index 06e17887..12995d27 100644 --- a/packages/wasi/lib/worker/wasi-host.ts +++ b/packages/wasi/lib/worker/wasi-host.ts @@ -77,6 +77,9 @@ export class WASIWorkerHost { fs: this.context.fs, isTTY: this.context.isTTY, }); + }).then((result) => { + this.worker?.terminate(); + return result; }); return this.result; diff --git a/sandbox/README.md b/sandbox/README.md index 1d88e6c8..6a178dd7 100644 --- a/sandbox/README.md +++ b/sandbox/README.md @@ -27,7 +27,8 @@ If you want to process files within your sandbox, you can include them in the file system by using the `run_fs` method. ```python -from runno import run_code, WASIFS, StringFile +from runno import run_fs, WASIFS, StringFile, WASITimestamps + fs: WASIFS = { "/program.py": StringFile( path="/program.py", @@ -63,3 +64,20 @@ await run_fs(runtime, "/program", fs) You can use the same technique to include pure python packages as dependencies. The interface for this is not super nice right now, but it's on the way. + +## Limiting Execution Time + +You can limit how much time is allocated for execution using an optional +`timeout` kwarg (measured in seconds). Like: + +``` +from runno import run_code + +code = "while True: pass" +result = await run_code("python", code, timeout=5) + +if result.result_type == "timeout": + print("Timed out.") +else: + print("Wow, it ran forever.") +``` diff --git a/sandbox/cli/bin/main.ts b/sandbox/cli/bin/main.ts index 41f75525..b9d4873d 100644 --- a/sandbox/cli/bin/main.ts +++ b/sandbox/cli/bin/main.ts @@ -10,6 +10,7 @@ import { CrashResult, Runtime, TerminatedResult, + TimeoutResult, } from "../lib/types.ts"; import { fetchWASIFS } from "../lib/main.ts"; import { extractTarGz } from "../lib/tar.ts"; @@ -32,14 +33,18 @@ const command = new Command() .description( `A CLI for running code in a sandbox environment, powered by Runno & WASI. Supports python, ruby, quickjs, php-cgi, clang, and clangpp. -Entry name is the name of the entrypoint in the base filesystem. +Entry path is the path of the entrypoint in the base filesystem. ` ) - .arguments(" ") + .arguments(" [entry-path:string]") .option( "-f, --filesystem ", "A tgz file to use as the base filesystem" ) + .option( + "-t --timeout ", + "The maximum amount of time to allow the code to run (in seconds)" + ) .option( "--filesystem-stdin", "Read the base filesystem from stdin as a tgz file (useful for piping)" @@ -56,17 +61,17 @@ Entry name is the name of the entrypoint in the base filesystem. filesystemStdin?: true; entryStdin?: true; json?: true; + timeout?: number; }, - ...args: [string, string] + ...args: [string, string | undefined] ) => { - const [runtimeString, entry] = args; + let [runtimeString, entry] = args; if (!isRuntime(runtimeString)) { throw new Error(`Unsupported runtime: ${runtimeString}`); } const runtime: Runtime = runtimeString; - // TODO: Read the entry file from stdin - + entry = entry ?? "/program"; const entryPath = entry.startsWith("/") ? entry : `/${entry}`; let fs: WASIFS = {}; @@ -95,7 +100,9 @@ Entry name is the name of the entrypoint in the base filesystem. }; } - const result = await runFS(runtime, entryPath, fs); + const result = await runFS(runtime, entryPath, fs, { + timeout: options.timeout, + }); if (options.json) { let jsonResult: JSONResult; if (result.resultType === "complete") { @@ -115,6 +122,10 @@ Entry name is the name of the entrypoint in the base filesystem. console.error(result.stderr); console.log(result.stdout); break; + case "timeout": + console.error("Timeout"); + Deno.exit(1); + break; case "crash": console.error(result.error); Deno.exit(1); @@ -173,4 +184,5 @@ export type StringWASIFile = { export type JSONResult = | (Omit & { fs: Base64WASIFS }) | CrashResult - | TerminatedResult; + | TerminatedResult + | TimeoutResult; diff --git a/sandbox/cli/deno.json b/sandbox/cli/deno.json index 779db6f9..a5aa65e4 100644 --- a/sandbox/cli/deno.json +++ b/sandbox/cli/deno.json @@ -17,7 +17,7 @@ "test": "deno test --allow-read", "cli": "deno --allow-net --allow-read src/main.ts", "bootstrap": "cp -R ../../langs .", - "compile": "deno compile --frozen --include langs --output runno bin/main.ts" + "compile": "deno compile --frozen --include langs --include lib/worker.ts --output runno bin/main.ts" }, "lock": { "path": "./deno.lock", diff --git a/sandbox/cli/deno.lock b/sandbox/cli/deno.lock index 677b65eb..2ffdcdd3 100644 --- a/sandbox/cli/deno.lock +++ b/sandbox/cli/deno.lock @@ -10,8 +10,8 @@ "jsr:@std/fmt@~1.0.2": "1.0.3", "jsr:@std/internal@^1.0.5": "1.0.5", "jsr:@std/io@0.225": "0.225.0", - "jsr:@std/streams@^1.0.7": "1.0.8", - "jsr:@std/tar@~0.1.3": "0.1.3", + "jsr:@std/streams@^1.0.8": "1.0.8", + "jsr:@std/tar@~0.1.3": "0.1.4", "jsr:@std/text@~1.0.7": "1.0.8", "npm:@runno/wasi@0.7": "0.7.0" }, @@ -65,8 +65,8 @@ "@std/streams@1.0.8": { "integrity": "b41332d93d2cf6a82fe4ac2153b930adf1a859392931e2a19d9fabfb6f154fb3" }, - "@std/tar@0.1.3": { - "integrity": "531270fc707b37ab9b5f051aa4943e7b16b86905e0398a4ebe062983b0c93115", + "@std/tar@0.1.4": { + "integrity": "1bc1f1f9bfd557e849b31d6521348fdf5848886d87c851f1f0f992d002fe0ff5", "dependencies": [ "jsr:@std/streams" ] diff --git a/sandbox/cli/lib/host.ts b/sandbox/cli/lib/host.ts new file mode 100644 index 00000000..0e032b2c --- /dev/null +++ b/sandbox/cli/lib/host.ts @@ -0,0 +1,128 @@ +import type { WASIExecutionResult, WASIContextOptions } from "@runno/wasi"; +import type { HostMessage, WorkerMessage } from "./worker.ts"; + +function sendMessage(worker: Worker, message: WorkerMessage) { + worker.postMessage(message); +} + +type WASIWorkerHostContext = Partial>; + +export class WASIWorkerHostKilledError extends Error {} + +export class WASIWorkerHost { + binaryURL: string; + + // 8kb should be big enough + stdinBuffer: SharedArrayBuffer = new SharedArrayBuffer(8 * 1024); + + context: WASIWorkerHostContext; + + result?: Promise; + worker?: Worker; + reject?: (reason?: unknown) => void; + + constructor(binaryURL: string, context: WASIWorkerHostContext) { + this.binaryURL = binaryURL; + this.context = context; + } + + start() { + if (this.result) { + throw new Error("WASIWorker Host can only be started once"); + } + + this.result = new Promise((resolve, reject) => { + this.reject = reject; + this.worker = new Worker(new URL("./worker.ts", import.meta.url), { + type: "module", + }); + + this.worker.addEventListener("message", (messageEvent) => { + const message: HostMessage = messageEvent.data; + switch (message.type) { + case "stdout": + this.context.stdout?.(message.text); + break; + case "stderr": + this.context.stderr?.(message.text); + break; + case "debug": + this.context.debug?.( + message.name, + message.args, + message.ret, + message.data + ); + break; + case "result": + resolve(message.result); + break; + case "crash": + reject(message.error); + break; + } + }); + + sendMessage(this.worker, { + target: "client", + type: "start", + binaryURL: this.binaryURL, + stdinBuffer: this.stdinBuffer, + + // Unfortunately can't just splat these because it includes types + // that can't be sent as a message. + args: this.context.args, + env: this.context.env, + fs: this.context.fs, + isTTY: this.context.isTTY, + }); + }).then((result) => { + this.worker?.terminate(); + return result; + }); + + return this.result; + } + + kill() { + if (!this.worker) { + throw new Error("WASIWorker has not started"); + } + this.worker.terminate(); + this.reject?.(new WASIWorkerHostKilledError("WASI Worker was killed")); + } + + async pushStdin(data: string) { + const view = new DataView(this.stdinBuffer); + + // Wait until the stdinbuffer is fully consumed at the other end + // before pushing more data on. + + // first four bytes (Int32) is the length of the text + while (view.getInt32(0) !== 0) { + // TODO: Switch to Atomics.waitAsync when supported by firefox + await new Promise((resolve) => setTimeout(resolve, 0)); + } + + // Store the encoded text offset by 4 bytes + const encodedText = new TextEncoder().encode(data); + const buffer = new Uint8Array(this.stdinBuffer, 4); + buffer.set(encodedText); + + // Store how long the text is in the first 4 bytes + view.setInt32(0, encodedText.byteLength); + Atomics.notify(new Int32Array(this.stdinBuffer), 0); + } + + async pushEOF() { + const view = new DataView(this.stdinBuffer); + + // TODO: Switch to Atomics.waitAsync when supported by firefox + while (view.getInt32(0) !== 0) { + await new Promise((resolve) => setTimeout(resolve, 0)); + } + + view.setInt32(0, -1); + Atomics.notify(new Int32Array(this.stdinBuffer), 0); + } +} diff --git a/sandbox/cli/lib/runtime.ts b/sandbox/cli/lib/runtime.ts index 144eabaa..4dc224fa 100644 --- a/sandbox/cli/lib/runtime.ts +++ b/sandbox/cli/lib/runtime.ts @@ -1,4 +1,8 @@ -import { WASIFS, WASI } from "@runno/wasi"; +import type { + WASIFS, + WASIExecutionResult, + WASIContextOptions, +} from "@runno/wasi"; import { Command, commandsForRuntime, @@ -6,11 +10,15 @@ import { } from "./commands.ts"; import { fetchWASIFS, makeBlobFromPath, makeRunnoError } from "./helpers.ts"; import { CompleteResult, RunResult, Runtime } from "./types.ts"; +import { WASIWorkerHost } from "./host.ts"; -export async function runCode( +export function runCode( runtime: Runtime, code: string, - stdin?: string + options: { + stdin?: string; + timeout?: number; + } = {} ): Promise { const fs: WASIFS = { "/program": { @@ -24,14 +32,18 @@ export async function runCode( }, }, }; - return runFS(runtime, "/program", fs, stdin); + + return runFS(runtime, "/program", fs, options); } export async function runFS( runtime: Runtime, entryPath: string, fs: WASIFS, - stdin?: string + options: { + stdin?: string; + timeout?: number; + } = {} ): Promise { const commands = commandsForRuntime(runtime, entryPath); @@ -40,6 +52,12 @@ export async function runFS( prepare = await headlessPrepareFS(commands.prepare ?? [], fs); fs = prepare.fs; } catch (e) { + if (e instanceof TimeoutError) { + return { + resultType: "timeout", + }; + } + return { resultType: "crash", error: makeRunnoError(e), @@ -61,31 +79,44 @@ export async function runFS( } } - let stdinBytes = new TextEncoder().encode(stdin ?? ""); + try { + const resultWithLimits = await _startWASIWithLimits( + binaryPath, + { + args: [run.binaryName, ...(run.args ?? [])], + env: run.env, + fs, + stdin: options.stdin, + stdout: (out: string) => { + prepare.stdout += out; + prepare.tty += out; + }, + stderr: (err: string) => { + prepare.stderr += err; + prepare.tty += err; + }, + }, + { + timeout: options.timeout, + } + ); - const result = await WASI.start(fetch(makeBlobFromPath(binaryPath)), { - args: [run.binaryName, ...(run.args ?? [])], - env: run.env, - fs, - stdin: (maxByteLength: number) => { - const chunk = stdinBytes.slice(0, maxByteLength); - stdinBytes = stdinBytes.slice(maxByteLength); - return new TextDecoder().decode(chunk); - }, - stdout: (out: string) => { - prepare.stdout += out; - prepare.tty += out; - }, - stderr: (err: string) => { - prepare.stderr += err; - prepare.tty += err; - }, - }); + prepare.fs = { ...fs, ...resultWithLimits.result.fs }; + prepare.exitCode = resultWithLimits.result.exitCode; - prepare.fs = { ...fs, ...result.fs }; - prepare.exitCode = result.exitCode; + return prepare; + } catch (e) { + if (e instanceof TimeoutError) { + return { + resultType: "timeout", + }; + } - return prepare; + return { + resultType: "crash", + error: makeRunnoError(e), + }; + } } type PrepareErrorData = { @@ -107,9 +138,10 @@ export class PrepareError extends Error { export async function headlessPrepareFS( prepareCommands: Command[], - fs: WASIFS + fs: WASIFS, + limits: Partial = {} ) { - let prepare: CompleteResult = { + const prepare: CompleteResult = { resultType: "complete", stdin: "", stdout: "", @@ -127,24 +159,31 @@ export async function headlessPrepareFS( prepare.fs = { ...prepare.fs, ...baseFS }; } - const result = await WASI.start(fetch(makeBlobFromPath(binaryPath)), { - args: [command.binaryName, ...(command.args ?? [])], - env: command.env, - fs: prepare.fs, - stdout: (out: string) => { - prepare.stdout += out; - prepare.tty += out; - }, - stderr: (err: string) => { - prepare.stderr += err; - prepare.tty += err; + const resultWithLimits = await _startWASIWithLimits( + binaryPath, + { + args: [command.binaryName, ...(command.args ?? [])], + env: command.env, + fs: prepare.fs, + stdout: (out: string) => { + prepare.stdout += out; + prepare.tty += out; + }, + stderr: (err: string) => { + prepare.stderr += err; + prepare.tty += err; + }, }, - }); + limits + ); - prepare.fs = result.fs; - prepare.exitCode = result.exitCode; + prepare.fs = resultWithLimits.result.fs; + prepare.exitCode = resultWithLimits.result.exitCode; - if (result.exitCode !== 0) { + // Consume the timeout + limits.timeout = resultWithLimits.remainingLimits.timeout; + + if (resultWithLimits.result.exitCode !== 0) { // If a prepare step fails then we stop. throw new PrepareError( "Prepare step returned a non-zero exit code", @@ -155,3 +194,60 @@ export async function headlessPrepareFS( return prepare; } + +type Limits = { + timeout: number; +}; + +type WASIExecutionResultWithLimits = { + result: WASIExecutionResult; + remainingLimits: Partial; +}; + +type Context = Omit & { stdin: string }; + +async function _startWASIWithLimits( + binaryPath: string, + context: Partial, + limits: Partial +): Promise { + const startTime = performance.now(); + const workerHost = new WASIWorkerHost(makeBlobFromPath(binaryPath), context); + + if (context.stdin) { + workerHost.pushStdin(context.stdin); + } + + let result: WASIExecutionResult | "timeout"; + if (limits.timeout) { + const timeoutPromise: Promise<"timeout"> = new Promise((resolve) => + setTimeout(() => resolve("timeout"), limits.timeout! * 1000) + ); + + result = await Promise.race([workerHost.start(), timeoutPromise]); + } else { + result = await workerHost.start(); + } + + if (result === "timeout") { + workerHost.kill(); + throw new TimeoutError(); + } + + const endTime = performance.now(); + + return { + result, + remainingLimits: { + timeout: limits.timeout + ? limits.timeout - (endTime - startTime) / 1000 + : undefined, + }, + }; +} + +class TimeoutError extends Error { + constructor() { + super("Execution timed out"); + } +} diff --git a/sandbox/cli/lib/types.ts b/sandbox/cli/lib/types.ts index ff7d6fec..04d268af 100644 --- a/sandbox/cli/lib/types.ts +++ b/sandbox/cli/lib/types.ts @@ -40,7 +40,15 @@ export type TerminatedResult = { resultType: "terminated"; }; -export type RunResult = CompleteResult | CrashResult | TerminatedResult; +export type TimeoutResult = { + resultType: "timeout"; +}; + +export type RunResult = + | CompleteResult + | CrashResult + | TerminatedResult + | TimeoutResult; export type RuntimeMethods = { showControls: () => void; diff --git a/sandbox/cli/lib/worker.ts b/sandbox/cli/lib/worker.ts new file mode 100644 index 00000000..72a136b5 --- /dev/null +++ b/sandbox/cli/lib/worker.ts @@ -0,0 +1,191 @@ +import { + WASI, + WASIContextOptions, + WASIContext, + WASIExecutionResult, +} from "@runno/wasi"; + +type WorkerWASIContext = Partial< + Omit +>; + +type StartWorkerMessage = { + target: "client"; + type: "start"; + binaryURL: string; + stdinBuffer: SharedArrayBuffer; +} & WorkerWASIContext; + +export type WorkerMessage = StartWorkerMessage; + +type StdoutHostMessage = { + target: "host"; + type: "stdout"; + text: string; +}; + +type StderrHostMessage = { + target: "host"; + type: "stderr"; + text: string; +}; + +type DebugHostMessage = { + target: "host"; + type: "debug"; + name: string; + args: string[]; + ret: number; + data: { [key: string]: any }[]; +}; + +type ResultHostMessage = { + target: "host"; + type: "result"; + result: WASIExecutionResult; +}; + +type CrashHostMessage = { + target: "host"; + type: "crash"; + error: { + message: string; + type: string; + }; +}; + +export type HostMessage = + | StdoutHostMessage + | StderrHostMessage + | DebugHostMessage + | ResultHostMessage + | CrashHostMessage; + +(globalThis as any).onmessage = async (ev: MessageEvent) => { + const data = ev.data as WorkerMessage; + + switch (data.type) { + case "start": + try { + const result = await start(data.binaryURL, data.stdinBuffer, data); + sendMessage({ + target: "host", + type: "result", + result, + }); + } catch (e) { + let error; + if (e instanceof Error) { + error = { + message: e.message, + type: e.constructor.name, + }; + } else { + error = { + message: `unknown error - ${e}`, + type: "Unknown", + }; + } + sendMessage({ + target: "host", + type: "crash", + error, + }); + } + + break; + } +}; + +function sendMessage(message: HostMessage) { + (globalThis as any).postMessage(message); +} + +function start( + binaryURL: string, + stdinBuffer: SharedArrayBuffer, + context: WorkerWASIContext +) { + return WASI.start( + fetch(binaryURL), + new WASIContext({ + ...context, + stdout: sendStdout, + stderr: sendStderr, + stdin: (maxByteLength) => getStdin(maxByteLength, stdinBuffer), + debug: sendDebug, + }) + ); +} + +function sendStdout(out: string) { + sendMessage({ + target: "host", + type: "stdout", + text: out, + }); +} + +function sendStderr(err: string) { + sendMessage({ + target: "host", + type: "stderr", + text: err, + }); +} + +function sendDebug( + name: string, + args: string[], + ret: number, + data: { [key: string]: any }[] +) { + // this debug data comes through as part of a message + // we need to make sure it can be encoded by sendMessage + data = JSON.parse(JSON.stringify(data)); + sendMessage({ + target: "host", + type: "debug", + name, + args, + ret, + data, + }); + + // TODO: debugging WASI supports substituting a return value + // but it's hard to do async, so lets just always return + // the same value + return ret; +} + +function getStdin( + maxByteLength: number, + stdinBuffer: SharedArrayBuffer +): string | null { + // Wait until the integer at the start of the buffer has a length in it + Atomics.wait(new Int32Array(stdinBuffer), 0, 0); + + // First four bytes are a Int32 of how many bytes are in the buffer + const view = new DataView(stdinBuffer); + const numBytes = view.getInt32(0); + if (numBytes < 0) { + view.setInt32(0, 0); + return null; + } + + const buffer = new Uint8Array(stdinBuffer, 4, numBytes); + + // Decode the buffer into text, but only as much as was asked for + const returnValue = new TextDecoder().decode(buffer.slice(0, maxByteLength)); + + // Rewrite the buffer with the remaining bytes + const remaining = buffer.slice(maxByteLength, buffer.length); + view.setInt32(0, remaining.byteLength); + buffer.set(remaining); + + return returnValue; +} + +export default class WASIWorker { + // This is a decoy class to make the Deno type system happy +} diff --git a/sandbox/cli/tests/sandbox.test.ts b/sandbox/cli/tests/sandbox.test.ts index a0de5641..cb1fc821 100644 --- a/sandbox/cli/tests/sandbox.test.ts +++ b/sandbox/cli/tests/sandbox.test.ts @@ -77,3 +77,8 @@ int main() { assertEquals(result.stderr, ""); assertEquals(result.stdout, "Hello, World!\n"); }); + +Deno.test("a simple timeout example", async () => { + const result = await runCode("python", `while True:pass`, { timeout: 0.5 }); + assertEquals(result.resultType, "timeout"); +}); diff --git a/sandbox/pyproject.toml b/sandbox/pyproject.toml index 537f3b79..b333261b 100644 --- a/sandbox/pyproject.toml +++ b/sandbox/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "runno" -version = "0.1.2" +version = "0.2.0" description = "Run untrusted code inside the Runno WebAssembly sandbox." authors = [{ name = "Ben Taylor", email = "me@taybenlor.com" }] maintainers = [{ name = "Ben Taylor", email = "me@taybenlor.com" }] diff --git a/sandbox/runno/__init__.py b/sandbox/runno/__init__.py index 1aca12d8..315a85f4 100644 --- a/sandbox/runno/__init__.py +++ b/sandbox/runno/__init__.py @@ -9,6 +9,7 @@ StringFile, BinaryFile, WASIFS, + WASITimestamps, ) -__all__ = ["run_code", "run_fs", "StringFile", "BinaryFile", "WASIFS"] +__all__ = ["run_code", "run_fs", "StringFile", "BinaryFile", "WASIFS", "WASITimestamps"] diff --git a/sandbox/runno/main.py b/sandbox/runno/main.py index 7cdb8b7a..533a16e3 100644 --- a/sandbox/runno/main.py +++ b/sandbox/runno/main.py @@ -8,10 +8,12 @@ BaseFile, CompleteResult, CrashResult, + Options, RunnoError, Runtime, WASIFS, TerminatedResult, + TimeoutResult, WASIPath, StringFile, RunResult, @@ -19,7 +21,7 @@ ) -async def run_code(runtime: Runtime, code: str) -> RunResult: +async def run_code(runtime: Runtime, code: str, **kwargs: Options) -> RunResult: """ Run code in a Runno sandbox. @@ -37,10 +39,12 @@ async def run_code(runtime: Runtime, code: str) -> RunResult: content=code, ) } - return await run_fs(runtime, "/program", fs) + return await run_fs(runtime, "/program", fs, **kwargs) -async def run_fs(runtime: Runtime, entry_path: WASIPath, fs: WASIFS) -> RunResult: +async def run_fs( + runtime: Runtime, entry_path: WASIPath, fs: WASIFS, **kwargs: Options +) -> RunResult: """ Run code in a Runno sandbox with a custom filesystem. @@ -52,6 +56,18 @@ async def run_fs(runtime: Runtime, entry_path: WASIPath, fs: WASIFS) -> RunResul See the types module for more information on the WASIFS type. """ + try: + return await asyncio.wait_for( + _internal_run_fs(runtime, entry_path, fs, **kwargs), + timeout=kwargs.get("timeout", None), + ) + except asyncio.TimeoutError: + return TimeoutResult(result_type="timeout") + + +async def _internal_run_fs( + runtime: Runtime, entry_path: WASIPath, fs: WASIFS, **kwargs: Options +) -> RunResult: proc = await asyncio.create_subprocess_exec( "./runno", runtime, @@ -84,7 +100,9 @@ async def run_fs(runtime: Runtime, entry_path: WASIPath, fs: WASIFS) -> RunResul if exit_code != 0: raise RuntimeError( - f"Runno sandbox subprocess failed with exit code {exit_code}" + f"Runno sandbox subprocess failed with exit code {exit_code}", + stdout, + stderr, ) data = json.loads(stdout) diff --git a/sandbox/runno/types.py b/sandbox/runno/types.py index 4a43d00b..dbd2157e 100644 --- a/sandbox/runno/types.py +++ b/sandbox/runno/types.py @@ -1,5 +1,5 @@ import base64 -from typing import Literal, Union, Dict +from typing import Literal, Union, Dict, TypedDict from pydantic import BaseModel from datetime import datetime @@ -86,7 +86,29 @@ class TerminatedResult(BaseModel): result_type: Literal["terminated"] -type RunResult = Union[CompleteResult, CrashResult, TerminatedResult] +class TimeoutResult(BaseModel): + result_type: Literal["timeout"] -__all__ = ["StringFile", "BinaryFile", "WASIFS"] +type RunResult = Union[CompleteResult, CrashResult, TerminatedResult, TimeoutResult] + + +class Options(TypedDict): + timeout: float + + +__all__ = [ + "StringFile", + "BinaryFile", + "WASIFS", + "WASITimestamps", + "Options", + "RunnoError", + "Runtime", + "WASIPath", + "RunResult", + "TimeoutResult", + "TerminatedResult", + "CompleteResult", + "CrashResult", +] diff --git a/sandbox/tests/test_code.py b/sandbox/tests/test_code.py index 73716a0c..7e8d5e9c 100644 --- a/sandbox/tests/test_code.py +++ b/sandbox/tests/test_code.py @@ -61,3 +61,9 @@ async def test_simple_cpp(): assert result.stdout == "Hello, World!\n" assert result.stderr == "" assert result.exit_code == 0 + + +@pytest.mark.asyncio +async def test_timeout(): + result = await run_code("python", "while True: pass", timeout=0.1) + assert result.result_type == "timeout"