Skip to content

[NFC] Remove duplication across allocate jit memory windows for x64 and arm64 #129

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lucianopa-msft wants to merge 1 commit into
base: main
Choose a base branch
from
+54 −94
Open

[NFC] Remove duplication across allocate jit memory windows for x64 and arm64 #129

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 54 additions & 94 deletions src/injector_core/common.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -158,107 +158,14 @@ 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.
///
/// For AArch64, memory must be within ±128MB due to instruction encoding limits (e.g., B/BL).
/// 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 {
Expand All @@ -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.
Expand Down
Loading