use super::commitment::{leaf_hash, node_hash, StorageCommitment, MAX_COMMITMENT_KEY_COUNT};
use super::protocol::compute_audit_digest;
use crate::ant_protocol::XorName;
use serde::{Deserialize, Serialize};
pub const SMALL_TREE_FULL_AUDIT_FLOOR: u32 = 4;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct SubtreeLeaf {
pub key: XorName,
pub bytes_hash: [u8; 32],
pub nonced_hash: [u8; 32],
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct SubtreeProof {
pub leaves: Vec<SubtreeLeaf>,
pub sibling_cut_hashes: Vec<[u8; 32]>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct SubtreePath {
pub depth: u32,
pub slot: u32,
pub leaf_start: u32,
pub leaf_end: u32,
}
impl SubtreePath {
#[must_use]
pub fn real_leaf_count(&self) -> u32 {
self.leaf_end - self.leaf_start
}
}
#[must_use]
fn tree_depth(key_count: u32) -> Option<u32> {
if key_count == 0 || key_count > MAX_COMMITMENT_KEY_COUNT {
return None;
}
let rounded = key_count.checked_next_power_of_two()?;
Some(rounded.trailing_zeros())
}
#[must_use]
fn real_leaves_under(depth: u32, slot: u64, key_count: u32, total_depth: u32) -> u32 {
let levels_below = total_depth - depth;
let span = 1u64 << levels_below;
let start = slot.saturating_mul(span).min(u64::from(key_count));
let end = slot
.saturating_add(1)
.saturating_mul(span)
.min(u64::from(key_count));
u32::try_from(end - start).unwrap_or(0)
}
#[must_use]
fn sqrt_floor(key_count: u32) -> u32 {
let n = u64::from(key_count);
if n <= 1 {
return 1;
}
let mut x = n;
let mut y = x.div_ceil(2);
while y < x {
x = y;
y = (x + n / x) / 2;
}
let ceil = if x.saturating_mul(x) == n { x } else { x + 1 };
u32::try_from(ceil.max(1)).unwrap_or(u32::MAX)
}
#[must_use]
fn nonce_bit(nonce: &[u8; 32], index: u32) -> bool {
let byte = (index / 8) as usize;
let bit = 7 - (index % 8);
nonce.get(byte).is_some_and(|b| (b >> bit) & 1 == 1)
}
#[must_use]
pub fn select_subtree_path(nonce: &[u8; 32], key_count: u32) -> Option<SubtreePath> {
let total_depth = tree_depth(key_count)?;
if key_count <= SMALL_TREE_FULL_AUDIT_FLOOR {
return Some(SubtreePath {
depth: 0,
slot: 0,
leaf_start: 0,
leaf_end: key_count,
});
}
let floor = sqrt_floor(key_count);
let mut depth = 0u32;
let mut slot = 0u64;
while depth < total_depth {
let go_left = nonce_bit(nonce, depth);
let child_slot = slot * 2 + u64::from(!go_left);
let child_real = real_leaves_under(depth + 1, child_slot, key_count, total_depth);
if child_real < floor {
break; }
depth += 1;
slot = child_slot;
}
let span = 1u64 << (total_depth - depth);
let leaf_start =
u32::try_from(slot.saturating_mul(span).min(u64::from(key_count))).unwrap_or(key_count);
let leaf_end = u32::try_from(
slot.saturating_add(1)
.saturating_mul(span)
.min(u64::from(key_count)),
)
.unwrap_or(key_count);
Some(SubtreePath {
depth,
slot: u32::try_from(slot).unwrap_or(u32::MAX),
leaf_start,
leaf_end,
})
}
#[must_use]
pub fn select_spotcheck_indices(nonce: &[u8; 32], path: &SubtreePath, k: u32) -> Vec<u32> {
let n = path.real_leaf_count();
if n == 0 {
return Vec::new();
}
if n <= k {
return (0..n).collect();
}
let mut out: Vec<u32> = Vec::with_capacity(k as usize);
let mut counter: u32 = 0;
while u32::try_from(out.len()).unwrap_or(u32::MAX) < k {
let mut h = blake3::Hasher::new();
h.update(b"autonomi.ant.replication.audit_spotcheck.v1");
h.update(nonce);
h.update(&counter.to_le_bytes());
let digest = *h.finalize().as_bytes();
let mut word = [0u8; 4];
word.copy_from_slice(&digest[..4]);
let idx = u32::from_le_bytes(word) % n;
if !out.contains(&idx) {
out.push(idx);
}
counter = counter.wrapping_add(1);
if counter > k.saturating_mul(64) {
break;
}
}
let mut candidate: u32 = 0;
while u32::try_from(out.len()).unwrap_or(u32::MAX) < k && candidate < n {
if !out.contains(&candidate) {
out.push(candidate);
}
candidate += 1;
}
out
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StructureVerdict {
Valid,
Invalid(&'static str),
}
#[must_use]
pub fn verify_subtree_proof(
proof: &SubtreeProof,
nonce: &[u8; 32],
commitment: &StorageCommitment,
) -> StructureVerdict {
let Some(path) = select_subtree_path(nonce, commitment.key_count) else {
return StructureVerdict::Invalid("out-of-protocol key_count");
};
let expected_leaves = path.real_leaf_count() as usize;
if proof.leaves.len() != expected_leaves {
return StructureVerdict::Invalid("wrong leaf count");
}
if proof.sibling_cut_hashes.len() != path.depth as usize {
return StructureVerdict::Invalid("wrong cut-hash count");
}
for w in proof.leaves.windows(2) {
if let [a, b] = w {
if a.key >= b.key {
return StructureVerdict::Invalid("leaves not strictly ascending");
}
}
}
let Some(total_depth) = tree_depth(commitment.key_count) else {
return StructureVerdict::Invalid("out-of-protocol key_count");
};
let leaf_hashes: Vec<[u8; 32]> = proof
.leaves
.iter()
.map(|l| leaf_hash(&l.key, &l.bytes_hash))
.collect();
let levels_to_subtree_root = total_depth - path.depth;
let mut cur = fold_levels(leaf_hashes, levels_to_subtree_root);
let mut node_index = u64::from(path.slot);
for level_above in (0..path.depth).rev() {
let Some(sibling) = proof.sibling_cut_hashes.get(level_above as usize) else {
return StructureVerdict::Invalid("missing cut-hash");
};
cur = if node_index % 2 == 0 {
node_hash(&cur, sibling)
} else {
node_hash(sibling, &cur)
};
node_index /= 2;
}
if cur == commitment.root {
StructureVerdict::Valid
} else {
StructureVerdict::Invalid("root mismatch")
}
}
#[must_use]
fn fold_levels(mut level: Vec<[u8; 32]>, levels: u32) -> [u8; 32] {
if level.is_empty() {
return [0u8; 32];
}
for _ in 0..levels {
level = crate::replication::commitment::build_next_level(&level);
}
level.first().copied().unwrap_or([0u8; 32])
}
#[must_use]
pub fn nonced_leaf_hash(
nonce: &[u8; 32],
challenged_peer_id: &[u8; 32],
key: &XorName,
record_bytes: &[u8],
) -> [u8; 32] {
compute_audit_digest(nonce, challenged_peer_id, key, record_bytes)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum BuildProofError {
BadKeyCount,
MissingKey {
leaf_index: u32,
},
MissingBytes {
key: XorName,
},
}
pub fn build_subtree_proof(
tree: &super::commitment::MerkleTree,
nonce: &[u8; 32],
challenged_peer_id: &[u8; 32],
bytes_for: impl Fn(&XorName) -> Option<Vec<u8>>,
) -> Result<SubtreeProof, BuildProofError> {
let plan = subtree_plan(tree, nonce)?;
let mut leaves = Vec::with_capacity(plan.leaf_keys.len());
for key in &plan.leaf_keys {
let bytes = bytes_for(key).ok_or(BuildProofError::MissingBytes { key: *key })?;
leaves.push(subtree_leaf(nonce, challenged_peer_id, key, &bytes));
}
Ok(SubtreeProof {
leaves,
sibling_cut_hashes: plan.sibling_cut_hashes,
})
}
#[derive(Debug, Clone)]
pub struct SubtreePlan {
pub leaf_keys: Vec<XorName>,
pub sibling_cut_hashes: Vec<[u8; 32]>,
}
pub fn subtree_plan(
tree: &super::commitment::MerkleTree,
nonce: &[u8; 32],
) -> Result<SubtreePlan, BuildProofError> {
let key_count = tree.key_count();
let path = select_subtree_path(nonce, key_count).ok_or(BuildProofError::BadKeyCount)?;
let mut leaf_keys = Vec::with_capacity(path.real_leaf_count() as usize);
for idx in path.leaf_start..path.leaf_end {
let key = tree
.key_at(idx as usize)
.ok_or(BuildProofError::MissingKey { leaf_index: idx })?;
leaf_keys.push(key);
}
let total_depth = u32::try_from(tree.levels_count().saturating_sub(1)).unwrap_or(0);
let mut sibling_cut_hashes = Vec::with_capacity(path.depth as usize);
let mut slot = 0u64;
for d in 0..path.depth {
let go_left = nonce_bit(nonce, d);
let child = slot * 2 + u64::from(!go_left);
let sibling = child ^ 1;
let level_from_leaves = (total_depth - (d + 1)) as usize;
let chosen_hash = tree.node_at(level_from_leaves, child);
let sib_hash = tree
.node_at(level_from_leaves, sibling)
.or(chosen_hash)
.ok_or(BuildProofError::BadKeyCount)?;
sibling_cut_hashes.push(sib_hash);
slot = child;
}
Ok(SubtreePlan {
leaf_keys,
sibling_cut_hashes,
})
}
#[must_use]
pub fn subtree_leaf(
nonce: &[u8; 32],
challenged_peer_id: &[u8; 32],
key: &XorName,
bytes: &[u8],
) -> SubtreeLeaf {
SubtreeLeaf {
key: *key,
bytes_hash: *blake3::hash(bytes).as_bytes(),
nonced_hash: nonced_leaf_hash(nonce, challenged_peer_id, key, bytes),
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests {
use super::*;
use crate::replication::commitment::MerkleTree;
fn xn_u32(i: u32) -> XorName {
let mut k = [0u8; 32];
k[..4].copy_from_slice(&i.to_be_bytes()); k
}
fn nonce_of(seed: u8) -> [u8; 32] {
[seed; 32]
}
#[test]
fn sqrt_floor_is_exact_ceil() {
assert_eq!(sqrt_floor(1), 1);
assert_eq!(sqrt_floor(4), 2);
assert_eq!(sqrt_floor(5), 3); assert_eq!(sqrt_floor(9), 3);
assert_eq!(sqrt_floor(10), 4);
assert_eq!(sqrt_floor(100), 10);
assert_eq!(sqrt_floor(101), 11);
assert_eq!(sqrt_floor(1_000_000), 1000);
}
#[test]
fn real_leaves_under_root_is_all() {
let d = tree_depth(100).unwrap();
assert_eq!(real_leaves_under(0, 0, 100, d), 100);
}
#[test]
fn real_leaves_under_padding_slot_is_zero() {
let d = tree_depth(5).unwrap();
assert_eq!(d, 3);
assert_eq!(real_leaves_under(1, 0, 5, d), 4); assert_eq!(real_leaves_under(1, 1, 5, d), 1); assert_eq!(real_leaves_under(3, 7, 5, d), 0); assert_eq!(real_leaves_under(2, 3, 5, d), 0); }
#[test]
fn selection_never_empty_across_many_sizes_and_nonces() {
for n in [
5u32, 6, 7, 9, 13, 17, 33, 65, 100, 129, 333, 1000, 1024, 1025,
] {
let floor = sqrt_floor(n);
for seed in 0u8..=255 {
let path = select_subtree_path(&nonce_of(seed), n).unwrap();
assert!(
path.real_leaf_count() >= floor.min(n),
"n={n} seed={seed}: real={} < floor={floor}",
path.real_leaf_count()
);
assert!(
path.real_leaf_count() >= 1,
"n={n} seed={seed}: empty selection"
);
assert!(path.leaf_end <= n);
assert!(path.leaf_start < path.leaf_end);
}
}
}
#[test]
fn small_trees_select_whole_tree() {
for n in 1..=SMALL_TREE_FULL_AUDIT_FLOOR {
let path = select_subtree_path(&nonce_of(7), n).unwrap();
assert_eq!(path.depth, 0);
assert_eq!(path.leaf_start, 0);
assert_eq!(path.leaf_end, n);
}
}
#[test]
fn selection_is_deterministic() {
let n = 500;
let a = select_subtree_path(&nonce_of(42), n).unwrap();
let b = select_subtree_path(&nonce_of(42), n).unwrap();
assert_eq!(a, b);
}
#[test]
fn different_nonces_cover_different_branches_over_time() {
let n = 1024;
let mut starts = std::collections::HashSet::new();
for seed in 0u8..=255 {
let p = select_subtree_path(&nonce_of(seed), n).unwrap();
starts.insert(p.leaf_start);
}
assert!(
starts.len() > 4,
"nonce should spread selection: {}",
starts.len()
);
}
fn nonce_for_trial(i: u32) -> [u8; 32] {
let mut h = blake3::Hasher::new();
h.update(b"detection-sim-trial");
h.update(&i.to_le_bytes());
*h.finalize().as_bytes()
}
fn catch_rate(n: u32, deleted: &std::collections::HashSet<u32>, trials: u32) -> f64 {
let mut caught = 0u32;
for t in 0..trials {
let path = select_subtree_path(&nonce_for_trial(t), n).unwrap();
if (path.leaf_start..path.leaf_end).any(|i| deleted.contains(&i)) {
caught += 1;
}
}
f64::from(caught) / f64::from(trials)
}
#[test]
fn detection_uniform_fast_clustered_floor() {
let n = 1024u32; let del_count = n / 10;
let uniform: std::collections::HashSet<u32> =
(0..del_count).map(|i| (i * n / del_count) % n).collect();
let uniform_rate = catch_rate(n, &uniform, 256);
let clustered: std::collections::HashSet<u32> = (0..del_count).collect();
let clustered_rate = catch_rate(n, &clustered, 256);
assert!(
uniform_rate > 0.95,
"uniform deletions should be caught almost every audit, got {uniform_rate}"
);
assert!(
(0.04..=0.30).contains(&clustered_rate),
"clustered catch-rate should be near f=0.1, got {clustered_rate}"
);
assert!(
uniform_rate > clustered_rate * 2.0,
"uniform ({uniform_rate}) must be far easier to catch than clustered ({clustered_rate})"
);
}
#[test]
fn subtree_size_near_sqrt_for_balanced_tree() {
let n = 1024; let path = select_subtree_path(&nonce_of(3), n).unwrap();
assert!(path.real_leaf_count() >= 32);
assert!(
path.real_leaf_count() <= 64,
"got {}",
path.real_leaf_count()
);
}
fn chunk_bytes(key: &XorName) -> Vec<u8> {
let mut v = key.to_vec();
v.extend_from_slice(b"chunk-body");
v
}
fn entries_for(n: u32) -> Vec<(XorName, [u8; 32])> {
(0..n)
.map(|i| {
let key = xn_u32(i);
let bytes_hash = *blake3::hash(&chunk_bytes(&key)).as_bytes();
(key, bytes_hash)
})
.collect()
}
fn build_proof(
entries: &[(XorName, [u8; 32])],
nonce: &[u8; 32],
peer_id: &[u8; 32],
) -> (SubtreeProof, StorageCommitment) {
let tree = MerkleTree::build(entries.to_vec()).unwrap();
let key_count = tree.key_count();
let proof = build_subtree_proof(&tree, nonce, peer_id, |k| Some(chunk_bytes(k))).unwrap();
let commitment = fake_commitment(tree.root(), key_count, *peer_id);
(proof, commitment)
}
fn fake_commitment(root: [u8; 32], key_count: u32, peer: [u8; 32]) -> StorageCommitment {
StorageCommitment {
root,
key_count,
sender_peer_id: peer,
sender_public_key: vec![0u8; 1952],
signature: vec![0u8; 3293],
}
}
#[test]
fn honest_proof_verifies_at_many_sizes() {
let peer = [0xABu8; 32];
for n in [5u32, 8, 13, 17, 64, 100, 256, 1000] {
let entries = entries_for(n);
for seed in [1u8, 2, 7, 42, 200] {
let nonce = nonce_of(seed);
let (proof, commitment) = build_proof(&entries, &nonce, &peer);
assert_eq!(
verify_subtree_proof(&proof, &nonce, &commitment),
StructureVerdict::Valid,
"n={n} seed={seed}"
);
}
}
}
#[test]
fn honest_proof_verifies_for_every_size_and_nonce() {
let peer = [7u8; 32];
for n in 5u32..=600 {
let entries = entries_for(n);
for seed in 0u8..32 {
let nonce = nonce_of(seed.wrapping_mul(17).wrapping_add(3));
let (proof, commitment) = build_proof(&entries, &nonce, &peer);
assert_eq!(
verify_subtree_proof(&proof, &nonce, &commitment),
StructureVerdict::Valid,
"honest proof must verify at n={n} seed={seed}"
);
}
}
}
#[test]
fn tampered_leaf_breaks_root() {
let peer = [9u8; 32];
let entries = entries_for(100);
let nonce = nonce_of(5);
let (mut proof, commitment) = build_proof(&entries, &nonce, &peer);
proof.leaves[0].bytes_hash[0] ^= 0x01;
assert!(matches!(
verify_subtree_proof(&proof, &nonce, &commitment),
StructureVerdict::Invalid(_)
));
}
#[test]
fn tampered_cut_hash_breaks_root() {
let peer = [9u8; 32];
let entries = entries_for(256);
let nonce = nonce_of(11);
let (mut proof, commitment) = build_proof(&entries, &nonce, &peer);
if let Some(c) = proof.sibling_cut_hashes.first_mut() {
c[0] ^= 0x01;
}
assert!(matches!(
verify_subtree_proof(&proof, &nonce, &commitment),
StructureVerdict::Invalid(_)
));
}
#[test]
fn wrong_leaf_count_rejected() {
let peer = [9u8; 32];
let entries = entries_for(100);
let nonce = nonce_of(5);
let (mut proof, commitment) = build_proof(&entries, &nonce, &peer);
proof.leaves.pop();
assert_eq!(
verify_subtree_proof(&proof, &nonce, &commitment),
StructureVerdict::Invalid("wrong leaf count")
);
}
#[test]
fn non_ascending_leaves_rejected() {
let peer = [9u8; 32];
let entries = entries_for(100);
let nonce = nonce_of(5);
let (mut proof, commitment) = build_proof(&entries, &nonce, &peer);
if proof.leaves.len() >= 2 {
proof.leaves.swap(0, 1);
}
assert!(matches!(
verify_subtree_proof(&proof, &nonce, &commitment),
StructureVerdict::Invalid(_)
));
}
#[test]
fn spotcheck_indices_in_range_and_distinct() {
let n = 1024;
let nonce = nonce_of(3);
let path = select_subtree_path(&nonce, n).unwrap();
let k = 8;
let idxs = select_spotcheck_indices(&nonce, &path, k);
assert_eq!(
u32::try_from(idxs.len()).unwrap(),
k.min(path.real_leaf_count())
);
let mut seen = std::collections::HashSet::new();
for i in &idxs {
assert!(*i < path.real_leaf_count());
assert!(seen.insert(*i), "duplicate spot-check index {i}");
}
}
#[test]
fn build_proof_reports_missing_bytes() {
let entries = entries_for(100);
let tree = MerkleTree::build(entries).unwrap();
let nonce = nonce_of(5);
let path = select_subtree_path(&nonce, tree.key_count()).unwrap();
let victim = tree.key_at(path.leaf_start as usize).unwrap();
let err = build_subtree_proof(&tree, &nonce, &[1u8; 32], |k| {
if *k == victim {
None
} else {
Some(chunk_bytes(k))
}
})
.unwrap_err();
assert_eq!(err, BuildProofError::MissingBytes { key: victim });
}
#[test]
fn spotcheck_returns_all_when_subtree_small() {
let path = SubtreePath {
depth: 0,
slot: 0,
leaf_start: 0,
leaf_end: 3,
};
let idxs = select_spotcheck_indices(&nonce_of(1), &path, 8);
assert_eq!(idxs, vec![0, 1, 2]);
}
#[test]
fn spotcheck_always_yields_exactly_min_k_n_distinct_indices() {
for size in [1u32, 2, 3, 7, 8, 64, 1000] {
let path = SubtreePath {
depth: 0,
slot: 0,
leaf_start: 0,
leaf_end: size,
};
for k in [1u32, 3, 5, 8] {
for seed in 0..32u8 {
let nonce = nonce_of(seed);
let idxs = select_spotcheck_indices(&nonce, &path, k);
let expected = k.min(path.real_leaf_count()) as usize;
assert_eq!(
idxs.len(),
expected,
"size={size} k={k} seed={seed}: must yield exactly min(k, n)"
);
let mut seen = std::collections::HashSet::new();
for i in &idxs {
assert!(*i < path.real_leaf_count(), "index out of range");
assert!(seen.insert(*i), "duplicate index {i}");
}
assert_eq!(
idxs,
select_spotcheck_indices(&nonce, &path, k),
"selection must be deterministic"
);
}
}
}
}
#[test]
fn fabricated_nonced_hash_caught_by_spotcheck_probability() {
let peer = [1u8; 32];
let entries = entries_for(400);
let nonce = nonce_of(9);
let (mut proof, _commitment) = build_proof(&entries, &nonce, &peer);
proof.leaves[0].nonced_hash[0] ^= 0xFF;
let leaf = &proof.leaves[0];
let real_bytes = chunk_bytes(&leaf.key);
let expected = nonced_leaf_hash(&nonce, &peer, &leaf.key, &real_bytes);
assert_ne!(
leaf.nonced_hash, expected,
"fabricated nonced hash must differ from real"
);
}
#[test]
fn responder_cannot_substitute_a_different_branch() {
let peer = [0x5Au8; 32];
let n = 1024u32; let entries = entries_for(n);
let audit_nonce = nonce_of(7);
let audit_path = select_subtree_path(&audit_nonce, n).unwrap();
let mut decoy: Option<([u8; 32], SubtreePath)> = None;
for seed in 0u8..=255 {
let cand_nonce = nonce_of(seed);
let cand = select_subtree_path(&cand_nonce, n).unwrap();
if cand.depth == audit_path.depth
&& cand.slot != audit_path.slot
&& cand.real_leaf_count() == audit_path.real_leaf_count()
{
decoy = Some((cand_nonce, cand));
break;
}
}
let (decoy_nonce, decoy_path) =
decoy.expect("a same-depth, different-slot decoy branch must exist for n=1024");
assert_ne!(decoy_path.slot, audit_path.slot);
assert_eq!(decoy_path.depth, audit_path.depth);
assert_eq!(decoy_path.real_leaf_count(), audit_path.real_leaf_count());
let tree = MerkleTree::build(entries).unwrap();
let decoy_proof =
build_subtree_proof(&tree, &decoy_nonce, &peer, |k| Some(chunk_bytes(k))).unwrap();
let commitment = fake_commitment(tree.root(), n, peer);
assert_eq!(
verify_subtree_proof(&decoy_proof, &decoy_nonce, &commitment),
StructureVerdict::Valid,
"the decoy proof must be a genuinely valid proof for its own nonce"
);
let verdict = verify_subtree_proof(&decoy_proof, &audit_nonce, &commitment);
assert_eq!(
verdict,
StructureVerdict::Invalid("root mismatch"),
"substituting a different valid branch must be rejected by the root check, got {verdict:?}"
);
}
}