Add AArch64 guest exception callback

CHANGELOG.md

@@ -5,6 +5,8 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
 ## [Prerelease] - Unreleased
 
 ### Added
+* AArch64 guests can register an exception callback to update ELR and saved
+  general-purpose registers before Hyperlight aborts.
 * `SandboxBuilder`, the entry point for creating a sandbox. It gathers machine
   configuration, host functions, init data and memory mappings, then builds a
   `MultiUseSandbox` from a guest binary on disk, a guest binary in memory, or a

src/hyperlight_guest_bin/src/arch/aarch64/exception/handle.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // Copyright 2026 The Hyperlight Authors.
 use core::fmt::Write;
+use core::sync::atomic::{AtomicU64, Ordering};
 
 use hyperlight_common::arch::exn::{DataFault, DataFaultKind, Exception, decode_syndrome};
 use hyperlight_common::vmem::{
@@ -10,10 +11,50 @@ use hyperlight_guest::error::ErrorCode;
 use hyperlight_guest::exit::write_abort;
 use hyperlight_guest::layout::{MAIN_STACK_LIMIT_GVA, MAIN_STACK_TOP_GVA};
 
-use super::super::mrs;
+use super::super::{mrs, msr};
 use super::types::*;
 use crate::HyperlightAbortWriter;
 
+/// Callback for synchronous AArch64 exceptions not handled by Hyperlight.
+///
+/// The arguments are the decoded exception, raw ESR_EL1, FAR_EL1, mutable
+/// ELR_EL1, and saved x0 through x30. Return `true` to resume execution.
+/// Call [`crate::paging::barrier::first_valid_same_ctx`] after mapping a page
+/// for a translation fault.
+pub type ExceptionHandler = fn(Exception, u64, u64, &mut u64, &mut [u64; 31]) -> bool;
+
+static HANDLER: AtomicU64 = AtomicU64::new(0);
+
+/// Register the callback for synchronous AArch64 exceptions.
+///
+/// A new callback replaces any existing callback.
+pub fn register_exception_handler(handler: ExceptionHandler) {
+    HANDLER.store(handler as usize as u64, Ordering::Release);
+}
+
+/// Remove the registered AArch64 exception callback.
+pub fn unregister_exception_handler() {
+    HANDLER.store(0, Ordering::Release);
+}
+
+fn try_handle_registered_exception(
+    exn: Exception,
+    esr: u64,
+    far: u64,
+    elr: &mut u64,
+    registers: &mut [u64; 31],
+) -> bool {
+    let handler = HANDLER.load(Ordering::Acquire);
+    if handler == 0 {
+        return false;
+    }
+
+    // SAFETY: HANDLER contains only function pointers stored by
+    // register_exception_handler.
+    let handler = unsafe { core::mem::transmute::<u64, ExceptionHandler>(handler) };
+    handler(exn, esr, far, elr, registers)
+}
+
 fn handle_stack_fault(far: u64) {
     // TODO: perhaps we should have a sanity check that the
     // stack grows only one page at a time, which should be
@@ -124,7 +165,7 @@ fn handle_internal_fault(exn: Exception, far: u64) -> bool {
 pub(super) extern "C" fn handle_exception(
     typ: ExceptionType,
     from: ExceptionFrom,
-    _regs: *mut ExceptionContext,
+    regs: *mut ExceptionContext,
 ) {
     let esr = unsafe { mrs!(ESR_EL1) };
     let far = unsafe { mrs!(FAR_EL1) };
@@ -134,6 +175,17 @@ pub(super) extern "C" fn handle_exception(
         if handle_internal_fault(exn, far) {
             return;
         }
+
+        let mut elr = unsafe { mrs!(ELR_EL1) };
+        // SAFETY: Exception entry passes its live, uniquely borrowed save area.
+        let registers = unsafe { &mut (*regs).x };
+        if try_handle_registered_exception(exn, esr, far, &mut elr, registers) {
+            // SAFETY: The callback supplies the resume address for eret.
+            unsafe {
+                msr!(ELR_EL1, elr);
+            }
+            return;
+        }
     }
 
     // Die with some diagnostic information

src/hyperlight_guest_bin/src/arch/aarch64/mod.rs

@@ -39,7 +39,7 @@ pub mod dispatch {
     );
 }
 
-mod exception;
+pub(crate) mod exception;
 
 macro_rules! msr {
     ($sysreg:ident, $expr:expr) => {

src/hyperlight_guest_bin/src/exception.rs

@@ -1,8 +1,18 @@
 // SPDX-License-Identifier: Apache-2.0
 // Copyright 2025 The Hyperlight Authors.
 
+#[cfg(target_arch = "x86_64")]
 pub mod arch {
     pub use crate::arch::context::Context;
     pub use crate::arch::exception::handle::HANDLERS;
     pub use crate::arch::machine::ExceptionInfo;
 }
+
+#[cfg(target_arch = "aarch64")]
+pub mod arch {
+    pub use hyperlight_common::arch::exn::Exception;
+
+    pub use crate::arch::exception::handle::{
+        ExceptionHandler, register_exception_handler, unregister_exception_handler,
+    };
+}

src/hyperlight_guest_bin/src/lib.rs

@@ -23,10 +23,7 @@ use hyperlight_guest::guest_handle::handle::GuestHandle;
 #[cfg_attr(target_arch = "x86_64", path = "arch/amd64/mod.rs")]
 #[cfg_attr(target_arch = "aarch64", path = "arch/aarch64/mod.rs")]
 mod arch;
-// temporarily expose the architecture-specific exception interface;
-// this should be replaced with something a bit more abstract in the
-// near future.
-#[cfg(target_arch = "x86_64")]
+// Temporarily expose architecture-specific exception interfaces.
 pub mod exception;
 pub mod guest_function {
     pub(super) mod call;