use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::OnceLock;
fn page_size() -> usize {
static PAGE_SIZE: OnceLock<usize> = OnceLock::new();
*PAGE_SIZE.get_or_init(|| {
#[cfg(miri)]
{
4096
}
#[cfg(not(miri))]
{
let v = unsafe { libc::sysconf(libc::_SC_PAGESIZE) };
if v > 0 {
v as usize
} else {
4096
}
}
})
}
#[derive(Debug)]
pub struct PinnedPageRange {
pub start: *const u8,
pub page_count: usize,
}
impl PinnedPageRange {
fn empty() -> Self {
Self {
start: std::ptr::null(),
page_count: 0,
}
}
}
impl Drop for PinnedPageRange {
fn drop(&mut self) {
if self.page_count == 0 || self.start.is_null() {
return;
}
let bytes = self.page_count * page_size();
let _ = unsafe { sys_munlock(self.start, bytes) };
}
}
pub fn pin_pages_for(buf: &[u8]) -> PinnedPageRange {
if buf.is_empty() {
return PinnedPageRange::empty();
}
let (start, page_count) = round_to_pages(buf.as_ptr() as usize, buf.len(), page_size());
let bytes = page_count * page_size();
let start_ptr = start as *const u8;
mlock_state().record_attempt();
let result = unsafe { sys_mlock_attempt(start_ptr, bytes) };
match result {
Ok(()) => PinnedPageRange {
start: start_ptr,
page_count,
},
Err(errno) => {
if errno == libc::EINVAL {
debug_assert!(
false,
"mlock returned EINVAL — should be unreachable from slice-fn API by construction (page_size={}, bytes={})",
page_size(),
bytes,
);
}
mlock_state().record_failure(errno, buf.len());
PinnedPageRange::empty()
}
}
}
fn round_to_pages(addr: usize, len: usize, page: usize) -> (usize, usize) {
debug_assert!(page.is_power_of_two(), "page size must be a power of two");
if len == 0 {
return (0, 0);
}
let mask = page - 1;
let start = addr & !mask;
let end = (addr + len + mask) & !mask;
(start, (end - start) / page)
}
struct MlockState {
attempts: AtomicUsize,
failure_count: AtomicUsize,
total_bytes_unlocked: AtomicUsize,
first_errno: OnceLock<i32>,
}
impl MlockState {
const fn new() -> Self {
Self {
attempts: AtomicUsize::new(0),
failure_count: AtomicUsize::new(0),
total_bytes_unlocked: AtomicUsize::new(0),
first_errno: OnceLock::new(),
}
}
fn record_attempt(&self) {
self.attempts.fetch_add(1, Ordering::Relaxed);
}
fn record_failure(&self, errno: i32, bytes: usize) {
self.failure_count.fetch_add(1, Ordering::Relaxed);
self.total_bytes_unlocked.fetch_add(bytes, Ordering::Relaxed);
let _ = self.first_errno.set(errno);
}
}
static MLOCK_STATE: OnceLock<MlockState> = OnceLock::new();
fn mlock_state() -> &'static MlockState {
MLOCK_STATE.get_or_init(MlockState::new)
}
pub fn report_at_exit() {
let Some(st) = MLOCK_STATE.get() else {
return;
};
let failures = st.failure_count.load(Ordering::Relaxed);
if failures == 0 {
return;
}
let attempts = st.attempts.load(Ordering::Relaxed);
let bytes = st.total_bytes_unlocked.load(Ordering::Relaxed);
let errno_name = st
.first_errno
.get()
.map(|&e| errno_to_name(e))
.unwrap_or("?");
eprintln!("warning: {failures} of {attempts} secret regions could not be locked");
eprintln!(" (first errno: {errno_name}, {bytes} bytes total); secret");
eprintln!(" data remains in heap and may be swappable.");
eprintln!("hint: set RLIMIT_MEMLOCK >= 64KiB or grant CAP_IPC_LOCK");
eprintln!(" to eliminate this warning.");
}
fn errno_to_name(errno: i32) -> &'static str {
match errno {
libc::EPERM => "EPERM",
libc::ENOMEM => "ENOMEM",
libc::EAGAIN => "EAGAIN",
libc::EINVAL => "EINVAL",
libc::ENOTSUP => "ENOTSUP",
_ => "UNKNOWN",
}
}
pub fn page_size_for_test() -> usize {
page_size()
}
pub fn failure_count_for_test() -> usize {
MLOCK_STATE
.get()
.map(|s| s.failure_count.load(Ordering::Relaxed))
.unwrap_or(0)
}
pub fn attempts_for_test() -> usize {
MLOCK_STATE
.get()
.map(|s| s.attempts.load(Ordering::Relaxed))
.unwrap_or(0)
}
pub fn first_errno_for_test() -> Option<i32> {
MLOCK_STATE.get().and_then(|s| s.first_errno.get().copied())
}
#[cfg(miri)]
unsafe fn sys_mlock_attempt(_addr: *const u8, _len: usize) -> Result<(), i32> {
Ok(())
}
#[cfg(miri)]
unsafe fn sys_munlock(_addr: *const u8, _len: usize) -> i32 {
0
}
#[cfg(all(not(miri), not(test)))]
unsafe fn sys_mlock_attempt(addr: *const u8, len: usize) -> Result<(), i32> {
let rc = unsafe { libc::mlock(addr as *const libc::c_void, len) };
if rc == 0 {
Ok(())
} else {
Err(last_os_errno())
}
}
#[cfg(all(not(miri), not(test)))]
unsafe fn sys_munlock(addr: *const u8, len: usize) -> i32 {
unsafe { libc::munlock(addr as *const libc::c_void, len) }
}
#[cfg(all(not(miri), test))]
unsafe fn sys_mlock_attempt(addr: *const u8, len: usize) -> Result<(), i32> {
match fail_mode::current() {
fail_mode::FailMode::Off => {
let rc = unsafe { libc::mlock(addr as *const libc::c_void, len) };
if rc == 0 {
Ok(())
} else {
Err(last_os_errno())
}
}
fail_mode::FailMode::EPerm => Err(libc::EPERM),
fail_mode::FailMode::ENoMem => Err(libc::ENOMEM),
fail_mode::FailMode::EInval => Err(libc::EINVAL),
}
}
#[cfg(all(not(miri), test))]
unsafe fn sys_munlock(addr: *const u8, len: usize) -> i32 {
unsafe { libc::munlock(addr as *const libc::c_void, len) }
}
#[cfg(not(miri))]
fn last_os_errno() -> i32 {
std::io::Error::last_os_error().raw_os_error().unwrap_or(0)
}
#[cfg(test)]
mod fail_mode {
use std::sync::OnceLock;
pub enum FailMode {
Off,
EPerm,
ENoMem,
EInval,
}
pub fn parse(s: &str) -> Option<FailMode> {
match s {
"off" => Some(FailMode::Off),
"eperm" => Some(FailMode::EPerm),
"enomem" => Some(FailMode::ENoMem),
"einval" => Some(FailMode::EInval),
_ => None,
}
}
static FAIL_MODE: OnceLock<FailMode> = OnceLock::new();
pub fn current() -> &'static FailMode {
FAIL_MODE.get_or_init(|| {
std::env::var("MNEMONIC_TEST_MLOCK_FAIL_MODE")
.ok()
.as_deref()
.and_then(parse)
.unwrap_or(FailMode::Off)
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn page_rounding_formula_single_page() {
let page = 4096;
let (start, count) = round_to_pages(0x1234, 64, page);
assert_eq!(start, 0x1000, "round down to page boundary");
assert_eq!(count, 1, "small buf fits in 1 page");
}
#[test]
fn page_rounding_formula_multi_page() {
let page = 4096;
let (start, count) = round_to_pages(0x1000, 2 * page + 1, page);
assert_eq!(start, 0x1000);
assert_eq!(count, 3, "2*page+1 spans 3 pages when starting page-aligned");
}
#[test]
fn page_rounding_formula_zero_length() {
let page = 4096;
let (start, count) = round_to_pages(0x1234, 0, page);
assert_eq!(start, 0, "zero-length yields empty range");
assert_eq!(count, 0);
}
#[test]
fn page_rounding_formula_exactly_one_page_aligned() {
let page = 4096;
let (start, count) = round_to_pages(0x2000, page, page);
assert_eq!(start, 0x2000, "already page-aligned");
assert_eq!(count, 1, "exactly-one-page buf at aligned address yields 1");
}
#[test]
fn mlockstate_record_failure_idempotent_on_first_errno() {
let st = MlockState::new();
st.record_failure(libc::EPERM, 64);
st.record_failure(libc::ENOMEM, 128);
assert_eq!(
st.first_errno.get().copied(),
Some(libc::EPERM),
"first_errno is set once and stays",
);
}
#[test]
fn mlockstate_record_failure_monotonic_on_counters() {
let st = MlockState::new();
st.record_attempt();
st.record_attempt();
st.record_attempt();
st.record_failure(libc::EPERM, 64);
st.record_failure(libc::ENOMEM, 128);
assert_eq!(st.attempts.load(Ordering::Relaxed), 3);
assert_eq!(st.failure_count.load(Ordering::Relaxed), 2);
assert_eq!(st.total_bytes_unlocked.load(Ordering::Relaxed), 192);
}
#[test]
fn g4_a_pin_and_zeroize_compose_without_panic() {
use zeroize::Zeroize;
let mut v: Vec<u8> = vec![0xAAu8; 64];
assert_eq!(v[0], 0xAA);
let pin = pin_pages_for(&v);
assert_eq!(pin.page_count, 1, "64-byte buf pins exactly one page");
v.zeroize();
assert_eq!(v.len(), 0, "zeroize clears Vec len after scrubbing");
drop(pin); }
#[test]
#[ignore = "subprocess: requires MNEMONIC_TEST_MLOCK_FAIL_MODE=eperm in env"]
fn g2_1_eperm_increments_failure_count() {
let buf = vec![0u8; 64];
let _pin = pin_pages_for(&buf);
assert!(
failure_count_for_test() > 0,
"FAIL_MODE=eperm must increment failure_count via record_failure",
);
assert_eq!(
first_errno_for_test(),
Some(libc::EPERM),
"first_errno must be EPERM after eperm injection",
);
}
#[test]
#[cfg(debug_assertions)]
#[ignore = "subprocess: requires MNEMONIC_TEST_MLOCK_FAIL_MODE=einval in env"]
#[should_panic(expected = "EINVAL")]
fn g2_3_einval_debug_panics() {
let buf = vec![0u8; 64];
let _pin = pin_pages_for(&buf);
}
#[test]
#[cfg(not(debug_assertions))]
#[ignore = "subprocess: requires MNEMONIC_TEST_MLOCK_FAIL_MODE=einval in env"]
fn g2_3_einval_release_soft_fails() {
let buf = vec![0u8; 64];
let _pin = pin_pages_for(&buf);
assert!(
failure_count_for_test() > 0,
"FAIL_MODE=einval in release build must soft-fail (increment failure_count, no panic)",
);
assert_eq!(
first_errno_for_test(),
Some(libc::EINVAL),
"first_errno must be EINVAL after einval injection in release build",
);
}
#[test]
#[ignore = "subprocess: requires MNEMONIC_TEST_MLOCK_FAIL_MODE=off + sufficient ulimit"]
fn g2_4_off_no_synthesized_failures() {
let buf = vec![0u8; 64];
let _pin = pin_pages_for(&buf);
assert_eq!(
failure_count_for_test(),
0,
"FAIL_MODE=off must not synthesize failures (test env requires ulimit -l >= 64KiB)",
);
}
}