#![cfg(all(feature = "encryption", feature = "test-faults"))]
use std::collections::BTreeSet;
use std::fs;
use std::process::Command;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering};
use std::sync::{Mutex, MutexGuard, OnceLock};
use armdb::armour::Db;
use armdb::format_migration::failpoints;
use armdb::{
CollectionMeta, Config, DbError, FormatAction, FormatMigration, NoHook, PageCipher, TxConfig,
};
use armour_core::GetType;
use tempfile::tempdir;
use zerocopy::{FromBytes, Immutable, IntoBytes};
const NAME: &str = "format_migration_encrypted";
const FAIL_NAME: &str = "format_migration_encrypted_cleanup";
const FAIL_ROWS: usize = 256;
const SENTINEL: u64 = 0x1122_3344_5566_7788;
static INSPECT_ROOT: OnceLock<Mutex<Option<std::path::PathBuf>>> = OnceLock::new();
static TEMP_VIOLATION: AtomicBool = AtomicBool::new(false);
static TEMP_PEAK: AtomicU64 = AtomicU64::new(0);
static PREFLIGHT_ROOT: OnceLock<Mutex<Option<std::path::PathBuf>>> = OnceLock::new();
static PREFLIGHT_TEMP_EXISTS: AtomicBool = AtomicBool::new(false);
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes)]
#[repr(C)]
struct V1(u32);
impl CollectionMeta for V1 {
type SelfId = [u8; 8];
const NAME: &'static str = NAME;
const VERSION: u16 = 1;
}
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes, PartialEq, Eq, Debug)]
#[repr(C)]
struct V2(u64);
fn widen(_: &[u8], old: &[u8], out: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
out.extend_from_slice(
&(u32::from_ne_bytes(old.try_into().expect("V1 width")) as u64).to_ne_bytes(),
);
Ok(FormatAction::Write)
}
const WIDEN: FormatMigration = FormatMigration::raw_checked(1, &V1::TYPE, &V2::TYPE, widen);
impl CollectionMeta for V2 {
type SelfId = [u8; 8];
const NAME: &'static str = NAME;
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[WIDEN];
}
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes, PartialEq, Eq, Debug)]
#[repr(C)]
struct RetryV2(u64);
fn widen_changed(_: &[u8], old: &[u8], out: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
let value = u32::from_ne_bytes(old.try_into().expect("V1 width"));
out.extend_from_slice(&(u64::from(value) + 10_000).to_ne_bytes());
Ok(FormatAction::Write)
}
const WIDEN_CHANGED: FormatMigration =
FormatMigration::raw_checked(1, &V1::TYPE, &RetryV2::TYPE, widen_changed);
impl CollectionMeta for RetryV2 {
type SelfId = [u8; 8];
const NAME: &'static str = NAME;
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[WIDEN_CHANGED];
}
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes, PartialEq, Eq, Debug)]
#[repr(C)]
struct ObservedV2(u64);
fn widen_observed(_: &[u8], old: &[u8], out: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
inspect_temporary_output();
out.extend_from_slice(
&(u32::from_ne_bytes(old.try_into().expect("V1 width")) as u64).to_ne_bytes(),
);
Ok(FormatAction::Write)
}
const WIDEN_OBSERVED: FormatMigration =
FormatMigration::raw_checked(1, &V1::TYPE, &ObservedV2::TYPE, widen_observed);
impl CollectionMeta for ObservedV2 {
type SelfId = [u8; 8];
const NAME: &'static str = NAME;
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[WIDEN_OBSERVED];
}
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes, PartialEq, Eq, Debug)]
#[repr(C)]
struct PreflightV2(u64);
fn widen_preflight(_: &[u8], old: &[u8], out: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
inspect_preflight_before_temporary_output();
out.extend_from_slice(
&(u32::from_ne_bytes(old.try_into().expect("V1 width")) as u64).to_ne_bytes(),
);
Ok(FormatAction::Write)
}
const PREFLIGHT_WIDEN: FormatMigration =
FormatMigration::raw_checked(1, &V1::TYPE, &PreflightV2::TYPE, widen_preflight);
impl CollectionMeta for PreflightV2 {
type SelfId = [u8; 8];
const NAME: &'static str = NAME;
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[PREFLIGHT_WIDEN];
}
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes, PartialEq, Eq, Debug)]
#[repr(C)]
struct DropV2(u32);
fn drop_row(_: &[u8], _: &[u8], _: &mut Vec<u8>) -> armdb::DbResult<FormatAction> {
Ok(FormatAction::Drop)
}
const DROP: FormatMigration = FormatMigration::raw_checked(1, &V1::TYPE, &DropV2::TYPE, drop_row);
impl CollectionMeta for DropV2 {
type SelfId = [u8; 8];
const NAME: &'static str = NAME;
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[DROP];
}
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes)]
#[repr(C)]
struct FailV1(u32);
impl CollectionMeta for FailV1 {
type SelfId = [u8; 8];
const NAME: &'static str = FAIL_NAME;
const VERSION: u16 = 1;
}
#[derive(Clone, Copy, FromBytes, GetType, Immutable, IntoBytes, PartialEq, Eq, Debug)]
#[repr(C)]
struct FailV2(u64);
static FAIL_CALLS: AtomicUsize = AtomicUsize::new(0);
static FAIL_AT: AtomicUsize = AtomicUsize::new(usize::MAX);
static FAILPOINT_LOCK: Mutex<()> = Mutex::new(());
struct FailpointGuard {
_lock: MutexGuard<'static, ()>,
}
impl FailpointGuard {
fn acquire() -> Self {
let lock = FAILPOINT_LOCK
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
reset_failpoints();
Self { _lock: lock }
}
}
impl Drop for FailpointGuard {
fn drop(&mut self) {
reset_failpoints();
}
}
fn reset_failpoints() {
FAIL_CALLS.store(0, Ordering::SeqCst);
FAIL_AT.store(usize::MAX, Ordering::SeqCst);
failpoints::CLEANUP_ERROR.store(failpoints::CLEANUP_NONE, Ordering::SeqCst);
TEMP_VIOLATION.store(false, Ordering::SeqCst);
TEMP_PEAK.store(0, Ordering::SeqCst);
PREFLIGHT_TEMP_EXISTS.store(false, Ordering::SeqCst);
if let Some(root) = INSPECT_ROOT.get() {
*root.lock().unwrap() = None;
}
if let Some(root) = PREFLIGHT_ROOT.get() {
*root.lock().unwrap() = None;
}
}
#[test]
fn encrypted_failpoint_guard_resets_observer_globals() {
let guard = FailpointGuard::acquire();
TEMP_VIOLATION.store(true, Ordering::SeqCst);
TEMP_PEAK.store(1, Ordering::SeqCst);
PREFLIGHT_TEMP_EXISTS.store(true, Ordering::SeqCst);
*INSPECT_ROOT
.get_or_init(|| Mutex::new(None))
.lock()
.unwrap() = Some(std::path::PathBuf::from("stale-inspect-root"));
*PREFLIGHT_ROOT
.get_or_init(|| Mutex::new(None))
.lock()
.unwrap() = Some(std::path::PathBuf::from("stale-preflight-root"));
drop(guard);
let _guard = FailpointGuard::acquire();
assert!(!TEMP_VIOLATION.load(Ordering::SeqCst));
assert_eq!(TEMP_PEAK.load(Ordering::SeqCst), 0);
assert!(!PREFLIGHT_TEMP_EXISTS.load(Ordering::SeqCst));
assert!(INSPECT_ROOT.get().unwrap().lock().unwrap().is_none());
assert!(PREFLIGHT_ROOT.get().unwrap().lock().unwrap().is_none());
}
fn widen_with_late_failure(
_: &[u8],
old: &[u8],
out: &mut Vec<u8>,
) -> armdb::DbResult<FormatAction> {
let call = FAIL_CALLS.fetch_add(1, Ordering::SeqCst);
if call == FAIL_AT.load(Ordering::SeqCst) {
return Err(DbError::Internal("injected encrypted output failure"));
}
out.extend_from_slice(
&(u32::from_ne_bytes(old.try_into().expect("FailV1 width")) as u64).to_ne_bytes(),
);
Ok(FormatAction::Write)
}
const FAIL_WIDEN: FormatMigration =
FormatMigration::raw_checked(1, &FailV1::TYPE, &FailV2::TYPE, widen_with_late_failure);
impl CollectionMeta for FailV2 {
type SelfId = [u8; 8];
const NAME: &'static str = FAIL_NAME;
const VERSION: u16 = 2;
const FORMAT_MIGRATIONS: &'static [FormatMigration] = &[FAIL_WIDEN];
}
fn config() -> Config {
let mut config = Config::test();
config.encryption_key = Some([0x51; 32]);
config.shard_count = 1;
config.write_buffer_size = 8192;
config.max_file_size = 8192;
config
}
fn data_ids(dir: &std::path::Path) -> Vec<u32> {
data_ids_in(&dir.join(NAME).join("shard_000"))
}
fn data_ids_in(dir: &std::path::Path) -> Vec<u32> {
let mut ids = fs::read_dir(dir)
.unwrap()
.filter_map(Result::ok)
.filter_map(|entry| {
entry
.file_name()
.to_str()?
.strip_suffix(".data")?
.parse()
.ok()
})
.collect::<Vec<u32>>();
ids.sort_unstable();
ids
}
fn child_at_stage(path: &std::path::Path, stage: u8) -> std::process::ExitStatus {
Command::new(std::env::current_exe().unwrap())
.arg("--exact")
.arg("format_migration_encrypted_peak_child")
.arg("--nocapture")
.env("ARMDB_ENCRYPTED_MIGRATION_CHILD", "1")
.env("ARMDB_ENCRYPTED_MIGRATION_DIR", path)
.env("ARMDB_ENCRYPTED_MIGRATION_STAGE", stage.to_string())
.status()
.unwrap()
}
fn migration_artifact_names(
entries: &[fs::DirEntry],
predicate: impl Fn(&str) -> bool,
) -> BTreeSet<String> {
entries
.iter()
.filter_map(|entry| {
let name = entry.file_name().to_string_lossy().into_owned();
predicate(&name).then_some(name)
})
.collect()
}
fn assert_no_plaintext_temporary_artifacts(root: &std::path::Path, plaintext: &[u8]) {
let mut pending = vec![root.to_path_buf()];
while let Some(dir) = pending.pop() {
for entry in fs::read_dir(dir).unwrap() {
let entry = entry.unwrap();
let path = entry.path();
if path.is_dir() {
pending.push(path);
continue;
}
let name = entry.file_name();
let name = name.to_string_lossy();
assert!(
!name.starts_with("spool_"),
"plaintext spool survived: {path:?}"
);
if name.ends_with(".data") || name.ends_with(".tmp") {
let bytes = fs::read(&path).unwrap();
assert!(
!bytes
.windows(plaintext.len())
.any(|window| window == plaintext),
"decoded plaintext survived in {path:?}"
);
}
}
}
}
fn inspect_preflight_before_temporary_output() {
let Some(root) = PREFLIGHT_ROOT
.get_or_init(|| Mutex::new(None))
.lock()
.unwrap()
.take()
else {
return;
};
PREFLIGHT_TEMP_EXISTS.store(
root.join(NAME).join("shard_000.mig").exists(),
Ordering::SeqCst,
);
}
fn inspect_temporary_output() {
let Some(root) = INSPECT_ROOT
.get()
.and_then(|root| root.lock().unwrap().clone())
else {
return;
};
let collection = root.join(NAME);
let temporary = [
collection.join(".format-stage-000"),
collection.join("shard_000.mig"),
]
.into_iter()
.find(|path| path.is_dir());
let Some(temporary) = temporary else {
return;
};
let Ok(entries) = fs::read_dir(temporary) else {
return;
};
let plaintext = (SENTINEL as u32 as u64).to_ne_bytes();
let mut bytes = 0;
for entry in entries.filter_map(Result::ok) {
let name = entry.file_name();
let name = name.to_string_lossy();
if name.starts_with("spool_") {
TEMP_VIOLATION.store(true, Ordering::SeqCst);
}
if let Ok(metadata) = entry.metadata() {
bytes += metadata.len();
}
if (name.ends_with(".data") || name.ends_with(".tmp"))
&& fs::read(entry.path()).is_ok_and(|contents| {
contents
.windows(plaintext.len())
.any(|window| window == plaintext)
})
{
TEMP_VIOLATION.store(true, Ordering::SeqCst);
}
}
TEMP_PEAK.fetch_max(bytes, Ordering::SeqCst);
}
fn force_only_encrypted_data_file_id(root: &std::path::Path, new_id: u32) {
let shard = root.join(NAME).join("shard_000");
let ids = data_ids(root);
assert_eq!(
ids.len(),
1,
"fixture must have exactly one source data file"
);
let old_id = ids[0];
let old_data = shard.join(format!("{old_id:06}.data"));
let old_tags = shard.join(format!("{old_id:06}.tags"));
let mut data = fs::read(&old_data).unwrap();
let tags = fs::read(&old_tags).unwrap();
let salt: [u8; 16] = fs::read(root.join(NAME).join("db.salt"))
.unwrap()
.try_into()
.unwrap();
let cipher = PageCipher::derive(&[0x51; 32], &salt, 0).unwrap();
let mut new_tags = Vec::with_capacity(tags.len());
for (page_index, page) in data.chunks_exact_mut(4096).enumerate() {
let old_tag: [u8; 16] = tags[page_index * 16..page_index * 16 + 16]
.try_into()
.unwrap();
cipher
.decrypt_page(old_id, page_index as u64, page, &old_tag)
.unwrap();
new_tags.extend_from_slice(
&cipher
.encrypt_page(new_id, page_index as u64, page)
.unwrap(),
);
}
fs::write(shard.join(format!("{new_id:06}.data")), data).unwrap();
fs::write(shard.join(format!("{new_id:06}.tags")), new_tags).unwrap();
fs::remove_file(old_data).unwrap();
fs::remove_file(old_tags).unwrap();
let old_hint = shard.join(format!("{old_id:06}.hint"));
if old_hint.exists() {
fs::remove_file(old_hint).unwrap();
}
}
#[test]
fn format_migration_encrypted_uses_fresh_file_ids_and_round_trips() {
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
for value in 0..1_000u32 {
tree.put(&(value as u64).to_be_bytes(), &V1(value)).unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
let previous_max = *data_ids(dir.path())
.last()
.expect("at least two source data files");
assert!(
previous_max >= 2,
"fixture must exercise multiple encrypted source ids"
);
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V2, 8, _>(config(), NoHook, &[])
.unwrap();
let output_ids = data_ids(dir.path());
assert!(
output_ids.iter().all(|id| *id > previous_max),
"migration must never reuse encrypted file nonces"
);
for value in 0..1_000u32 {
assert_eq!(
tree.get(&(value as u64).to_be_bytes()),
Some(V2(value as u64))
);
}
}
#[test]
fn format_migration_encrypted_file_id_exhaustion_creates_no_temporary_output() {
let _failpoints = FailpointGuard::acquire();
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
tree.put(&1u64.to_be_bytes(), &V1(7)).unwrap();
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
force_only_encrypted_data_file_id(dir.path(), u32::MAX);
PREFLIGHT_TEMP_EXISTS.store(false, Ordering::SeqCst);
*PREFLIGHT_ROOT
.get_or_init(|| Mutex::new(None))
.lock()
.unwrap() = Some(dir.path().to_path_buf());
let db = Db::open_test(dir.path()).unwrap();
let error = match db.open_zero_tree::<PreflightV2, 8, _>(config(), NoHook, &[]) {
Ok(_) => panic!("u32::MAX cannot reserve a fresh encrypted nonce"),
Err(error) => error,
};
assert!(matches!(
error,
DbError::FormatFileIdExhausted {
collection: NAME,
shard: 0,
needed: 1
}
));
assert!(
!PREFLIGHT_TEMP_EXISTS.load(Ordering::SeqCst),
"file-id exhaustion must be detected before .mig exists, not after cleanup"
);
assert!(
!dir.path().join(NAME).join("shard_000.mig").exists(),
"file-id preflight must run before creating .mig"
);
assert_no_plaintext_temporary_artifacts(dir.path(), &7u64.to_ne_bytes());
}
#[test]
fn format_migration_encrypted_all_drop_preserves_the_file_id_nonce_floor() {
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
tree.put(&1u64.to_be_bytes(), &V1(7)).unwrap();
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
assert_eq!(data_ids(dir.path()), vec![1]);
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<DropV2, 4, _>(config(), NoHook, &[])
.unwrap();
assert_eq!(tree.get(&1u64.to_be_bytes()), None);
tree.put(&2u64.to_be_bytes(), &DropV2(9)).unwrap();
tree.flush().unwrap();
let ids = data_ids(dir.path());
assert!(
ids.iter().all(|id| *id > 1),
"an all-Drop migration must not let the next encrypted write reuse file id 1: {ids:?}"
);
assert_eq!(tree.get(&2u64.to_be_bytes()), Some(DropV2(9)));
}
#[test]
fn format_migration_encrypted_reserves_the_post_open_active_file_id() {
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
tree.put(&1u64.to_be_bytes(), &V1(7)).unwrap();
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
force_only_encrypted_data_file_id(dir.path(), u32::MAX - 1);
let source = dir
.path()
.join(NAME)
.join("shard_000")
.join(format!("{:06}.data", u32::MAX - 1));
let before = fs::read(&source).unwrap();
let db = Db::open_test(dir.path()).unwrap();
let error = match db.open_zero_tree::<V2, 8, _>(config(), NoHook, &[]) {
Ok(_) => panic!("migration output at u32::MAX leaves no legal next active file id"),
Err(error) => error,
};
assert!(matches!(
error,
DbError::FormatFileIdExhausted {
collection: NAME,
shard: 0,
needed: 1
}
));
assert_eq!(
fs::read(source).unwrap(),
before,
"source must be untouched"
);
assert!(
!dir.path().join(NAME).join("shard_000.mig").exists(),
"post-open file-id reserve must be checked before .mig is created"
);
assert!(
!dir.path().join(NAME).join("shard_000.old").exists(),
"post-open file-id reserve failure must not begin the swap"
);
}
#[test]
fn format_migration_encrypted_near_max_uses_exact_output_file_count() {
const ROWS: u64 = 100;
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..ROWS {
tree.put(&key.to_be_bytes(), &V1(key as u32)).unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
force_only_encrypted_data_file_id(dir.path(), u32::MAX - 2);
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V2, 8, _>(config(), NoHook, &[])
.expect("100 small rows fit one output file and leave one successor id");
assert_eq!(
data_ids(dir.path()),
vec![u32::MAX - 1],
"file-id preflight must reserve the exact one-file output layout"
);
for key in 0..ROWS {
assert_eq!(tree.get(&key.to_be_bytes()), Some(V2(key)));
}
}
#[test]
fn format_migration_encrypted_discards_a_pre_admission_staging_directory() {
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
tree.put(&1u64.to_be_bytes(), &V1(7)).unwrap();
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
let staging = dir.path().join(NAME).join(".format-stage-000");
fs::create_dir(&staging).unwrap();
fs::write(staging.join("interrupted.stage"), b"discard me").unwrap();
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V2, 8, _>(config(), NoHook, &[])
.expect("pre-admission staging is disposable because the source is untouched");
assert_eq!(tree.get(&1u64.to_be_bytes()), Some(V2(7)));
assert!(!staging.exists());
}
#[test]
fn format_migration_encrypted_failed_output_is_cleaned_without_plaintext_artifacts() {
let _failpoints = FailpointGuard::acquire();
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<FailV1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..FAIL_ROWS as u64 {
tree.put(&key.to_be_bytes(), &FailV1(SENTINEL as u32))
.unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
FAIL_CALLS.store(0, Ordering::SeqCst);
FAIL_AT.store(128, Ordering::SeqCst);
let db = Db::open_test(dir.path()).unwrap();
let error = match db.open_zero_tree::<FailV2, 8, _>(config(), NoHook, &[]) {
Ok(_) => panic!("injected output-pass failure must abort migration"),
Err(error) => error,
};
assert!(matches!(error, DbError::FormatStepFailed { .. }));
assert!(
!dir.path().join(FAIL_NAME).join("shard_000.mig").exists(),
"ordinary failures before source rename must clean all output"
);
assert_no_plaintext_temporary_artifacts(dir.path(), &SENTINEL.to_ne_bytes());
FAIL_AT.store(usize::MAX, Ordering::SeqCst);
let tree = db
.open_zero_tree::<FailV2, 8, _>(config(), NoHook, &[])
.unwrap();
assert_eq!(
tree.get(&0u64.to_be_bytes()),
Some(FailV2(SENTINEL as u32 as u64))
);
}
fn failed_output_fixture() -> tempfile::TempDir {
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<FailV1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..FAIL_ROWS as u64 {
tree.put(&key.to_be_bytes(), &FailV1(SENTINEL as u32))
.unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
FAIL_CALLS.store(0, Ordering::SeqCst);
FAIL_AT.store(128, Ordering::SeqCst);
dir
}
#[test]
fn format_migration_encrypted_propagates_pre_swap_cleanup_remove_failure() {
let _failpoints = FailpointGuard::acquire();
let dir = failed_output_fixture();
failpoints::CLEANUP_ERROR.store(failpoints::CLEANUP_REMOVE_ERROR, Ordering::SeqCst);
let db = Db::open_test(dir.path()).unwrap();
let error = match db.open_zero_tree::<FailV2, 8, _>(config(), NoHook, &[]) {
Ok(_) => panic!("injected cleanup remove failure must reach the caller"),
Err(error) => error,
};
assert!(
error
.to_string()
.contains("injected format migration cleanup remove failure"),
"cleanup error was lost behind the transform error: {error}"
);
assert!(
dir.path().join(FAIL_NAME).join("shard_000.mig").exists(),
"failed removal must leave discoverable cleanup state for retry"
);
assert_no_plaintext_temporary_artifacts(dir.path(), &SENTINEL.to_ne_bytes());
failpoints::CLEANUP_ERROR.store(failpoints::CLEANUP_NONE, Ordering::SeqCst);
FAIL_AT.store(usize::MAX, Ordering::SeqCst);
let tree = db
.open_zero_tree::<FailV2, 8, _>(config(), NoHook, &[])
.unwrap();
assert_eq!(
tree.get(&0u64.to_be_bytes()),
Some(FailV2(SENTINEL as u32 as u64))
);
}
#[test]
fn format_migration_encrypted_propagates_pre_swap_cleanup_parent_sync_failure() {
let _failpoints = FailpointGuard::acquire();
let dir = failed_output_fixture();
failpoints::CLEANUP_ERROR.store(failpoints::CLEANUP_SYNC_ERROR, Ordering::SeqCst);
let db = Db::open_test(dir.path()).unwrap();
let error = match db.open_zero_tree::<FailV2, 8, _>(config(), NoHook, &[]) {
Ok(_) => panic!("injected cleanup parent-sync failure must reach the caller"),
Err(error) => error,
};
assert!(
error
.to_string()
.contains("injected format migration cleanup parent sync failure"),
"cleanup durability error was lost behind the transform error: {error}"
);
assert!(
!dir.path().join(FAIL_NAME).join("shard_000.mig").exists(),
"sync failure is injected only after the temporary tree is removed"
);
assert_no_plaintext_temporary_artifacts(dir.path(), &SENTINEL.to_ne_bytes());
failpoints::CLEANUP_ERROR.store(failpoints::CLEANUP_NONE, Ordering::SeqCst);
FAIL_AT.store(usize::MAX, Ordering::SeqCst);
let tree = db
.open_zero_tree::<FailV2, 8, _>(config(), NoHook, &[])
.unwrap();
assert_eq!(
tree.get(&0u64.to_be_bytes()),
Some(FailV2(SENTINEL as u32 as u64))
);
}
#[test]
fn format_migration_encrypted_streaming_temporary_state_is_ciphertext_only() {
let _failpoints = FailpointGuard::acquire();
const ROWS: u64 = 1_000;
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..ROWS {
tree.put(&key.to_be_bytes(), &V1(SENTINEL as u32)).unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
TEMP_VIOLATION.store(false, Ordering::SeqCst);
TEMP_PEAK.store(0, Ordering::SeqCst);
*INSPECT_ROOT
.get_or_init(|| Mutex::new(None))
.lock()
.unwrap() = Some(dir.path().to_path_buf());
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<ObservedV2, 8, _>(config(), NoHook, &[])
.unwrap();
*INSPECT_ROOT.get().unwrap().lock().unwrap() = None;
assert_eq!(
tree.get(&0u64.to_be_bytes()),
Some(ObservedV2(SENTINEL as u32 as u64))
);
assert!(
TEMP_PEAK.load(Ordering::SeqCst) > 0,
"test must observe in-flight output"
);
assert!(
!TEMP_VIOLATION.load(Ordering::SeqCst),
"in-flight encrypted migration exposed a plaintext spool or value"
);
}
#[test]
fn format_migration_encrypted_peak_child() {
if std::env::var_os("ARMDB_ENCRYPTED_MIGRATION_CHILD").is_none() {
return;
}
let path = std::env::var("ARMDB_ENCRYPTED_MIGRATION_DIR").unwrap();
let stage = std::env::var("ARMDB_ENCRYPTED_MIGRATION_STAGE")
.unwrap()
.parse()
.unwrap();
failpoints::STAGE.store(stage, Ordering::SeqCst);
let db = Db::open_test(path).unwrap();
let _ = db.open_zero_tree::<V2, 8, _>(config(), NoHook, &[]);
unreachable!("selected encrypted migration stage must hard-exit");
}
#[test]
fn format_migration_encrypted_peak_disk_is_bounded_during_stage_finalize() {
const ROWS: u64 = 1_000;
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..ROWS {
tree.put(&key.to_be_bytes(), &V1(SENTINEL as u32)).unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
assert_eq!(
child_at_stage(dir.path(), failpoints::DURING_STAGE_FINALIZE).code(),
Some(9)
);
let migration = dir.path().join(NAME).join("shard_000.mig");
let entries = fs::read_dir(&migration)
.unwrap()
.collect::<Result<Vec<_>, _>>()
.unwrap();
let stage_data = migration_artifact_names(&entries, |name| {
name.starts_with("stage_") && name.ends_with(".data")
});
let final_data = migration_artifact_names(&entries, |name| {
!name.starts_with("stage_") && name.ends_with(".data")
});
let stage_tags = migration_artifact_names(&entries, |name| {
name.starts_with("stage_") && name.ends_with(".tags")
});
let final_tags = migration_artifact_names(&entries, |name| {
!name.starts_with("stage_") && name.ends_with(".tags")
});
let stage_hints = migration_artifact_names(&entries, |name| {
name.starts_with("stage_") && name.ends_with(".hint")
});
let final_hints = migration_artifact_names(&entries, |name| {
!name.starts_with("stage_") && name.ends_with(".hint")
});
assert!(!stage_data.is_empty(), "snapshot must observe staged files");
assert!(
!final_data.is_empty(),
"snapshot must observe finalized files"
);
let logical_bytes = ROWS * 32; let expected_files = logical_bytes.div_ceil(config().max_file_size) as usize;
assert_eq!(expected_files, 4, "fixture layout drifted");
let data = stage_data
.union(&final_data)
.cloned()
.collect::<BTreeSet<_>>();
assert_eq!(
data.len(),
expected_files,
"staged plus finalized files must be exactly one output file set: {data:?}"
);
let expected_tags = data
.iter()
.map(|name| format!("{}.tags", name.strip_suffix(".data").unwrap()))
.collect::<BTreeSet<_>>();
assert_eq!(
stage_tags
.union(&final_tags)
.cloned()
.collect::<BTreeSet<_>>(),
expected_tags,
"every in-flight data artifact must retain its matching tag artifact"
);
assert!(stage_hints.is_empty(), "staged files must not have hints");
let expected_final_hints = final_data
.iter()
.map(|name| format!("{}.hint", name.strip_suffix(".data").unwrap()))
.collect::<BTreeSet<_>>();
assert_eq!(
final_hints, expected_final_hints,
"hints must exist only for finalized encrypted output"
);
const OUTPUT_SET_BYTES: u64 = 32_768;
let physical_data_bytes = entries
.iter()
.filter(|entry| entry.file_name().to_string_lossy().ends_with(".data"))
.map(|entry| entry.metadata().unwrap().len())
.sum::<u64>();
assert!(physical_data_bytes >= logical_bytes);
assert!(
physical_data_bytes <= OUTPUT_SET_BYTES,
"in-flight conversion duplicated encrypted stage data: {physical_data_bytes} > {OUTPUT_SET_BYTES}"
);
let tag_set_bytes = OUTPUT_SET_BYTES / 4096 * 16;
let hint_set_bytes = logical_bytes + expected_files as u64 * 36;
let complete_output_set_bytes = OUTPUT_SET_BYTES + tag_set_bytes + hint_set_bytes;
let tree_bytes = entries
.iter()
.map(|entry| entry.metadata().unwrap().len())
.sum::<u64>();
assert!(
tree_bytes <= complete_output_set_bytes,
"in-flight migration tree exceeded one complete encrypted output set: {tree_bytes} > {complete_output_set_bytes}"
);
assert_no_plaintext_temporary_artifacts(&migration, &(SENTINEL as u32 as u64).to_ne_bytes());
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V2, 8, _>(config(), NoHook, &[])
.unwrap();
assert_eq!(
tree.get(&0u64.to_be_bytes()),
Some(V2(SENTINEL as u32 as u64))
);
}
#[test]
fn format_migration_encrypted_retry_burns_failed_output_ids_before_resealing_changed_bytes() {
const ROWS: u64 = 1_000;
for stage in [
failpoints::DURING_STAGE_FINALIZE,
failpoints::BEFORE_SOURCE_RENAME,
] {
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..ROWS {
tree.put(&key.to_be_bytes(), &V1(key as u32)).unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
assert_eq!(child_at_stage(dir.path(), stage).code(), Some(9));
let migration = dir.path().join(NAME).join("shard_000.mig");
let attempted_ids = data_ids_in(&migration);
assert!(
!attempted_ids.is_empty(),
"stage {stage} must have sealed at least one target-domain file"
);
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<RetryV2, 8, _>(config(), NoHook, &[])
.expect("retry with changed migration bytes");
let retry_ids = data_ids(dir.path());
assert!(
retry_ids
.iter()
.all(|id| *id > *attempted_ids.last().unwrap()),
"stage {stage} reused a burned encrypted output id: attempted={attempted_ids:?}, retry={retry_ids:?}"
);
assert_eq!(tree.get(&0u64.to_be_bytes()), Some(RetryV2(10_000)));
assert_eq!(
tree.get(&(ROWS - 1).to_be_bytes()),
Some(RetryV2(ROWS - 1 + 10_000))
);
}
}
#[test]
fn normal_encrypted_rotation_stays_above_ids_burned_by_a_failed_migration() {
const ROWS: u64 = 1_000;
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..ROWS {
tree.put(&key.to_be_bytes(), &V1(key as u32)).unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
assert_eq!(
child_at_stage(dir.path(), failpoints::BEFORE_SOURCE_RENAME).code(),
Some(9)
);
let migration = dir.path().join(NAME).join("shard_000.mig");
let burned = *data_ids_in(&migration)
.last()
.expect("failed migration must have sealed output");
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.expect("old collection remains recoverable before the swap");
for key in ROWS..ROWS + 1_000 {
tree.put(&key.to_be_bytes(), &V1(key as u32)).unwrap();
}
tree.flush().unwrap();
assert!(
data_ids(dir.path()).iter().any(|id| *id > burned),
"normal rotation reused or failed to advance beyond burned migration id {burned}: {:?}",
data_ids(dir.path())
);
}
#[test]
fn format_migration_encrypted_final_output_has_data_tags_and_hints() {
const ROWS: u64 = 1_000;
let dir = tempdir().unwrap();
{
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
for key in 0..ROWS {
tree.put(&key.to_be_bytes(), &V1(SENTINEL as u32)).unwrap();
}
tree.flush().unwrap();
drop(tree);
db.close().unwrap();
}
assert_eq!(
child_at_stage(dir.path(), failpoints::AFTER_FORMAT_VERSION).code(),
Some(9)
);
let migration = dir.path().join(NAME).join("shard_000.mig");
let entries = fs::read_dir(&migration)
.unwrap()
.collect::<Result<Vec<_>, _>>()
.unwrap();
assert!(
entries.iter().all(|entry| {
let name = entry.file_name();
let name = name.to_string_lossy();
name == "format.v"
|| name.ends_with(".data")
|| name.ends_with(".tags")
|| name.ends_with(".hint")
}),
"temporary migration state must contain only the final encrypted output set"
);
let data = entries
.iter()
.filter(|entry| entry.file_name().to_string_lossy().ends_with(".data"))
.collect::<Vec<_>>();
let tags = entries
.iter()
.filter(|entry| entry.file_name().to_string_lossy().ends_with(".tags"))
.collect::<Vec<_>>();
let hints = entries
.iter()
.filter(|entry| entry.file_name().to_string_lossy().ends_with(".hint"))
.collect::<Vec<_>>();
let logical_bytes = ROWS * 32; let expected_files = logical_bytes.div_ceil(config().max_file_size) as usize;
assert_eq!(data.len(), expected_files);
assert_eq!(tags.len(), expected_files);
assert_eq!(hints.len(), expected_files);
let physical_data_bytes = data
.iter()
.map(|entry| entry.metadata().unwrap().len())
.sum::<u64>();
assert!(physical_data_bytes >= logical_bytes);
assert!(
physical_data_bytes <= logical_bytes + expected_files as u64 * 4095,
"page padding is the only permitted encrypted data overhead"
);
let tag_bytes = tags
.iter()
.map(|entry| entry.metadata().unwrap().len())
.sum::<u64>();
assert_eq!(tag_bytes, physical_data_bytes / 4096 * 16);
assert_no_plaintext_temporary_artifacts(&migration, &(SENTINEL as u32 as u64).to_ne_bytes());
let db = Db::open_test(dir.path()).unwrap();
let tree = db
.open_zero_tree::<V2, 8, _>(config(), NoHook, &[])
.unwrap();
assert_eq!(
tree.get(&0u64.to_be_bytes()),
Some(V2(SENTINEL as u32 as u64))
);
}
#[test]
fn format_migration_encrypted_pending_txlog_segments_remain_byte_identical() {
let dir = tempdir().unwrap();
let tx_config = TxConfig {
log_encryption_key: Some([0x73; 32]),
..TxConfig::default()
};
{
let db = Arc::new(Db::open_with_tx_config(dir.path(), tx_config).unwrap());
let tree = db
.open_zero_tree::<V1, 4, _>(config(), NoHook, &[])
.unwrap();
Db::tx_debug_skip_apply(true);
db.tx()
.put(&tree, &1u64.to_be_bytes(), &V1(9))
.unwrap()
.durable()
.commit()
.unwrap();
Db::tx_debug_skip_apply(false);
drop(tree);
db.close().unwrap();
}
let snapshot = |root: &std::path::Path| {
let mut paths = fs::read_dir(root.join("txlog"))
.unwrap()
.collect::<Result<Vec<_>, _>>()
.unwrap();
paths.sort_by_key(std::fs::DirEntry::file_name);
paths
.into_iter()
.map(|entry| (entry.file_name(), fs::read(entry.path()).unwrap()))
.collect::<Vec<_>>()
};
let before = snapshot(dir.path());
let db = Arc::new(Db::open_with_tx_config(dir.path(), tx_config).unwrap());
db.migrate_format::<V2>(&config())
.expect("install the encrypted target format before managed replay");
assert_eq!(
snapshot(dir.path()),
before,
"explicit encrypted migration must not reseal txlog segments"
);
let tree = db
.open_zero_tree::<V2, 8, _>(config(), NoHook, &[])
.unwrap();
assert_eq!(tree.get(&1u64.to_be_bytes()), Some(V2(9)));
assert_eq!(
snapshot(dir.path()),
before,
"migration/replay must never reseal txlog segments"
);
}