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}