substrate/hook/
memory.rs

1use crate::error::{Result, SubstrateError};
2use libc::{c_void, mmap, mprotect, munmap, sysconf, MAP_ANON, MAP_FAILED, MAP_PRIVATE, PROT_EXEC, PROT_READ, PROT_WRITE, _SC_PAGESIZE};
3use std::ptr;
4
5pub struct ProtectedMemory {
6    address: *mut u8,
7    width: usize,
8}
9
10impl ProtectedMemory {
11    pub unsafe fn new(data: *mut u8, size: usize) -> Result<Self> {
12        if data.is_null() || size == 0 {
13            return Err(SubstrateError::NullPointer);
14        }
15
16        let page_size = sysconf(_SC_PAGESIZE) as usize;
17        let base = (data as usize / page_size) * page_size;
18        let width = ((data as usize + size - 1) / page_size + 1) * page_size - base;
19        let address = base as *mut u8;
20
21        let result = mprotect(
22            address as *mut c_void,
23            width,
24            PROT_READ | PROT_WRITE | PROT_EXEC,
25        );
26
27        if result == -1 {
28            let err = std::io::Error::last_os_error();
29            return Err(SubstrateError::MemoryProtection(format!("{}", err)));
30        }
31
32        Ok(Self { address, width })
33    }
34}
35
36impl Drop for ProtectedMemory {
37    fn drop(&mut self) {
38        unsafe {
39            mprotect(
40                self.address as *mut c_void,
41                self.width,
42                PROT_READ | PROT_EXEC,
43            );
44
45            #[cfg(any(target_arch = "arm", target_arch = "aarch64"))]
46            {
47                extern "C" {
48                    fn __clear_cache(begin: *mut c_void, end: *mut c_void);
49                }
50                __clear_cache(
51                    self.address as *mut c_void,
52                    self.address.add(self.width) as *mut c_void,
53                );
54            }
55        }
56    }
57}
58
59pub unsafe fn allocate_trampoline(size: usize) -> Result<*mut u8> {
60    let ptr = mmap(
61        ptr::null_mut(),
62        size,
63        PROT_READ | PROT_WRITE,
64        MAP_ANON | MAP_PRIVATE,
65        -1,
66        0,
67    );
68
69    if ptr == MAP_FAILED {
70        let err = std::io::Error::last_os_error();
71        return Err(SubstrateError::MemoryMap(format!("{}", err)));
72    }
73
74    Ok(ptr as *mut u8)
75}
76
77pub unsafe fn make_executable(ptr: *mut u8, size: usize) -> Result<()> {
78    let result = mprotect(ptr as *mut c_void, size, PROT_READ | PROT_EXEC);
79
80    if result == -1 {
81        let err = std::io::Error::last_os_error();
82        return Err(SubstrateError::MemoryProtection(format!("{}", err)));
83    }
84
85    Ok(())
86}
87
88pub unsafe fn free_trampoline(ptr: *mut u8, size: usize) {
89    munmap(ptr as *mut c_void, size);
90}