Support AOT rewriting for AArch64 Linux guests on macOS

litebox_common_linux/src/loader.rs

@@ -72,6 +72,8 @@ struct TrampolineInfo {
 /// The magic number used to identify the LiteBox trampoline.
 /// This must match `TRAMPOLINE_MAGIC` in `litebox_syscall_rewriter`.
 const TRAMPOLINE_MAGIC: u64 = u64::from_le_bytes(*b"LITEBOX0");
+/// This must match `litebox_syscall_rewriter::TRAMPOLINE_FILE_ALIGNMENT`.
+const TRAMPOLINE_FILE_ALIGNMENT: u64 = 4096;
 
 /// Trampoline header for 64-bit: 8 (magic) + 8 (file_offset) + 8 (vaddr) + 8 (size) = 32 bytes
 #[repr(C, packed)]
@@ -357,8 +359,8 @@ impl ElfParsedFile {
             return Ok(());
         }
 
-        // Verify the file offset is page-aligned (as required by the rewriter)
-        if !file_offset.is_multiple_of(PAGE_SIZE as u64) {
+        // Verify the rewriter-defined file alignment.
+        if !file_offset.is_multiple_of(TRAMPOLINE_FILE_ALIGNMENT) {
             return Err(ElfParseError::BadTrampoline);
         }
 
@@ -577,6 +579,10 @@ impl ElfParsedFile {
             info.brk = info.brk.max(trampoline_end);
             return Ok(());
         }
+        debug_assert!(
+            trampoline.file_offset.is_multiple_of(PAGE_SIZE as u64),
+            "non-populating loaders map the trampoline directly from its file offset"
+        );
         mapper
             .map_file(
                 trampoline_start,

litebox_runner_linux_on_macos_userland/Cargo.toml

@@ -15,6 +15,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 
 [target.'cfg(all(target_os = "macos", target_arch = "aarch64"))'.dev-dependencies]
 tempfile = "3"
+litebox_syscall_rewriter = { path = "../litebox_syscall_rewriter", version = "0.1.0", default-features = false }
 
 [lints]
 workspace = true

litebox_runner_linux_on_macos_userland/tests/loader.rs

@@ -3,10 +3,11 @@
 
 #![cfg(all(target_os = "macos", target_arch = "aarch64"))]
 
+use litebox_syscall_rewriter::{RewriteOptions, TargetHost, hook_syscalls_in_elf_with_options};
 use std::{path::Path, process::Command};
 
 // Prebuilt AArch64 Linux programs with their dynamic loader and glibc.
-fn run_program(name: &str) {
+fn run_program(name: &str, aot: bool) {
     let fixtures = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/test-bins");
     let directory = tempfile::tempdir().unwrap();
     let root = directory.path().join("root");
@@ -18,6 +19,16 @@ fn run_program(name: &str) {
         let path = root.join(destination);
         std::fs::create_dir_all(path.parent().unwrap()).unwrap();
         std::fs::copy(fixtures.join(source), &path).unwrap();
+        if aot {
+            let original = std::fs::read(&path).unwrap();
+            let rewritten = hook_syscalls_in_elf_with_options(
+                &original,
+                None,
+                RewriteOptions::new(TargetHost::MacOs, true),
+            )
+            .unwrap_or_else(|error| panic!("rewriting {source}: {error}"));
+            std::fs::write(&path, rewritten).unwrap();
+        }
     }
     let archive = directory.path().join("root.tar");
     let tar = Command::new("tar")
@@ -54,7 +65,7 @@ fn run_program(name: &str) {
         assert_eq!(
             output.status.code(),
             Some(0),
-            "{name} (from_tar={from_tar}): {}\nstdout: {}\nstderr: {}",
+            "{name} (aot={aot}, from_tar={from_tar}): {}\nstdout: {}\nstderr: {}",
             output.status,
             String::from_utf8_lossy(&output.stdout),
             String::from_utf8_lossy(&output.stderr)
@@ -64,5 +75,10 @@ fn run_program(name: &str) {
 
 #[test]
 fn test_load_exec_dynamic() {
-    run_program("hello_world_dyn");
+    run_program("hello_world_dyn", false);
+}
+
+#[test]
+fn test_syscall_rewriter() {
+    run_program("hello_world_dyn", true);
 }

litebox_runner_linux_on_macos_userland/tests/runner/gates.rs

@@ -2,6 +2,8 @@
 // Licensed under the MIT license.
 
 use super::*;
+use litebox_syscall_rewriter::{RewriteOptions, TargetHost, hook_syscalls_in_elf_with_options};
+
 // Exercises TLS, syscalls, x18 spills, SP writeback, and ADR/BLR through x18.
 const X18: &[u32] = &[
     0xd2824692, 0xd2822229, 0xd51bd049, 0xd2801588, 0xd4000001, 0xd53bd04a, 0xeb0a013f, 0x54000221,
@@ -33,10 +35,20 @@ fn guest_signal_return_restores_x18_and_vector_state() {
     );
 }
 
-#[test]
-fn x18_gates_preserve_registers_and_branch_targets() {
+fn run_x18_fixture(aot: bool) {
     let fixture = Fixture::new();
-    std::fs::write(fixture.0.join("program"), elf(X18)).unwrap();
+    let original = elf(X18);
+    let code = if aot {
+        hook_syscalls_in_elf_with_options(
+            &original,
+            None,
+            RewriteOptions::new(TargetHost::MacOs, false),
+        )
+        .unwrap()
+    } else {
+        original
+    };
+    std::fs::write(fixture.0.join("program"), code).unwrap();
     let output = fixture.run(&[]);
     assert_eq!(
         output.status.code(),
@@ -45,3 +57,13 @@ fn x18_gates_preserve_registers_and_branch_targets() {
         String::from_utf8_lossy(&output.stderr)
     );
 }
+
+#[test]
+fn runtime_x18_gates_preserve_registers_and_branch_targets() {
+    run_x18_fixture(false);
+}
+
+#[test]
+fn aot_x18_gates_preserve_registers_and_branch_targets() {
+    run_x18_fixture(true);
+}

litebox_syscall_rewriter/src/lib.rs

@@ -127,6 +127,9 @@ const BUN_FOOTER_MARKER: &[u8] = b"\n---- Bun! ----\n";
 /// This is checked by the loader to verify that the trampoline is valid.
 pub const TRAMPOLINE_MAGIC: &[u8; 8] = b"LITEBOX0";
 
+/// Required file alignment of the appended trampoline payload.
+pub const TRAMPOLINE_FILE_ALIGNMENT: usize = 0x1000;
+
 /// Host operating system for AArch64 guest rewriting.
 #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
 pub enum TargetHost {
@@ -929,12 +932,21 @@ fn append_trampoline_footer(
     header_vaddr: u64,
     align_trampoline_size: bool,
 ) {
-    let remain = out.len() % 0x1000;
-    out.extend_from_slice(&vec![0; if remain == 0 { 0 } else { 0x1000 - remain }]);
+    let remain = out.len() % TRAMPOLINE_FILE_ALIGNMENT;
+    out.extend_from_slice(&vec![
+        0;
+        if remain == 0 {
+            0
+        } else {
+            TRAMPOLINE_FILE_ALIGNMENT - remain
+        }
+    ]);
 
     let trampoline_file_offset = out.len() as u64;
     if align_trampoline_size {
-        let trampoline_size = trampoline_data.len().next_multiple_of(0x1000);
+        let trampoline_size = trampoline_data
+            .len()
+            .next_multiple_of(TRAMPOLINE_FILE_ALIGNMENT);
         trampoline_data.extend_from_slice(&vec![0; trampoline_size - trampoline_data.len()]);
     }
     let trampoline_size = trampoline_data.len();
@@ -1129,7 +1141,10 @@ fn is_already_hooked(input_binary: &[u8], arch: Arch) -> bool {
         return true;
     }
 
-    if file_offset % 0x1000 != 0 {
+    if usize::try_from(file_offset)
+        .ok()
+        .is_none_or(|offset| !offset.is_multiple_of(TRAMPOLINE_FILE_ALIGNMENT))
+    {
         return false;
     }
     if vaddr % 0x1000 != 0 {