Skip to main content

bwx/
locked.rs

1use zeroize::Zeroize as _;
2
3const LEN: usize = 4096;
4
5static MLOCK_WORKS: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6
7/// RAII guard around `mlock`/`munlock`. `munlock` can spuriously fail with
8/// `ENOMEM` on musl under `RLIMIT_MEMLOCK` pressure (common in CI
9/// containers); pages are released on process exit anyway, so unlock is
10/// best-effort and never panics on drop.
11struct MlockGuard {
12    ptr: *mut core::ffi::c_void,
13    len: usize,
14}
15
16// The guard only tracks an address + length owned for the lifetime of the
17// owning `FixedVec`; safe to move across threads.
18unsafe impl Send for MlockGuard {}
19unsafe impl Sync for MlockGuard {}
20
21impl Drop for MlockGuard {
22    fn drop(&mut self) {
23        // SAFETY: (ptr, len) came from a successful `mlock` call on a
24        // `Box<FixedVec>` that is still live (guard is dropped before the
25        // box).
26        let _ = unsafe { rustix::mm::munlock(self.ptr, self.len) };
27    }
28}
29
30fn try_mlock(ptr: *const u8, len: usize) -> rustix::io::Result<MlockGuard> {
31    // rustix takes *mut c_void to match the POSIX signature, even though
32    // mlock doesn't mutate.
33    let p = ptr.cast::<core::ffi::c_void>().cast_mut();
34    // SAFETY: `ptr` points to a live allocation of at least `len` bytes
35    // owned by the caller.
36    unsafe { rustix::mm::mlock(p, len) }?;
37    Ok(MlockGuard { ptr: p, len })
38}
39
40pub struct FixedVec<const N: usize> {
41    data: [u8; N],
42    len: usize,
43}
44
45impl<const N: usize> FixedVec<N> {
46    fn new() -> Self {
47        Self {
48            data: [0u8; N],
49            len: 0,
50        }
51    }
52
53    const fn capacity() -> usize {
54        N
55    }
56
57    fn as_ptr(&self) -> *const u8 {
58        self.data.as_ptr()
59    }
60
61    fn as_slice(&self) -> &[u8] {
62        &self.data[..self.len]
63    }
64
65    fn as_mut_slice(&mut self) -> &mut [u8] {
66        &mut self.data[..self.len]
67    }
68
69    fn truncate(&mut self, len: usize) {
70        if len < self.len {
71            self.len = len;
72        }
73    }
74
75    fn extend(&mut self, it: impl Iterator<Item = u8>) {
76        for b in it {
77            assert!(self.len < N, "FixedVec capacity exceeded");
78            self.data[self.len] = b;
79            self.len += 1;
80        }
81    }
82}
83
84impl<const N: usize> Drop for FixedVec<N> {
85    fn drop(&mut self) {
86        self.data[..self.len].zeroize();
87    }
88}
89
90pub struct Vec {
91    data: Box<FixedVec<LEN>>,
92    _lock: Option<MlockGuard>,
93}
94
95impl Default for Vec {
96    fn default() -> Self {
97        let data = Box::new(FixedVec::<LEN>::new());
98        let lock = match MLOCK_WORKS.get() {
99            Some(true) => {
100                try_mlock(data.as_ptr(), FixedVec::<LEN>::capacity()).ok()
101            }
102            Some(false) => None,
103            None => {
104                match try_mlock(data.as_ptr(), FixedVec::<LEN>::capacity()) {
105                    Ok(lock) => {
106                        let _ = MLOCK_WORKS.set(true);
107                        Some(lock)
108                    }
109                    Err(e) => {
110                        if MLOCK_WORKS.set(false).is_ok() {
111                            eprintln!("failed to lock memory region: {e}");
112                        }
113                        None
114                    }
115                }
116            }
117        };
118        Self { data, _lock: lock }
119    }
120}
121
122impl Vec {
123    pub fn new() -> Self {
124        Self::default()
125    }
126
127    pub fn data(&self) -> &[u8] {
128        self.data.as_slice()
129    }
130
131    pub fn data_mut(&mut self) -> &mut [u8] {
132        self.data.as_mut_slice()
133    }
134
135    pub fn zero(&mut self) {
136        self.truncate(0);
137        self.data.extend(std::iter::repeat_n(0, LEN));
138    }
139
140    pub fn extend(&mut self, it: impl Iterator<Item = u8>) {
141        self.data.extend(it);
142    }
143
144    pub fn truncate(&mut self, len: usize) {
145        self.data.truncate(len);
146    }
147}
148
149impl Drop for Vec {
150    fn drop(&mut self) {
151        self.zero();
152        self.data.as_mut_slice().zeroize();
153    }
154}
155
156impl Clone for Vec {
157    fn clone(&self) -> Self {
158        let mut new_vec = Self::new();
159        new_vec.extend(self.data().iter().copied());
160        new_vec
161    }
162}
163
164#[derive(Clone)]
165pub struct Password {
166    password: Vec,
167}
168
169impl Password {
170    pub fn new(password: Vec) -> Self {
171        Self { password }
172    }
173
174    pub fn password(&self) -> &[u8] {
175        self.password.data()
176    }
177}
178
179#[derive(Clone)]
180pub struct Keys {
181    keys: Vec,
182}
183
184impl Keys {
185    pub fn new(keys: Vec) -> Self {
186        Self { keys }
187    }
188
189    pub fn enc_key(&self) -> &[u8] {
190        &self.keys.data()[0..32]
191    }
192
193    pub fn mac_key(&self) -> &[u8] {
194        &self.keys.data()[32..64]
195    }
196
197    /// Full 64-byte `enc_key` || `mac_key` buffer.
198    pub fn as_bytes(&self) -> &[u8] {
199        &self.keys.data()[0..64]
200    }
201}
202
203#[derive(Clone)]
204pub struct PasswordHash {
205    hash: Vec,
206}
207
208impl PasswordHash {
209    pub fn new(hash: Vec) -> Self {
210        Self { hash }
211    }
212
213    pub fn hash(&self) -> &[u8] {
214        self.hash.data()
215    }
216}
217
218#[derive(Clone)]
219pub struct PrivateKey {
220    private_key: Vec,
221}
222
223impl PrivateKey {
224    pub fn new(private_key: Vec) -> Self {
225        Self { private_key }
226    }
227
228    pub fn private_key(&self) -> &[u8] {
229        self.private_key.data()
230    }
231}
232
233#[derive(Clone)]
234pub struct ApiKey {
235    client_id: Password,
236    client_secret: Password,
237}
238
239impl ApiKey {
240    pub fn new(client_id: Password, client_secret: Password) -> Self {
241        Self {
242            client_id,
243            client_secret,
244        }
245    }
246
247    pub fn client_id(&self) -> &[u8] {
248        self.client_id.password()
249    }
250
251    pub fn client_secret(&self) -> &[u8] {
252        self.client_secret.password()
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::FixedVec;
259
260    #[test]
261    fn push_len_and_slice() {
262        let mut v = FixedVec::<8>::new();
263        v.extend([1u8, 2, 3, 4].into_iter());
264        assert_eq!(v.as_slice().len(), 4);
265        assert_eq!(v.as_slice(), &[1, 2, 3, 4]);
266    }
267
268    #[test]
269    fn truncate_and_clear() {
270        let mut v = FixedVec::<8>::new();
271        v.extend([1u8, 2, 3, 4].into_iter());
272        v.truncate(0);
273        assert!(v.as_slice().is_empty());
274        assert_eq!(v.data[..4], [1, 2, 3, 4]);
275    }
276
277    #[test]
278    #[should_panic(expected = "FixedVec capacity exceeded")]
279    fn push_past_capacity_panics() {
280        let mut v = FixedVec::<2>::new();
281        v.extend([1u8, 2, 3].into_iter());
282    }
283
284    #[test]
285    fn fixed_vec_drop_zeros_written_bytes() {
286        // FixedVec::Drop must zeroize the written region. The memory cannot
287        // be observed after drop, so the Drop body's zeroize call is invoked
288        // manually here and the internal `data` array is checked.
289        let mut v = FixedVec::<8>::new();
290        v.extend([0xaa_u8, 0xbb, 0xcc, 0xdd].into_iter());
291        assert_eq!(v.data[..4], [0xaa, 0xbb, 0xcc, 0xdd]);
292        {
293            use zeroize::Zeroize as _;
294            v.data[..v.len].zeroize();
295        }
296        assert_eq!(v.data[..4], [0, 0, 0, 0]);
297    }
298
299    #[test]
300    fn locked_vec_extend_and_data() {
301        let mut v = super::Vec::new();
302        v.extend([1_u8, 2, 3, 4].iter().copied());
303        assert_eq!(v.data(), &[1, 2, 3, 4]);
304    }
305
306    #[test]
307    fn locked_vec_zero_fills_and_exposes_full_slice() {
308        let mut v = super::Vec::new();
309        v.extend([9_u8; 16].iter().copied());
310        v.zero();
311        // After zero(), the logical slice covers full capacity and reads
312        // as all zeros; previous contents must not be visible.
313        assert_eq!(v.data().len(), super::LEN);
314        assert!(v.data().iter().all(|b| *b == 0));
315    }
316
317    #[test]
318    fn locked_vec_truncate_shrinks_visible_slice() {
319        let mut v = super::Vec::new();
320        v.extend((0_u8..32).chain(std::iter::repeat_n(0, 0)));
321        assert_eq!(v.data().len(), 32);
322        v.truncate(8);
323        assert_eq!(v.data(), &(0_u8..8).collect::<std::vec::Vec<_>>()[..]);
324    }
325
326    #[test]
327    fn locked_vec_clone_is_independent() {
328        let mut original = super::Vec::new();
329        original.extend([1_u8, 2, 3, 4].iter().copied());
330        let copy = original.clone();
331        assert_eq!(copy.data(), &[1, 2, 3, 4]);
332        original.data_mut()[0] = 99;
333        assert_eq!(copy.data(), &[1, 2, 3, 4]);
334    }
335
336    #[test]
337    fn keys_exposes_enc_mac_split() {
338        let mut buf = super::Vec::new();
339        buf.extend((0_u8..64).collect::<std::vec::Vec<_>>().into_iter());
340        let k = super::Keys::new(buf);
341        assert_eq!(k.enc_key().len(), 32);
342        assert_eq!(k.mac_key().len(), 32);
343        assert_eq!(k.as_bytes().len(), 64);
344        assert_eq!(k.enc_key()[0], 0);
345        assert_eq!(k.enc_key()[31], 31);
346        assert_eq!(k.mac_key()[0], 32);
347        assert_eq!(k.mac_key()[31], 63);
348    }
349}