//! Trampoline generator for carrying context with function pointer. //! //! This makes use of the `mm0` register to pass the context as an implicit "parameter" because `mm0` is //! not used to pass parameters and is almost never used by modern compilers. It's still better to call //! `get_context()` as early as possible in the callee function though, as a good practice. //! //! Variadic functions are not supported because `rax` is used by the trampoline code. use crate::loader::CodeMemory; use crate::vm::Ctx; use std::fmt; use std::{mem, slice}; lazy_static! { /// Reads the context pointer from `mm0`. /// /// This function generates code at runtime since `asm!` macro is not yet stable. static ref GET_CONTEXT: extern "C" fn () -> *const CallContext = { static CODE: &'static [u8] = &[ 0x48, 0x0f, 0x7e, 0xc0, // movq %mm0, %rax 0xc3, // retq ]; let mut mem = CodeMemory::new(4096); mem[..CODE.len()].copy_from_slice(CODE); mem.make_executable(); let ptr = mem.as_ptr(); mem::forget(mem); unsafe { mem::transmute(ptr) } }; } /// An opaque type for pointers to a callable memory location. pub enum CallTarget {} /// An opaque type for context pointers. pub enum CallContext {} /// An opaque type for generated trampolines' call entries. pub enum Trampoline {} /// Trampoline Buffer Builder. pub struct TrampolineBufferBuilder { code: Vec, offsets: Vec, } /// Trampoline Buffer. pub struct TrampolineBuffer { code: CodeMemory, offsets: Vec, } fn value_to_bytes(ptr: &T) -> &[u8] { unsafe { slice::from_raw_parts(ptr as *const T as *const u8, mem::size_of::()) } } /// Calls `GET_CONTEXT` and returns the current context. pub fn get_context() -> *const CallContext { GET_CONTEXT() } impl TrampolineBufferBuilder { pub fn new() -> TrampolineBufferBuilder { TrampolineBufferBuilder { code: vec![], offsets: vec![], } } /// Adds a context trampoline. /// /// This generates a transparent trampoline function that forwards any call to `target` with /// unmodified params/returns. When called from the trampoline, `target` will have access to /// the `context` specified here through `get_context()`. /// /// Note that since `rax` is overwritten internally, variadic functions are not supported as `target`. pub fn add_context_trampoline( &mut self, target: *const CallTarget, context: *const CallContext, ) -> usize { let idx = self.offsets.len(); self.offsets.push(self.code.len()); self.code.extend_from_slice(&[ 0x48, 0xb8, // movabsq ?, %rax ]); self.code.extend_from_slice(value_to_bytes(&context)); self.code.extend_from_slice(&[ 0x48, 0x0f, 0x6e, 0xc0, // movq %rax, %mm0 ]); self.code.extend_from_slice(&[ 0x48, 0xb8, // movabsq ?, %rax ]); self.code.extend_from_slice(value_to_bytes(&target)); self.code.extend_from_slice(&[ 0xff, 0xe0, // jmpq *%rax ]); idx } pub fn add_context_rsp_state_preserving_trampoline( &mut self, target: unsafe extern "C" fn(&mut Ctx, *const CallContext, *const u64), context: *const CallContext, ) -> usize { let idx = self.offsets.len(); self.offsets.push(self.code.len()); self.code.extend_from_slice(&[ 0x53, // push %rbx 0x41, 0x54, // push %r12 0x41, 0x55, // push %r13 0x41, 0x56, // push %r14 0x41, 0x57, // push %r15 ]); self.code.extend_from_slice(&[ 0x48, 0xbe, // movabsq ?, %rsi ]); self.code.extend_from_slice(value_to_bytes(&context)); self.code.extend_from_slice(&[ 0x48, 0x89, 0xe2, // mov %rsp, %rdx ]); self.code.extend_from_slice(&[ 0x48, 0xb8, // movabsq ?, %rax ]); self.code.extend_from_slice(value_to_bytes(&target)); self.code.extend_from_slice(&[ 0xff, 0xd0, // callq *%rax ]); self.code.extend_from_slice(&[ 0x48, 0x81, 0xc4, // add ?, %rsp ]); self.code.extend_from_slice(value_to_bytes(&40i32)); // 5 * 8 self.code.extend_from_slice(&[ 0xc3, //retq ]); idx } /// Adds a callinfo trampoline. /// /// This generates a trampoline function that collects `num_params` parameters into an array /// and passes the array into `target` as the second argument when called. The first argument /// of `target` is the `context` specified here. /// /// Note that non-integer parameters/variadic functions are not supported. pub fn add_callinfo_trampoline( &mut self, target: unsafe extern "C" fn(*const CallContext, *const u64) -> u64, context: *const CallContext, num_params: u32, ) -> usize { let idx = self.offsets.len(); self.offsets.push(self.code.len()); let mut stack_offset: u32 = num_params.checked_mul(8).unwrap(); if stack_offset % 16 == 0 { stack_offset += 8; } self.code.extend_from_slice(&[0x48, 0x81, 0xec]); // sub ?, %rsp self.code.extend_from_slice(value_to_bytes(&stack_offset)); for i in 0..num_params { match i { 0..=5 => { // mov %?, ?(%rsp) let prefix: &[u8] = match i { 0 => &[0x48, 0x89, 0xbc, 0x24], // rdi 1 => &[0x48, 0x89, 0xb4, 0x24], // rsi 2 => &[0x48, 0x89, 0x94, 0x24], // rdx 3 => &[0x48, 0x89, 0x8c, 0x24], // rcx 4 => &[0x4c, 0x89, 0x84, 0x24], // r8 5 => &[0x4c, 0x89, 0x8c, 0x24], // r9 _ => unreachable!(), }; self.code.extend_from_slice(prefix); self.code.extend_from_slice(value_to_bytes(&(i * 8u32))); } _ => { self.code.extend_from_slice(&[ 0x48, 0x8b, 0x84, 0x24, // mov ?(%rsp), %rax ]); self.code.extend_from_slice(value_to_bytes( &((i - 6) * 8u32 + stack_offset + 8/* ret addr */), )); // mov %rax, ?(%rsp) self.code.extend_from_slice(&[0x48, 0x89, 0x84, 0x24]); self.code.extend_from_slice(value_to_bytes(&(i * 8u32))); } } } self.code.extend_from_slice(&[ 0x48, 0xbf, // movabsq ?, %rdi ]); self.code.extend_from_slice(value_to_bytes(&context)); self.code.extend_from_slice(&[ 0x48, 0x89, 0xe6, // mov %rsp, %rsi ]); self.code.extend_from_slice(&[ 0x48, 0xb8, // movabsq ?, %rax ]); self.code.extend_from_slice(value_to_bytes(&target)); self.code.extend_from_slice(&[ 0xff, 0xd0, // callq *%rax ]); self.code.extend_from_slice(&[ 0x48, 0x81, 0xc4, // add ?, %rsp ]); self.code.extend_from_slice(value_to_bytes(&stack_offset)); self.code.extend_from_slice(&[ 0xc3, //retq ]); idx } /// Consumes the builder and builds the trampoline buffer. pub fn build(self) -> TrampolineBuffer { get_context(); // ensure lazy initialization is completed let mut code = CodeMemory::new(self.code.len()); code[..self.code.len()].copy_from_slice(&self.code); code.make_executable(); TrampolineBuffer { code, offsets: self.offsets, } } } impl TrampolineBuffer { /// Returns the trampoline pointer at index `idx`. pub fn get_trampoline(&self, idx: usize) -> *const Trampoline { &self.code[self.offsets[idx]] as *const u8 as *const Trampoline } } impl fmt::Debug for TrampolineBuffer { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "TrampolineBuffer {{}}") } } #[cfg(test)] mod tests { use super::*; #[test] fn test_context_trampoline() { struct TestContext { value: i32, } extern "C" fn do_add(a: i32, b: f32) -> f32 { let ctx = unsafe { &*(get_context() as *const TestContext) }; a as f32 + b + ctx.value as f32 } let mut builder = TrampolineBufferBuilder::new(); let ctx = TestContext { value: 3 }; let idx = builder.add_context_trampoline( do_add as usize as *const _, &ctx as *const TestContext as *const _, ); let buf = builder.build(); let t = buf.get_trampoline(idx); let ret = unsafe { mem::transmute::<_, extern "C" fn(i32, f32) -> f32>(t)(1, 2.0) as i32 }; assert_eq!(ret, 6); } #[test] fn test_callinfo_trampoline() { struct TestContext { value: i32, } unsafe extern "C" fn do_add(ctx: *const CallContext, args: *const u64) -> u64 { let ctx = &*(ctx as *const TestContext); let args: &[u64] = slice::from_raw_parts(args, 8); (args.iter().map(|x| *x as i32).fold(0, |a, b| a + b) + ctx.value) as u64 } let mut builder = TrampolineBufferBuilder::new(); let ctx = TestContext { value: 100 }; let idx = builder.add_callinfo_trampoline(do_add, &ctx as *const TestContext as *const _, 8); let buf = builder.build(); let t = buf.get_trampoline(idx); let ret = unsafe { mem::transmute::<_, extern "C" fn(i32, i32, i32, i32, i32, i32, i32, i32) -> i32>(t)( 1, 2, 3, 4, 5, 6, 7, 8, ) as i32 }; assert_eq!(ret, 136); } }