diff --git a/examples/python-agent-driver/hl_pydriver.c b/examples/python-agent-driver/hl_pydriver.c index 5cc3929..0781702 100644 --- a/examples/python-agent-driver/hl_pydriver.c +++ b/examples/python-agent-driver/hl_pydriver.c @@ -35,6 +35,7 @@ #include #include #include +#include #include /* -- Minimal FlatBuffers reader for the incoming FunctionCall --------- @@ -137,14 +138,69 @@ static inline void wrmsr_fsbase(uint64_t v) /* -- Core work ------------------------------------------------------- */ +static void report_exit_code(int code) +{ + char req[128]; + int n = snprintf(req, sizeof(req), + "{\"name\":\"__hl_exit\",\"args\":{\"code\":%d}}", code); + int fd = open("/dev/hcall", O_RDWR); + if (fd < 0) + return; + write(fd, req, n); + char resp[128]; + read(fd, resp, sizeof(resp)); + close(fd); +} + +static int run_code_with_exceptions(const char *code) +{ + PyObject *m = PyImport_AddModule("__main__"); + if (!m) return 1; + PyObject *d = PyModule_GetDict(m); + if (!d) return 1; + + PyObject *result = PyRun_String(code, Py_file_input, d, d); + if (result) { + Py_DECREF(result); + return 0; + } + + if (PyErr_ExceptionMatches(PyExc_SystemExit)) { + PyObject *type, *value, *tb; + PyErr_Fetch(&type, &value, &tb); + PyErr_NormalizeException(&type, &value, &tb); + int exit_code = 1; + if (value) { + PyObject *ca = PyObject_GetAttrString(value, "code"); + if (ca) { + if (PyLong_Check(ca)) { + exit_code = (int)PyLong_AsLong(ca); + } else if (ca == Py_None) { + exit_code = 0; + } else { + PyObject *s = PyObject_Str(ca); + if (s) { + const char *msg = PyUnicode_AsUTF8(s); + if (msg) + fprintf(stderr, "%s\n", msg); + Py_DECREF(s); + } + } + Py_DECREF(ca); + } + } + Py_XDECREF(type); + Py_XDECREF(value); + Py_XDECREF(tb); + return exit_code; + } + + PyErr_Print(); + return 1; +} + static void py_run_user_code(const uint8_t *fc, size_t fc_len) { - /* Restore Python's FS_BASE. Between calls, the kernel-side - * dispatch_prepare (and possibly Hyperlight's own segment - * restore) can leave FS_BASE at a kernel-default value, which - * breaks Python's TLS-reliant code (PyGILState, pthread local - * state, …) the moment we touch it. g_py_fsbase was captured - * at the end of the one-time init below. */ if (g_py_fsbase) wrmsr_fsbase(g_py_fsbase); @@ -153,23 +209,27 @@ static void py_run_user_code(const uint8_t *fc, size_t fc_len) if (!code) return; - /* PyRun_SimpleString wants a NUL-terminated C string; the FC - * hlstring isn't NUL-terminated in its on-the-wire form. Small - * copy into a stack/heap buffer to terminate. */ - if (code_len < 4096) { - char buf[4096]; - memcpy(buf, code, code_len); - buf[code_len] = '\0'; - PyRun_SimpleString(buf); + char stack_buf[4096]; + char *buf; + if (code_len < sizeof(stack_buf)) { + memcpy(stack_buf, code, code_len); + stack_buf[code_len] = '\0'; + buf = stack_buf; } else { - char *buf = malloc(code_len + 1); + buf = malloc(code_len + 1); if (!buf) return; memcpy(buf, code, code_len); buf[code_len] = '\0'; - PyRun_SimpleString(buf); - free(buf); } + + int exit_code = run_code_with_exceptions(buf); + + if (buf != stack_buf) + free(buf); + + if (exit_code != 0) + report_exit_code(exit_code); } static void py_initialize_once(void) @@ -202,6 +262,24 @@ static void py_initialize_once(void) " importlib.import_module(_mod)\n" " except Exception as _e:\n" " sys.stderr.write(f'warn: preload {_mod} failed: {_e}\\n')\n"); + + /* Monkey-patch time.sleep to call the host via /dev/hcall. + * Unikraft's cooperative scheduler on Hyperlight has no timer + * interrupt, so the kernel-level nanosleep is a no-op. This + * routes the sleep to the host thread which actually blocks. */ + PyRun_SimpleString( + "import time as _hl_time\n" + "def _hl_sleep(secs):\n" + " if secs <= 0:\n" + " return\n" + " import json\n" + " fd = open('/dev/hcall', 'r+b', buffering=0)\n" + " fd.write(json.dumps(" + "{'name':'__hl_sleep','args':{'ns':int(secs*1e9)}}).encode())\n" + " fd.read()\n" + " fd.close()\n" + "_hl_time.sleep = _hl_sleep\n" + "del _hl_time, _hl_sleep\n"); } /* -- Entry points --------------------------------------------------- */ diff --git a/host/src/bin/pyhl.rs b/host/src/bin/pyhl.rs index 49c1ec1..6df59d0 100644 --- a/host/src/bin/pyhl.rs +++ b/host/src/bin/pyhl.rs @@ -504,14 +504,19 @@ fn cmd_run(args: RunArgs) -> Result<()> { full }; + sandbox.reset_exit_code(); let t_call = Instant::now(); let _: () = sandbox.call_named("run", payload)?; let call_ms = t_call.elapsed().as_secs_f64() * 1000.0; + let exit_code = sandbox.last_exit_code(); if args.verbose { eprintln!( - "[pyhl] run {i}/{total} restore={restore_ms:.1}ms call={call_ms:.1}ms (hermetic)" + "[pyhl] run {i}/{total} restore={restore_ms:.1}ms call={call_ms:.1}ms exit={exit_code} (hermetic)" ); } + if exit_code != 0 { + std::process::exit(exit_code); + } } Ok(()) diff --git a/host/src/lib.rs b/host/src/lib.rs index ef1fcef..0759ada 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -67,6 +67,7 @@ use hyperlight_host::sandbox::SandboxConfiguration; use hyperlight_host::{GuestBinary, HostFunctions, MultiUseSandbox, UninitializedSandbox}; use std::collections::HashMap; use std::path::Path; +use std::sync::atomic::{AtomicI32, Ordering}; use std::sync::Arc; use std::time::Duration; @@ -791,6 +792,25 @@ fn build_tools( Ok(Some(registry)) } +/// Register internal tools (`__hl_exit`, `__hl_sleep`) on a tool registry. +/// These are plumbing used by the guest driver (`hl_pydriver.c`) and are +/// always present regardless of user-supplied tools or preopens. +fn register_internal_tools(tools: &mut ToolRegistry, exit_code: &Arc) { + let ec = exit_code.clone(); + tools.register("__hl_exit", move |args| { + let code = args["code"].as_i64().unwrap_or(1) as i32; + ec.store(code, Ordering::Relaxed); + Ok(serde_json::json!({})) + }); + tools.register("__hl_sleep", |args| { + let ns = args["ns"].as_u64().unwrap_or(0); + if ns > 0 { + std::thread::sleep(std::time::Duration::from_nanos(ns)); + } + Ok(serde_json::json!({})) + }); +} + /// Routes incoming fs_* tool calls to the matching `FsSandbox` by /// matching the guest-supplied path against each preopen's guest path. #[derive(Clone)] @@ -1041,6 +1061,7 @@ pub struct Sandbox { /// Snapshot restore unmaps all non-snapshot regions. file_mapping_path: Option, file_mapping_base: u64, + exit_code: Arc, } /// Where the initrd comes from — either a file (zero-copy `map_file_cow`) @@ -1213,17 +1234,16 @@ impl Sandbox { let mut usbox = UninitializedSandbox::new(env, Some(config.sandbox_config()))?; - let tools = build_tools(tools, preopens)?; - - if let Some(tools) = tools { - let tools = Arc::new(tools); - let tools_ref = tools.clone(); - usbox.register_host_function("__dispatch", move |payload: Vec| -> Vec { - tools_ref.dispatch(&payload) - })?; - } + let exit_code = Arc::new(AtomicI32::new(0)); + let mut tools = build_tools(tools, preopens)?.unwrap_or_default(); + register_internal_tools(&mut tools, &exit_code); + let tools = Arc::new(tools); + let tools_ref = tools.clone(); + usbox.register_host_function("__dispatch", move |payload: Vec| -> Vec { + tools_ref.dispatch(&payload) + })?; - Self::finish_evolve(usbox, None, 0) + Self::finish_evolve(usbox, None, 0, exit_code) } /// Low-level: boot with a zero-copy mapped initrd file. Prefer the builder. @@ -1263,24 +1283,28 @@ impl Sandbox { usbox.map_file_cow(path, INITRD_MAP_BASE, Some("initrd"))?; } - let tools = build_tools(tools, preopens)?; - - // Register tool dispatch if needed - if let Some(tools) = tools { - let tools = Arc::new(tools); - let tools_ref = tools.clone(); - usbox.register_host_function("__dispatch", move |payload: Vec| -> Vec { - tools_ref.dispatch(&payload) - })?; - } + let exit_code = Arc::new(AtomicI32::new(0)); + let mut tools = build_tools(tools, preopens)?.unwrap_or_default(); + register_internal_tools(&mut tools, &exit_code); + let tools = Arc::new(tools); + let tools_ref = tools.clone(); + usbox.register_host_function("__dispatch", move |payload: Vec| -> Vec { + tools_ref.dispatch(&payload) + })?; - Self::finish_evolve(usbox, initrd_path.map(|p| p.to_path_buf()), INITRD_MAP_BASE) + Self::finish_evolve( + usbox, + initrd_path.map(|p| p.to_path_buf()), + INITRD_MAP_BASE, + exit_code, + ) } fn finish_evolve( usbox: UninitializedSandbox, file_mapping_path: Option, file_mapping_base: u64, + exit_code: Arc, ) -> Result { let mut inner = usbox.evolve()?; let snapshot = inner.snapshot().ok(); @@ -1289,6 +1313,7 @@ impl Sandbox { snapshot, file_mapping_path, file_mapping_base, + exit_code, }) } @@ -1339,6 +1364,18 @@ impl Sandbox { Ok(self.inner.call(func_name, args)?) } + /// Read the exit code reported by the guest via `__hl_exit`. + /// Defaults to 0 (success) if the guest never called it. + pub fn last_exit_code(&self) -> i32 { + self.exit_code.load(Ordering::Relaxed) + } + + /// Reset the stored exit code to 0. Call before each guest + /// invocation so a previous non-zero code doesn't leak. + pub fn reset_exit_code(&self) { + self.exit_code.store(0, Ordering::Relaxed); + } + /// Take a new snapshot of the current guest state. /// /// Useful for the "snapshot after one-time warm-up" pattern: call @@ -1405,21 +1442,16 @@ impl Sandbox { let loaded = Snapshot::from_file_unchecked(path.as_ref())?; let arc = Arc::new(loaded); + let exit_code = Arc::new(AtomicI32::new(0)); + let mut tools = build_tools(None, preopens)?.unwrap_or_default(); + register_internal_tools(&mut tools, &exit_code); + let tools = Arc::new(tools); + let tools_ref = tools.clone(); + let mut host_funcs = HostFunctions::default(); - if !preopens.is_empty() { - if let Some(tools) = build_tools(None, preopens)? { - let tools = Arc::new(tools); - let tools_ref = tools.clone(); - host_funcs - .register_host_function("__dispatch", move |payload: Vec| -> Vec { - tools_ref.dispatch(&payload) - })?; - } - } else { - host_funcs.register_host_function("__dispatch", |_payload: Vec| -> Vec { - Vec::new() - })?; - } + host_funcs.register_host_function("__dispatch", move |payload: Vec| -> Vec { + tools_ref.dispatch(&payload) + })?; let inner = MultiUseSandbox::from_snapshot(arc.clone(), host_funcs, None)?; @@ -1428,6 +1460,7 @@ impl Sandbox { snapshot: Some(arc), file_mapping_path: None, file_mapping_base: 0, + exit_code, }) } } diff --git a/host/src/pyhl.rs b/host/src/pyhl.rs index 69cf762..df52d88 100644 --- a/host/src/pyhl.rs +++ b/host/src/pyhl.rs @@ -197,6 +197,7 @@ pub fn install(opts: &InstallOptions<'_>) -> Result { pub struct RunTiming { pub restore_ms: f64, pub call_ms: f64, + pub exit_code: i32, } /// A pyhl runtime backed by a warmed-up snapshot. Cheap to keep around; @@ -244,10 +245,12 @@ impl Runtime { t.restore_ms = tr.elapsed().as_secs_f64() * 1000.0; } self.first_run = false; + self.sandbox.reset_exit_code(); let tc = Instant::now(); let _: () = self.sandbox.call_named("run", code.to_string())?; t.call_ms = tc.elapsed().as_secs_f64() * 1000.0; + t.exit_code = self.sandbox.last_exit_code(); Ok(t) }