unc_vm_engine/universal/
code_memory.rs

1// This file contains code from external sources.
2// Attributions: https://github.com/wasmerio/wasmer/blob/master/ATTRIBUTIONS.md
3
4//! Memory management for executable code.
5use rustix::mm::{self, MapFlags, MprotectFlags, ProtFlags};
6use std::sync::Arc;
7use unc_vm_compiler::CompileError;
8
9/// The optimal alignment for functions.
10///
11/// On x86-64, this is 16 since it's what the optimizations assume.
12/// When we add support for other architectures, we should also figure out their
13/// optimal alignment values.
14pub(crate) const ARCH_FUNCTION_ALIGNMENT: u16 = 16;
15
16/// The optimal alignment for data.
17///
18pub(crate) const DATA_SECTION_ALIGNMENT: u16 = 64;
19
20fn round_up(size: usize, multiple: usize) -> usize {
21    debug_assert!(multiple.is_power_of_two());
22    (size + (multiple - 1)) & !(multiple - 1)
23}
24
25pub struct CodeMemoryWriter<'a> {
26    memory: &'a mut CodeMemory,
27    offset: usize,
28}
29
30impl<'a> CodeMemoryWriter<'a> {
31    /// Write the contents from the provided buffer into the location of `self.memory` aligned to
32    /// provided `alignment`.
33    ///
34    /// The `alignment` actually used may be greater than the spepcified value. This is relevant,
35    /// for example, when calling this function after a sequence of [`Self::write_executable`]
36    /// calls.
37    ///
38    /// Returns the position within the mapping at which the buffer was written.
39    pub fn write_data(&mut self, mut alignment: u16, input: &[u8]) -> Result<usize, CompileError> {
40        if self.offset == self.memory.executable_end {
41            alignment = u16::try_from(rustix::param::page_size()).expect("page size > u16::MAX");
42        }
43        self.write_inner(alignment, input)
44    }
45
46    /// Write the executable code from the provided buffer into the executable portion of
47    /// `self.memory`.
48    ///
49    /// All executable parts must be written out before `self.write_data` is called for the first
50    /// time.
51    ///
52    /// Returns the position within the mapping at which the buffer was written.
53    pub fn write_executable(
54        &mut self,
55        alignment: u16,
56        input: &[u8],
57    ) -> Result<usize, CompileError> {
58        assert_eq!(
59            self.memory.executable_end, self.offset,
60            "may not interleave executable and data in the same map"
61        );
62        let result = self.write_inner(alignment, input);
63        self.memory.executable_end = self.offset;
64        result
65    }
66
67    fn write_inner(&mut self, alignment: u16, input: &[u8]) -> Result<usize, CompileError> {
68        let entry_offset = self.offset;
69        let aligned_offset = round_up(entry_offset, usize::from(alignment));
70        let final_offset = aligned_offset + input.len();
71        let out_buffer = self.memory.as_slice_mut();
72        // Fill out the padding with zeroes, if only to make sure there are no gadgets in there.
73        out_buffer
74            .get_mut(entry_offset..aligned_offset)
75            .ok_or_else(|| CompileError::Resource("out of code memory space".into()))?
76            .fill(0);
77        out_buffer
78            .get_mut(aligned_offset..final_offset)
79            .ok_or_else(|| CompileError::Resource("out of code memory space".into()))?
80            .copy_from_slice(input);
81        self.offset = final_offset;
82        Ok(aligned_offset)
83    }
84
85    /// The current position of the writer.
86    pub fn position(&self) -> usize {
87        self.offset
88    }
89}
90
91/// Mappings to regions of memory storing the executable JIT code.
92pub struct CodeMemory {
93    /// Where to return this memory to when dropped.
94    source_pool: Option<Arc<crossbeam_queue::ArrayQueue<Self>>>,
95
96    /// The mapping
97    map: *mut u8,
98
99    /// Mapping size
100    size: usize,
101
102    /// Addresses `0..executable_end` contain executable memory.
103    ///
104    /// In a populated buffer rounding this up to the next page will give the address of the
105    /// read-write data portion of this memory.
106    executable_end: usize,
107}
108
109impl CodeMemory {
110    fn create(size: usize) -> rustix::io::Result<Self> {
111        // Make sure callers don’t pass in a 0-sized map request. That is most likely a bug.
112        assert!(size != 0);
113        let size = round_up(size, rustix::param::page_size());
114        let map = unsafe {
115            mm::mmap_anonymous(
116                std::ptr::null_mut(),
117                size,
118                ProtFlags::WRITE | ProtFlags::READ,
119                MapFlags::SHARED,
120            )?
121        };
122        Ok(Self { source_pool: None, map: map.cast(), executable_end: 0, size })
123    }
124
125    fn as_slice_mut(&mut self) -> &mut [u8] {
126        unsafe {
127            // SAFETY: We have made sure that this is the only reference to the memory region by
128            // requiring a mutable self reference.
129            std::slice::from_raw_parts_mut(self.map, self.size)
130        }
131    }
132
133    /// Ensure this CodeMemory is at least of the requested size.
134    ///
135    /// This will invalidate any data previously written into the mapping if the mapping needs to
136    /// be resized.
137    pub fn resize(mut self, size: usize) -> rustix::io::Result<Self> {
138        if self.size < size {
139            // Ideally we would use mremap, but see
140            // https://bugzilla.kernel.org/show_bug.cgi?id=8691
141            let source_pool = unsafe {
142                mm::munmap(self.map.cast(), self.size)?;
143                let source_pool = self.source_pool.take();
144                std::mem::forget(self);
145                source_pool
146            };
147            Self::create(size).map(|mut m| {
148                m.source_pool = source_pool;
149                m
150            })
151        } else {
152            self.executable_end = 0;
153            Ok(self)
154        }
155    }
156
157    /// Write to this code memory from the beginning of the mapping.
158    ///
159    /// # Safety
160    ///
161    /// At the time this method is called, there should remain no dangling readable/executable
162    /// references to this `CodeMemory`, for the original code memory that those references point
163    /// to are invalidated as soon as this method is invoked.
164    pub unsafe fn writer(&mut self) -> CodeMemoryWriter<'_> {
165        self.executable_end = 0;
166        CodeMemoryWriter { memory: self, offset: 0 }
167    }
168
169    /// Publish the specified number of bytes as executable code.
170    ///
171    /// # Safety
172    ///
173    /// Calling this requires that no mutable references to the code memory remain.
174    pub unsafe fn publish(&mut self) -> Result<(), CompileError> {
175        mm::mprotect(
176            self.map.cast(),
177            self.executable_end,
178            MprotectFlags::EXEC | MprotectFlags::READ,
179        )
180        .map_err(|e| {
181            CompileError::Resource(format!("could not make code memory executable: {}", e))
182        })
183    }
184
185    /// Remap the offset into an absolute address within a read-execute mapping.
186    ///
187    /// Offset must not exceed `isize::MAX`.
188    pub unsafe fn executable_address(&self, offset: usize) -> *const u8 {
189        // TODO: encapsulate offsets so that this `offset` is guaranteed to be sound.
190        debug_assert!(offset <= isize::MAX as usize);
191        self.map.offset(offset as isize)
192    }
193
194    /// Remap the offset into an absolute address within a read-write mapping.
195    ///
196    /// Offset must not exceed `isize::MAX`.
197    pub unsafe fn writable_address(&self, offset: usize) -> *mut u8 {
198        // TODO: encapsulate offsets so that this `offset` is guaranteed to be sound.
199        debug_assert!(offset <= isize::MAX as usize);
200        self.map.offset(offset as isize)
201    }
202}
203
204impl Drop for CodeMemory {
205    fn drop(&mut self) {
206        if let Some(source_pool) = self.source_pool.take() {
207            unsafe {
208                let result = mm::mprotect(
209                    self.map.cast(),
210                    self.size,
211                    MprotectFlags::WRITE | MprotectFlags::READ,
212                );
213                if let Err(e) = result {
214                    panic!(
215                        "could not mprotect mapping before returning it to the memory pool: \
216                         map={:?}, size={:?}, error={}",
217                        self.map, self.size, e
218                    );
219                }
220            }
221            drop(source_pool.push(Self {
222                source_pool: None,
223                map: self.map,
224                size: self.size,
225                executable_end: 0,
226            }));
227        } else {
228            unsafe {
229                if let Err(e) = mm::munmap(self.map.cast(), self.size) {
230                    tracing::error!(
231                        target: "unc_vm",
232                        message="could not unmap mapping",
233                        map=?self.map, size=self.size, error=%e
234                    );
235                }
236            }
237        }
238    }
239}
240
241unsafe impl Send for CodeMemory {}
242
243/// The pool of preallocated memory maps for storing the code.
244///
245/// This pool cannot grow and will only allow up to a number of code mappings that were specified
246/// at construction time.
247///
248/// However it is possible for the mappings inside to grow to accomodate larger code.
249#[derive(Clone)]
250pub struct LimitedMemoryPool {
251    pool: Arc<crossbeam_queue::ArrayQueue<CodeMemory>>,
252}
253
254impl LimitedMemoryPool {
255    /// Create a new pool with `count` mappings initialized to `default_memory_size` each.
256    pub fn new(count: usize, default_memory_size: usize) -> rustix::io::Result<Self> {
257        let pool = Arc::new(crossbeam_queue::ArrayQueue::new(count));
258        let this = Self { pool };
259        for _ in 0..count {
260            this.pool
261                .push(CodeMemory::create(default_memory_size)?)
262                .unwrap_or_else(|_| panic!("ArrayQueue could not accomodate {count} memories!"));
263        }
264        Ok(this)
265    }
266
267    /// Get a memory mapping, at least `size` bytes large.
268    pub fn get(&self, size: usize) -> rustix::io::Result<CodeMemory> {
269        let mut memory = self.pool.pop().ok_or(rustix::io::Errno::NOMEM)?;
270        memory.source_pool = Some(Arc::clone(&self.pool));
271        if memory.size < size {
272            Ok(memory.resize(size)?)
273        } else {
274            Ok(memory)
275        }
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::CodeMemory;
282    fn _assert() {
283        fn _assert_send<T: Send>() {}
284        _assert_send::<CodeMemory>();
285    }
286}
unc_vm_engine/universal/code_memory.rs

unc_vm_engine/universal/
code_memory.rs
