Remove allocations in panic handler by formatting panic message

using a new FixedStringBuf type backed by a static mut array.
Adds a test to verify that StackOverflow no longer occurs on OOM panic

Signed-off-by: adamperlin <[email protected]>

Author: adamperlin

src/hyperlight_common/src/fixed_buf.rs

@@ -0,0 +1,63 @@
+use core::fmt;
+use core::result::Result;
+use core::ffi::{CStr, FromBytesUntilNulError};
+
+pub struct FixedStringBuf<'a> {
+    pub buf: &'a mut [u8],
+    pub pos: usize,
+}
+
+impl<'a> fmt::Write for FixedStringBuf<'a> {
+    fn write_str(&mut self, s: &str) -> fmt::Result {
+        // we always reserve 1 byte for the null terminator, 
+        // as the buffer must be convertible to a CStr.
+        let buf_end = self.buf.len() - 1;
+        if self.pos + s.len() > buf_end {
+            return Err(fmt::Error)
+        }
+
+        self.buf[self.pos..self.pos + s.len()].copy_from_slice(s.as_bytes());
+        self.pos += s.len();
+        Ok(())
+    }
+}
+
+impl <'a> FixedStringBuf<'a> {
+    pub fn as_str(&self) -> Result<&str, core::str::Utf8Error> {
+        core::str::from_utf8(&self.buf[..self.pos])
+    }
+
+    pub fn new(buf: &'a mut [u8]) -> FixedStringBuf<'a> {
+        assert!(buf.len() > 0);
+        FixedStringBuf {
+            buf,
+            pos: 0,
+        }
+    }
+
+    pub fn as_c_str(&mut self) -> Result<&CStr, FromBytesUntilNulError> {
+        // null terminate buffer. 
+        // we are guaranteed to have enough space since we always reserve one extra
+        // byte for null in write_str
+        self.buf[self.pos] = 0;
+        CStr::from_bytes_until_nul(&self.buf[..self.pos + 1])
+    }
+}
+
+mod test {
+    use super::{FixedStringBuf};
+    use core::fmt::Write;
+    use core::fmt;
+    #[test]
+    fn test_fixed_buf() {
+        let mut bs = [0; 21];
+        let mut buf = FixedStringBuf::new(&mut bs);
+        write!(&mut buf, "{}", "0123456789").expect("Failed to write to FixedBuf");
+        write!(&mut buf, "{}", "0123456789").expect("Failed to write to FixedBuf");
+        assert_eq!(buf.as_str().unwrap(), "01234567890123456789");
+        assert_eq!(buf.pos, 20);
+
+        let res = write!(&mut buf, "10");
+        assert_eq!(res, Err(fmt::Error));
+    }
+}

src/hyperlight_common/src/lib.rs

@@ -42,3 +42,5 @@ pub mod outb;
 
 /// cbindgen:ignore
 pub mod resource;
+
+pub mod fixed_buf;

src/hyperlight_guest_bin/src/lib.rs

@@ -18,8 +18,9 @@ limitations under the License.
 // === Dependencies ===
 extern crate alloc;
 
-use alloc::string::ToString;
+use core::fmt::{Write};
 
+use alloc::string::ToString;
 use buddy_system_allocator::LockedHeap;
 #[cfg(target_arch = "x86_64")]
 use exceptions::{gdt::load_gdt, idtr::load_idt};
@@ -30,11 +31,12 @@ use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::mem::HyperlightPEB;
 #[cfg(feature = "mem_profile")]
 use hyperlight_common::outb::OutBAction;
+use hyperlight_common::fixed_buf::{FixedStringBuf};
 use hyperlight_guest::exit::{abort_with_code_and_message, halt};
 use hyperlight_guest::guest_handle::handle::GuestHandle;
 use hyperlight_guest_tracing::{trace, trace_function};
 use log::LevelFilter;
-use spin::Once;
+use spin::{Once};
 
 // === Modules ===
 #[cfg(target_arch = "x86_64")]
@@ -139,7 +141,33 @@ 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();
+    _panic_handler(info)
+}
+
+
+static mut PANIC_MSG: [u8; 512] = [0u8; 512];
+
+#[allow(static_mut_refs)]
+static mut PANIC_BUF: FixedStringBuf = FixedStringBuf{
+    buf: unsafe { &mut PANIC_MSG },
+    pos: 0,
+};
+
+
+#[inline(always)]
+#[allow(static_mut_refs)]
+fn _panic_handler(info: &core::panic::PanicInfo) -> ! {
+    let panic_buf: &mut FixedStringBuf = unsafe { &mut PANIC_BUF };
+    write!(panic_buf, "{}", info).expect("panic: message format failed");
+
+    // create a CStr from the underlying array in panic_buf.
+    // Note that we do NOT use CString here to avoid allocating
+    let c_string = panic_buf.as_c_str().expect("panic: failed to convert to CStr");
+    unsafe { abort_with_code_and_message(&[ErrorCode::UnknownError as u8], c_string.as_ptr()) }
+}
+
+fn _panic_handler_old(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());
 
@@ -213,4 +241,4 @@ pub extern "C" fn entrypoint(peb_address: u64, seed: u64, ops: u64, max_log_leve
     });
 
     halt();
-}
+}

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();
+    println!("{:?}", res);
+
+    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,9 +29,11 @@ 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;
+use core::sync::atomic::AtomicPtr;
 
 use hyperlight_common::flatbuffer_wrappers::function_call::{FunctionCall, FunctionCallType};
 use hyperlight_common::flatbuffer_wrappers::function_types::{
@@ -507,6 +509,18 @@ 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);
+    }
+
+    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 +1037,14 @@ pub extern "C" fn hyperlight_main() {
     );
     register_function(call_malloc_def);
 
+    let call_malloc_and_panic_def = GuestFunctionDefinition::new(
+        "ExhaustHeap".to_string(),
+        Vec::new(),
+        ReturnType::Int,
+        exhaust_heap as usize,
+    );
+    register_function(call_malloc_and_panic_def);
+
     let malloc_and_free_def = GuestFunctionDefinition::new(
         "MallocAndFree".to_string(),
         Vec::from(&[ParameterType::Int]),