closure_ffi/
jit_alloc.rs

1//! Abstractions around allocators that provide dual-mapped memory with XOR protection rules (one RW
2//! view and one RX view) suitable for emitting code at runtime.
3//!
4//! Meant to be an abstraction over the `jit-allocator` crate's API so that it can be swapped with
5//! user-provided allocators.
6//!
7//! See the [`JitAlloc`] trait for more information.
8
9/// Anonymous error that may be returned by [`JitAlloc`] implementations when [`JitAlloc::alloc`] or
10/// [`JitAlloc::release`] fail.
11#[derive(Debug)]
12pub struct JitAllocError;
13
14/// Values to use with [`JitAlloc::protect_jit_memory`].
15#[derive(Debug, Clone, Copy, Eq, PartialEq)]
16pub enum ProtectJitAccess {
17    /// Protect JIT memory with Read+Write permissions.
18    ReadWrite = 0,
19    /// Protect JIT memory with Read+Execute permissions.
20    ReadExecute = 1,
21}
22
23/// Generic allocator providing virtual memory suitable for emitting code at runtime.
24///
25/// The API is meant to be a thin abstraction over the `jit-allocator` crate's API, to allow it
26/// to be swapped with other allocators.
27pub trait JitAlloc {
28    /// Allocates `size` bytes in the executable memory region.
29    /// Returns two pointers. One points to Read-Execute mapping and another to Read-Write mapping.
30    /// All code writes *must* go to the Read-Write mapping.
31    fn alloc(&self, size: usize) -> Result<(*const u8, *mut u8), JitAllocError>;
32
33    /// Releases the memory allocated by `alloc`.
34    ///
35    /// # Safety
36    /// - `rx_ptr` must have been returned from `alloc`
37    /// - `rx_ptr` must have been allocated from this allocator
38    /// - `rx_ptr` must not have been passed to `release` before
39    /// - `rx_ptr` must point to read-execute part of memory returned from `alloc`.
40    unsafe fn release(&self, rx_ptr: *const u8) -> Result<(), JitAllocError>;
41
42    /// On hardened architectures with `MAP_JIT`-like memory flags, set the access for the current
43    /// thread.
44    ///
45    /// This is expected to be a no-op on most platforms, but should be called before writing
46    /// or executing JIT memory.
47    ///
48    /// # Safety
49    ///
50    /// - `ptr` must point at least `size` bytes of readable memory.
51    unsafe fn protect_jit_memory(ptr: *const u8, size: usize, access: ProtectJitAccess);
52
53    /// Flushes the instruction cache for (at least) the given slice of executable memory. Should be
54    /// called after the JIT memory is ready to be executed.
55    ///
56    /// On architectures with shared data/instruction caches, like x86_64, this is a no-op.
57    ///
58    /// # Safety
59    /// - `rx_ptr` must point at least `size` bytes of Read-Execute memory.
60    unsafe fn flush_instruction_cache(rx_ptr: *const u8, size: usize);
61}
62
63impl<J: JitAlloc> JitAlloc for &J {
64    fn alloc(&self, size: usize) -> Result<(*const u8, *mut u8), JitAllocError> {
65        (*self).alloc(size)
66    }
67
68    unsafe fn release(&self, rx_ptr: *const u8) -> Result<(), JitAllocError> {
69        (*self).release(rx_ptr)
70    }
71
72    #[inline(always)]
73    unsafe fn flush_instruction_cache(rx_ptr: *const u8, size: usize) {
74        J::flush_instruction_cache(rx_ptr, size);
75    }
76
77    #[inline(always)]
78    unsafe fn protect_jit_memory(ptr: *const u8, size: usize, access: ProtectJitAccess) {
79        J::protect_jit_memory(ptr, size, access);
80    }
81}
82
83#[cfg(any(test, feature = "bundled_jit_alloc"))]
84mod bundled_jit_alloc {
85    use jit_allocator::JitAllocator;
86
87    use super::*;
88
89    #[inline(always)]
90    fn convert_access(access: ProtectJitAccess) -> jit_allocator::ProtectJitAccess {
91        match access {
92            ProtectJitAccess::ReadExecute => jit_allocator::ProtectJitAccess::ReadExecute,
93            ProtectJitAccess::ReadWrite => jit_allocator::ProtectJitAccess::ReadWrite,
94        }
95    }
96
97    impl JitAlloc for core::cell::RefCell<JitAllocator> {
98        fn alloc(&self, size: usize) -> Result<(*const u8, *mut u8), JitAllocError> {
99            self.borrow_mut().alloc(size).map_err(|_| JitAllocError)
100        }
101
102        unsafe fn release(&self, rx_ptr: *const u8) -> Result<(), JitAllocError> {
103            self.borrow_mut().release(rx_ptr).map_err(|_| JitAllocError)
104        }
105
106        #[inline(always)]
107        unsafe fn flush_instruction_cache(rx_ptr: *const u8, size: usize) {
108            jit_allocator::flush_instruction_cache(rx_ptr, size);
109        }
110
111        #[inline(always)]
112        unsafe fn protect_jit_memory(_ptr: *const u8, _size: usize, access: ProtectJitAccess) {
113            jit_allocator::protect_jit_memory(convert_access(access));
114        }
115    }
116
117    #[cfg(not(feature = "no_std"))]
118    impl JitAlloc for std::sync::RwLock<JitAllocator> {
119        fn alloc(&self, size: usize) -> Result<(*const u8, *mut u8), JitAllocError> {
120            self.write().unwrap().alloc(size).map_err(|_| JitAllocError)
121        }
122
123        unsafe fn release(&self, rx_ptr: *const u8) -> Result<(), JitAllocError> {
124            self.write().unwrap().release(rx_ptr).map_err(|_| JitAllocError)
125        }
126
127        #[inline(always)]
128        unsafe fn flush_instruction_cache(rx_ptr: *const u8, size: usize) {
129            jit_allocator::flush_instruction_cache(rx_ptr, size);
130        }
131
132        #[inline(always)]
133        unsafe fn protect_jit_memory(_ptr: *const u8, _size: usize, access: ProtectJitAccess) {
134            jit_allocator::protect_jit_memory(convert_access(access));
135        }
136    }
137
138    #[cfg(not(feature = "no_std"))]
139    impl JitAlloc for std::sync::Mutex<JitAllocator> {
140        fn alloc(&self, size: usize) -> Result<(*const u8, *mut u8), JitAllocError> {
141            self.lock().unwrap().alloc(size).map_err(|_| JitAllocError)
142        }
143
144        unsafe fn release(&self, rx_ptr: *const u8) -> Result<(), JitAllocError> {
145            self.lock().unwrap().release(rx_ptr).map_err(|_| JitAllocError)
146        }
147
148        #[inline(always)]
149        unsafe fn flush_instruction_cache(rx_ptr: *const u8, size: usize) {
150            jit_allocator::flush_instruction_cache(rx_ptr, size);
151        }
152
153        #[inline(always)]
154        unsafe fn protect_jit_memory(_ptr: *const u8, _size: usize, access: ProtectJitAccess) {
155            jit_allocator::protect_jit_memory(convert_access(access));
156        }
157    }
158
159    #[cfg(feature = "no_std")]
160    static GLOBAL_JIT_ALLOC: spin::Mutex<Option<alloc::boxed::Box<JitAllocator>>> =
161        spin::Mutex::new(None);
162    #[cfg(not(feature = "no_std"))]
163    static GLOBAL_JIT_ALLOC: std::sync::Mutex<Option<Box<JitAllocator>>> =
164        std::sync::Mutex::new(None);
165
166    /// The default, global JIT allocator.
167    ///
168    /// This is currently implemented as a ZST deffering to a static [`jit_allocator::JitAllocator`]
169    /// behind a [`std::sync::Mutex`] (or a [`spin::Mutex`] under no_std).
170    ///
171    /// [`spin::Mutex`]: https://docs.rs/spin/0.9/spin/type.Mutex.html
172    #[derive(Default, Clone, Copy)]
173    pub struct GlobalJitAlloc;
174
175    impl GlobalJitAlloc {
176        fn use_alloc<T>(&self, action: impl FnOnce(&mut JitAllocator) -> T) -> T {
177            #[cfg(feature = "no_std")]
178            let mut maybe_alloc = GLOBAL_JIT_ALLOC.lock();
179            #[cfg(not(feature = "no_std"))]
180            let mut maybe_alloc = GLOBAL_JIT_ALLOC.lock().unwrap();
181
182            let alloc = maybe_alloc.get_or_insert_with(|| JitAllocator::new(Default::default()));
183            action(alloc)
184        }
185    }
186
187    impl JitAlloc for GlobalJitAlloc {
188        fn alloc(&self, size: usize) -> Result<(*const u8, *mut u8), JitAllocError> {
189            self.use_alloc(|a| a.alloc(size)).map_err(|_| JitAllocError)
190        }
191
192        unsafe fn release(&self, rx_ptr: *const u8) -> Result<(), JitAllocError> {
193            self.use_alloc(|a| a.release(rx_ptr)).map_err(|_| JitAllocError)
194        }
195
196        #[inline(always)]
197        unsafe fn flush_instruction_cache(rx_ptr: *const u8, size: usize) {
198            jit_allocator::flush_instruction_cache(rx_ptr, size);
199        }
200
201        #[inline(always)]
202        unsafe fn protect_jit_memory(_ptr: *const u8, _size: usize, access: ProtectJitAccess) {
203            jit_allocator::protect_jit_memory(convert_access(access));
204        }
205    }
206
207    #[cfg(not(feature = "no_std"))]
208    mod thread_jit_alloc {
209        use core::{cell::UnsafeCell, marker::PhantomData};
210
211        use jit_allocator::JitAllocator;
212
213        #[allow(unused_imports)]
214        use super::*;
215
216        thread_local! {
217            static THREAD_JIT_ALLOC: UnsafeCell<Box<JitAllocator>> =
218                UnsafeCell::new(JitAllocator::new(Default::default()));
219        }
220
221        /// Marker type providing access to a thread-local JIT allocator.
222        ///
223        /// Unlike [`GlobalJitAlloc`], this allocator is neither [`Send`] nor [`Sync`].
224        #[derive(Default, Clone)]
225        pub struct ThreadJitAlloc(PhantomData<*mut ()>);
226
227        impl JitAlloc for ThreadJitAlloc {
228            fn alloc(&self, size: usize) -> Result<(*const u8, *mut u8), JitAllocError> {
229                THREAD_JIT_ALLOC
230                    .with(|a| unsafe { &mut *a.get() }.alloc(size))
231                    .map_err(|_| JitAllocError)
232            }
233
234            unsafe fn release(&self, rx_ptr: *const u8) -> Result<(), JitAllocError> {
235                THREAD_JIT_ALLOC
236                    .with(|a| unsafe { &mut *a.get() }.release(rx_ptr))
237                    .map_err(|_| JitAllocError)
238            }
239
240            #[inline(always)]
241            unsafe fn flush_instruction_cache(rx_ptr: *const u8, size: usize) {
242                jit_allocator::flush_instruction_cache(rx_ptr, size);
243            }
244
245            #[inline(always)]
246            unsafe fn protect_jit_memory(_ptr: *const u8, _size: usize, access: ProtectJitAccess) {
247                jit_allocator::protect_jit_memory(convert_access(access));
248            }
249        }
250    }
251    #[cfg(not(feature = "no_std"))]
252    pub use thread_jit_alloc::*;
253}
254#[cfg(any(test, feature = "bundled_jit_alloc"))]
255pub use bundled_jit_alloc::*;
closure_ffi/jit_alloc.rs

closure_ffi/
jit_alloc.rs
