#![cfg(any(test, feature = "failpoints"))]
use mnemo::{KdfParams, Memory, MemoryType, Mnemo, MnemoConfig};
use tempfile::tempdir;
fn fast_cfg(dimensions: usize) -> MnemoConfig {
MnemoConfig {
dimensions,
kdf: KdfParams::fast(),
..Default::default()
}
}
fn mem(i: usize) -> Memory {
Memory::new(
format!("m-{i}"),
MemoryType::Semantic,
vec![i as f32, 1.0, 2.0, 3.0],
)
}
fn read_bytes(p: &std::path::Path) -> Vec<u8> {
std::fs::read(p).unwrap()
}
fn write_bytes(p: &std::path::Path, b: &[u8]) {
std::fs::write(p, b).unwrap()
}
#[test]
fn sector_0_zeroed_heals_from_wal() {
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = Mnemo::create(&path, "pw", fast_cfg(4)).unwrap();
for i in 0..6 {
db.remember(mem(i)).unwrap();
}
db.close().unwrap();
drop(db);
let mut bytes = read_bytes(&path);
for b in bytes.iter_mut().take(512) {
*b = 0;
}
write_bytes(&path, &bytes);
let db = Mnemo::open(&path, "pw").expect("zeroed sector-0 must heal");
assert_eq!(db.len(), 6);
}
#[test]
fn sector_0_random_garbage_heals_from_wal() {
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = Mnemo::create(&path, "pw", fast_cfg(4)).unwrap();
for i in 0..8 {
db.remember(mem(i)).unwrap();
}
db.close().unwrap();
drop(db);
let mut bytes = read_bytes(&path);
for (i, b) in bytes.iter_mut().take(512).enumerate() {
*b = (i as u8).wrapping_mul(37).wrapping_add(13);
}
write_bytes(&path, &bytes);
let db = Mnemo::open(&path, "pw").expect("garbaged sector-0 must heal");
assert_eq!(db.len(), 8);
}
#[test]
fn header_magic_and_version_bytes_corrupted_heal() {
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = Mnemo::create(&path, "pw", fast_cfg(4)).unwrap();
for i in 0..5 {
db.remember(mem(i)).unwrap();
}
db.close().unwrap();
drop(db);
let mut bytes = read_bytes(&path);
for b in bytes.iter_mut().take(6) {
*b ^= 0xFF;
}
write_bytes(&path, &bytes);
let db = Mnemo::open(&path, "pw").expect("corrupted magic/version must heal");
assert_eq!(db.len(), 5);
}
#[test]
fn header_seal_fields_stale_report_tampering() {
use mnemo::MnemoError;
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = Mnemo::create(&path, "pw", fast_cfg(4)).unwrap();
for i in 0..4 {
db.remember(mem(i)).unwrap();
}
db.close().unwrap();
drop(db);
let mut bytes = read_bytes(&path);
bytes[260] ^= 0x01;
write_bytes(&path, &bytes);
match Mnemo::open(&path, "pw") {
Err(MnemoError::HeaderTampered) => {}
Err(other) => panic!("expected HeaderTampered, got {other:?}"),
Ok(_) => panic!("expected HeaderTampered; silent heal would erase evidence"),
}
}
#[test]
fn header_sealed_field_altered_with_recomputed_crc_reports_tampering() {
use mnemo::MnemoError;
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = Mnemo::create(&path, "pw", fast_cfg(4)).unwrap();
for i in 0..4 {
db.remember(mem(i)).unwrap();
}
db.close().unwrap();
drop(db);
let mut bytes = read_bytes(&path);
bytes[42] ^= 0x01;
let new_crc = mnemo_crc32_adapter(&bytes[0..238]).to_le_bytes();
bytes[238..242].copy_from_slice(&new_crc);
write_bytes(&path, &bytes);
match Mnemo::open(&path, "pw") {
Err(MnemoError::HeaderTampered) => {}
Err(other) => panic!("expected HeaderTampered, got {other:?}"),
Ok(_) => panic!("CRC-recomputed sealed-field tamper slipped past the seal"),
}
}
fn mnemo_crc32_adapter(data: &[u8]) -> u32 {
let mut c = !0u32;
for &b in data {
c ^= b as u32;
for _ in 0..8 {
let mask = (c & 1).wrapping_neg();
c = (c >> 1) ^ (0xEDB8_8320 & mask);
}
}
!c
}
#[test]
fn relocated_txn_larger_than_64_pages_heals_torn_page_0() {
let dir = tempdir().unwrap();
let path = dir.path().join("big.mnemo");
let mut db = Mnemo::create(&path, "pw", fast_cfg(2)).unwrap();
for i in 0..25_000 {
db.remember(Memory::new(
format!("rec-{i}"),
MemoryType::Semantic,
vec![((i % 101) + 1) as f32, ((i % 97) + 1) as f32],
))
.unwrap();
}
db.flush().unwrap();
let wal_pages_before = db.stats().unwrap().wal_pages;
db.close().unwrap();
drop(db);
assert!(
wal_pages_before > 64,
"the WAL must have grown past 64 pages for this test to be meaningful \
(observed {wal_pages_before} pages)"
);
let mut bytes = read_bytes(&path);
for b in bytes.iter_mut().take(512) {
*b = 0;
}
write_bytes(&path, &bytes);
let mut db = Mnemo::open(&path, "pw")
.expect("torn page 0 with the WAL at the file tail must heal via the streaming scan");
let stats = db.stats().unwrap();
assert_eq!(stats.memories, 25_000);
assert_eq!(stats.wal_pages, wal_pages_before);
}
fn disarm() {
mnemo::__failpoints::set_writes_until_fail(-1);
}
fn seed_up_to_relocation(path: &std::path::Path) -> Mnemo {
let mut db = Mnemo::create(path, "pw", fast_cfg(2)).unwrap();
for i in 0..1200 {
db.remember(Memory::new(
format!("seed-{i}"),
MemoryType::Semantic,
vec![((i % 7) + 1) as f32, ((i % 11) + 1) as f32],
))
.unwrap();
}
db.flush().unwrap();
db
}
#[test]
fn relocation_crash_recovers_to_pre_or_post_state() {
for fail_at in 0i64..32 {
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = seed_up_to_relocation(&path);
let pre = db.len();
for i in 0..20_000 {
db.remember(Memory::new(
format!("post-{i}"),
MemoryType::Semantic,
vec![((i % 13) + 1) as f32, ((i % 17) + 1) as f32],
))
.unwrap();
}
let post = pre + 20_000;
mnemo::__failpoints::set_writes_until_fail(fail_at);
let flush_result = db.flush();
disarm();
drop(db);
let db = Mnemo::open(&path, "pw")
.unwrap_or_else(|e| panic!("reopen after fail_at={fail_at} failed: {e:?}"));
let observed = db.len();
assert!(
observed == pre || observed == post,
"fail_at={fail_at}: observed={observed}, pre={pre}, post={post}, \
flush_result={flush_result:?}"
);
}
}
#[test]
fn rekey_crash_recovers_to_old_or_new_passphrase() {
for fail_at in 0i64..16 {
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = Mnemo::create(&path, "old-pw", fast_cfg(4)).unwrap();
for i in 0..5 {
db.remember(mem(i)).unwrap();
}
db.flush().unwrap();
mnemo::__failpoints::set_writes_until_fail(fail_at);
let _ = db.rekey("new-pw", KdfParams::fast());
disarm();
drop(db);
let opened = match Mnemo::open(&path, "new-pw") {
Ok(h) => Some(("new", h)),
Err(_) => Mnemo::open(&path, "old-pw").ok().map(|h| ("old", h)),
};
let (which, db) = opened
.unwrap_or_else(|| panic!("rekey crash at fail_at={fail_at} left no openable file"));
assert_eq!(db.len(), 5, "rekey crash ({which}) lost memories");
}
}
#[test]
fn rekey_crash_on_never_flushed_file_keeps_dek() {
for fail_at in 0i64..8 {
let dir = tempdir().unwrap();
let path = dir.path().join("t.mnemo");
let mut db = Mnemo::create(&path, "old-pw", fast_cfg(4)).unwrap();
mnemo::__failpoints::set_writes_until_fail(fail_at);
let _ = db.rekey("new-pw", KdfParams::fast());
disarm();
drop(db);
let opened = Mnemo::open(&path, "new-pw").or_else(|_| Mnemo::open(&path, "old-pw"));
assert!(
opened.is_ok(),
"never-flushed rekey crash at fail_at={fail_at} lost the DEK"
);
}
}