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/src/exit.rs

@@ -67,6 +67,14 @@ pub unsafe fn abort_with_code_and_message(code: &[u8], message_ptr: *const c_cha
     }
 }
 
+/// This function exists to give the guest more manual control
+/// over the abort sequence. For example, in `hyperlight_guest_bin`'s panic handler,
+/// we have a message of unknown length that we want to stream
+/// to the host, which requires sending the message in chunks
+pub fn write_abort(code: &[u8]) {
+    outb(OutBAction::Abort as u16, code);
+}
+
 /// OUT bytes to the host through multiple exits.
 #[hyperlight_guest_tracing::trace_function]
 pub(crate) fn outb(port: u16, data: &[u8]) {

src/hyperlight_guest_bin/src/lib.rs

@@ -18,7 +18,7 @@ limitations under the License.
 // === Dependencies ===
 extern crate alloc;
 
-use alloc::string::ToString;
+use core::fmt::Write;
 
 use buddy_system_allocator::LockedHeap;
 #[cfg(target_arch = "x86_64")]
@@ -30,7 +30,7 @@ use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::mem::HyperlightPEB;
 #[cfg(feature = "mem_profile")]
 use hyperlight_common::outb::OutBAction;
-use hyperlight_guest::exit::{abort_with_code_and_message, halt};
+use hyperlight_guest::exit::{halt, write_abort};
 use hyperlight_guest::guest_handle::handle::GuestHandle;
 use hyperlight_guest_tracing::{trace, trace_function};
 use log::LevelFilter;
@@ -139,11 +139,37 @@ 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)
+}
+
+/// A writer that sends all output to the hyperlight host
+/// using output ports. This allows us to not impose a
+/// buffering limit on error message size on the guest end,
+/// though one exists for the host.
+struct HyperlightAbortWriter;
+impl core::fmt::Write for HyperlightAbortWriter {
+    fn write_str(&mut self, s: &str) -> core::fmt::Result {
+        write_abort(s.as_bytes());
+        Ok(())
+    }
+}
+
+#[inline(always)]
+fn _panic_handler(info: &core::panic::PanicInfo) -> ! {
+    let mut w = HyperlightAbortWriter;
+
+    // begin abort sequence by writing the error code
+    write_abort(&[ErrorCode::UnknownError as u8]);
+
+    let write_res = write!(w, "{}", info);
+    if write_res.is_err() {
+        write_abort("panic: message format failed".as_bytes());
+    }
 
-    unsafe { abort_with_code_and_message(&[ErrorCode::UnknownError as u8], c_string.as_ptr()) }
+    // write abort terminator to finish the abort
+    // and signal to the host that the message can now be read
+    write_abort(&[0xFF]);
+    unreachable!();
 }
 
 // === Entrypoint ===

src/hyperlight_host/tests/integration_test.rs

@@ -531,6 +531,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;
@@ -508,6 +509,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() {
@@ -1108,6 +1126,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]),