From ba5c7c28766b465f4ed035035b6aa7946f0832b6 Mon Sep 17 00:00:00 2001 From: Evian-Zhang Date: Fri, 26 Jul 2024 15:10:36 +0800 Subject: [PATCH] Complete macOS aarch64 support --- .github/workflows/ci.yml | 1 + Cargo.lock | 10 ++++ Cargo.toml | 3 ++ src/arch/aarch64.rs | 3 +- src/code_manipulate.rs | 17 ++----- src/lib.rs | 11 +---- src/os/linux.rs | 32 +++++------- src/os/macos.rs | 103 ++++++++++++++++++++++++++++++++++++++- 8 files changed, 134 insertions(+), 46 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9c6c0a3..8aac2a3 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,6 +17,7 @@ jobs: matrix: os: [ ubuntu-latest, + macos-latest, ] steps: - uses: actions/checkout@v4 diff --git a/Cargo.lock b/Cargo.lock index 06b92ca..e7f3f18 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -42,6 +42,15 @@ version = "0.2.155" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c" +[[package]] +name = "mach2" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19b955cdeb2a02b9117f121ce63aa52d08ade45de53e48fe6a38b39c10f6f709" +dependencies = [ + "libc", +] + [[package]] name = "memchr" version = "2.7.4" @@ -117,6 +126,7 @@ name = "static-keys" version = "0.2.0" dependencies = [ "libc", + "mach2", "trybuild", ] diff --git a/Cargo.toml b/Cargo.toml index 7b53f1d..19ae06e 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/arch/aarch64.rs b/src/arch/aarch64.rs index 0bfd8ed..49cdf51 100644 --- a/src/arch/aarch64.rs +++ b/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] } diff --git a/src/code_manipulate.rs b/src/code_manipulate.rs index 39ef426..c9a0575 100644 --- a/src/code_manipulate.rs +++ b/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(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(_addr: *mut core::ffi::c_void, _data: &[u8; L]) {} } diff --git a/src/lib.rs b/src/lib.rs index 95bb79c..2c0e74b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -318,16 +318,7 @@ unsafe fn jump_entry_update(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 ---------------------------- diff --git a/src/os/linux.rs b/src/os/linux.rs index ee3f0be..3dbe59c 100644 --- a/src/os/linux.rs +++ b/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(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"); } diff --git a/src/os/macos.rs b/src/os/macos.rs index 29eaac6..8204d3e 100644 --- a/src/os/macos.rs +++ b/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(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"); + } + } +}