gaia_jit/
lib.rs

1#![warn(missing_docs)]
2
3//! JIT memory management for Gaia.
4//!
5//! This crate provides a safe abstraction for allocating executable memory
6//! and writing machine code at runtime.
7
8use std::{io, ptr::NonNull, slice};
9
10mod os;
11
12/// A buffer for JIT code execution.
13///
14/// # Example
15///
16/// ```rust
17/// use gaia_jit::JitMemory;
18///
19/// let mut jit = JitMemory::new(4096).unwrap();
20/// // Write some machine code...
21/// jit.write(&[0x90, 0xc3]).unwrap(); // nop; ret
22/// jit.make_executable().unwrap();
23///
24/// unsafe {
25///     let f: unsafe extern "C" fn() = jit.as_fn();
26///     f();
27/// }
28/// ```
29pub struct JitMemory {
30    ptr: NonNull<u8>,
31    len: usize,
32    capacity: usize,
33}
34
35impl JitMemory {
36    /// Allocate a new JIT memory block with the given capacity.
37    pub fn new(capacity: usize) -> io::Result<Self> {
38        let page_size = os::get_page_size();
39        let rounded_capacity = (capacity + page_size - 1) & !(page_size - 1);
40
41        let ptr = os::alloc_rw(rounded_capacity)?;
42
43        Ok(Self { ptr, len: 0, capacity: rounded_capacity })
44    }
45
46    /// Write code into the buffer.
47    ///
48    /// This method automatically ensures the memory is writable before writing,
49    /// and restores the previous protection after if possible.
50    pub fn write(&mut self, code: &[u8]) -> io::Result<()> {
51        if self.len + code.len() > self.capacity {
52            return Err(io::Error::new(io::ErrorKind::OutOfMemory, "Not enough capacity in JIT buffer"));
53        }
54
55        unsafe {
56            os::write_protect_scoped(self.ptr.as_ptr(), self.capacity, || {
57                std::ptr::copy_nonoverlapping(code.as_ptr(), self.ptr.as_ptr().add(self.len), code.len());
58            })?;
59        }
60
61        self.len += code.len();
62        Ok(())
63    }
64
65    /// Make the memory executable and return a pointer to the start.
66    pub fn make_executable(&self) -> io::Result<*const u8> {
67        unsafe {
68            os::make_executable(self.ptr.as_ptr(), self.capacity)?;
69        }
70        Ok(self.ptr.as_ptr())
71    }
72
73    /// Make the memory readable and writable again (e.g. for updates).
74    pub fn make_writable(&self) -> io::Result<()> {
75        unsafe {
76            os::make_writable(self.ptr.as_ptr(), self.capacity)?;
77        }
78        Ok(())
79    }
80
81    /// Get the current length of the code in the buffer.
82    pub fn len(&self) -> usize {
83        self.len
84    }
85
86    /// Get the capacity of the buffer.
87    pub fn capacity(&self) -> usize {
88        self.capacity
89    }
90
91    /// Is the buffer empty?
92    pub fn is_empty(&self) -> bool {
93        self.len == 0
94    }
95
96    /// As a slice of bytes (read-only).
97    pub fn as_slice(&self) -> &[u8] {
98        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
99    }
100
101    /// Cast the JIT memory to a function pointer of type `S`.
102    pub unsafe fn as_fn<S>(&self) -> S {
103        unsafe { std::mem::transmute_copy(&self.ptr.as_ptr()) }
104    }
105}
106
107impl Drop for JitMemory {
108    fn drop(&mut self) {
109        unsafe {
110            os::dealloc(self.ptr.as_ptr(), self.capacity);
111        }
112    }
113}
114
115unsafe impl Send for JitMemory {}
116unsafe impl Sync for JitMemory {}
gaia_jit/lib.rs

gaia_jit/
lib.rs
