1use std::sync::atomic::{AtomicUsize, Ordering};
14use std::sync::{Mutex, PoisonError};
15
16pub const ANCHOR_SYMBOL: &str = "__frust_hotpatch_anchor";
19
20pub const UNKNOWN_IMAGE: u32 = u32::MAX;
22
23static ANCHOR: AtomicUsize = AtomicUsize::new(0);
24
25pub fn set_anchor(address: usize) {
33 ANCHOR.store(address, Ordering::Release);
34}
35
36#[doc(hidden)]
40pub fn aslr_reference() -> usize {
41 ANCHOR.load(Ordering::Acquire)
42}
43
44#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46pub(crate) struct Image {
47 pub(crate) index: u32,
48 pub(crate) start: usize,
49 pub(crate) end: usize,
50 pub(crate) slide: usize,
52}
53
54static IMAGES: Mutex<Vec<Image>> = Mutex::new(Vec::new());
55
56fn images() -> std::sync::MutexGuard<'static, Vec<Image>> {
57 IMAGES.lock().unwrap_or_else(PoisonError::into_inner)
58}
59
60pub(crate) fn register_base(start: usize, end: usize, slide: usize) {
62 let mut images = images();
63 images.retain(|image| image.index != 0);
64 images.insert(
65 0,
66 Image {
67 index: 0,
68 start,
69 end,
70 slide,
71 },
72 );
73}
74
75pub(crate) fn register_patch(start: usize, end: usize, slide: usize) -> u32 {
77 let mut images = images();
78 let index = images.iter().map(|image| image.index).max().unwrap_or(0) + 1;
79 images.push(Image {
80 index,
81 start,
82 end,
83 slide,
84 });
85 index
86}
87
88pub(crate) fn classify(address: u64) -> (u32, u64) {
91 let found = images()
92 .iter()
93 .copied()
94 .find(|image| usize::try_from(address).is_ok_and(|a| image.start <= a && a < image.end));
95 match found {
96 Some(image) => (image.index, address.wrapping_sub(image.slide as u64)),
97 None => (UNKNOWN_IMAGE, address),
98 }
99}
100
101pub(crate) fn image_range_containing(address: usize) -> Option<(usize, usize)> {
104 platform::image_containing(address).map(|(lo, hi, _)| (lo, hi))
105}
106
107pub(crate) fn image_slide_containing(address: usize) -> Option<usize> {
111 platform::image_containing(address).map(|(_, _, slide)| slide)
112}
113
114#[cfg(test)]
115pub(crate) fn reset_for_test() {
116 ANCHOR.store(0, Ordering::Release);
117 images().clear();
118}
119
120#[cfg(test)]
121pub(crate) fn register_for_test(index: u32, start: usize, end: usize, slide: usize) {
122 let mut images = images();
123 images.retain(|image| image.index != index);
124 images.push(Image {
125 index,
126 start,
127 end,
128 slide,
129 });
130}
131
132#[cfg(any(target_os = "linux", target_os = "android"))]
133mod platform {
134 use std::ffi::c_void;
135
136 struct Probe {
137 address: usize,
138 found: Option<(usize, usize, usize)>,
139 }
140
141 unsafe extern "C" fn visit(
142 info: *mut libc::dl_phdr_info,
143 _size: libc::size_t,
144 data: *mut c_void,
145 ) -> libc::c_int {
146 let (info, probe) = unsafe { (&*info, &mut *data.cast::<Probe>()) };
149 if info.dlpi_phdr.is_null() {
150 return 0;
151 }
152 let headers =
154 unsafe { std::slice::from_raw_parts(info.dlpi_phdr, usize::from(info.dlpi_phnum)) };
155 let base = info.dlpi_addr as usize;
156 let (mut lo, mut hi) = (usize::MAX, 0usize);
157 for header in headers.iter().filter(|h| h.p_type == libc::PT_LOAD) {
158 let start = base.wrapping_add(header.p_vaddr as usize);
159 lo = lo.min(start);
160 hi = hi.max(start.wrapping_add(header.p_memsz as usize));
161 }
162 if lo <= probe.address && probe.address < hi {
163 probe.found = Some((lo, hi, base));
164 return 1;
165 }
166 0
167 }
168
169 pub(super) fn image_containing(address: usize) -> Option<(usize, usize, usize)> {
170 #[cfg(target_pointer_width = "64")]
173 let address = address & 0x00FF_FFFF_FFFF_FFFF;
174 let mut probe = Probe {
175 address,
176 found: None,
177 };
178 unsafe {
180 libc::dl_iterate_phdr(Some(visit), std::ptr::from_mut(&mut probe).cast::<c_void>());
181 }
182 probe.found
183 }
184}
185
186#[cfg(all(target_vendor = "apple", target_pointer_width = "64"))]
187mod platform {
188 unsafe extern "C" {
191 fn _dyld_image_count() -> u32;
192 fn _dyld_get_image_header(index: u32) -> *const MachHeader64;
193 fn _dyld_get_image_vmaddr_slide(index: u32) -> isize;
194 }
195
196 #[repr(C)]
197 struct MachHeader64 {
198 magic: u32,
199 cpu_type: i32,
200 cpu_subtype: i32,
201 file_type: u32,
202 command_count: u32,
203 commands_size: u32,
204 flags: u32,
205 reserved: u32,
206 }
207
208 #[repr(C)]
209 struct LoadCommand {
210 cmd: u32,
211 size: u32,
212 }
213
214 #[repr(C)]
215 struct Segment64 {
216 cmd: u32,
217 size: u32,
218 name: [u8; 16],
219 vm_address: u64,
220 vm_size: u64,
221 }
222
223 const LC_SEGMENT_64: u32 = 0x19;
224
225 pub(super) fn image_containing(address: usize) -> Option<(usize, usize, usize)> {
226 let count = unsafe { _dyld_image_count() };
230 for index in 0..count {
231 let (header, slide) = unsafe {
233 (
234 _dyld_get_image_header(index),
235 _dyld_get_image_vmaddr_slide(index) as usize,
236 )
237 };
238 if header.is_null() {
239 continue;
240 }
241 let range = unsafe { segments_extent(header, slide) };
243 if let Some((lo, hi)) = range
244 && lo <= address
245 && address < hi
246 {
247 return Some((lo, hi, slide));
248 }
249 }
250 None
251 }
252
253 unsafe fn segments_extent(header: *const MachHeader64, slide: usize) -> Option<(usize, usize)> {
259 let (count, mut cursor) = unsafe { ((*header).command_count, header.add(1).cast::<u8>()) };
262 let (mut lo, mut hi) = (usize::MAX, 0usize);
263 for _ in 0..count {
264 let command = unsafe { &*cursor.cast::<LoadCommand>() };
266 if command.cmd == LC_SEGMENT_64 {
267 let segment = unsafe { &*cursor.cast::<Segment64>() };
269 if !segment.name.starts_with(b"__PAGEZERO") && segment.vm_size != 0 {
270 let start = (segment.vm_address as usize).wrapping_add(slide);
271 lo = lo.min(start);
272 hi = hi.max(start.wrapping_add(segment.vm_size as usize));
273 }
274 }
275 if command.size == 0 {
276 break;
277 }
278 cursor = unsafe { cursor.add(command.size as usize) };
280 }
281 (lo < hi).then_some((lo, hi))
282 }
283}
284
285#[cfg(windows)]
286mod platform {
287 use std::ffi::c_void;
288
289 unsafe extern "system" {
290 fn GetModuleHandleExW(flags: u32, address: *const c_void, module: *mut *mut c_void) -> i32;
291 }
292
293 const FROM_ADDRESS: u32 = 0x4;
294 const UNCHANGED_REFCOUNT: u32 = 0x2;
295
296 pub(super) fn image_containing(address: usize) -> Option<(usize, usize, usize)> {
297 let mut module: *mut c_void = std::ptr::null_mut();
298 let ok = unsafe {
301 GetModuleHandleExW(
302 FROM_ADDRESS | UNCHANGED_REFCOUNT,
303 address as *const c_void,
304 &mut module,
305 )
306 };
307 if ok == 0 || module.is_null() {
308 return None;
309 }
310 let base = module as usize;
311 let size = unsafe {
315 let e_lfanew = std::ptr::read_unaligned((base + 0x3C) as *const u32) as usize;
316 std::ptr::read_unaligned((base + e_lfanew + 24 + 56) as *const u32) as usize
317 };
318 let preferred = unsafe {
321 let e_lfanew = std::ptr::read_unaligned((base + 0x3C) as *const u32) as usize;
322 std::ptr::read_unaligned((base + e_lfanew + 24 + 24) as *const u64) as usize
323 };
324 Some((base, base + size, base.wrapping_sub(preferred)))
325 }
326}
327
328#[cfg(not(any(
329 target_os = "linux",
330 target_os = "android",
331 all(target_vendor = "apple", target_pointer_width = "64"),
332 windows
333)))]
334mod platform {
335 pub(super) fn image_containing(_address: usize) -> Option<(usize, usize, usize)> {
336 None
337 }
338}
339
340#[cfg(test)]
341mod tests {
342 use super::*;
343
344 fn here() {}
345
346 #[test]
347 fn anchor_reads_zero_until_set_and_is_not_cached_while_unset() {
348 let _serial = crate::patch::test_support::serial();
349 reset_for_test();
350 assert_eq!(aslr_reference(), 0);
351 assert_eq!(aslr_reference(), 0);
352 set_anchor(0x1234);
353 assert_eq!(aslr_reference(), 0x1234);
354 reset_for_test();
355 }
356
357 #[test]
358 fn classify_names_the_image_and_subtracts_its_slide() {
359 let _serial = crate::patch::test_support::serial();
360 reset_for_test();
361 register_for_test(0, 0x1000, 0x2000, 0x1000);
362 register_for_test(1, 0x8000, 0x9000, 0x7000);
363 assert_eq!(classify(0x1800), (0, 0x800));
364 assert_eq!(classify(0x8100), (1, 0x1100));
365 assert_eq!(classify(0x9000), (UNKNOWN_IMAGE, 0x9000));
366 reset_for_test();
367 }
368
369 #[test]
370 fn patch_images_are_numbered_in_load_order() {
371 let _serial = crate::patch::test_support::serial();
372 reset_for_test();
373 register_base(0x1000, 0x2000, 0);
374 assert_eq!(register_patch(0x4000, 0x5000, 0), 1);
375 assert_eq!(register_patch(0x6000, 0x7000, 0), 2);
376 assert_eq!(classify(0x6100).0, 2);
377 reset_for_test();
378 }
379
380 #[test]
381 #[cfg(any(
382 target_os = "linux",
383 target_os = "android",
384 all(target_vendor = "apple", target_pointer_width = "64"),
385 windows
386 ))]
387 fn the_running_executable_has_a_slide_no_larger_than_its_start() {
388 let address = here as *const () as usize;
389 let slide = image_slide_containing(address).expect("a slide for this test binary");
390 let (lo, _) = image_range_containing(address).expect("a range");
391 assert!(slide <= lo, "slide {slide:#x} above image start {lo:#x}");
392 }
393
394 #[test]
395 #[cfg(any(
396 target_os = "linux",
397 target_os = "android",
398 all(target_vendor = "apple", target_pointer_width = "64"),
399 windows
400 ))]
401 fn the_running_executable_has_a_range_containing_its_own_code() {
402 let address = here as *const () as usize;
403 let (lo, hi) = image_range_containing(address).expect("a range for this test binary");
404 assert!(lo <= address && address < hi);
405 }
406}