1use alloc::sync::Arc;
4
5use chacha20poly1305::aead::AeadInOut;
6use chacha20poly1305::{ChaCha20Poly1305, KeyInit, Nonce};
7use rand::{TryRng, rngs::SysRng};
8use sha2::{Digest, Sha256};
9use zeroize::Zeroize;
10
11use crate::{Fixed, RevealSecret};
12
13pub(crate) const PREKEY_LEN: usize = 16 * 1024;
15
16pub(crate) const TAG_LEN: usize = 16;
18
19pub(crate) const SALT_LEN: usize = 32;
21
22#[repr(C, align(16384))]
26struct Block([u8; PREKEY_LEN]);
27
28#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub struct LockError {
42 os_code: Option<i32>,
43 page_size: Option<usize>,
44}
45
46impl LockError {
47 fn from_region(err: ®ion::Error) -> Self {
48 let os_code = match err {
49 region::Error::SystemCall(code) | region::Error::MachCall(code) => Some(*code),
50 _ => None,
51 };
52 Self {
53 os_code,
54 page_size: None,
55 }
56 }
57
58 pub fn raw_os_error(&self) -> Option<i32> {
61 self.os_code
62 }
63
64 pub fn page_too_large(&self) -> Option<usize> {
67 self.page_size
68 }
69}
70
71impl core::fmt::Display for LockError {
72 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
73 match (self.os_code, self.page_size) {
74 (_, Some(size)) => write!(
75 f,
76 "the sealed-memory pre-key was not locked: the {size}-byte page would be \
77 shared with other memory"
78 ),
79 (Some(code), None) => write!(
80 f,
81 "the sealed-memory pre-key could not be locked in RAM (OS error {code})"
82 ),
83 (None, None) => f.write_str("the sealed-memory pre-key could not be locked in RAM"),
84 }
85 }
86}
87
88impl std::error::Error for LockError {}
89
90pub struct Prekey {
122 block: alloc::boxed::Box<Block>,
124 lock: Option<region::LockGuard>,
125 lock_error: Option<LockError>,
126}
127
128impl Prekey {
129 pub fn new() -> Arc<Self> {
138 let mut block = alloc::boxed::Box::new(Block([0u8; PREKEY_LEN]));
139 let lock = lock_block(&block, region::page::size());
142 SysRng
143 .try_fill_bytes(&mut block.0)
144 .expect("SysRng failure is a program error");
145 Arc::new(Self::from_parts(block, lock))
146 }
147
148 fn from_parts(
149 block: alloc::boxed::Box<Block>,
150 lock: Result<region::LockGuard, LockError>,
151 ) -> Self {
152 let (lock, lock_error) = match lock {
153 Ok(guard) => (Some(guard), None),
154 Err(err) => (None, Some(err)),
155 };
156 Self {
157 block,
158 lock,
159 lock_error,
160 }
161 }
162
163 pub fn lock_status(&self) -> Result<(), LockError> {
165 match self.lock_error {
166 None => Ok(()),
167 Some(err) => Err(err),
168 }
169 }
170
171 #[inline]
173 fn sealing_key(&self, salt: &[u8; SALT_LEN]) -> Fixed<[u8; 32]> {
174 Fixed::new_with(|key: &mut [u8; 32]| {
175 let mut hasher = Sha256::new();
176 hasher.update(salt);
177 hasher.update(&self.block.0);
178 hasher.finalize_into(key.into());
179 })
180 }
181
182 #[inline]
183 fn cipher(&self, salt: &[u8; SALT_LEN]) -> ChaCha20Poly1305 {
184 self.sealing_key(salt)
185 .with_secret(|key| ChaCha20Poly1305::new(key.into()))
186 }
187
188 pub(crate) fn fresh_salt() -> [u8; SALT_LEN] {
191 let mut rng = SysRng;
192 let words: [u64; SALT_LEN / 8] = core::array::from_fn(|_| {
193 rng.try_next_u64()
194 .expect("SysRng failure is a program error")
195 });
196 core::array::from_fn(|i| words[i / 8].to_le_bytes()[i % 8])
197 }
198
199 pub(crate) fn seal_in_place(&self, salt: &[u8; SALT_LEN], buf: &mut [u8]) -> [u8; TAG_LEN] {
204 self.cipher(salt)
205 .encrypt_inout_detached(&Nonce::default(), &[], buf.into())
206 .expect("ChaCha20-Poly1305 rejects only inputs over 256 GiB")
207 .into()
208 }
209
210 pub(crate) fn open_in_place(&self, salt: &[u8; SALT_LEN], tag: &[u8; TAG_LEN], buf: &mut [u8]) {
215 let opened = self.cipher(salt).decrypt_inout_detached(
216 &Nonce::default(),
217 &[],
218 buf.into(),
219 tag.into(),
220 );
221 if opened.is_err() {
222 buf.zeroize();
223 panic!("sealed memory was modified or corrupted");
224 }
225 }
226}
227
228fn lock_block(block: &Block, page_size: usize) -> Result<region::LockGuard, LockError> {
231 if page_size > PREKEY_LEN {
232 return Err(LockError {
233 os_code: None,
234 page_size: Some(page_size),
235 });
236 }
237 region::lock(block.0.as_ptr(), PREKEY_LEN).map_err(|err| LockError::from_region(&err))
238}
239
240const SCRUB_LEN: usize = 32 * 1024;
242
243#[inline(never)]
254pub(crate) fn scrub_stack() {
255 let mut area = [0u64; SCRUB_LEN / 8];
256 area.zeroize();
257 core::hint::black_box(&area);
258}
259
260#[inline(always)]
267pub(crate) fn scrubbed<R>(f: impl FnOnce() -> R) -> R {
268 let result = std::panic::catch_unwind(core::panic::AssertUnwindSafe(f));
269 scrub_stack();
270 match result {
271 Ok(value) => value,
272 Err(payload) => std::panic::resume_unwind(payload),
273 }
274}
275
276impl Prekey {
277 fn wipe_and_unlock(&mut self) {
280 self.block.0.zeroize();
281 drop(self.lock.take());
282 }
283}
284
285impl Drop for Prekey {
286 fn drop(&mut self) {
287 self.wipe_and_unlock();
288 }
289}
290
291impl core::fmt::Debug for Prekey {
293 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
294 f.write_str("[REDACTED]")
295 }
296}
297
298#[cfg(test)]
299pub(crate) mod tests {
300 use super::*;
301
302 pub(crate) fn quota() -> std::sync::MutexGuard<'static, ()> {
306 static QUOTA: std::sync::Mutex<()> = std::sync::Mutex::new(());
307 QUOTA
308 .lock()
309 .unwrap_or_else(std::sync::PoisonError::into_inner)
310 }
311
312 fn random_block() -> alloc::boxed::Box<Block> {
313 let mut block = alloc::boxed::Box::new(Block([0u8; PREKEY_LEN]));
314 SysRng.try_fill_bytes(&mut block.0).unwrap();
315 block
316 }
317
318 #[test]
319 fn block_is_page_aligned_and_fills_whole_pages() {
320 let block = random_block();
321 let addr = block.0.as_ptr() as usize;
322 assert_eq!(addr % 16384, 0);
323 assert_eq!(core::mem::size_of::<Block>(), PREKEY_LEN);
324 let page = region::page::size();
325 if page <= PREKEY_LEN {
326 assert_eq!(PREKEY_LEN % page, 0);
327 } else {
328 let _q = quota();
329 assert_eq!(
330 Prekey::new().lock_status().unwrap_err().page_too_large(),
331 Some(page)
332 );
333 }
334 }
335
336 #[test]
337 fn a_failed_lock_is_reported_and_sealing_still_works() {
338 let _q = quota();
339 let prekey = Arc::new(Prekey::from_parts(
340 random_block(),
341 Err(LockError::from_region(®ion::Error::SystemCall(1453))),
342 ));
343 let err = prekey.lock_status().unwrap_err();
344 assert_eq!(err.raw_os_error(), Some(1453));
345 assert_eq!(err.page_too_large(), None);
346 assert_eq!(
347 format!("{err:?}"),
348 "LockError { os_code: Some(1453), page_size: None }"
349 );
350 assert_eq!(
351 err.to_string(),
352 "the sealed-memory pre-key could not be locked in RAM (OS error 1453)"
353 );
354
355 let salt = Prekey::fresh_salt();
356 let mut buf = *b"still sealed";
357 let tag = prekey.seal_in_place(&salt, &mut buf);
358 assert_ne!(&buf, b"still sealed");
359 prekey.open_in_place(&salt, &tag, &mut buf);
360 assert_eq!(&buf, b"still sealed");
361 }
362
363 #[test]
364 fn a_non_os_lock_failure_has_no_code() {
365 let prekey = Prekey::from_parts(
366 random_block(),
367 Err(LockError::from_region(®ion::Error::InvalidParameter(
368 "size",
369 ))),
370 );
371 let err = prekey.lock_status().unwrap_err();
372 assert_eq!(err.raw_os_error(), None);
373 assert_eq!(
374 format!("{err:?}"),
375 "LockError { os_code: None, page_size: None }"
376 );
377 }
378
379 #[test]
380 fn wipe_and_unlock_zeroes_the_prekey_and_releases_the_lock() {
381 let _q = quota();
382 let lock = region::lock(core::ptr::null::<u8>(), 0).map_err(|_| ());
383 assert!(lock.is_err(), "a zero-length lock is rejected");
384 let block = random_block();
385 let guard = region::lock(block.0.as_ptr(), PREKEY_LEN);
386 let locked = guard.is_ok();
389 let mut prekey = Prekey::from_parts(block, guard.map_err(|e| LockError::from_region(&e)));
390 assert!(prekey.block.0.iter().any(|&b| b != 0));
391 assert_eq!(prekey.lock.is_some(), locked);
392 prekey.wipe_and_unlock();
393 assert!(prekey.block.0.iter().all(|&b| b == 0));
394 assert!(prekey.lock.is_none());
395 }
396
397 #[test]
404 fn wipe_and_unlock_returns_the_quota() {
405 let _q = quota();
406 let mut held = Vec::new();
407 while held.len() < 1024 {
408 let block = random_block();
409 let lock = region::lock(block.0.as_ptr(), PREKEY_LEN);
410 let refused = lock.is_err();
411 held.push(Prekey::from_parts(
412 block,
413 lock.map_err(|e| LockError::from_region(&e)),
414 ));
415 if refused {
416 break;
417 }
418 }
419 let locked = held.iter().filter(|p| p.lock.is_some()).count();
420 if locked == held.len() {
421 return; }
423 for p in &mut held {
424 p.wipe_and_unlock();
425 }
426 let relocked: Vec<_> = (0..locked)
427 .map(|_| {
428 let block = random_block();
429 let lock = region::lock(block.0.as_ptr(), PREKEY_LEN);
430 Prekey::from_parts(block, lock.map_err(|e| LockError::from_region(&e)))
431 })
432 .collect();
433 let ok = relocked.iter().filter(|p| p.lock.is_some()).count();
434 assert!(
435 ok * 2 >= locked,
436 "only {ok} of {locked} relocked after unlocking"
437 );
438 drop(held);
439 }
440
441 #[test]
442 fn the_sealing_key_depends_on_every_prekey_bit() {
443 let _q = quota();
444 let mut prekey = Arc::try_unwrap(Prekey::new()).unwrap();
445 let salt = Prekey::fresh_salt();
446 let before = prekey.sealing_key(&salt);
447 prekey.block.0[PREKEY_LEN - 1] ^= 1;
448 let after = prekey.sealing_key(&salt);
449 assert_ne!(before.with_secret(|k| *k), after.with_secret(|k| *k));
450 }
451
452 #[test]
453 fn the_sealing_key_is_sha256_of_salt_then_prekey() {
454 let _q = quota();
455 let prekey = Prekey::new();
456 let salt = Prekey::fresh_salt();
457 let mut hasher = Sha256::new();
458 hasher.update(salt);
459 hasher.update(prekey.block.0);
460 let expected: [u8; 32] = hasher.finalize().into();
461 assert_eq!(prekey.sealing_key(&salt).with_secret(|k| *k), expected);
462 }
463
464 #[test]
465 fn every_salt_is_fresh() {
466 assert_ne!(Prekey::fresh_salt(), Prekey::fresh_salt());
467 }
468
469 #[test]
470 #[should_panic(expected = "sealed memory was modified or corrupted")]
471 fn a_wrong_tag_panics() {
472 let _q = quota();
473 let prekey = Prekey::new();
474 let salt = Prekey::fresh_salt();
475 let mut buf = [1u8; 8];
476 let mut tag = prekey.seal_in_place(&salt, &mut buf);
477 tag[0] ^= 1;
478 prekey.open_in_place(&salt, &tag, &mut buf);
479 }
480
481 #[test]
482 fn a_failed_open_wipes_the_buffer_before_panicking() {
483 let _q = quota();
484 let prekey = Prekey::new();
485 let salt = Prekey::fresh_salt();
486 let mut buf = [1u8; 8];
487 let tag = prekey.seal_in_place(&salt, &mut buf);
488 buf[0] ^= 1;
489 let result = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| {
490 prekey.open_in_place(&salt, &tag, &mut buf);
491 }));
492 assert!(result.is_err());
493 assert_eq!(buf, [0u8; 8]);
494 }
495
496 #[test]
497 fn a_page_larger_than_the_prekey_is_refused_without_locking() {
498 let block = random_block();
499 let Err(refused) = lock_block(&block, 65536) else {
501 panic!("a page larger than the pre-key must be refused");
502 };
503 assert_eq!(refused.page_too_large(), Some(65536));
504 assert_eq!(refused.raw_os_error(), None);
505 assert_eq!(
506 refused.to_string(),
507 "the sealed-memory pre-key was not locked: the 65536-byte page would be shared \
508 with other memory"
509 );
510 let prekey = Prekey::from_parts(block, Err(refused));
511 assert!(prekey.lock.is_none());
512 assert_eq!(
513 prekey.lock_status().unwrap_err().page_too_large(),
514 Some(65536)
515 );
516 }
517
518 #[test]
519 fn a_page_no_larger_than_the_prekey_is_locked() {
520 let _q = quota();
521 let block = random_block();
522 let page = region::page::size().min(PREKEY_LEN);
523 assert!(lock_block(&block, page).is_ok());
524 }
525
526 #[test]
527 fn debug_is_redacted() {
528 let _q = quota();
529 assert_eq!(format!("{:?}", Prekey::new()), "[REDACTED]");
530 }
531}