juicebox_asm/
rt.rs

1// SPDX-License-Identifier: MIT
2//
3// Copyright (c) 2023, Johannes Stoelp <dev@memzero.de>
4
5//! Simple `mmap`ed runtime.
6//!
7//! This runtime supports adding code to executable pages and turn the added code into user
8//! specified function pointer.
9
10#[cfg(not(target_os = "linux"))]
11compile_error!("This runtime is only supported on linux");
12
13mod perf {
14    use std::fs;
15    use std::io::Write;
16
17    /// Provide support for the simple [perf jit interface][perf-jit].
18    ///
19    /// This allows a simple (static) jit runtime to generate meta data describing the generated
20    /// functions, which is used during post-processing by `perf report` to symbolize addresses
21    /// captured while executing jitted code.
22    ///
23    /// By the nature of this format, this can not be used for dynamic jit runtimes, which reuses
24    /// memory which previously contained jitted code.
25    ///
26    /// [perf-jit]: https://elixir.bootlin.com/linux/v6.6.6/source/tools/perf/Documentation/jit-interface.txt
27    pub(super) struct PerfMap {
28        file: std::fs::File,
29    }
30
31    impl PerfMap {
32        /// Create an empty perf map file.
33        pub(super) fn new() -> Self {
34            let name = format!("/tmp/perf-{}.map", unsafe { libc::getpid() });
35            let file = fs::OpenOptions::new()
36                .truncate(true)
37                .create(true)
38                .write(true)
39                .open(&name)
40                .unwrap_or_else(|_| panic!("Failed to open perf map file {}", &name));
41
42            PerfMap { file }
43        }
44
45        /// Add an entry to the perf map file.
46        pub(super) fn add_entry(&mut self, start: usize, len: usize) {
47            // Each line has the following format, fields separated with spaces:
48            //   START SIZE NAME
49            //
50            // START and SIZE are hex numbers without 0x.
51            // NAME is the rest of the line, so it could contain special characters.
52            writeln!(self.file, "{:x} {:x} jitfn_{:x}", start, len, start)
53                .expect("Failed to write PerfMap entry");
54        }
55    }
56}
57
58/// A simple `mmap`ed runtime with executable pages.
59pub struct Runtime {
60    buf: *mut u8,
61    len: usize,
62    idx: usize,
63    perf: Option<perf::PerfMap>,
64}
65
66impl Runtime {
67    /// Create a new [Runtime].
68    ///
69    /// # Panics
70    ///
71    /// Panics if the `mmap` call fails.
72    pub fn new() -> Runtime {
73        // Allocate a single page.
74        let len = 4096;
75        let buf = unsafe {
76            libc::mmap(
77                std::ptr::null_mut(),
78                len,
79                libc::PROT_NONE,
80                libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
81                0, /* fd */
82                0, /* off */
83            ) as *mut u8
84        };
85        assert_ne!(
86            buf.cast(),
87            libc::MAP_FAILED,
88            "Failed to mmap runtime code page"
89        );
90
91        Runtime {
92            buf,
93            len,
94            idx: 0,
95            perf: None,
96        }
97    }
98
99    /// Create a new [Runtime] which also generates static perf metat data.
100    ///
101    /// For each function added to the [Runtime], an entry will be generated in the
102    /// `/tmp/perf-<PID>.map` file, which `perf report` uses to symbolicate unknown addresses.
103    /// This is applicable for static runtimes only.
104    ///
105    /// # Panics
106    ///
107    /// Panics if the `mmap` call fails.
108    pub fn with_profile() -> Runtime {
109        let mut rt = Runtime::new();
110        rt.perf = Some(perf::PerfMap::new());
111        rt
112    }
113
114    /// Add the block of `code` to the runtime and a get function pointer of type `F`.
115    ///
116    /// # Panics
117    ///
118    /// Panics if the `code` does not fit on the `mmap`ed pages or is empty.
119    ///
120    /// # Safety
121    ///
122    /// The code added must fulfill the ABI of the specified function `F` and the returned function
123    /// pointer is only valid until the [`Runtime`] is dropped.
124    ///
125    /// # Examples
126    ///
127    /// ```
128    /// let mut rt = juicebox_asm::Runtime::new();
129    ///
130    /// let code = [ 0x90 /* nop */, 0xc3 /* ret */ ];
131    /// let nop = unsafe { rt.add_code::<extern "C" fn()>(&code) };
132    ///
133    /// nop();
134    /// ```
135    pub unsafe fn add_code<F>(&mut self, code: impl AsRef<[u8]>) -> F {
136        // Get pointer to start of next free byte.
137        assert!(self.idx < self.len, "Runtime code page full");
138        let fn_start = self.buf.add(self.idx);
139
140        // Copy over code.
141        let code = code.as_ref();
142        assert!(!code.is_empty(), "Adding empty code not supported");
143        assert!(
144            code.len() <= (self.len - self.idx),
145            "Code does not fit on the runtime code page"
146        );
147        self.unprotect();
148        unsafe { std::ptr::copy_nonoverlapping(code.as_ptr(), fn_start, code.len()) };
149        self.protect();
150
151        // Increment index to next free byte.
152        self.idx += code.len();
153
154        // Add perf map entry.
155        if let Some(map) = &mut self.perf {
156            map.add_entry(fn_start as usize, code.len());
157        }
158
159        // Return function to newly added code.
160        unsafe { Self::as_fn::<F>(fn_start) }
161    }
162
163    /// Disassemble the code currently added to the runtime, using
164    /// [`ndisasm`](https://nasm.us/index.php) and print it to _stdout_. If
165    /// `ndisasm` is not available on the system this prints a warning and
166    /// becomes a nop.
167    ///
168    /// # Panics
169    ///
170    /// Panics if anything goes wrong with spawning, writing to or reading from
171    /// the `ndisasm` child process.
172    pub fn disasm(&self) {
173        assert!(self.idx <= self.len);
174        crate::disasm::disasm(unsafe { core::slice::from_raw_parts(self.buf, self.idx) });
175    }
176
177    /// Reinterpret the block of code pointed to by `fn_start` as `F`.
178    #[inline]
179    unsafe fn as_fn<F>(fn_start: *mut u8) -> F {
180        unsafe { std::mem::transmute_copy(&fn_start) }
181    }
182
183    /// Add write protection the underlying code page(s).
184    ///
185    /// # Panics
186    ///
187    /// Panics if the `mprotect` call fails.
188    fn protect(&mut self) {
189        unsafe {
190            // Remove write permissions from code page and allow to read-execute from it.
191            let ret = libc::mprotect(self.buf.cast(), self.len, libc::PROT_READ | libc::PROT_EXEC);
192            assert_eq!(ret, 0, "Failed to RX mprotect runtime code page");
193        }
194    }
195
196    /// Remove write protection the underlying code page(s).
197    ///
198    /// # Panics
199    ///
200    /// Panics if the `mprotect` call fails.
201    fn unprotect(&mut self) {
202        unsafe {
203            // Add write permissions to code page.
204            let ret = libc::mprotect(self.buf.cast(), self.len, libc::PROT_WRITE);
205            assert_eq!(ret, 0, "Failed to W mprotect runtime code page");
206        }
207    }
208}
209
210impl Drop for Runtime {
211    /// Unmaps the code page. This invalidates all the function pointer returned by
212    /// [`Runtime::add_code`].
213    fn drop(&mut self) {
214        unsafe {
215            let ret = libc::munmap(self.buf.cast(), self.len);
216            assert_eq!(ret, 0, "Failed to munmap runtime");
217        }
218    }
219}
220
221#[cfg(test)]
222mod test {
223    use super::*;
224
225    #[test]
226    fn test_code_max_size() {
227        let mut rt = Runtime::new();
228        let code = [0u8; 4096];
229        unsafe {
230            rt.add_code::<extern "C" fn()>(code);
231        }
232    }
233
234    #[test]
235    #[should_panic]
236    fn test_code_max_size_plus_1() {
237        let mut rt = Runtime::new();
238        let code = [0u8; 4097];
239        unsafe {
240            rt.add_code::<extern "C" fn()>(code);
241        }
242    }
243
244    #[test]
245    #[should_panic]
246    fn test_code_max_size_plus_1_2() {
247        let mut rt = Runtime::new();
248        let code = [0u8; 4096];
249        unsafe {
250            rt.add_code::<extern "C" fn()>(code);
251        }
252
253        let code = [0u8; 1];
254        unsafe {
255            rt.add_code::<extern "C" fn()>(code);
256        }
257    }
258
259    #[test]
260    #[should_panic]
261    fn test_empty_code() {
262        let mut rt = Runtime::new();
263        let code = [0u8; 0];
264        unsafe {
265            rt.add_code::<extern "C" fn()>(code);
266        }
267    }
268}