locked_vec/
locked_vec.rs

1use std::io;
2use std::thread;
3use std::time::Duration;
4
5/// A minimal safe wrapper around a locked Vec<u8>.
6///
7/// Design notes:
8/// - Owns its buffer and never changes capacity after lock (no push/reserve API exposed).
9/// - Attempts `mlock` on construction; gracefully handles Unsupported platforms.
10/// - Zeroizes the content before drop, then calls `munlock` if it was locked.
11/// - On Linux, it also attempts `madvise(MADV_DONTDUMP)` as a best-effort hint.
12///
13/// This is an example for demonstration and is not a hardened secret container.
14/// For production, prefer a battle-tested type from a dedicated crate and enforce
15/// stricter invariants (pinning, poison on failure, etc).
16struct LockedVec {
17    buf: Vec<u8>,
18    locked: bool,
19}
20
21impl LockedVec {
22    /// Construct a new LockedVec with a fixed length.
23    ///
24    /// Returns Ok even when `mlock` is Unsupported on this platform/build.
25    /// Returns Err for other OS errors (e.g., resource limits).
26    pub fn new(len: usize) -> io::Result<Self> {
27        // Allocate a zeroed buffer. We won't change capacity after locking.
28        let buf = vec![0u8; len];
29
30        // Attempt to lock pages. Treat Unsupported as a non-fatal condition.
31        let ptr = buf.as_ptr() as *const std::os::raw::c_void;
32        let locked = match unsafe { os_memlock::mlock(ptr, buf.len()) } {
33            Ok(()) => true,
34            Err(e) if e.kind() == io::ErrorKind::Unsupported => {
35                eprintln!(
36                    "os-memlock: mlock unsupported on this platform/build; continuing unlocked"
37                );
38                false
39            }
40            Err(e) => return Err(e),
41        };
42
43        // Best-effort: on Linux, advise dump exclusion.
44        #[cfg(target_os = "linux")]
45        {
46            let mut_ptr = buf.as_mut_ptr() as *mut std::os::raw::c_void;
47            match unsafe { os_memlock::madvise_dontdump(mut_ptr, buf.len()) } {
48                Ok(()) => (),
49                Err(e) if e.kind() == io::ErrorKind::Unsupported => {
50                    // Not supported on this target; ignore gracefully.
51                }
52                Err(e) => {
53                    // Non-fatal: just log the error.
54                    eprintln!("os-memlock: madvise(MADV_DONTDUMP) failed: {e}");
55                }
56            }
57        }
58
59        Ok(Self { buf, locked })
60    }
61
62    /// Length in bytes.
63    pub fn len(&self) -> usize {
64        self.buf.len()
65    }
66
67    /// Get a mutable slice view of the buffer.
68    ///
69    /// Caller may write secrets here. The buffer will be zeroized on Drop
70    /// and `munlock` will be called if the buffer was locked.
71    pub fn as_mut_slice(&mut self) -> &mut [u8] {
72        &mut self.buf
73    }
74
75    /// Whether `mlock` succeeded for this buffer.
76    pub fn is_locked(&self) -> bool {
77        self.locked
78    }
79}
80
81impl Drop for LockedVec {
82    fn drop(&mut self) {
83        // Zeroize contents while still locked (if locked).
84        for b in &mut self.buf {
85            *b = 0;
86        }
87
88        if self.locked {
89            let ptr = self.buf.as_ptr() as *const std::os::raw::c_void;
90            let len = self.buf.len();
91            match unsafe { os_memlock::munlock(ptr, len) } {
92                Ok(()) => (),
93                Err(e) if e.kind() == io::ErrorKind::Unsupported => {
94                    // Not supported on this target; ignore.
95                }
96                Err(e) => {
97                    // We can't do much in Drop; log error.
98                    eprintln!("os-memlock: munlock failed: {e}");
99                }
100            }
101        }
102    }
103}
104
105fn main() -> io::Result<()> {
106    const LEN: usize = 4096;
107
108    let mut secrets = LockedVec::new(LEN)?;
109    println!(
110        "LockedVec allocated: {} bytes, locked: {}",
111        secrets.len(),
112        secrets.is_locked()
113    );
114
115    // Write a demo secret (for example purposes).
116    // In real usage, manage keys carefully and avoid unnecessary copies.
117    let s = secrets.as_mut_slice();
118    let demo = b"super-secret-demo";
119    if s.len() >= demo.len() {
120        s[..demo.len()].copy_from_slice(demo);
121    }
122
123    // Simulate work while the buffer is (ideally) locked in memory.
124    println!("Working with secret data (simulated)...");
125    thread::sleep(Duration::from_millis(200));
126
127    // Drop occurs at end of scope: zeroize then munlock (if locked).
128    println!("Example complete; dropping LockedVec (zeroize + munlock).");
129    Ok(())
130}