Complete macOS aarch64 support

.github/workflows/ci.yml

@@ -17,6 +17,7 @@ jobs:
       matrix:
         os: [
           ubuntu-latest,
+          macos-latest,
         ]
     steps:
     - uses: actions/checkout@v4

Cargo.toml

@@ -15,5 +15,8 @@ categories = ["rust-patterns", "no-std"]
 [target.'cfg(target_os = "linux")'.dependencies]
 libc = { version = "0.2", default-features = false }
 
+[target.'cfg(target_os = "macos")'.dependencies]
+mach2 = "0.4"
+
 [dev-dependencies]
 trybuild = "1"

src/arch/aarch64.rs

@@ -15,8 +15,7 @@ pub fn arch_jump_entry_instruction(
         JumpLabelType::Jmp => {
             // Note that aarch64 only supports relative address within +/-128MB.
             // In current implementation, this assumption is always hold.
-            let relative_addr =
-                (jump_entry.target_addr() - (jump_entry.code_addr() + ARCH_JUMP_INS_LENGTH)) as u32;
+            let relative_addr = (jump_entry.target_addr() - jump_entry.code_addr()) as u32;
             let [a, b, c, d] = (relative_addr / 4).to_ne_bytes();
             [a, b, c, d | 0b00010100]
         }

src/code_manipulate.rs

@@ -2,30 +2,19 @@
 //!
 //! Since we need to make the code region writable and restore it during jump entry update,
 //! we need to provide utility functions here.
-//!
-//! For `std` environment, we can directly use [`NixCodeManipulator`] here, which utilizes
-//! [`nix`] to manipulate memory protection with `mprotect`. For `no_std` environment, there
-//! are either no memory protection mechanism or complicated memory protections, so implement
-//! it you self. :)
 
 /// Manipulate memory protection in code region.
 pub trait CodeManipulator {
-    /// Mark the code region starting at `addr` with `length` writable.
+    /// Write `data` as code instruction to `addr`.
     ///
     /// The `addr` is not aligned, you need to align it you self. The length is not too long, usually
     /// 5 bytes.
-    unsafe fn mark_code_region_writable(addr: *const core::ffi::c_void, length: usize) -> Self;
-    /// Restore the code region protection after the instruction has been updated.
-    unsafe fn restore_code_region_protect(&self);
+    unsafe fn write_code<const L: usize>(addr: *mut core::ffi::c_void, data: &[u8; L]);
 }
 
 /// Dummy code manipulator. Do nothing. Used to declare a dummy static key which is never modified
 pub(crate) struct DummyCodeManipulator;
 
 impl CodeManipulator for DummyCodeManipulator {
-    unsafe fn mark_code_region_writable(_addr: *const core::ffi::c_void, _length: usize) -> Self {
-        Self
-    }
-
-    unsafe fn restore_code_region_protect(&self) {}
+    unsafe fn write_code<const L: usize>(_addr: *mut core::ffi::c_void, _data: &[u8; L]) {}
 }

src/lib.rs

@@ -318,16 +318,7 @@ unsafe fn jump_entry_update<M: CodeManipulator>(jump_entry: &JumpEntry, enabled:
     };
     let code_bytes = arch::arch_jump_entry_instruction(jump_label_type, jump_entry);
 
-    let manipulator = M::mark_code_region_writable(
-        jump_entry.code_addr() as *const _,
-        arch::ARCH_JUMP_INS_LENGTH,
-    );
-    core::ptr::copy_nonoverlapping(
-        code_bytes.as_ptr(),
-        jump_entry.code_addr() as usize as *mut u8,
-        arch::ARCH_JUMP_INS_LENGTH,
-    );
-    manipulator.restore_code_region_protect();
+    M::write_code(jump_entry.code_addr() as *mut _, &code_bytes);
 }
 
 // ---------------------------- Use ----------------------------

src/os/linux.rs

@@ -25,20 +25,17 @@ extern "Rust" {
 }
 
 /// Arch-specific [`CodeManipulator`] using [`libc`] with `mprotect`.
-pub struct ArchCodeManipulator {
-    /// Aligned addr
-    addr: *mut core::ffi::c_void,
-    /// Aligned length
-    length: usize,
-}
+pub struct ArchCodeManipulator;
 
 impl CodeManipulator for ArchCodeManipulator {
-    unsafe fn mark_code_region_writable(addr: *const core::ffi::c_void, length: usize) -> Self {
+    /// Due to limitation of Linux, we cannot get the original memory protection flags easily
+    /// without parsing `/proc/[pid]/maps`. As a result, we just make the code region non-writable.
+    unsafe fn write_code<const L: usize>(addr: *mut core::ffi::c_void, data: &[u8; L]) {
         // TODO: page_size can be initialized once
         let page_size = unsafe { libc::sysconf(libc::_SC_PAGESIZE) as usize };
         let aligned_addr_val = (addr as usize) / page_size * page_size;
         let aligned_addr = aligned_addr_val as *mut core::ffi::c_void;
-        let aligned_length = if (addr as usize) + length - aligned_addr_val > page_size {
+        let aligned_length = if (addr as usize) + L - aligned_addr_val > page_size {
             page_size * 2
         } else {
             page_size
@@ -53,17 +50,14 @@ impl CodeManipulator for ArchCodeManipulator {
         if res != 0 {
             panic!("Unable to make code region writable");
         }
-        Self {
-            addr: aligned_addr,
-            length: aligned_length,
-        }
-    }
-
-    /// Due to limitation of Linux, we cannot get the original memory protection flags easily
-    /// without parsing `/proc/[pid]/maps`. As a result, we just make the code region non-writable.
-    unsafe fn restore_code_region_protect(&self) {
-        let res =
-            unsafe { libc::mprotect(self.addr, self.length, libc::PROT_READ | libc::PROT_EXEC) };
+        core::ptr::copy_nonoverlapping(data.as_ptr(), addr.cast(), L);
+        let res = unsafe {
+            libc::mprotect(
+                aligned_addr,
+                aligned_length,
+                libc::PROT_READ | libc::PROT_EXEC,
+            )
+        };
         if res != 0 {
             panic!("Unable to restore code region to non-writable");
         }

src/os/macos.rs

@@ -1,6 +1,6 @@
 //! macOS-specific implementations
 
-use crate::JumpEntry;
+use crate::{code_manipulate::CodeManipulator, JumpEntry};
 
 // See https://developer.apple.com/library/archive/documentation/DeveloperTools/Reference/Assembler/040-Assembler_Directives/asm_directives.html#//apple_ref/doc/uid/TP30000823-CJBIFBJG
 /// Name and attribute of section storing jump entries
@@ -21,3 +21,104 @@ extern "Rust" {
     #[link_name = "\x01section$end$__DATA$__static_keys"]
     pub static mut JUMP_ENTRY_STOP: JumpEntry;
 }
+
+extern "C" {
+    // libkern/OSCacheControl.h
+    // void	sys_dcache_flush( void *start, size_t len) __OSX_AVAILABLE_STARTING(__MAC_10_5, __IPHONE_2_0);
+    fn sys_dcache_flush(start: *mut core::ffi::c_void, len: usize);
+
+    // libkern/OSCacheControl.h
+    // void	sys_icache_invalidate( void *start, size_t len) __OSX_AVAILABLE_STARTING(__MAC_10_5, __IPHONE_2_0);
+    fn sys_icache_invalidate(start: *mut core::ffi::c_void, len: usize);
+}
+
+/// Arch-specific [`CodeManipulator`] using `mach_vm_remap` to remap the code page
+/// to a writable page, and remap back to bypass the W xor X rule.
+pub struct ArchCodeManipulator;
+
+// From mach/vm_statistics.h
+/// Return address of target data, rather than base of page
+const VM_FLAGS_RETURN_DATA_ADDR: i32 = 0x00100000;
+
+impl CodeManipulator for ArchCodeManipulator {
+    // See https://stackoverflow.com/a/76552040/10005095
+    unsafe fn write_code<const L: usize>(addr: *mut core::ffi::c_void, data: &[u8; L]) {
+        let mut remap_addr = 0;
+        let mut cur_prot = 0;
+        let mut max_prot = 0;
+        let self_task = mach2::traps::mach_task_self();
+        let length = L as u64;
+
+        // 1. Remap the page somewhere else
+        let ret = mach2::vm::mach_vm_remap(
+            self_task,
+            &mut remap_addr,
+            length,
+            0,
+            mach2::vm_statistics::VM_FLAGS_ANYWHERE | VM_FLAGS_RETURN_DATA_ADDR,
+            self_task,
+            addr as u64,
+            0,
+            &mut cur_prot,
+            &mut max_prot,
+            mach2::vm_inherit::VM_INHERIT_NONE,
+        );
+        if ret != mach2::kern_return::KERN_SUCCESS {
+            panic!("mach_vm_remap to new failed");
+        }
+
+        // 2. Reprotect the page to rw- (needs VM_PROT_COPY because the max protection is currently r-x)
+        let ret = mach2::vm::mach_vm_protect(
+            self_task,
+            remap_addr,
+            length,
+            0,
+            mach2::vm_prot::VM_PROT_READ
+                | mach2::vm_prot::VM_PROT_WRITE
+                | mach2::vm_prot::VM_PROT_COPY,
+        );
+        if ret != mach2::kern_return::KERN_SUCCESS {
+            panic!("mach_vm_protect to write failed");
+        }
+
+        // 3. Write the changes
+        core::ptr::copy_nonoverlapping(data.as_ptr(), remap_addr as *mut _, L);
+
+        // 4. Flush the data cache
+        sys_dcache_flush(addr, L);
+
+        // 5. Reprotect the page to r-x
+        let ret = mach2::vm::mach_vm_protect(
+            self_task,
+            remap_addr,
+            length,
+            0,
+            mach2::vm_prot::VM_PROT_READ | mach2::vm_prot::VM_PROT_EXECUTE,
+        );
+        if ret != mach2::kern_return::KERN_SUCCESS {
+            panic!("mach_vm_protect to execute failed");
+        }
+
+        // 6. Invalidate the instruction cache
+        sys_icache_invalidate(addr, L);
+
+        // 7. Remap the page back over the original
+        let mut origin_addr = addr as u64;
+        let ret = mach2::vm::mach_vm_remap(
+            self_task,
+            &mut origin_addr,
+            length,
+            0,
+            mach2::vm_statistics::VM_FLAGS_OVERWRITE | VM_FLAGS_RETURN_DATA_ADDR,
+            self_task,
+            remap_addr,
+            0,
+            &mut cur_prot,
+            &mut max_prot,
+            mach2::vm_inherit::VM_INHERIT_NONE,
+        );
+        if ret != mach2::kern_return::KERN_SUCCESS {
+            panic!("mach_vm_remap to origin failed");
+        }
+    }
+}