Merge pull request #20 from hyperlight-dev/fix/exit-code-and-sleep

Author: danbugs

examples/python-agent-driver/hl_pydriver.c

@@ -35,6 +35,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <unistd.h>
+#include <fcntl.h>
 #include <sys/syscall.h>
 
 /* -- 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 --------------------------------------------------- */

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(())

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<AtomicI32>) {
+    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<std::path::PathBuf>,
     file_mapping_base: u64,
+    exit_code: Arc<AtomicI32>,
 }
 
 /// 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<u8>| -> Vec<u8> {
-                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<u8>| -> Vec<u8> {
+            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<u8>| -> Vec<u8> {
-                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<u8>| -> Vec<u8> {
+            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<std::path::PathBuf>,
         file_mapping_base: u64,
+        exit_code: Arc<AtomicI32>,
     ) -> Result<Self> {
         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<u8>| -> Vec<u8> {
-                        tools_ref.dispatch(&payload)
-                    })?;
-            }
-        } else {
-            host_funcs.register_host_function("__dispatch", |_payload: Vec<u8>| -> Vec<u8> {
-                Vec::new()
-            })?;
-        }
+        host_funcs.register_host_function("__dispatch", move |payload: Vec<u8>| -> Vec<u8> {
+            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,
         })
     }
 }

host/src/pyhl.rs

@@ -197,6 +197,7 @@ pub fn install(opts: &InstallOptions<'_>) -> Result<InstallReport> {
 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)
     }