From f1578c8737c0487e30a2998e9b89af49f7e28735 Mon Sep 17 00:00:00 2001 From: Luciano Passos Almeida Date: Fri, 27 Mar 2026 19:08:50 -0300 Subject: [PATCH] [NFC] Remove duplication across allocate jit memory for x64 and arm64 --- src/injector_core/common.rs | 148 +++++++++++++----------------------- 1 file changed, 54 insertions(+), 94 deletions(-) diff --git a/src/injector_core/common.rs b/src/injector_core/common.rs index fc6d3b4..74c0f35 100644 --- a/src/injector_core/common.rs +++ b/src/injector_core/common.rs @@ -158,6 +158,7 @@ fn allocate_jit_memory_unix(_src: &FuncPtrInternal, code_size: usize) -> *mut u8 ptr as *mut u8 } } + // See https://github.com/microsoft/injectorppforrust/issues/84 /// Allocate executable JIT memory on Windows platforms. /// @@ -165,100 +166,6 @@ fn allocate_jit_memory_unix(_src: &FuncPtrInternal, code_size: usize) -> *mut u8 /// For x86_64, memory must be within ±2GB for `jmp rel32` instructions. #[cfg(target_os = "windows")] fn allocate_jit_memory_windows(_src: &FuncPtrInternal, code_size: usize) -> *mut u8 { - #[cfg(target_arch = "aarch64")] - { - let max_range: u64 = 0x8000000; // ±128MB - let original_addr = _src.as_ptr() as u64; - let page_size = unsafe { get_page_size() as u64 }; - - // Search outward from the function address to find the CLOSEST free page. - let mut offset: u64 = 0; - while offset <= max_range { - for &dir in &[1i64, -1i64] { - let hint = if dir > 0 { - original_addr.checked_add(offset) - } else if offset > 0 { - original_addr.checked_sub(offset) - } else { - continue; - }; - - let Some(hint_addr) = hint else { continue }; - - let ptr = unsafe { - VirtualAlloc( - hint_addr as *mut c_void, - code_size, - MEM_COMMIT | MEM_RESERVE, - PAGE_EXECUTE_READWRITE, - ) - }; - if !ptr.is_null() { - let allocated = ptr as u64; - let diff = allocated.abs_diff(original_addr); - if diff <= max_range { - return ptr as *mut u8; - } else { - unsafe { - VirtualFree(ptr, 0, MEM_RELEASE); - } - } - } - } - offset += page_size; - } - - panic!("Failed to allocate executable memory within ±128MB of original function address on AArch64 Windows"); - } - - #[cfg(target_arch = "x86_64")] - { - let max_range: u64 = 0x8000_0000; // ±2GB - let original_addr = _src.as_ptr() as u64; - let page_size = unsafe { get_page_size() as u64 }; - - // Search outward from the function address to find the CLOSEST free page. - // This avoids allocating far from the function (e.g., in/near stack memory), - // which could disrupt the stack guard page and cause STATUS_STACK_OVERFLOW. - let mut offset: u64 = 0; - while offset <= max_range { - for &dir in &[1i64, -1i64] { - let hint = if dir > 0 { - original_addr.checked_add(offset) - } else if offset > 0 { - original_addr.checked_sub(offset) - } else { - continue; // Already tried offset=0 with dir=1 - }; - - let Some(hint_addr) = hint else { continue }; - - let ptr = unsafe { - VirtualAlloc( - hint_addr as *mut c_void, - code_size, - MEM_COMMIT | MEM_RESERVE, - PAGE_EXECUTE_READWRITE, - ) - }; - if !ptr.is_null() { - let allocated = ptr as u64; - let diff = allocated.abs_diff(original_addr); - if diff <= max_range { - return ptr as *mut u8; - } else { - unsafe { - VirtualFree(ptr, 0, MEM_RELEASE); - } - } - } - } - offset += page_size; - } - - panic!("Failed to allocate executable memory within ±2GB of original function address on x86_64 Windows"); - } - #[cfg(all(not(target_arch = "x86_64"), not(target_arch = "aarch64")))] { let ptr = unsafe { @@ -276,6 +183,59 @@ fn allocate_jit_memory_windows(_src: &FuncPtrInternal, code_size: usize) -> *mut ptr as *mut u8 } + + #[cfg(target_arch = "aarch64")] + let max_range: u64 = 0x8000000; // ±128MB + + #[cfg(target_arch = "x86_64")] + let max_range: u64 = 0x8000_0000; // ±2GB + + let original_addr = _src.as_ptr() as u64; + let page_size = unsafe { get_page_size() as u64 }; + + // Search outward from the function address to find the CLOSEST free page. + // This avoids allocating far from the function (e.g., in/near stack memory), + // which could disrupt the stack guard page and cause STATUS_STACK_OVERFLOW. + let mut offset: u64 = 0; + while offset <= max_range { + for &dir in &[1i64, -1i64] { + let hint = if dir > 0 { + original_addr.checked_add(offset) + } else if offset > 0 { + original_addr.checked_sub(offset) + } else { + continue; // Already tried offset=0 with dir=1 + }; + + let Some(hint_addr) = hint else { continue }; + + let ptr = unsafe { + VirtualAlloc( + hint_addr as *mut c_void, + code_size, + MEM_COMMIT | MEM_RESERVE, + PAGE_EXECUTE_READWRITE, + ) + }; + if !ptr.is_null() { + let allocated = ptr as u64; + let diff = allocated.abs_diff(original_addr); + if diff <= max_range { + return ptr as *mut u8; + } else { + unsafe { + VirtualFree(ptr, 0, MEM_RELEASE); + } + } + } + } + offset += page_size; + } + + panic!( + "Failed to allocate JIT memory within ±{max_range} bytes of source on {} arch", + std::env::consts::ARCH + ); } /// Unsafely reads `len` bytes from `ptr` and returns them as a Vec.