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#[cfg_attr(
40 any(
41 target_arch = "x86",
42 target_arch = "x86_64",
43 all(target_arch = "aarch64", target_vendor = "apple"),
44 ),
45 repr(C, align(16384))
46)]
47#[cfg_attr(
48 not(any(
49 target_arch = "x86",
50 target_arch = "x86_64",
51 all(target_arch = "aarch64", target_vendor = "apple"),
52 )),
53 repr(C, align(65536))
54)]
55struct Block([u8; PREKEY_LEN]);
56
57const BLOCK_ALIGN: usize = core::mem::align_of::<Block>();
59
60const _: () = assert!(BLOCK_ALIGN >= PREKEY_LEN);
63
64#[derive(Clone, Copy, Debug, PartialEq, Eq)]
78pub struct LockError {
79 os_code: Option<i32>,
80 page_size: Option<usize>,
81}
82
83impl LockError {
84 fn from_region(err: ®ion::Error) -> Self {
85 let os_code = match err {
86 region::Error::SystemCall(code) | region::Error::MachCall(code) => Some(*code),
87 _ => None,
88 };
89 Self {
90 os_code,
91 page_size: None,
92 }
93 }
94
95 pub fn raw_os_error(&self) -> Option<i32> {
98 self.os_code
99 }
100
101 pub fn page_too_large(&self) -> Option<usize> {
106 self.page_size
107 }
108}
109
110impl core::fmt::Display for LockError {
111 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
112 match (self.os_code, self.page_size) {
113 (_, Some(size)) => write!(
114 f,
115 "the sealed-memory pre-key was not locked: the {size}-byte page would be \
116 shared with other memory"
117 ),
118 (Some(code), None) => write!(
119 f,
120 "the sealed-memory pre-key could not be locked in RAM (OS error {code})"
121 ),
122 (None, None) => f.write_str("the sealed-memory pre-key could not be locked in RAM"),
123 }
124 }
125}
126
127impl std::error::Error for LockError {}
128
129pub struct Prekey {
161 block: alloc::boxed::Box<Block>,
163 lock: Option<region::LockGuard>,
164 lock_error: Option<LockError>,
165}
166
167impl Prekey {
168 pub fn new() -> Arc<Self> {
185 let mut block = alloc::boxed::Box::new(Block([0u8; PREKEY_LEN]));
186 let lock = lock_block(&block, region::page::size());
189 SysRng
190 .try_fill_bytes(&mut block.0)
191 .expect("SysRng failure is a program error");
192 Arc::new(Self::from_parts(block, lock))
193 }
194
195 fn from_parts(
196 block: alloc::boxed::Box<Block>,
197 lock: Result<region::LockGuard, LockError>,
198 ) -> Self {
199 let (lock, lock_error) = match lock {
200 Ok(guard) => (Some(guard), None),
201 Err(err) => (None, Some(err)),
202 };
203 Self {
204 block,
205 lock,
206 lock_error,
207 }
208 }
209
210 pub fn lock_status(&self) -> Result<(), LockError> {
212 match self.lock_error {
213 None => Ok(()),
214 Some(err) => Err(err),
215 }
216 }
217
218 #[inline]
220 fn sealing_key(&self, salt: &[u8; SALT_LEN]) -> Fixed<[u8; 32]> {
221 Fixed::new_with(|key: &mut [u8; 32]| {
222 let mut hasher = Sha256::new();
223 hasher.update(salt);
224 hasher.update(&self.block.0);
225 hasher.finalize_into(key.into());
226 })
227 }
228
229 #[inline]
230 fn cipher(&self, salt: &[u8; SALT_LEN]) -> ChaCha20Poly1305 {
231 self.sealing_key(salt)
232 .with_secret(|key| ChaCha20Poly1305::new(key.into()))
233 }
234
235 pub(crate) fn fresh_salt() -> [u8; SALT_LEN] {
238 let mut rng = SysRng;
239 let words: [u64; SALT_LEN / 8] = core::array::from_fn(|_| {
240 rng.try_next_u64()
241 .expect("SysRng failure is a program error")
242 });
243 core::array::from_fn(|i| words[i / 8].to_le_bytes()[i % 8])
244 }
245
246 pub(crate) fn seal_in_place(&self, salt: &[u8; SALT_LEN], buf: &mut [u8]) -> [u8; TAG_LEN] {
251 self.cipher(salt)
252 .encrypt_inout_detached(&Nonce::default(), &[], buf.into())
253 .expect("ChaCha20-Poly1305 rejects only inputs over 256 GiB")
254 .into()
255 }
256
257 pub(crate) fn open_in_place(&self, salt: &[u8; SALT_LEN], tag: &[u8; TAG_LEN], buf: &mut [u8]) {
262 let opened = self.cipher(salt).decrypt_inout_detached(
263 &Nonce::default(),
264 &[],
265 buf.into(),
266 tag.into(),
267 );
268 if opened.is_err() {
269 buf.zeroize();
270 panic!("sealed memory was modified or corrupted");
271 }
272 }
273}
274
275fn lock_block(block: &Block, page_size: usize) -> Result<region::LockGuard, LockError> {
278 if page_size > BLOCK_ALIGN {
279 return Err(LockError {
280 os_code: None,
281 page_size: Some(page_size),
282 });
283 }
284 region::lock(block.0.as_ptr(), core::mem::size_of::<Block>())
285 .map_err(|err| LockError::from_region(&err))
286}
287
288const SCRUB_LEN: usize = 32 * 1024;
290
291#[inline(never)]
302pub(crate) fn scrub_stack() {
303 let mut area = [0u64; SCRUB_LEN / 8];
304 area.zeroize();
305 core::hint::black_box(&area);
306}
307
308#[inline(always)]
315pub(crate) fn scrubbed<R>(f: impl FnOnce() -> R) -> R {
316 let result = std::panic::catch_unwind(core::panic::AssertUnwindSafe(f));
317 scrub_stack();
318 match result {
319 Ok(value) => value,
320 Err(payload) => std::panic::resume_unwind(payload),
321 }
322}
323
324impl Prekey {
325 fn wipe_and_unlock(&mut self) {
328 self.block.0.zeroize();
329 drop(self.lock.take());
330 }
331}
332
333impl Drop for Prekey {
334 fn drop(&mut self) {
335 self.wipe_and_unlock();
336 }
337}
338
339impl core::fmt::Debug for Prekey {
341 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
342 f.write_str("[REDACTED]")
343 }
344}
345
346#[cfg(test)]
347pub(crate) mod tests {
348 use super::*;
349
350 pub(crate) fn quota() -> std::sync::MutexGuard<'static, ()> {
354 static QUOTA: std::sync::Mutex<()> = std::sync::Mutex::new(());
355 QUOTA
356 .lock()
357 .unwrap_or_else(std::sync::PoisonError::into_inner)
358 }
359
360 fn random_block() -> alloc::boxed::Box<Block> {
361 let mut block = alloc::boxed::Box::new(Block([0u8; PREKEY_LEN]));
362 SysRng.try_fill_bytes(&mut block.0).unwrap();
363 block
364 }
365
366 #[test]
367 fn block_is_page_aligned_and_fills_whole_pages() {
368 let block = random_block();
369 let addr = block.0.as_ptr() as usize;
370 assert_eq!(addr % BLOCK_ALIGN, 0);
371 let page = region::page::size();
372 if page <= BLOCK_ALIGN {
373 assert_eq!(addr % page, 0);
374 assert_eq!(core::mem::size_of::<Block>() % page, 0);
375 } else {
376 let _q = quota();
377 assert_eq!(
378 Prekey::new().lock_status().unwrap_err().page_too_large(),
379 Some(page)
380 );
381 }
382 }
383
384 #[test]
385 fn a_failed_lock_is_reported_and_sealing_still_works() {
386 let _q = quota();
387 let prekey = Arc::new(Prekey::from_parts(
388 random_block(),
389 Err(LockError::from_region(®ion::Error::SystemCall(1453))),
390 ));
391 let err = prekey.lock_status().unwrap_err();
392 assert_eq!(err.raw_os_error(), Some(1453));
393 assert_eq!(err.page_too_large(), None);
394 assert_eq!(
395 format!("{err:?}"),
396 "LockError { os_code: Some(1453), page_size: None }"
397 );
398 assert_eq!(
399 err.to_string(),
400 "the sealed-memory pre-key could not be locked in RAM (OS error 1453)"
401 );
402
403 let salt = Prekey::fresh_salt();
404 let mut buf = *b"still sealed";
405 let tag = prekey.seal_in_place(&salt, &mut buf);
406 assert_ne!(&buf, b"still sealed");
407 prekey.open_in_place(&salt, &tag, &mut buf);
408 assert_eq!(&buf, b"still sealed");
409 }
410
411 #[test]
412 fn a_non_os_lock_failure_has_no_code() {
413 let prekey = Prekey::from_parts(
414 random_block(),
415 Err(LockError::from_region(®ion::Error::InvalidParameter(
416 "size",
417 ))),
418 );
419 let err = prekey.lock_status().unwrap_err();
420 assert_eq!(err.raw_os_error(), None);
421 assert_eq!(
422 format!("{err:?}"),
423 "LockError { os_code: None, page_size: None }"
424 );
425 }
426
427 #[test]
428 fn wipe_and_unlock_zeroes_the_prekey_and_releases_the_lock() {
429 let _q = quota();
430 let lock = region::lock(core::ptr::null::<u8>(), 0).map_err(|_| ());
431 assert!(lock.is_err(), "a zero-length lock is rejected");
432 let block = random_block();
433 let guard = region::lock(block.0.as_ptr(), PREKEY_LEN);
434 let locked = guard.is_ok();
437 let mut prekey = Prekey::from_parts(block, guard.map_err(|e| LockError::from_region(&e)));
438 assert!(prekey.block.0.iter().any(|&b| b != 0));
439 assert_eq!(prekey.lock.is_some(), locked);
440 prekey.wipe_and_unlock();
441 assert!(prekey.block.0.iter().all(|&b| b == 0));
442 assert!(prekey.lock.is_none());
443 }
444
445 #[test]
452 fn wipe_and_unlock_returns_the_quota() {
453 let _q = quota();
454 let mut held = Vec::new();
455 while held.len() < 1024 {
456 let block = random_block();
457 let lock = region::lock(block.0.as_ptr(), PREKEY_LEN);
458 let refused = lock.is_err();
459 held.push(Prekey::from_parts(
460 block,
461 lock.map_err(|e| LockError::from_region(&e)),
462 ));
463 if refused {
464 break;
465 }
466 }
467 let locked = held.iter().filter(|p| p.lock.is_some()).count();
468 if locked == held.len() {
469 return; }
471 for p in &mut held {
472 p.wipe_and_unlock();
473 }
474 let relocked: Vec<_> = (0..locked)
475 .map(|_| {
476 let block = random_block();
477 let lock = region::lock(block.0.as_ptr(), PREKEY_LEN);
478 Prekey::from_parts(block, lock.map_err(|e| LockError::from_region(&e)))
479 })
480 .collect();
481 let ok = relocked.iter().filter(|p| p.lock.is_some()).count();
482 assert!(
483 ok * 2 >= locked,
484 "only {ok} of {locked} relocked after unlocking"
485 );
486 drop(held);
487 }
488
489 #[test]
490 fn the_sealing_key_depends_on_every_prekey_bit() {
491 let _q = quota();
492 let mut prekey = Arc::try_unwrap(Prekey::new()).unwrap();
493 let salt = Prekey::fresh_salt();
494 let before = prekey.sealing_key(&salt);
495 prekey.block.0[PREKEY_LEN - 1] ^= 1;
496 let after = prekey.sealing_key(&salt);
497 assert_ne!(before.with_secret(|k| *k), after.with_secret(|k| *k));
498 }
499
500 #[test]
501 fn the_sealing_key_is_sha256_of_salt_then_prekey() {
502 let _q = quota();
503 let prekey = Prekey::new();
504 let salt = Prekey::fresh_salt();
505 let mut hasher = Sha256::new();
506 hasher.update(salt);
507 hasher.update(prekey.block.0);
508 let expected: [u8; 32] = hasher.finalize().into();
509 assert_eq!(prekey.sealing_key(&salt).with_secret(|k| *k), expected);
510 }
511
512 #[test]
513 fn every_salt_is_fresh() {
514 assert_ne!(Prekey::fresh_salt(), Prekey::fresh_salt());
515 }
516
517 #[test]
518 #[should_panic(expected = "sealed memory was modified or corrupted")]
519 fn a_wrong_tag_panics() {
520 let _q = quota();
521 let prekey = Prekey::new();
522 let salt = Prekey::fresh_salt();
523 let mut buf = [1u8; 8];
524 let mut tag = prekey.seal_in_place(&salt, &mut buf);
525 tag[0] ^= 1;
526 prekey.open_in_place(&salt, &tag, &mut buf);
527 }
528
529 #[test]
530 fn a_failed_open_wipes_the_buffer_before_panicking() {
531 let _q = quota();
532 let prekey = Prekey::new();
533 let salt = Prekey::fresh_salt();
534 let mut buf = [1u8; 8];
535 let tag = prekey.seal_in_place(&salt, &mut buf);
536 buf[0] ^= 1;
537 let result = std::panic::catch_unwind(core::panic::AssertUnwindSafe(|| {
538 prekey.open_in_place(&salt, &tag, &mut buf);
539 }));
540 assert!(result.is_err());
541 assert_eq!(buf, [0u8; 8]);
542 }
543
544 #[test]
545 fn a_page_larger_than_the_prekey_is_refused_without_locking() {
546 let block = random_block();
547 let too_large = BLOCK_ALIGN * 2;
550 let Err(refused) = lock_block(&block, too_large) else {
551 panic!("a page larger than the block's alignment must be refused");
552 };
553 assert_eq!(refused.page_too_large(), Some(too_large));
554 assert_eq!(refused.raw_os_error(), None);
555 assert_eq!(
556 refused.to_string(),
557 format!(
558 "the sealed-memory pre-key was not locked: the {too_large}-byte page would \
559 be shared with other memory"
560 )
561 );
562 let prekey = Prekey::from_parts(block, Err(refused));
563 assert!(prekey.lock.is_none());
564 assert_eq!(
565 prekey.lock_status().unwrap_err().page_too_large(),
566 Some(too_large)
567 );
568 }
569
570 #[test]
571 fn a_page_no_larger_than_the_block_alignment_is_locked() {
572 let page = region::page::size();
573 if page > BLOCK_ALIGN {
574 return;
577 }
578 let _q = quota();
579 let block = random_block();
580 assert!(lock_block(&block, page).is_ok());
581 }
582
583 #[test]
584 fn debug_is_redacted() {
585 let _q = quota();
586 assert_eq!(format!("{:?}", Prekey::new()), "[REDACTED]");
587 }
588}