flex_alloc_secure/
stack.rs

1//! Support for stack-allocated types which employ memory locking.
2
3use core::any::type_name;
4use core::fmt;
5use core::mem::MaybeUninit;
6use core::ptr;
7use core::slice;
8
9use const_default::ConstDefault;
10use rand_core::RngCore;
11use zeroize::{DefaultIsZeroes, Zeroize};
12
13use crate::{
14    alloc::{lock_pages, unlock_pages, UNINIT_ALLOC_BYTE},
15    bytes::FillBytes,
16    protect::SecureRef,
17};
18
19/// A stack-allocated value protected by locking its virtual memory page in physical memory.
20#[repr(align(4096))]
21#[cfg_attr(all(target_arch = "aarch64", target_os = "macos"), repr(align(16384)))]
22pub struct Secured<T: Copy>(MaybeUninit<T>);
23
24impl<T: Copy> Secured<T> {
25    /// For an existing `Secured` instance, fill with the default value
26    /// of `T` and call the closure `f` with a mutable reference.
27    pub fn borrow_default<F, R>(&mut self, f: F) -> R
28    where
29        F: FnOnce(SecureRef<&mut T>) -> R,
30        T: Default,
31    {
32        let lock = SecuredGuard::new(&mut self.0);
33        lock.0.write(T::default());
34        unsafe { lock.eval_inited(f) }
35    }
36
37    /// For an existing `Secured` instance, fill with a set of random
38    /// bytes, then return the result of calling the closure `f` with
39    /// a mutable reference.
40    pub fn borrow_random<F, R>(&mut self, rng: impl RngCore, f: F) -> R
41    where
42        F: FnOnce(SecureRef<&mut T>) -> R,
43        T: FillBytes,
44    {
45        let mut lock = SecuredGuard::new(&mut self.0);
46        lock.fill_random(rng);
47        unsafe { lock.eval_inited(f) }
48    }
49
50    /// For an existing `Secured` instance, fill with the default value
51    /// of `T` and call the closure `f` with a mutable reference.
52    pub fn borrow_take<F, R>(&mut self, take: &mut T, f: F) -> R
53    where
54        F: FnOnce(SecureRef<&mut T>) -> R,
55        T: DefaultIsZeroes,
56    {
57        let lock = SecuredGuard::new(&mut self.0);
58        lock.0.write(*take);
59        take.zeroize();
60        unsafe { lock.eval_inited(f) }
61    }
62
63    /// For an existing `Secured` instance, call a closure `f` with a
64    /// mutable reference to an uninitialized `T`.
65    pub fn borrow_uninit<F, R>(&mut self, f: F) -> R
66    where
67        F: FnOnce(SecureRef<&mut MaybeUninit<T>>) -> R,
68    {
69        let mut lock = SecuredGuard::new(&mut self.0);
70        lock.fill_bytes(UNINIT_ALLOC_BYTE);
71        f(SecureRef::new_mut(lock.0))
72    }
73
74    /// Fill a `Secured` with the default value of `T` and return the result
75    /// of calling the closure `f` with a mutable reference.
76    pub fn default<F, R>(f: F) -> R
77    where
78        F: FnOnce(SecureRef<&mut T>) -> R,
79        T: Default,
80    {
81        let mut slf = Secured::DEFAULT;
82        slf.borrow_default(f)
83    }
84
85    /// Fill a `Secured` with a set of random bytes, then return the
86    /// result of calling the closure `f` with a mutable reference.
87    pub fn random<F, R>(rng: impl RngCore, f: F) -> R
88    where
89        F: FnOnce(SecureRef<&mut T>) -> R,
90        T: FillBytes,
91    {
92        let mut slf = Secured::DEFAULT;
93        slf.borrow_random(rng, f)
94    }
95
96    /// Fill a `Secured` by coping an existing value of type `T`,
97    /// and zeroize the original copy. Return the result of calling the
98    /// closure `f` with a mutable reference.
99    pub fn take<F, R>(take: &mut T, f: F) -> R
100    where
101        F: FnOnce(SecureRef<&mut T>) -> R,
102        T: DefaultIsZeroes,
103    {
104        let mut slf = Secured::DEFAULT;
105        slf.borrow_take(take, f)
106    }
107
108    /// Call the closure `f` with a mutable reference to an uninitialized
109    /// `T`.
110    pub fn uninit<F, R>(f: F) -> R
111    where
112        F: FnOnce(SecureRef<&mut MaybeUninit<T>>) -> R,
113    {
114        let mut slf = Secured::DEFAULT;
115        slf.borrow_uninit(f)
116    }
117}
118
119impl<const N: usize> Secured<[u8; N]> {
120    /// For an existing `Secured` instance, call a closure with a mutable
121    /// reference to an array of bytes. The values of the bytes are
122    /// initialized to a standard indicator value.
123    pub fn borrow_bytes<F, R>(&mut self, f: F) -> R
124    where
125        F: FnOnce(SecureRef<&mut [u8; N]>) -> R,
126    {
127        let mut lock = SecuredGuard::new(&mut self.0);
128        lock.fill_bytes(UNINIT_ALLOC_BYTE);
129        unsafe { lock.eval_inited(f) }
130    }
131
132    /// Call the closure `f` with a mutable reference to an array of
133    /// bytes. The values of the bytes are initialized to a standard
134    /// indicator value.
135    pub fn bytes<F, R>(f: F) -> R
136    where
137        F: FnOnce(SecureRef<&mut [u8; N]>) -> R,
138    {
139        let mut slf = Secured::DEFAULT;
140        slf.borrow_bytes(f)
141    }
142}
143
144impl<T: Copy> fmt::Debug for Secured<T> {
145    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
146        f.write_fmt(format_args!("Secured<{}>", type_name::<T>()))
147    }
148}
149
150impl<T: Copy> ConstDefault for Secured<T> {
151    const DEFAULT: Self = Secured(MaybeUninit::<T>::uninit());
152}
153
154impl<T: Copy> Default for Secured<T> {
155    fn default() -> Self {
156        Self::DEFAULT
157    }
158}
159
160struct SecuredGuard<'a, T>(&'a mut MaybeUninit<T>);
161
162impl<'a, T> SecuredGuard<'a, T> {
163    pub fn new(data: &'a mut MaybeUninit<T>) -> Self {
164        lock_pages(data.as_mut_ptr().cast(), size_of::<T>()).expect("Error locking stack memory");
165        Self(data)
166    }
167
168    #[inline]
169    // SAFETY: `self.0` must be initialized prior to calling.
170    pub unsafe fn eval_inited<R>(self, f: impl FnOnce(SecureRef<&mut T>) -> R) -> R {
171        struct Dropper<'d, D>(&'d mut D);
172
173        impl<D> Drop for Dropper<'_, D> {
174            fn drop(&mut self) {
175                unsafe {
176                    ptr::drop_in_place(self.0);
177                }
178            }
179        }
180
181        let drop = Dropper(self.0.assume_init_mut());
182        f(SecureRef::new_mut(drop.0))
183    }
184}
185
186unsafe impl<T> FillBytes for SecuredGuard<'_, T> {
187    fn as_bytes_mut(&mut self) -> &mut [u8] {
188        let len: usize = size_of_val(self.0);
189        unsafe { slice::from_raw_parts_mut(self.0 as *mut MaybeUninit<T> as *mut u8, len) }
190    }
191}
192
193impl<T> Drop for SecuredGuard<'_, T> {
194    fn drop(&mut self) {
195        self.0.zeroize();
196        match unlock_pages(self.0.as_mut_ptr().cast(), size_of::<T>()) {
197            Ok(_) => (),
198            Err(_) => {
199                if !std::thread::panicking() {
200                    panic!("Error unlocking memory");
201                }
202            }
203        };
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use const_default::ConstDefault;
210    use rand_core::OsRng;
211
212    use super::Secured;
213    use crate::alloc::UNINIT_ALLOC_BYTE;
214
215    #[test]
216    fn secured_default() {
217        let mut sec = Secured::<usize>::DEFAULT;
218        #[cfg_attr(miri, allow(unused))]
219        let ptr = sec.borrow_default(|mut b| {
220            assert_eq!(&*b, &0);
221            *b = 99usize;
222            &*b as *const usize
223        });
224        // ensure the value is zeroized after use
225        #[cfg(not(miri))]
226        assert_eq!(unsafe { *ptr }, 0usize);
227
228        Secured::<[u8; 10]>::default(|r| {
229            assert_eq!(&*r, &[0; 10]);
230        });
231    }
232
233    #[test]
234    fn secured_random() {
235        // The comparisons below could spuriously fail, with a low probability.
236
237        Secured::<[u8; 10]>::random(OsRng, |r| {
238            assert_ne!(&*r, &[0u8; 10]);
239        });
240
241        let mut sec = Secured::<[u8; 20]>::DEFAULT;
242        sec.borrow_random(OsRng, |r| {
243            assert_ne!(&*r, &[0u8; 20]);
244        });
245    }
246
247    #[test]
248    fn secured_take() {
249        let mut value = 99usize;
250
251        Secured::take(&mut value, |v| {
252            assert_eq!(&*v, &99);
253        });
254        // ensure the value is zeroized when taken
255        assert_eq!(value, 0);
256    }
257
258    #[test]
259    fn secured_uninit() {
260        Secured::<[u8; 10]>::bytes(|r| {
261            assert_eq!(&*r, &[UNINIT_ALLOC_BYTE; 10]);
262        });
263        Secured::<u32>::uninit(|m| {
264            let val = unsafe { m.assume_init() };
265            assert_eq!(val.to_ne_bytes(), [UNINIT_ALLOC_BYTE; 4]);
266        });
267    }
268}
flex_alloc_secure/stack.rs

flex_alloc_secure/
stack.rs
