diff --git a/host/Cargo.lock b/host/Cargo.lock index 0e78b0b..41938cc 100644 --- a/host/Cargo.lock +++ b/host/Cargo.lock @@ -636,7 +636,7 @@ dependencies = [ [[package]] name = "hyperlight-unikraft-host" -version = "0.5.0" +version = "0.6.0" dependencies = [ "anyhow", "base64", diff --git a/host/Cargo.toml b/host/Cargo.toml index e06326b..f69f6a2 100644 --- a/host/Cargo.toml +++ b/host/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "hyperlight-unikraft-host" -version = "0.5.0" +version = "0.6.0" edition = "2021" description = "Embedded Hyperlight host for running Unikraft unikernels" license = "MIT OR Apache-2.0" diff --git a/host/src/lib.rs b/host/src/lib.rs index 6456587..427e4c5 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -68,7 +68,7 @@ use std::collections::{HashMap, HashSet}; use std::net::IpAddr; use std::path::Path; use std::sync::atomic::{AtomicI32, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Condvar, Mutex}; use std::time::Duration; /// Magic header for cmdline embedded in initrd: "HLCMDLN\0" @@ -108,6 +108,42 @@ const MAX_DISPATCH_PAYLOAD: usize = 64 * 1024 * 1024; /// Cap for `__hl_sleep` duration to prevent unbounded host-thread blocking (60 s). const MAX_SLEEP_NS: u64 = 60_000_000_000; +/// Shared cancellation primitive for `__hl_sleep`. Calling +/// [`SleepCancel::cancel`] wakes up any in-progress sleep immediately so +/// the host function returns and the hypervisor execution loop can detect +/// the pending cancellation. +#[derive(Clone)] +pub struct SleepCancel(Arc<(Mutex, Condvar)>); + +impl SleepCancel { + fn new() -> Self { + Self(Arc::new((Mutex::new(false), Condvar::new()))) + } + + /// Wake any in-progress `__hl_sleep` immediately. + pub fn cancel(&self) { + let (lock, cvar) = &*self.0; + *lock.lock().unwrap() = true; + cvar.notify_all(); + } + + /// Reset so the next guest call can sleep normally. + pub fn reset(&self) { + *self.0 .0.lock().unwrap() = false; + } + + fn wait(&self, dur: Duration) { + let (lock, cvar) = &*self.0; + let guard = lock.lock().unwrap(); + if *guard { + return; + } + // wait_timeout_while handles spurious wakeups by re-checking the + // predicate; we only return early when actually cancelled. + let _ = cvar.wait_timeout_while(guard, dur, |cancelled| !*cancelled); + } +} + /// Cap for `fs_list` directory entries to prevent host OOM on huge directories. const MAX_DIR_ENTRIES: usize = 100_000; @@ -957,6 +993,7 @@ fn build_tools( fn register_internal_tools( tools: &mut ToolRegistry, exit_code: &Arc, + sleep_cancel: &SleepCancel, network: Option<&NetworkPolicy>, listen_ports: Option<&ListenPorts>, ) -> Option>> { @@ -966,10 +1003,11 @@ fn register_internal_tools( ec.store(code, Ordering::Relaxed); Ok(serde_json::json!({})) }); - tools.register("__hl_sleep", |args| { + let sc = sleep_cancel.clone(); + tools.register("__hl_sleep", move |args| { let ns = args["ns"].as_u64().unwrap_or(0).min(MAX_SLEEP_NS); if ns > 0 { - std::thread::sleep(std::time::Duration::from_nanos(ns)); + sc.wait(Duration::from_nanos(ns)); } Ok(serde_json::json!({})) }); @@ -982,7 +1020,6 @@ fn register_internal_tools( use socket2::{Domain, Protocol, SockAddr, Socket, Type}; use std::net::SocketAddr; -use std::sync::Mutex; struct HostSocket { socket: Socket, @@ -1940,6 +1977,8 @@ pub struct Sandbox { /// Shared socket table — cleared on [`Sandbox::restore`] so that /// host-side fds don't leak across guest restore cycles. socket_table: Option>>, + /// Cancellation token for in-progress `__hl_sleep` host calls. + sleep_cancel: SleepCancel, } /// Where the initrd comes from — either a file (zero-copy `map_file_cow`) @@ -2148,15 +2187,17 @@ impl Sandbox { let mut usbox = UninitializedSandbox::new(env, Some(config.sandbox_config()))?; let exit_code = Arc::new(AtomicI32::new(0)); + let sleep_cancel = SleepCancel::new(); let mut tools = build_tools(tools, preopens)?.unwrap_or_default(); - let socket_table = register_internal_tools(&mut tools, &exit_code, network, listen_ports); + let socket_table = + register_internal_tools(&mut tools, &exit_code, &sleep_cancel, network, listen_ports); 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, exit_code, socket_table) + Self::finish_evolve(usbox, None, 0, exit_code, sleep_cancel, socket_table) } /// Low-level: boot with a zero-copy mapped initrd file. Prefer the builder. @@ -2200,8 +2241,10 @@ impl Sandbox { } let exit_code = Arc::new(AtomicI32::new(0)); + let sleep_cancel = SleepCancel::new(); let mut tools = build_tools(tools, preopens)?.unwrap_or_default(); - let socket_table = register_internal_tools(&mut tools, &exit_code, network, listen_ports); + let socket_table = + register_internal_tools(&mut tools, &exit_code, &sleep_cancel, network, listen_ports); let tools = Arc::new(tools); let tools_ref = tools.clone(); usbox.register_host_function("__dispatch", move |payload: Vec| -> Vec { @@ -2213,6 +2256,7 @@ impl Sandbox { initrd_path.map(|p| p.to_path_buf()), INITRD_MAP_BASE, exit_code, + sleep_cancel, socket_table, ) } @@ -2222,6 +2266,7 @@ impl Sandbox { file_mapping_path: Option, file_mapping_base: u64, exit_code: Arc, + sleep_cancel: SleepCancel, socket_table: Option>>, ) -> Result { let mut inner = usbox.evolve()?; @@ -2233,6 +2278,7 @@ impl Sandbox { file_mapping_base, exit_code, socket_table, + sleep_cancel, }) } @@ -2299,6 +2345,18 @@ impl Sandbox { self.exit_code.store(0, Ordering::Relaxed); } + /// Obtain a handle that can interrupt a running guest call from + /// another thread. See [`hyperlight_host::hypervisor::InterruptHandle`]. + pub fn interrupt_handle(&self) -> Arc { + self.inner.interrupt_handle() + } + + /// Obtain the sleep-cancellation token. Call `.cancel()` on it to + /// wake any in-progress `__hl_sleep` immediately. + pub fn sleep_cancel(&self) -> SleepCancel { + self.sleep_cancel.clone() + } + /// Take a new snapshot of the current guest state. /// /// Useful for the "snapshot after one-time warm-up" pattern: call @@ -2412,8 +2470,10 @@ impl Sandbox { let arc = Arc::new(loaded); let exit_code = Arc::new(AtomicI32::new(0)); + let sleep_cancel = SleepCancel::new(); let mut tools = build_tools(None, preopens)?.unwrap_or_default(); - let socket_table = register_internal_tools(&mut tools, &exit_code, network, listen_ports); + let socket_table = + register_internal_tools(&mut tools, &exit_code, &sleep_cancel, network, listen_ports); let tools = Arc::new(tools); let tools_ref = tools.clone(); @@ -2436,6 +2496,7 @@ impl Sandbox { file_mapping_base: INITRD_MAP_BASE, exit_code, socket_table, + sleep_cancel, }) } } @@ -3159,10 +3220,12 @@ mod tests { fn net_tools_registered_with_blocklist() { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); + let sc = SleepCancel::new(); let bl = BlockList::from_hosts(&["1.2.3.4"]).unwrap(); register_internal_tools( &mut tools, &exit_code, + &sc, Some(&NetworkPolicy::BlockList(bl)), None, ); @@ -3176,7 +3239,8 @@ mod tests { fn net_tools_not_registered_without_policy() { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); - register_internal_tools(&mut tools, &exit_code, None, None); + let sc = SleepCancel::new(); + register_internal_tools(&mut tools, &exit_code, &sc, None, None); let req = br#"{"name":"net_socket","args":{"family":2,"type":1}}"#; let resp = tools.dispatch(req); let s = std::str::from_utf8(&resp).unwrap(); @@ -3187,7 +3251,14 @@ mod tests { fn net_tools_registered_with_allow_all() { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); - register_internal_tools(&mut tools, &exit_code, Some(&NetworkPolicy::AllowAll), None); + let sc = SleepCancel::new(); + register_internal_tools( + &mut tools, + &exit_code, + &sc, + Some(&NetworkPolicy::AllowAll), + None, + ); let req = br#"{"name":"net_socket","args":{"family":2,"type":1}}"#; let resp = tools.dispatch(req); let s = std::str::from_utf8(&resp).unwrap(); @@ -3213,7 +3284,14 @@ mod tests { fn net_bind_denied_without_listen_ports() { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); - register_internal_tools(&mut tools, &exit_code, Some(&NetworkPolicy::AllowAll), None); + let sc = SleepCancel::new(); + register_internal_tools( + &mut tools, + &exit_code, + &sc, + Some(&NetworkPolicy::AllowAll), + None, + ); // Create a socket first let req = br#"{"name":"net_socket","args":{"family":2,"type":1}}"#; let resp = tools.dispatch(req); @@ -3235,10 +3313,12 @@ mod tests { fn net_bind_allowed_with_matching_port() { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); + let sc = SleepCancel::new(); let lp = ListenPorts::from_ports([8080]); register_internal_tools( &mut tools, &exit_code, + &sc, Some(&NetworkPolicy::AllowAll), Some(&lp), ); @@ -3259,10 +3339,12 @@ mod tests { fn net_bind_denied_with_wrong_port() { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); + let sc = SleepCancel::new(); let lp = ListenPorts::from_ports([8080]); register_internal_tools( &mut tools, &exit_code, + &sc, Some(&NetworkPolicy::AllowAll), Some(&lp), ); @@ -3302,7 +3384,8 @@ mod tests { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); - register_internal_tools(&mut tools, &exit_code, None, None); + let sc = SleepCancel::new(); + register_internal_tools(&mut tools, &exit_code, &sc, None, None); let req = br#"{"name":"__hl_sleep","args":{"ns":0}}"#; let resp = tools.dispatch(req); @@ -3310,6 +3393,35 @@ mod tests { assert!(!s.contains("\"error\""), "sleep(0) should succeed: {s}"); } + #[test] + fn test_sleep_cancel_wakes_immediately() { + let mut tools = ToolRegistry::new(); + let exit_code = Arc::new(AtomicI32::new(0)); + let sc = SleepCancel::new(); + register_internal_tools(&mut tools, &exit_code, &sc, None, None); + + let sc2 = sc.clone(); + let handle = std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(200)); + sc2.cancel(); + }); + + let start = std::time::Instant::now(); + let req = br#"{"name":"__hl_sleep","args":{"ns":60000000000}}"#; + let resp = tools.dispatch(req); + let elapsed = start.elapsed(); + handle.join().unwrap(); + sc.reset(); + + let s = std::str::from_utf8(&resp).unwrap(); + assert!(!s.contains("\"error\""), "sleep should succeed: {s}"); + assert!( + elapsed.as_secs() < 5, + "cancelled sleep should wake promptly, took {:.1}s", + elapsed.as_secs_f64() + ); + } + #[test] fn net_getsockopt_returns_correct_type_for_dgram() { let mut reg = ToolRegistry::new(); @@ -3467,9 +3579,15 @@ mod tests { fn net_socket_has_default_timeout() { let mut tools = ToolRegistry::new(); let exit_code = Arc::new(AtomicI32::new(0)); - let table = - register_internal_tools(&mut tools, &exit_code, Some(&NetworkPolicy::AllowAll), None) - .expect("network tools should be registered"); + let sc = SleepCancel::new(); + let table = register_internal_tools( + &mut tools, + &exit_code, + &sc, + Some(&NetworkPolicy::AllowAll), + None, + ) + .expect("network tools should be registered"); let req = br#"{"name":"net_socket","args":{"family":2,"type":1}}"#; let resp = tools.dispatch(req); diff --git a/host/src/pyhl.rs b/host/src/pyhl.rs index 8f2d02b..154bcd7 100644 --- a/host/src/pyhl.rs +++ b/host/src/pyhl.rs @@ -44,6 +44,7 @@ use anyhow::{anyhow, bail, Context, Result}; use std::fs; use std::path::{Path, PathBuf}; +use std::sync::Arc; use std::time::Instant; use crate::{Preopen, Sandbox}; @@ -312,6 +313,73 @@ impl Runtime { Ok(t) } + /// Like [`run_code`](Self::run_code), but kills the guest if it + /// exceeds `timeout`. Returns an error when the guest is + /// interrupted. The runtime is left in a usable state — the next + /// `run_code*` call will restore from the snapshot automatically. + pub fn run_code_with_timeout( + &mut self, + code: &str, + timeout: std::time::Duration, + ) -> Result { + let mut t = RunTiming::default(); + if !self.first_run { + let tr = Instant::now(); + self.sandbox.restore()?; + t.restore_ms = tr.elapsed().as_secs_f64() * 1000.0; + } + self.first_run = false; + self.sandbox.reset_exit_code(); + + let handle = self.sandbox.interrupt_handle(); + let sleep_cancel = self.sandbox.sleep_cancel(); + let done = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let done_clone = done.clone(); + + let timer = std::thread::spawn(move || { + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if done_clone.load(std::sync::atomic::Ordering::Relaxed) { + return false; + } + std::thread::sleep(std::time::Duration::from_millis(10)); + } + if !done_clone.load(std::sync::atomic::Ordering::Relaxed) { + sleep_cancel.cancel(); + handle.kill(); + return true; + } + false + }); + + let tc = Instant::now(); + let call_result: Result<()> = self.sandbox.call_named("run", code.to_string()); + t.call_ms = tc.elapsed().as_secs_f64() * 1000.0; + + done.store(true, std::sync::atomic::Ordering::Relaxed); + let timed_out = timer.join().unwrap_or(false); + + // Always reset so a subsequent guest call can sleep normally, even + // if the timer fired right as the call completed (race where + // timed_out=true but call_result=Ok). + self.sandbox.sleep_cancel.reset(); + + match call_result { + Ok(()) => { + t.exit_code = self.sandbox.last_exit_code(); + Ok(t) + } + Err(_) if timed_out => { + self.sandbox.restore()?; + Err(anyhow!( + "execution timed out after {:.1}s", + timeout.as_secs_f64() + )) + } + Err(e) => Err(e), + } + } + /// Convenience: read a file and run its contents. pub fn run_script(&mut self, path: &Path) -> Result { let code = diff --git a/host/tests/pyhl_runtime.rs b/host/tests/pyhl_runtime.rs index 3e46886..66f372d 100644 --- a/host/tests/pyhl_runtime.rs +++ b/host/tests/pyhl_runtime.rs @@ -339,6 +339,80 @@ fn runtime_network_disabled_by_default() { ); } +// --------------------------------------------------------------------------- +// Timeout enforcement +// --------------------------------------------------------------------------- + +#[test] +fn runtime_timeout_kills_busy_spin() { + let Some((_home, mut rt)) = setup() else { + return; + }; + let start = std::time::Instant::now(); + let result = rt.run_code_with_timeout("while True: pass", std::time::Duration::from_secs(2)); + let elapsed = start.elapsed(); + assert!(result.is_err(), "busy spin should be killed"); + let err = result.unwrap_err().to_string(); + assert!( + err.contains("timed out"), + "error should mention timeout, got: {err}" + ); + eprintln!("busy spin killed in {:.1}s", elapsed.as_secs_f64()); + assert!( + elapsed.as_secs() < 10, + "busy spin should be killed promptly, took {:.1}s", + elapsed.as_secs_f64() + ); +} + +#[test] +fn runtime_timeout_kills_time_sleep() { + let Some((_home, mut rt)) = setup() else { + return; + }; + let start = std::time::Instant::now(); + let result = rt.run_code_with_timeout( + "import time; time.sleep(120)", + std::time::Duration::from_secs(2), + ); + let elapsed = start.elapsed(); + assert!(result.is_err(), "sleeping code should be killed"); + eprintln!("time.sleep killed in {:.1}s", elapsed.as_secs_f64()); + assert!( + elapsed.as_secs() < 10, + "time.sleep should be killed promptly, took {:.1}s", + elapsed.as_secs_f64() + ); +} + +#[test] +fn runtime_timeout_allows_fast_code() { + let Some((_home, mut rt)) = setup() else { + return; + }; + let timing = rt + .run_code_with_timeout("print('fast')", std::time::Duration::from_secs(10)) + .unwrap(); + assert_eq!(timing.exit_code, 0); +} + +#[test] +fn runtime_usable_after_timeout() { + let Some((_home, mut rt)) = setup() else { + return; + }; + // First call: times out + let result = rt.run_code_with_timeout( + "import time; time.sleep(120)", + std::time::Duration::from_secs(2), + ); + assert!(result.is_err()); + + // Second call: should work normally + let timing = rt.run_code("print('recovered')").unwrap(); + assert_eq!(timing.exit_code, 0); +} + // --------------------------------------------------------------------------- // Helpers // ---------------------------------------------------------------------------