1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78
use fnoption::FnOption; use hide::{hide_mem, hide_ptr}; /// Calls a closure and overwrites its stack on return. /// /// This function calls `clear_stack` after calling the passed closure, /// taking care to prevent either of them being inlined, so the stack /// used by the closure will be overwritten with zeros (as long as a /// large enough number of `pages` is used). /// /// For technical reasons, this function can be used only with `Fn` or /// `FnMut`. If all you have is a `FnOnce`, use the auxiliary function /// `clear_stack_on_return_fnonce` instead. /// /// # Example /// /// ``` /// # use clear_on_drop::clear_stack_on_return; /// # fn encrypt(input: &[u8]) -> Vec<u8> { input.to_owned() } /// let input = b"abc"; /// let result = clear_stack_on_return(1, || encrypt(input)); /// ``` #[inline] pub fn clear_stack_on_return<F, R>(pages: usize, mut f: F) -> R where F: FnMut() -> R, { let _clear = ClearStackOnDrop { pages: pages }; // Do not inline f to make sure clear_stack uses the same stack space. hide_ptr::<&mut FnMut() -> R>(&mut f)() } /// Calls a closure and overwrites its stack on return. /// /// This function is a variant of `clear_stack_on_return` which also /// accepts `FnOnce`, at the cost of being slightly slower. /// /// # Example /// /// ``` /// # use clear_on_drop::clear_stack_on_return_fnonce; /// # fn encrypt(input: Vec<u8>) -> Vec<u8> { input } /// let input = vec![97, 98, 99]; /// let result = clear_stack_on_return_fnonce(1, || encrypt(input)); /// ``` #[inline] pub fn clear_stack_on_return_fnonce<F, R>(pages: usize, f: F) -> R where F: FnOnce() -> R, { let mut f = FnOption::new(f); clear_stack_on_return(pages, || f.call_mut()).unwrap() } struct ClearStackOnDrop { pages: usize, } impl Drop for ClearStackOnDrop { #[inline] fn drop(&mut self) { // Do not inline clear_stack. hide_ptr::<fn(usize)>(clear_stack)(self.pages); } } /// Overwrites a few pages of stack. /// /// This function will overwrite `pages` 4096-byte blocks of the stack /// with zeros. pub fn clear_stack(pages: usize) { if pages > 0 { let mut buf = [0u8; 4096]; hide_mem(&mut buf); // prevent moving after recursive call clear_stack(pages - 1); hide_mem(&mut buf); // prevent reuse of stack space for call } }