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}