#![cfg(feature = "encryption")]
use mongreldb_core::query::{Condition, Query};
use mongreldb_core::schema::{ColumnDef, ColumnFlags, Schema, TypeId};
use mongreldb_core::{read_header, Table, Value};
use std::path::{Path, PathBuf};
use tempfile::tempdir;
const CANARY_MIN: &[u8] = b"aaa-leak-canary-minimum-bound";
const CANARY_MAX: &[u8] = b"zzz-leak-canary-maximum-bound";
const N: i64 = 200_000;
fn schema() -> Schema {
Schema {
schema_id: 71,
columns: vec![
ColumnDef {
id: 1,
name: "id".into(),
ty: TypeId::Int64,
flags: ColumnFlags::empty().with(ColumnFlags::PRIMARY_KEY),
default_value: None,
},
ColumnDef {
id: 2,
name: "cost".into(),
ty: TypeId::Float64,
flags: ColumnFlags::empty(),
default_value: None,
},
ColumnDef {
id: 3,
name: "tag".into(),
ty: TypeId::Bytes,
flags: ColumnFlags::empty(),
default_value: None,
},
],
indexes: vec![],
colocation: vec![],
constraints: Default::default(),
clustered: false,
}
}
fn rows() -> Vec<Vec<(u16, Value)>> {
(0..N)
.map(|i| {
let tag: Vec<u8> = if i == 0 {
CANARY_MIN.to_vec()
} else if i == N - 1 {
CANARY_MAX.to_vec()
} else {
format!("tag-{i:07}").into_bytes()
};
vec![
(1, Value::Int64(i)),
(2, Value::Float64(i as f64)),
(3, Value::Bytes(tag)),
]
})
.collect()
}
fn run_file(dir: &Path) -> PathBuf {
let runs = dir.join("_runs");
let mut files: Vec<PathBuf> = std::fs::read_dir(&runs)
.unwrap()
.flatten()
.map(|e| e.path())
.filter(|p| p.extension().map(|e| e == "sr").unwrap_or(false))
.collect();
assert_eq!(files.len(), 1, "expected exactly one run in {runs:?}");
files.pop().unwrap()
}
fn range_query(db: &mut Table, lo: f64, hi: f64) -> Vec<i64> {
let q = Query::new().and(Condition::RangeF64 {
column_id: 2,
lo,
lo_inclusive: true,
hi,
hi_inclusive: true,
});
let mut ids: Vec<i64> = db
.query(&q)
.unwrap()
.iter()
.map(|r| match r.columns.get(&1) {
Some(Value::Int64(v)) => *v,
other => panic!("expected Int64 id, got {other:?}"),
})
.collect();
ids.sort_unstable();
ids
}
#[test]
fn encrypted_range_results_match_plaintext() {
let enc_dir = tempdir().unwrap();
let plain_dir = tempdir().unwrap();
let mut enc = Table::create_encrypted(enc_dir.path(), schema(), 1, "pw").unwrap();
let mut plain = Table::create(plain_dir.path(), schema(), 1).unwrap();
enc.bulk_load(rows()).unwrap();
plain.bulk_load(rows()).unwrap();
for (lo, hi) in [
(0.0, 100.0), (65_000.0, 66_000.0), ((N - 50) as f64, (N + 100) as f64), ] {
let got = range_query(&mut enc, lo, hi);
let want = range_query(&mut plain, lo, hi);
assert_eq!(got, want, "encrypted != plaintext for [{lo}, {hi}]");
assert!(!want.is_empty(), "test range [{lo}, {hi}] must match rows");
}
}
#[test]
fn encrypted_stats_envelope_leaks_no_plaintext_bounds() {
let dir = tempdir().unwrap();
let mut db = Table::create_encrypted(dir.path(), schema(), 1, "pw").unwrap();
db.bulk_load(rows()).unwrap();
let path = run_file(dir.path());
let header = read_header(&path).unwrap();
assert_ne!(
header.encrypted_stats_offset, 0,
"encrypted run must carry the encrypted stats envelope"
);
assert_ne!(header.encrypted_stats_len, 0);
let bytes = std::fs::read(&path).unwrap();
for canary in [CANARY_MIN, CANARY_MAX] {
assert!(
!bytes.windows(canary.len()).any(|w| w == canary),
"plaintext page bound leaked into the run file: {:?}",
String::from_utf8_lossy(canary)
);
}
}
#[test]
fn tampered_stats_envelope_is_rejected() {
let dir = tempdir().unwrap();
{
let mut db = Table::create_encrypted(dir.path(), schema(), 1, "pw").unwrap();
db.bulk_load(rows()).unwrap();
}
let path = run_file(dir.path());
let header = read_header(&path).unwrap();
let mut bytes = std::fs::read(&path).unwrap();
let target = header.encrypted_stats_offset as usize + 4;
bytes[target] ^= 0xFF;
let foot = header.footer_offset as usize;
let checksum = {
use sha2::{Digest, Sha256};
Sha256::digest(&bytes[..foot])
};
bytes[foot + 16..foot + 48].copy_from_slice(&checksum);
std::fs::write(&path, &bytes).unwrap();
let err = Table::open_encrypted(dir.path(), "pw");
assert!(
err.is_err(),
"open must reject a run whose stats envelope fails authentication"
);
}