Add test and fix alignement

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

Author: jsturtevant

src/hyperlight_guest_bin/src/exceptions/handler.rs

@@ -21,46 +21,171 @@ use hyperlight_common::flatbuffer_wrappers::guest_error::ErrorCode;
 use hyperlight_common::outb::Exception;
 use hyperlight_guest::exit::abort_with_code_and_message;
 
+/// Error type for handler installation failures
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum InstallError {
+    /// Exception vector must be in range 0-30 (architecture-defined exceptions only)
+    InvalidVector,
+    /// A handler is already installed for this vector
+    HandlerAlreadyInstalled,
+}
+
+/// Exception information pushed onto the stack by the Hyperlight exception handler.
+///
 /// See AMD64 Architecture Programmer's Manual, Volume 2
 ///     §8.9.3 Interrupt Stack Frame, pp. 283--284
 ///       Figure 8-14: Long-Mode Stack After Interrupt---Same Privilege,
 ///       Figure 8-15: Long-Mode Stack After Interrupt---Higher Privilege
-/// Subject to the proviso that we push a dummy error code of 0 for exceptions
-/// for which the processor does not provide one
+/// Note: For exceptions that don't provide an error code, we push a dummy value of 0.
 #[repr(C)]
 pub struct ExceptionInfo {
+    /// Error code provided by the processor (or 0 if not applicable).
     pub error_code: u64,
+    /// Instruction pointer at the time of the exception.
     pub rip: u64,
+    /// Code segment selector.
     pub cs: u64,
+    /// CPU flags register.
     pub rflags: u64,
+    /// Stack pointer at the time of the exception.
     pub rsp: u64,
+    /// Stack segment selector.
     pub ss: u64,
 }
 const _: () = assert!(core::mem::offset_of!(ExceptionInfo, rip) == 8);
 const _: () = assert!(core::mem::offset_of!(ExceptionInfo, rsp) == 32);
 
+/// Saved CPU context pushed onto the stack by exception entry code.
+///
+/// This structure contains all the saved CPU state needed to resume execution
+/// after handling an exception. It includes segment registers, floating-point state,
+/// and general-purpose registers.
 #[repr(C)]
-/// Saved context, pushed onto the stack by exception entry code
 pub struct Context {
-    /// in order: gs, fs, es
-    pub segments: [u64; 3],
+    /// Segment registers in order: GS, FS, ES, DS.
+    pub segments: [u64; 4],
+    /// FPU/SSE state saved via FXSAVE instruction (512 bytes).
     pub fxsave: [u8; 512],
-    pub ds: u64,
-    /// no `rsp`, since the processor saved it
-    /// `rax` is at the top, `r15` the bottom
+    /// General-purpose registers (RAX through R15, excluding RSP).
+    ///
+    /// The stack pointer (RSP) is not included here since it's saved
+    /// by the processor in the `ExceptionInfo` structure.
+    /// RAX is at index 0, R15 is at index 14.
     pub gprs: [u64; 15],
+    /// Padding to ensure 16-byte alignment when combined with ExceptionInfo.
+    padding: [u64; 1],
 }
-const _: () = assert!(size_of::<Context>() == 152 + 512);
+const _: () = assert!(size_of::<Context>() == 32 + 512 + 120 + 8);
+const _: () = assert!((size_of::<Context>() + size_of::<ExceptionInfo>()) % 16 == 0);
 
-// TODO: This will eventually need to end up in a per-thread context,
-// when there are threads.
-pub static HANDLERS: [core::sync::atomic::AtomicU64; 31] =
+/// Array of installed exception handlers for vectors 0-30.
+///
+/// TODO: This will eventually need to be part of a per-thread context when threading is implemented.
+static HANDLERS: [core::sync::atomic::AtomicU64; 31] =
     [const { core::sync::atomic::AtomicU64::new(0) }; 31];
-pub type HandlerT = fn(n: u64, info: *mut ExceptionInfo, ctx: *mut Context, pf_addr: u64) -> bool;
 
-/// Exception handler
+/// Exception handler function type.
+///
+/// Handlers receive mutable pointers to the exception information and CPU context,
+/// allowing direct access and modification of exception state.
+///
+/// # Parameters
+/// * `exception_number` - Exception vector number (0-30)
+/// * `exception_info` - Mutable pointer to exception information (instruction pointer, error code, etc.)
+/// * `context` - Mutable pointer to saved CPU context (registers, FPU state, etc.)
+/// * `page_fault_address` - Page fault address (only valid for page fault exceptions)
+///
+/// # Returns
+/// * `true` - Suppress the default abort behavior and continue execution
+/// * `false` - Allow the default abort to occur
+///
+/// # Safety
+/// This function type uses raw mutable pointers. Handlers must ensure:
+/// - Pointers are valid for the duration of the handler
+/// - Any modifications to exception state maintain system integrity
+/// - Modified values are valid for CPU state (e.g., valid instruction pointers, aligned stack pointers)
+pub type ExceptionHandler = fn(
+    exception_number: u64,
+    exception_info: *mut ExceptionInfo,
+    context: *mut Context,
+    page_fault_address: u64,
+) -> bool;
+
+/// Install a custom exception handler for a specific vector.
+///
+/// # Arguments
+/// * `vector` - Exception vector (0-30). Must be an architecture-defined exception.
+/// * `handler` - The handler function to invoke when this exception occurs.
+///
+/// # Returns
+/// * `Ok(())` if the handler was successfully installed.
+/// * `Err(InstallError::InvalidVector)` if `vector >= 31`.
+/// * `Err(InstallError::HandlerAlreadyInstalled)` if a handler is already registered for this vector.
+///
+/// # Example
+/// ```ignore
+/// fn my_exception_handler(
+///     exception_number: u64,
+///     exception_info: *mut ExceptionInfo,
+///     context: *mut Context,
+///     page_fault_address: u64,
+/// ) -> bool {
+///     unsafe {
+///         // Read the faulting instruction pointer
+///         let faulting_rip = core::ptr::read_volatile(&(*exception_info).rip);
+///         
+///         // Save the original RIP to R9 register
+///         core::ptr::write_volatile(&mut (*context).gprs[8], faulting_rip);
+///         
+///         // Skip past the faulting instruction
+///         core::ptr::write_volatile(&mut (*exception_info).rip, faulting_rip + 2);
+///     }
+///     
+///     true // Return true to suppress abort and continue execution
+/// }
+///
+/// install_handler(3, my_exception_handler)?;  // Install for INT3 (breakpoint)
+/// ```
+pub fn install_handler(vector: u8, handler: ExceptionHandler) -> Result<(), InstallError> {
+    if vector >= 31 {
+        return Err(InstallError::InvalidVector);
+    }
+
+    // Use compare_exchange to atomically check and set, preventing races
+    match HANDLERS[vector as usize].compare_exchange(
+        0,
+        handler as usize as u64,
+        core::sync::atomic::Ordering::AcqRel,
+        core::sync::atomic::Ordering::Acquire,
+    ) {
+        Ok(_) => Ok(()),
+        Err(_) => Err(InstallError::HandlerAlreadyInstalled),
+    }
+}
+
+/// Remove a custom exception handler for a specific vector.
+///
+/// # Arguments
+/// * `vector` - Exception vector (0-30).
+///
+/// # Returns
+/// * `Ok(())` if the handler was successfully removed (or was already uninstalled).
+/// * `Err(InstallError::InvalidVector)` if `vector >= 31`.
+pub fn uninstall_handler(vector: u8) -> Result<(), InstallError> {
+    if vector >= 31 {
+        return Err(InstallError::InvalidVector);
+    }
+
+    HANDLERS[vector as usize].store(0, core::sync::atomic::Ordering::Release);
+    Ok(())
+}
+
+/// Internal exception handler invoked by the low-level exception entry code.
+///
+/// This function is called from assembly when an exception occurs. It checks for
+/// registered user handlers and either invokes them or aborts with an error message.
 #[unsafe(no_mangle)]
-pub extern "C" fn hl_exception_handler(
+pub(crate) extern "C" fn hl_exception_handler(
     stack_pointer: u64,
     exception_number: u64,
     page_fault_address: u64,
@@ -82,20 +207,18 @@ pub extern "C" fn hl_exception_handler(
         exception_number, saved_rip, page_fault_address, error_code, stack_pointer
     );
 
-    // We don't presently have any need for user-defined interrupts,
-    // so we only support handlers for the architecture-defined
-    // vectors (0-31)
+    // Check for registered user handlers (only for architecture-defined vectors 0-30)
     if exception_number < 31 {
         let handler =
             HANDLERS[exception_number as usize].load(core::sync::atomic::Ordering::Acquire);
-        if handler != 0
-            && unsafe {
-                core::mem::transmute::<u64, fn(u64, *mut ExceptionInfo, *mut Context, u64) -> bool>(
-                    handler,
-                )(exception_number, exn_info, ctx, page_fault_address)
-            }
-        {
-            return;
+        if handler != 0 {
+            unsafe {
+                let handler = core::mem::transmute::<u64, ExceptionHandler>(handler);
+                if handler(exception_number, exn_info, ctx, page_fault_address) {
+                    return;
+                }
+                // Handler returned false, fall through to abort
+            };
         }
     }
 

src/hyperlight_guest_bin/src/exceptions/interrupt_entry.rs

@@ -51,6 +51,8 @@ unsafe extern "C" {
 macro_rules! context_save {
     () => {
         concat!(
+            // Push padding to match Context struct (8 bytes)
+            "    push 0\n",
             // Save general-purpose registers
             "    push rax\n",
             "    push rbx\n",
@@ -67,10 +69,6 @@ macro_rules! context_save {
             "    push r13\n",
             "    push r14\n",
             "    push r15\n",
-            // Save one of the segment registers to get 16-byte alignment for
-            // FXSAVE. TODO: consider packing the segment registers
-            "    mov rax, ds\n",
-            "    push rax\n",
             // Save floating-point/SSE registers
             // TODO: Don't do this unconditionally: get the exn
             //       handlers compiled without sse
@@ -79,7 +77,9 @@ macro_rules! context_save {
             "    sub rsp, 512\n",
             "    mov rax, rsp\n",
             "    fxsave [rax]\n",
-            // Save the rest of the segment registers
+            // Save all segment registers
+            "    mov rax, ds\n",
+            "    push rax\n",
             "    mov rax, es\n",
             "    push rax\n",
             "    mov rax, fs\n",
@@ -93,20 +93,19 @@ macro_rules! context_save {
 macro_rules! context_restore {
     () => {
         concat!(
-            // Restore most segment registers
+            // Restore all segment registers
             "    pop rax\n",
             "    mov gs, rax\n",
             "    pop rax\n",
             "    mov fs, rax\n",
             "    pop rax\n",
             "    mov es, rax\n",
+            "    pop rax\n",
+            "    mov ds, rax\n",
             // Restore floating-point/SSE registers
             "    mov rax, rsp\n",
             "    fxrstor [rax]\n",
             "    add rsp, 512\n",
-            // Restore the last segment register
-            "    pop rax\n",
-            "    mov ds, rax\n",
             // Restore general-purpose registers
             "    pop r15\n",
             "    pop r14\n",
@@ -123,6 +122,8 @@ macro_rules! context_restore {
             "    pop rcx\n",
             "    pop rbx\n",
             "    pop rax\n",
+            // Skip padding (8 bytes)
+            "    add rsp, 8\n",
         )
     };
 }
@@ -178,6 +179,43 @@ macro_rules! generate_exceptions {
 //     mov rdx, 0
 //     jmp _do_excp_common
 // ```
+//
+// Stack layout after context_save!() (from high to low addresses):
+// ```
+// +------------------+ <-- Higher addresses
+// | SS               | (Pushed by CPU on exception)
+// | RSP              | (Pushed by CPU on exception)
+// | RFLAGS           | (Pushed by CPU on exception)
+// | CS               | (Pushed by CPU on exception)
+// | RIP              | (Pushed by CPU on exception)
+// | Error Code       | (Pushed by CPU or by handler)
+// +------------------+ <-- ExceptionInfo struct starts here
+// | Padding (8)      | (Pushed by context_save!)
+// +------------------+
+// | R15              | gprs[14]
+// | R14              | gprs[13]
+// | R13              | gprs[12]
+// | R12              | gprs[11]
+// | R11              | gprs[10]
+// | R10              | gprs[9]
+// | R9               | gprs[8]
+// | R8               | gprs[7]
+// | RBP              | gprs[6]
+// | RDI              | gprs[5]
+// | RSI              | gprs[4]
+// | RDX              | gprs[3]
+// | RCX              | gprs[2]
+// | RBX              | gprs[1]
+// | RAX              | gprs[0] (15 GPRs total, 120 bytes)
+// +------------------+
+// | FXSAVE area      | (512 bytes for FPU/SSE state)
+// +------------------+
+// | GS               | segments[3]
+// | FS               | segments[2]
+// | ES               | segments[1]
+// | DS               | segments[0] (4 segment registers, 32 bytes)
+// +------------------+ <-- Context struct starts here
+// ```
 macro_rules! generate_excp {
     ($num:expr) => {
         concat!(

src/hyperlight_host/tests/common/mod.rs

@@ -14,9 +14,9 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 use hyperlight_host::func::HostFunction;
-use hyperlight_host::{GuestBinary, MultiUseSandbox, Result, UninitializedSandbox};
 #[cfg(gdb)]
 use hyperlight_host::sandbox::config::DebugInfo;
+use hyperlight_host::{GuestBinary, MultiUseSandbox, Result, UninitializedSandbox};
 use hyperlight_testing::{c_simple_guest_as_string, simple_guest_as_string};
 
 /// Returns a rust/c simpleguest depending on environment variable GUEST.
@@ -37,13 +37,13 @@ pub fn new_uninit_rust() -> Result<UninitializedSandbox> {
         let mut cfg = SandboxConfiguration::default();
         let debug_info = DebugInfo { port: 8080 };
         cfg.set_guest_debug_info(debug_info);
-        
+
         UninitializedSandbox::new(
             GuestBinary::FilePath(simple_guest_as_string().unwrap()),
             Some(cfg),
         )
     }
-    
+
     #[cfg(not(gdb))]
     UninitializedSandbox::new(
         GuestBinary::FilePath(simple_guest_as_string().unwrap()),