thunk.rs

extern crate memmap;

#[cfg(target_pointer_width = "32")]
const MAGIC_VALUE: usize = 0xdedbeef0;
#[cfg(target_pointer_width = "32")]
const WORD_SIZE: usize = 4;
#[cfg(target_pointer_width = "64")]
const MAGIC_VALUE: usize = 0xd0e0a0d0b0e0e0f0;
#[cfg(target_pointer_width = "64")]
const WORD_SIZE: usize = 8;

type Foreign = ();
type Callback = unsafe extern fn(f: *mut Foreign) -> u32;

#[inline(never)]
unsafe extern fn thunk_template(f: *mut Foreign) -> u32 {
	let cb: *const *mut FnMut(*mut Foreign) -> u32 = MAGIC_VALUE as *const _;
	(**cb)(f)
}
#[inline(never)]
unsafe extern fn thunk_template_end() {}

struct Thunk {
	mmap: memmap::Mmap,
	_func: Box<Box<FnMut(*mut Foreign) -> u32>>,
}

impl Thunk {
	fn new<F: FnMut(&mut Foreign) -> u32 + 'static>(mut f: F) -> Thunk {
		let boxed: Box<FnMut(*mut Foreign) -> u32> = Box::new(move |ptr| {
			unsafe { f(&mut *ptr) }
		});
		let double_boxed = Box::new(boxed);

		Thunk {
			mmap: unsafe { make_map(&*double_boxed as *const _ as *const _) },
			_func: double_boxed,
		}
	}

	unsafe fn as_callback(&self) -> Callback {
		std::mem::transmute(self.mmap.ptr())
	}
}

fn replace(slice: &mut [u8], from: &[u8], to: &[u8]) {
	assert!(from.len() == to.len());
	for i in 0..slice.len() - from.len() {
		let subslice = &mut slice[i..i+from.len()];
		if subslice == from {
			subslice.copy_from_slice(to);
		}
	}
}

unsafe fn usize_slice(inp: &usize) -> &[u8] {
	std::slice::from_raw_parts(inp as *const _ as *const _, std::mem::size_of::<usize>())
}

unsafe fn make_map(pointer: *const Box<Fn(*mut Foreign) -> u32>) -> memmap::Mmap {
	let len = (thunk_template_end as usize) - (thunk_template as usize);
	let slice = std::slice::from_raw_parts(thunk_template as *const u8, len);

	let mut map = memmap::Mmap::anonymous(len, memmap::Protection::ReadWrite).unwrap();
	map.as_mut_slice().copy_from_slice(slice);
	replace(map.as_mut_slice(), usize_slice(&MAGIC_VALUE), usize_slice(&(pointer as usize)));
	replace(map.as_mut_slice(), usize_slice(&(MAGIC_VALUE + WORD_SIZE)), usize_slice(&(pointer as usize + WORD_SIZE)));

	map.set_protection(memmap::Protection::ReadExecute).unwrap();
	map
}

// ----->8-----
// client code

// used to know that the closure's context is not leaked
struct DropAlert(i32);
impl DropAlert {
	fn say_hi(&self) {
		println!("Hi from DropAlert({})", self.0);
	}
}
impl Drop for DropAlert {
	fn drop(&mut self) {
		println!("DropAlert({}) is dropped", self.0);
	}
}

// main function
fn main() {
	println!("MAGIC_VALUE = {:x}", MAGIC_VALUE);
	println!("thunk_template = {:p}", thunk_template as *const ());

	let mut param = ();
	println!("param = {:p}\n", &param);

	{ // inner scope: Thunk is dropped at the end
		let da = DropAlert(7);
		let thunk = Thunk::new(move |_| {
			println!("Hello from the inside");
			da.say_hi();
			5
		});
		let thunk_fn: Callback = unsafe { thunk.as_callback() };
		println!("closure thunk = {:p}", thunk_fn);
		println!("ret = {}", unsafe { thunk_fn(&mut param) });
	}

	println!("goodbye");
}

/* Output example ----
MAGIC_VALUE = d0e0a0d0b0e0e0f0
thunk_template = 0x4014f0
param = 0x22fbc0

closure thunk = 0x2e0000
Hello from the inside
Hi from DropAlert(7)
ret = 5
DropAlert(7) is dropped
goodbye
*/