diff --git a/src/isolation/stack.rs b/src/isolation/stack.rs index 29e73106a..f0b87ec8b 100644 --- a/src/isolation/stack.rs +++ b/src/isolation/stack.rs @@ -17,17 +17,33 @@ extern "C" { pub struct StackContext { addr: *mut u8, - size: usize, - thread_context: ThreadContext, + thread_context: Box>, } -struct ThreadContext { - target: fn(T) -> R, +pub struct ThreadContext { + target: fn(&mut ThreadContext, T) -> R, param: Option, ret: Option, jmp_buffer: [c_int; SETJMP_BUFFER_LEN], } +impl ThreadContext { + /// do_yield swaps the current state with the saved one. + /// + /// This is marked as safe because we assume that: + /// - A ThreadContext can only be constructed from within this module. + /// - This function is never inlined so that the compiler consider the whole ThreadContext as possibly modified after calling this. + #[inline(never)] // prevent compiler from looking into do_yield when performing optimizations. + pub fn do_yield(&mut self) { + unsafe { + let mut yield_point = self.jmp_buffer; + if setjmp(&mut self.jmp_buffer) == 0 { + longjmp(&mut yield_point, 1); + } + } + } +} + struct PthreadAttr { inner: pthread_attr_t, } @@ -56,36 +72,46 @@ impl Drop for PthreadAttr { } impl StackContext { - pub unsafe fn new(addr: *mut u8, size: usize, f: fn(T) -> R, param: T) -> StackContext { + /// Creates a new stack context. + /// + /// This function is unsafe because `addr` and `size` must be ensured by the user to be valid. + pub unsafe fn new( + addr: *mut u8, + f: fn(&mut ThreadContext, T) -> R, + param: T, + ) -> StackContext { extern "C" fn run_context(ctx: *mut c_void) -> *mut c_void { unsafe { let ctx = &mut *(ctx as *mut ThreadContext); if setjmp(&mut ctx.jmp_buffer) != 0 { - ctx.ret = Some((ctx.target)(ctx.param.take().unwrap())); + let target = ctx.target; + let param = ctx.param.take().unwrap(); + let ret = (target)(ctx, param); + ctx.ret = Some(ret); + ctx.do_yield(); } } ::std::ptr::null_mut() } let mut attr = PthreadAttr::new(); - if pthread_attr_setstack(&mut attr.inner, addr as *mut c_void, size as size_t) != 0 { - panic!( - "pthread_attr_setstack failed, addr = {:?}, size = {}", - addr, size - ); + + // The value 65536 is chosen arbitrarily. Any reasonable value for stack size works here. + if pthread_attr_setstack(&mut attr.inner, addr as *mut c_void, 65536) != 0 { + panic!("pthread_attr_setstack failed, addr = {:?}", addr); } - let mut thread_ctx = ThreadContext { + let mut thread_ctx = Box::new(ThreadContext { target: f, param: Some(param), ret: None, jmp_buffer: ::std::mem::uninitialized(), - }; + }); let mut pthread_handle: pthread_t = ::std::mem::uninitialized(); if pthread_create( &mut pthread_handle, &mut attr.inner, run_context::, - &mut thread_ctx as *mut ThreadContext as *mut c_void, + thread_ctx.as_mut() as *mut ThreadContext as *mut c_void, ) != 0 { panic!("pthread_create failed"); @@ -93,10 +119,20 @@ impl StackContext { pthread_join(pthread_handle, ::std::ptr::null_mut()); StackContext { addr: addr, - size: size, thread_context: thread_ctx, } } + + /// Continues execution of the current stack context. + /// + /// This function is unsafe because nested stack context is not supported and may cause UB. + pub unsafe fn next(mut self) -> Result { + self.thread_context.do_yield(); + match self.thread_context.ret { + Some(x) => Ok(x), + None => Err(self), + } + } } #[cfg(test)] @@ -105,16 +141,31 @@ mod tests { #[test] fn test_stack_context() { - fn some_fn(_: ()) {} + fn some_fn(ctx: &mut ThreadContext<&mut i32, ()>, out: &mut i32) { + for i in 0..100i32 { + *out += i; + ctx.do_yield(); + } + } let mut buf: Vec = Vec::with_capacity(65536); unsafe { buf.set_len(65536) }; - let ctx = unsafe { + + let mut sum: i32 = 0; + let mut ctx = Some(unsafe { StackContext::new( ((buf.as_mut_ptr().offset(buf.len() as isize) as usize) & (!4095usize)) as _, - buf.len(), some_fn, - (), + &mut sum, ) - }; + }); + loop { + match unsafe { ctx.take().unwrap().next() } { + Ok(x) => break, + Err(x) => { + ctx = Some(x); + } + } + } + assert_eq!(sum, 4950); } }