gaia-jit 0.0.4

JIT execution primitives for Gaia project
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

use std::{io, ptr::NonNull, slice};

mod os;

/// A buffer for JIT code execution.
///
/// # Example
///
/// ```rust
/// use gaia_jit::JitMemory;
///
/// let mut jit = JitMemory::new(4096).unwrap();
/// // Write some machine code...
/// jit.write(&[0x90, 0xc3]).unwrap(); // nop; ret
/// jit.make_executable().unwrap();
///
/// unsafe {
///     let f: unsafe extern "C" fn() = jit.as_fn();
///     f();
/// }
/// ```
pub struct JitMemory {
    ptr: NonNull<u8>,
    len: usize,
    capacity: usize,
}

impl JitMemory {
    /// Allocate a new JIT memory block with the given capacity.
    pub fn new(capacity: usize) -> io::Result<Self> {
        let page_size = os::get_page_size();
        let rounded_capacity = (capacity + page_size - 1) & !(page_size - 1);

        let ptr = os::alloc_rw(rounded_capacity)?;

        Ok(Self {
            ptr,
            len: 0,
            capacity: rounded_capacity,
        })
    }

    /// Write code into the buffer.
    ///
    /// This method automatically ensures the memory is writable before writing,
    /// and restores the previous protection after if possible.
    pub fn write(&mut self, code: &[u8]) -> io::Result<()> {
        if self.len + code.len() > self.capacity {
            return Err(io::Error::new(
                io::ErrorKind::OutOfMemory,
                "Not enough capacity in JIT buffer",
            ));
        }

        unsafe {
            os::write_protect_scoped(self.ptr.as_ptr(), self.capacity, || {
                std::ptr::copy_nonoverlapping(
                    code.as_ptr(),
                    self.ptr.as_ptr().add(self.len),
                    code.len(),
                );
            })?;
        }

        self.len += code.len();
        Ok(())
    }

    /// Make the memory executable and return a pointer to the start.
    pub fn make_executable(&self) -> io::Result<*const u8> {
        unsafe {
            os::make_executable(self.ptr.as_ptr(), self.capacity)?;
        }
        Ok(self.ptr.as_ptr())
    }

    /// Make the memory readable and writable again (e.g. for updates).
    pub fn make_writable(&self) -> io::Result<()> {
        unsafe {
            os::make_writable(self.ptr.as_ptr(), self.capacity)?;
        }
        Ok(())
    }

    /// Get the current length of the code in the buffer.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Get the capacity of the buffer.
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Is the buffer empty?
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// As a slice of bytes (read-only).
    pub fn as_slice(&self) -> &[u8] {
        unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
    }

    /// Cast the JIT memory to a function pointer of type `S`.
    pub unsafe fn as_fn<S>(&self) -> S {
        unsafe { std::mem::transmute_copy(&self.ptr.as_ptr()) }
    }
}

impl Drop for JitMemory {
    fn drop(&mut self) {
        unsafe {
            os::dealloc(self.ptr.as_ptr(), self.capacity);
        }
    }
}

unsafe impl Send for JitMemory {}
unsafe impl Sync for JitMemory {}