Don't allocate in exception handler (#1125)

Signed-off-by: James Sturtevant <jsturtevant@gmail.com>

Author: jsturtevant

src/hyperlight_guest_bin/src/exceptions/handler.rs

@@ -14,12 +14,13 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use alloc::format;
-use core::ffi::c_char;
+use core::fmt::Write;
 
 use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::outb::Exception;
-use hyperlight_guest::exit::abort_with_code_and_message;
+use hyperlight_guest::exit::write_abort;
+
+use crate::HyperlightAbortWriter;
 
 /// Exception information pushed onto the stack by the CPU during an excpection.
 ///
@@ -125,15 +126,6 @@ pub(crate) extern "C" fn hl_exception_handler(
     let saved_rip = unsafe { (&raw const (*exn_info).rip).read_volatile() };
     let error_code = unsafe { (&raw const (*exn_info).error_code).read_volatile() };
 
-    let msg = format!(
-        "Exception vector: {:#}\n\
-         Faulting Instruction: {:#x}\n\
-         Page Fault Address: {:#x}\n\
-         Error code: {:#x}\n\
-         Stack Pointer: {:#x}",
-        exception_number, saved_rip, page_fault_address, error_code, stack_pointer
-    );
-
     // Check for registered user handlers (only for architecture-defined vectors 0-30)
     if exception_number < 31 {
         let handler =
@@ -149,10 +141,24 @@ pub(crate) extern "C" fn hl_exception_handler(
         }
     }
 
-    unsafe {
-        abort_with_code_and_message(
-            &[ErrorCode::GuestError as u8, exception as u8],
-            msg.as_ptr() as *const c_char,
-        );
+    // begin abort sequence by writing the error code
+    let mut w = HyperlightAbortWriter;
+    write_abort(&[ErrorCode::GuestError as u8, exception as u8]);
+    let write_res = write!(
+        w,
+        "Exception vector: {}\n\
+         Faulting Instruction: {:#x}\n\
+         Page Fault Address: {:#x}\n\
+         Error code: {:#x}\n\
+         Stack Pointer: {:#x}",
+        exception_number, saved_rip, page_fault_address, error_code, stack_pointer
+    );
+    if write_res.is_err() {
+        write_abort("exception message format failed".as_bytes());
     }
+
+    write_abort(&[0xFF]);
+    // At this point, write_abort with the 0xFF terminator is expected to terminate guest execution,
+    // so control should never reach beyond this call.
+    unreachable!();
 }

src/hyperlight_host/tests/integration_test.rs

@@ -1668,6 +1668,44 @@ fn exception_handler_installation_and_validation() {
     assert_eq!(count, 2, "Handler should have been called twice");
 }
 
+/// Tests that an exception can be properly handled even when the heap is exhausted.
+/// The guest function fills the heap completely, then triggers a ud2 exception.
+/// This validates that the exception handling path does not require heap allocations.
+#[test]
+fn fill_heap_and_cause_exception() {
+    let mut sandbox: MultiUseSandbox = new_uninit_rust().unwrap().evolve().unwrap();
+    let result = sandbox.call::<()>("FillHeapAndCauseException", ());
+
+    // The call should fail with an exception error since there's no handler installed
+    assert!(result.is_err(), "Expected an error from ud2 exception");
+
+    let err = result.unwrap_err();
+    match &err {
+        HyperlightError::GuestAborted(code, message) => {
+            assert_eq!(*code, ErrorCode::GuestError as u8, "Full error: {:?}", err);
+
+            // Verify the message was properly formatted (proves no-allocation path worked)
+            // Exception vector 6 is #UD (Invalid Opcode from ud2 instruction)
+            assert!(
+                message.contains("Exception vector: 6"),
+                "Message should contain 'Exception vector: 6'\nFull error: {:?}",
+                err
+            );
+            assert!(
+                message.contains("Faulting Instruction:"),
+                "Message should contain 'Faulting Instruction:'\nFull error: {:?}",
+                err
+            );
+            assert!(
+                message.contains("Stack Pointer:"),
+                "Message should contain 'Stack Pointer:'\nFull error: {:?}",
+                err
+            );
+        }
+        _ => panic!("Expected GuestAborted error, got: {:?}", err),
+    }
+}
+
 /// This test is "likely" to catch a race condition where WHvCancelRunVirtualProcessor runs halfway, then the partition is deleted (by drop calling WHvDeletePartition),
 /// and WHvCancelRunVirtualProcessor continues, and tries to access freed memory.
 ///

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

@@ -336,6 +336,19 @@ fn call_malloc(size: i32) -> i32 {
     size
 }
 
+#[guest_function("FillHeapAndCauseException")]
+fn fill_heap_and_cause_exception() {
+    let layout: Layout = Layout::new::<u8>();
+    let mut ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
+    while !ptr.is_null() {
+        black_box(ptr);
+        ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
+    }
+
+    // trigger an undefined instruction exception
+    unsafe { core::arch::asm!("ud2") };
+}
+
 #[guest_function("ExhaustHeap")]
 fn exhaust_heap() {
     let layout: Layout = Layout::new::<u8>();