Remove Allocations from Panic Handler

flake.nix

@@ -107,7 +107,7 @@
           pname = "hyperlight";
           version = "0.0.0";
           src = lib.cleanSource ./.;
-          cargoHash = "sha256-mNKnsaSKVz4khzWO7VhmN0cR+Ed5ML7fD1PJJCeQQ6E=";
+          cargoHash = "sha256-hoeJEBdxaoyLlhQQ4X4Wk5X1QVtQ7RRQYaxkiGg8rWA=";
 
           nativeBuildInputs = [
             azure-cli

src/hyperlight_guest_bin/Cargo.toml

@@ -27,6 +27,7 @@ hyperlight-guest-tracing = { workspace = true, default-features = false }
 buddy_system_allocator = "0.11.0"
 log = { version = "0.4", default-features = false }
 spin = "0.10.0"
+heapless = "0.8.0"
 
 [lints]
 workspace = true

src/hyperlight_guest_bin/src/lib.rs

@@ -18,14 +18,16 @@ limitations under the License.
 // === Dependencies ===
 extern crate alloc;
 
-use alloc::string::ToString;
+use core::ffi::CStr;
+use core::fmt::Write;
 
 use buddy_system_allocator::LockedHeap;
 #[cfg(target_arch = "x86_64")]
 use exceptions::{gdt::load_gdt, idtr::load_idt};
 use guest_function::call::dispatch_function;
 use guest_function::register::GuestFunctionRegister;
 use guest_logger::init_logger;
+use heapless::String;
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::mem::HyperlightPEB;
 #[cfg(feature = "mem_profile")]
@@ -139,11 +141,39 @@ pub static mut OS_PAGE_SIZE: u32 = 0;
 // to satisfy the clippy when cfg == test
 #[allow(dead_code)]
 fn panic(info: &core::panic::PanicInfo) -> ! {
-    let msg = info.to_string();
-    let c_string = alloc::ffi::CString::new(msg)
-        .unwrap_or_else(|_| alloc::ffi::CString::new("panic (invalid utf8)").unwrap());
+    _panic_handler(info)
+}
+
+#[inline(always)]
+fn _panic_handler(info: &core::panic::PanicInfo) -> ! {
+    // stack allocate a 512-byte message buffer.
+    let mut panic_buf = String::<512>::new();
+    let write_res = write!(panic_buf, "{}\0", info);
+    if write_res.is_err() {
+        unsafe {
+            abort_with_code_and_message(
+                &[ErrorCode::UnknownError as u8],
+                c"panic: message format failed (limit: 512 bytes)".as_ptr(),
+            )
+        }
+    }
 
-    unsafe { abort_with_code_and_message(&[ErrorCode::UnknownError as u8], c_string.as_ptr()) }
+    let c_str_res = CStr::from_bytes_with_nul(panic_buf.as_bytes());
+    if c_str_res.is_err() {
+        unsafe {
+            abort_with_code_and_message(
+                &[ErrorCode::UnknownError as u8],
+                c"panic: failed to convert to CString".as_ptr(),
+            )
+        }
+    }
+
+    unsafe {
+        abort_with_code_and_message(
+            &[ErrorCode::UnknownError as u8],
+            c_str_res.unwrap().as_ptr(),
+        )
+    }
 }
 
 // === Entrypoint ===

src/hyperlight_host/tests/integration_test.rs

@@ -533,6 +533,36 @@ fn guest_malloc_abort() {
     ));
 }
 
+#[test]
+fn guest_panic_no_alloc() {
+    let heap_size = 0x4000;
+
+    let mut cfg = SandboxConfiguration::default();
+    cfg.set_heap_size(heap_size);
+    let uninit = UninitializedSandbox::new(
+        GuestBinary::FilePath(simple_guest_as_string().unwrap()),
+        Some(cfg),
+    )
+    .unwrap();
+    let mut sbox: MultiUseSandbox = uninit.evolve().unwrap();
+
+    let res = sbox
+        .call::<i32>(
+            "ExhaustHeap", // uses the rust allocator to allocate small blocks on the heap until OOM
+            (),
+        )
+        .unwrap_err();
+
+    if let HyperlightError::StackOverflow() = res {
+        panic!("panic on OOM caused stack overflow, this implies allocation in panic handler");
+    }
+
+    assert!(matches!(
+        res,
+        HyperlightError::GuestAborted(code, msg) if code == ErrorCode::UnknownError as u8 && msg.contains("memory allocation of ") && msg.contains("bytes failed")
+    ));
+}
+
 // Tests libc alloca
 #[test]
 fn dynamic_stack_allocate_c_guest() {

src/tests/rust_guests/simpleguest/src/main.rs

@@ -29,6 +29,7 @@ use alloc::boxed::Box;
 use alloc::string::ToString;
 use alloc::vec::Vec;
 use alloc::{format, vec};
+use core::alloc::Layout;
 use core::ffi::c_char;
 use core::hint::black_box;
 use core::ptr::write_volatile;
@@ -507,6 +508,23 @@ fn call_malloc(function_call: &FunctionCall) -> Result<Vec<u8>> {
     }
 }
 
+#[hyperlight_guest_tracing::trace_function]
+unsafe fn exhaust_heap(_: &FunctionCall) -> ! {
+    let layout: Layout = Layout::new::<u8>();
+    let mut ptr = alloc::alloc::alloc_zeroed(layout);
+    while !ptr.is_null() {
+        black_box(ptr);
+        ptr = alloc::alloc::alloc_zeroed(layout);
+    }
+
+    // after alloc::alloc_zeroed failure (null return when called in loop above)
+    // allocate a Vec to ensure OOM panic
+    let vec = Vec::<i32>::with_capacity(1);
+    black_box(vec);
+
+    panic!("function should have panicked before due to OOM")
+}
+
 #[hyperlight_guest_tracing::trace_function]
 fn malloc_and_free(function_call: &FunctionCall) -> Result<Vec<u8>> {
     if let ParameterValue::Int(size) = function_call.parameters.clone().unwrap()[0].clone() {
@@ -1023,6 +1041,14 @@ pub extern "C" fn hyperlight_main() {
     );
     register_function(call_malloc_def);
 
+    let exhaust_heap_def = GuestFunctionDefinition::new(
+        "ExhaustHeap".to_string(),
+        Vec::new(),
+        ReturnType::Int,
+        exhaust_heap as usize,
+    );
+    register_function(exhaust_heap_def);
+
     let malloc_and_free_def = GuestFunctionDefinition::new(
         "MallocAndFree".to_string(),
         Vec::from(&[ParameterType::Int]),