diff options
| -rw-r--r-- | examples/add.rs | 4 | ||||
| -rw-r--r-- | examples/fib.rs | 4 | ||||
| -rw-r--r-- | src/lib.rs | 4 | ||||
| -rw-r--r-- | src/rt.rs | 67 |
4 files changed, 54 insertions, 25 deletions
diff --git a/examples/add.rs b/examples/add.rs index ef010b6..8cd1984 100644 --- a/examples/add.rs +++ b/examples/add.rs @@ -22,8 +22,8 @@ fn main() { let code = asm.into_code(); std::fs::write("jit.asm", &code).unwrap(); - let rt = Runtime::new(&code); - let add42 = unsafe { rt.as_fn::<extern "C" fn(u32) -> u32>() }; + let mut rt = Runtime::new(); + let add42 = unsafe { rt.add_code::<extern "C" fn(u32) -> u32>(code) }; let res = add42(5); assert_eq!(res, 47); diff --git a/examples/fib.rs b/examples/fib.rs index 7acbb50..534dc13 100644 --- a/examples/fib.rs +++ b/examples/fib.rs @@ -62,8 +62,8 @@ fn main() { std::fs::write("jit.asm", &code).unwrap(); // Move code into executable page and get function pointer to it. - let rt = Runtime::new(&code); - let fib = unsafe { rt.as_fn::<extern "C" fn(u64) -> u64>() }; + let mut rt = Runtime::new(); + let fib = unsafe { rt.add_code::<extern "C" fn(u64) -> u64>(code) }; for n in 0..15 { let fib_jit = fib(n); @@ -61,8 +61,8 @@ //! asm.ret(); //! //! // Move code into executable page and get function pointer to it. -//! let rt = Runtime::new(&asm.into_code()); -//! let fib = unsafe { rt.as_fn::<extern "C" fn(u64) -> u64>() }; +//! let mut rt = Runtime::new(); +//! let fib = unsafe { rt.add_code::<extern "C" fn(u64) -> u64>(&asm.into_code()) }; //! //! for n in 0..15 { //! let fib_jit = fib(n); @@ -1,59 +1,88 @@ //! A simple runtime which can be used to execute emitted instructions. -use core::ffi::c_void; +use core::slice; use nix::sys::mman::{mmap, mprotect, munmap, MapFlags, ProtFlags}; /// A simple `mmap`ed runtime with executable pages. pub struct Runtime { - buf: *mut c_void, + buf: *mut u8, len: usize, + idx: usize, } impl Runtime { /// Create a new [Runtime]. - pub fn new(code: impl AsRef<[u8]>) -> Runtime { + pub fn new() -> Runtime { // Allocate a single page. let len = core::num::NonZeroUsize::new(4096).unwrap(); let buf = unsafe { mmap( None, len, - ProtFlags::PROT_WRITE, + ProtFlags::PROT_WRITE | ProtFlags::PROT_READ | ProtFlags::PROT_EXEC, MapFlags::MAP_PRIVATE | MapFlags::MAP_ANONYMOUS, 0, /* fd */ 0, /* off */ ) - .unwrap() + .unwrap() as *mut u8 }; - { - // Copy over code. - let code = code.as_ref(); - assert!(code.len() < len.get()); - unsafe { std::ptr::copy_nonoverlapping(code.as_ptr(), buf.cast(), len.get()) }; - } - unsafe { - // Remove write permissions from code buffer and allow to read-execute from it. - mprotect(buf, len.get(), ProtFlags::PROT_READ | ProtFlags::PROT_EXEC) - .expect("Failed to RX mprotect Runtime code buffer"); - } Runtime { buf, len: len.get(), + idx: 0, } } + /// Write protect the underlying code page(s). + pub fn protect(&mut self) { + unsafe { + // Remove write permissions from code buffer and allow to read-execute from it. + mprotect( + self.buf.cast(), + self.len, + ProtFlags::PROT_READ | ProtFlags::PROT_EXEC, + ) + .expect("Failed to RX mprotect Runtime code buffer"); + } + } + + /// Add block of code to the runtime and get function pointer back. + pub unsafe fn add_code<F>(&mut self, code: impl AsRef<[u8]>) -> F { + // Get pointer to start of next free byte. + assert!(self.idx < self.len); + let fn_start = self.buf.add(self.idx); + + // Copy over code. + let code = code.as_ref(); + assert!(code.len() < (self.len - self.idx)); + unsafe { std::ptr::copy_nonoverlapping(code.as_ptr(), fn_start, code.len()) }; + + // Increment index to next free byte. + self.idx += code.len(); + + // Return function to newly added code. + Self::as_fn::<F>(fn_start) + } + /// Reinterpret the block of code as `F`. #[inline] - pub unsafe fn as_fn<F>(&self) -> F { - unsafe { std::mem::transmute_copy(&self.buf) } + unsafe fn as_fn<F>(fn_start: *mut u8) -> F { + unsafe { std::mem::transmute_copy(&fn_start) } + } + + /// Dump the currently added code to a file called `jit.asm`. The disassembly can be inspected + /// as `ndisasm -b 64 jit.asm`. + pub fn dump(&self) { + let code = unsafe { slice::from_raw_parts(self.buf, self.idx) }; + std::fs::write("jit.asm", code).unwrap(); } } impl Drop for Runtime { fn drop(&mut self) { unsafe { - munmap(self.buf, self.len).expect("Failed to munmap Runtime"); + munmap(self.buf.cast(), self.len).expect("Failed to munmap Runtime"); } } } |