1#![no_std]
2
3use core::alloc::{GlobalAlloc, Layout};
4
5pub struct ZeroizingAllocator<Alloc: GlobalAlloc>(pub Alloc);
6
7unsafe fn zero(ptr: *mut u8, size: usize) {
8 for i in 0..size {
9 core::ptr::write_volatile(ptr.offset(i as isize), 0);
10 }
11 core::sync::atomic::compiler_fence(core::sync::atomic::Ordering::SeqCst);
12}
13
14unsafe impl<T> GlobalAlloc for ZeroizingAllocator<T>
15where
16 T: GlobalAlloc,
17{
18 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
19 self.0.alloc(layout)
20 }
21
22 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
23 zero(ptr, layout.size());
24 #[cfg(not(feature = "leaky"))]
25 self.0.dealloc(ptr, layout);
26 }
27
28 unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
29 self.0.alloc_zeroed(layout)
30 }
31}
32
33#[cfg(all(feature = "leaky", test))]
34mod test {
35 extern crate std;
36 use std::vec::Vec;
37
38 #[global_allocator]
39 static ALLOC: super::ZeroizingAllocator<std::alloc::System> =
40 super::ZeroizingAllocator(std::alloc::System);
41
42 #[test]
43 fn test() {
44 let mut a = Vec::with_capacity(2);
45 a.push(0xde);
46 a.push(0xad);
47 let ptr1: *const u8 = &a[0];
48 a.push(0xbe);
49 a.push(0xef);
50 let ptr2: *const u8 = &a[0];
51 assert_eq!(&[0xde, 0xad, 0xbe, 0xef], &a[..]);
52 assert_eq!(unsafe { ptr1.as_ref() }, Some(&0));
53 drop(a);
54 assert_eq!(unsafe { ptr2.as_ref() }, Some(&0));
55 }
56}