Skip to main content

commonware_storage/bmt/
mod.rs

1//! Stateless Binary Merkle Tree (BMT).
2//!
3//! The Binary Merkle Tree is constructed level-by-level. The first level consists of position-hashed leaf digests.
4//! On each additional level, pairs of nodes are hashed from the previous level (if a level contains an odd
5//! number of nodes, the last node is duplicated). The finalized root of the tree incorporates the leaf count
6//! to prevent proof malleability: `root = hash(leaf_count || tree_root)`.
7//!
8//! For example, given three leaves A, B, and C, the tree is constructed as follows:
9//!
10//! ```text
11//!     Level 2 (tree_root):  [hash(hash(hash(0,A),hash(1,B)),hash(hash(2,C),hash(2,C)))]
12//!     Level 1:              [hash(hash(0,A),hash(1,B)),hash(hash(2,C),hash(2,C))]
13//!     Level 0 (leaves):     [hash(0,A),hash(1,B),hash(2,C)]
14//!     Finalized root:       hash(3 || tree_root)
15//! ```
16//!
17//! A proof for one or more leaves is generated by collecting the siblings needed to reconstruct the root.
18//! An external process can then use this proof (with some trusted root) to verify that the leaves
19//! are part of the tree.
20//!
21//! # Example
22//!
23//! ```rust
24//! use commonware_storage::bmt::{Builder, Tree};
25//! use commonware_cryptography::{Sha256, sha256::Digest, Hasher as _};
26//!
27//! // Create transactions and compute their digests
28//! let txs = [b"tx1", b"tx2", b"tx3", b"tx4"];
29//! let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
30//!
31//! // Build a Merkle Tree from the digests
32//! let mut builder = Builder::<Sha256>::new(digests.len());
33//! for digest in &digests {
34//!    builder.add(digest);
35//! }
36//! let tree = builder.build();
37//! let root = tree.root();
38//!
39//! // Generate a proof for leaf at index 1
40//! let proof = tree.proof(1).unwrap();
41//! assert!(proof.verify_element_inclusion::<Sha256>(&digests[1], 1, &root).is_ok());
42//! ```
43
44use alloc::{
45    collections::btree_set::BTreeSet,
46    string::{String, ToString},
47    vec,
48    vec::Vec,
49};
50use bytes::{Buf, BufMut};
51use commonware_codec::{EncodeSize, Read, ReadExt, ReadRangeExt, Write};
52use commonware_cryptography::{Digest, Hasher};
53use commonware_utils::{non_empty_vec, vec::NonEmptyVec};
54use thiserror::Error;
55
56/// There should never be more than 32 sibling levels in a proof. Because
57/// [Proof::leaf_count] is a `u32`, a tree can have at most `u32::MAX` leaves,
58/// which requires at most `u32::BITS` sibling hashes per proven item.
59pub const MAX_LEVELS: usize = u32::BITS as usize;
60
61/// Errors that can occur when working with a Binary Merkle Tree (BMT).
62#[derive(Error, Debug)]
63pub enum Error {
64    #[error("invalid position: {0}")]
65    InvalidPosition(u32),
66    #[error("invalid proof: {0} != {1}")]
67    InvalidProof(String, String),
68    #[error("no leaves")]
69    NoLeaves,
70    #[error("unaligned proof")]
71    UnalignedProof,
72    #[error("duplicate position: {0}")]
73    DuplicatePosition(u32),
74}
75
76/// Constructor for a Binary Merkle Tree (BMT).
77pub struct Builder<H: Hasher> {
78    leaves: Vec<H::Digest>,
79}
80
81impl<H: Hasher> Builder<H> {
82    /// Creates a new Binary Merkle Tree builder.
83    pub fn new(leaves: usize) -> Self {
84        Self {
85            leaves: Vec::with_capacity(leaves),
86        }
87    }
88
89    /// Adds a leaf to the Binary Merkle Tree.
90    ///
91    /// When added, the leaf is hashed with its position.
92    ///
93    /// # Panics
94    ///
95    /// Panics if the tree already has `u32::MAX` leaves.
96    pub fn add(&mut self, leaf: &H::Digest) -> u32 {
97        // The count after this add must fit in Proof::leaf_count, a u32, so the
98        // maximum position is u32::MAX - 1.
99        let position: u32 = self.leaves.len().try_into().expect("too many leaves");
100        assert!(position < u32::MAX, "too many leaves");
101
102        let digest = H::hash(&[&position.to_be_bytes(), leaf.as_ref()]);
103        self.leaves.push(digest);
104        position
105    }
106
107    /// Builds the Binary Merkle Tree.
108    ///
109    /// It is valid to build a tree with no leaves, in which case
110    /// just an "empty" node is included (no leaves will be provable).
111    pub fn build(self) -> Tree<H::Digest> {
112        Tree::new::<H>(self.leaves)
113    }
114}
115
116/// Constructed Binary Merkle Tree (BMT).
117#[derive(Clone, Debug)]
118pub struct Tree<D: Digest> {
119    /// Records whether the tree is empty.
120    empty: bool,
121
122    /// The digests at each level of the tree (from leaves to root).
123    levels: NonEmptyVec<NonEmptyVec<D>>,
124
125    /// The finalized root digest, which incorporates the leaf count.
126    ///
127    /// This is computed as `H(leaf_count || tree_root)` to prevent
128    /// proof malleability where proofs that declare different leaf
129    /// counts could verify against the same root.
130    root: D,
131}
132
133impl<D: Digest> Tree<D> {
134    /// Builds a Merkle Tree from a slice of position-hashed leaf digests.
135    fn new<H: Hasher<Digest = D>>(mut leaves: Vec<D>) -> Self {
136        // If no leaves, add an empty node.
137        //
138        // Because this node only includes a position, there is no way a valid proof
139        // can be generated that references it.
140        let mut empty = false;
141        let leaf_count = leaves.len() as u32;
142        if leaves.is_empty() {
143            leaves.push(H::hash(&[]));
144            empty = true;
145        }
146
147        // Create the first level
148        let mut levels = non_empty_vec![non_empty_vec![@leaves]];
149
150        // Construct the tree level-by-level
151        let mut current_level = levels.last();
152        while !current_level.is_singleton() {
153            let mut next_level = Vec::with_capacity(current_level.len().get().div_ceil(2));
154
155            // Process four nodes (two sibling pairs) at a time, duplicating an unpaired
156            // trailing node. Hashing both pairs together lets the underlying hasher
157            // interleave independent messages (see `Hasher::hash_pair`). A trailing
158            // group with a single pair falls back to a single hash.
159            for group in current_level.chunks(4) {
160                match group {
161                    [a, b, c, d] => {
162                        let (left, right) =
163                            H::hash_pair(&[a.as_ref(), b.as_ref()], &[c.as_ref(), d.as_ref()]);
164                        next_level.push(left);
165                        next_level.push(right);
166                    }
167                    [a, b, c] => {
168                        let (left, right) =
169                            H::hash_pair(&[a.as_ref(), b.as_ref()], &[c.as_ref(), c.as_ref()]);
170                        next_level.push(left);
171                        next_level.push(right);
172                    }
173                    [a, b] => next_level.push(H::hash(&[a.as_ref(), b.as_ref()])),
174                    [a] => next_level.push(H::hash(&[a.as_ref(), a.as_ref()])),
175                    _ => unreachable!("chunks(4) yields at most 4 elements"),
176                }
177            }
178
179            // Add the computed level to the tree
180            levels.push(non_empty_vec![@next_level]);
181            current_level = levels.last();
182        }
183
184        // Compute the finalized root: H(leaf_count || tree_root)
185        // This binds the root to the tree size, preventing malleability attacks.
186        let tree_root = levels.last().first();
187        let root = H::hash(&[&leaf_count.to_be_bytes(), tree_root.as_ref()]);
188
189        Self {
190            empty,
191            levels,
192            root,
193        }
194    }
195
196    /// Returns the finalized root of the tree.
197    ///
198    /// The root incorporates the leaf count via `H(leaf_count || tree_root)`,
199    /// which prevents proof malleability attacks where different tree sizes
200    /// could produce valid proofs for the same root.
201    pub const fn root(&self) -> D {
202        self.root
203    }
204
205    /// Generates a Merkle proof for the leaf at `position`.
206    ///
207    /// This is a single-element multi-proof, which includes the minimal siblings
208    /// needed to reconstruct the root.
209    pub fn proof(&self, position: u32) -> Result<Proof<D>, Error> {
210        self.multi_proof(core::iter::once(position))
211    }
212
213    /// Generates a Merkle range proof for a contiguous set of leaves from `start`
214    /// to `end` (inclusive).
215    ///
216    /// The proof contains the minimal set of sibling digests needed to reconstruct
217    /// the root for all elements in the range. This is more efficient than individual
218    /// proofs when proving multiple consecutive elements.
219    pub fn range_proof(&self, start: u32, end: u32) -> Result<Proof<D>, Error> {
220        // For empty trees, return an empty proof
221        if self.empty {
222            if start == 0 && end == 0 {
223                return Ok(Proof::default());
224            }
225            return Err(Error::InvalidPosition(start));
226        }
227
228        // Validate range bounds
229        if start > end {
230            return Err(Error::InvalidPosition(start));
231        }
232        let leaf_count = self.levels.first().len().get() as u32;
233        if start >= leaf_count {
234            return Err(Error::InvalidPosition(start));
235        }
236        if end >= leaf_count {
237            return Err(Error::InvalidPosition(end));
238        }
239
240        // Compute required siblings without enumerating every leaf in the range.
241        let sibling_positions = siblings_required_for_range_proof(leaf_count, start, end)?;
242        let siblings: Vec<D> = sibling_positions
243            .iter()
244            .map(|&(level, index)| self.levels[level][index])
245            .collect();
246
247        Ok(Proof {
248            leaf_count,
249            siblings,
250        })
251    }
252
253    /// Generates a Merkle proof for multiple non-contiguous leaves at the given `positions`.
254    ///
255    /// The proof contains the minimal set of sibling digests needed to reconstruct
256    /// the root for all elements at the specified positions. This is more efficient
257    /// than individual proofs when proving multiple elements because shared siblings
258    /// are deduplicated.
259    ///
260    /// Positions are sorted internally; duplicate positions will return an error.
261    pub fn multi_proof<I, P>(&self, positions: I) -> Result<Proof<D>, Error>
262    where
263        I: IntoIterator<Item = P>,
264        P: core::borrow::Borrow<u32>,
265    {
266        let mut positions = positions.into_iter().peekable();
267
268        // Handle empty positions first - can't prove zero elements
269        let first = *positions.peek().ok_or(Error::NoLeaves)?.borrow();
270
271        // Handle empty tree case
272        if self.empty {
273            return Err(Error::InvalidPosition(first));
274        }
275
276        let leaf_count = self.levels.first().len().get() as u32;
277
278        // Get required sibling positions (this validates positions and checks for duplicates)
279        let sibling_positions =
280            siblings_required_for_multi_proof(leaf_count, positions.map(|p| *p.borrow()))?;
281
282        // Collect sibling digests in order
283        let siblings: Vec<D> = sibling_positions
284            .iter()
285            .map(|&(level, index)| self.levels[level][index])
286            .collect();
287
288        Ok(Proof {
289            leaf_count,
290            siblings,
291        })
292    }
293}
294
295/// A Merkle proof for multiple non-contiguous leaves in a Binary Merkle Tree.
296///
297/// This proof type is more space-efficient than generating individual proofs
298/// for each leaf because sibling nodes that are shared between multiple paths
299/// are deduplicated.
300///
301/// The proof contains the leaf count and sibling digests required for verification.
302/// The leaf count is incorporated into the root hash during finalization, so
303/// modifying it will cause verification to fail (preventing malleability attacks).
304#[derive(Clone, Debug, Eq, PartialEq)]
305pub struct Proof<D: Digest> {
306    /// The number of leaves in the tree.
307    ///
308    /// This value is incorporated into the root hash during finalization,
309    /// so modifying it will cause verification to fail (prevents malleability).
310    pub leaf_count: u32,
311
312    /// The deduplicated sibling digests required to verify all elements,
313    /// ordered by their position in the tree (level-major, then index within level).
314    pub siblings: Vec<D>,
315}
316
317impl<D: Digest> Default for Proof<D> {
318    fn default() -> Self {
319        Self {
320            leaf_count: 0,
321            siblings: Vec::new(),
322        }
323    }
324}
325
326impl<D: Digest> Write for Proof<D> {
327    fn write(&self, writer: &mut impl BufMut) {
328        self.leaf_count.write(writer);
329        self.siblings.write(writer);
330    }
331}
332
333impl<D: Digest> Read for Proof<D> {
334    /// The maximum number of items being proven.
335    ///
336    /// The upper bound on sibling hashes is derived as `max_items * MAX_LEVELS`.
337    type Cfg = usize;
338
339    fn read_cfg(
340        reader: &mut impl Buf,
341        max_items: &Self::Cfg,
342    ) -> Result<Self, commonware_codec::Error> {
343        let leaf_count = u32::read(reader)?;
344        let max_siblings = max_items.saturating_mul(MAX_LEVELS);
345        let siblings = Vec::<D>::read_range(reader, ..=max_siblings)?;
346        Ok(Self {
347            leaf_count,
348            siblings,
349        })
350    }
351}
352
353impl<D: Digest> EncodeSize for Proof<D> {
354    fn encode_size(&self) -> usize {
355        self.leaf_count.encode_size() + self.siblings.encode_size()
356    }
357}
358
359#[cfg(feature = "arbitrary")]
360impl<D: Digest> arbitrary::Arbitrary<'_> for Proof<D>
361where
362    D: for<'a> arbitrary::Arbitrary<'a>,
363{
364    fn arbitrary(u: &mut arbitrary::Unstructured<'_>) -> arbitrary::Result<Self> {
365        Ok(Self {
366            leaf_count: u.arbitrary()?,
367            siblings: u.arbitrary()?,
368        })
369    }
370}
371
372/// Returns the number of levels in a tree with `leaf_count` leaves.
373/// A tree with 1 leaf has 1 level, a tree with 2 leaves has 2 levels, etc.
374const fn levels_in_tree(leaf_count: u32) -> usize {
375    (u32::BITS - (leaf_count.saturating_sub(1)).leading_zeros() + 1) as usize
376}
377
378/// Returns the sorted, deduplicated positions of siblings required to prove
379/// inclusion of leaves at the given positions.
380///
381/// Each position in the result is encoded as `(level, index)` where level 0 is the leaf level.
382fn siblings_required_for_multi_proof(
383    leaf_count: u32,
384    positions: impl IntoIterator<Item = u32>,
385) -> Result<BTreeSet<(usize, usize)>, Error> {
386    // Validate positions and check for duplicates.
387    let mut current = BTreeSet::new();
388    for pos in positions {
389        if pos >= leaf_count {
390            return Err(Error::InvalidPosition(pos));
391        }
392        if !current.insert(pos as usize) {
393            return Err(Error::DuplicatePosition(pos));
394        }
395    }
396
397    if current.is_empty() {
398        return Err(Error::NoLeaves);
399    }
400
401    // Track positions we can compute at each level and record missing siblings.
402    // This keeps the work proportional to the number of positions, not the tree size.
403    let mut sibling_positions = BTreeSet::new();
404    let levels_count = levels_in_tree(leaf_count);
405    let mut level_size = leaf_count as usize;
406    for level in 0..levels_count - 1 {
407        for &index in &current {
408            let sibling_index = if index.is_multiple_of(2) {
409                if index + 1 < level_size {
410                    index + 1
411                } else {
412                    index
413                }
414            } else {
415                index - 1
416            };
417
418            if sibling_index != index && !current.contains(&sibling_index) {
419                sibling_positions.insert((level, sibling_index));
420            }
421        }
422
423        current = current.iter().map(|idx| idx / 2).collect();
424        level_size = level_size.div_ceil(2);
425    }
426
427    Ok(sibling_positions)
428}
429
430/// Returns the sorted, deduplicated positions of siblings required to prove
431/// inclusion of a contiguous range of leaves from `start` to `end` (inclusive).
432fn siblings_required_for_range_proof(
433    leaf_count: u32,
434    start: u32,
435    end: u32,
436) -> Result<BTreeSet<(usize, usize)>, Error> {
437    if leaf_count == 0 {
438        return Err(Error::NoLeaves);
439    }
440    if start > end {
441        return Err(Error::InvalidPosition(start));
442    }
443    if start >= leaf_count {
444        return Err(Error::InvalidPosition(start));
445    }
446    if end >= leaf_count {
447        return Err(Error::InvalidPosition(end));
448    }
449
450    let mut sibling_positions = BTreeSet::new();
451    let levels_count = levels_in_tree(leaf_count);
452    let mut level_start = start as usize;
453    let mut level_end = end as usize;
454    let mut level_size = leaf_count as usize;
455
456    for level in 0..levels_count - 1 {
457        if !level_start.is_multiple_of(2) {
458            sibling_positions.insert((level, level_start - 1));
459        }
460        if level_end.is_multiple_of(2) {
461            let right = level_end + 1;
462            if right < level_size {
463                sibling_positions.insert((level, right));
464            }
465        }
466
467        level_start /= 2;
468        level_end /= 2;
469        level_size = level_size.div_ceil(2);
470    }
471
472    Ok(sibling_positions)
473}
474
475impl<D: Digest> Proof<D> {
476    /// Verifies that a given `leaf` at `position` is included in a Binary Merkle Tree
477    /// with `root`.
478    ///
479    /// The proof consists of sibling hashes stored from the leaf up to the root. At each
480    /// level, if the current node is a left child (even index), the sibling is combined
481    /// to the right; if it is a right child (odd index), the sibling is combined to the
482    /// left.
483    ///
484    /// The `leaf_count` stored in the proof is incorporated into the finalized root
485    /// computation, so any modification to it will cause verification to fail.
486    pub fn verify_element_inclusion<H: Hasher<Digest = D>>(
487        &self,
488        leaf: &D,
489        mut position: u32,
490        root: &D,
491    ) -> Result<(), Error> {
492        // Validate position
493        if position >= self.leaf_count {
494            return Err(Error::InvalidPosition(position));
495        }
496
497        // Compute the position-hashed leaf
498        let mut computed = H::hash(&[&position.to_be_bytes(), leaf.as_ref()]);
499
500        // Track level size to handle odd-sized levels
501        let mut level_size = self.leaf_count as usize;
502        let mut sibling_iter = self.siblings.iter();
503
504        // Traverse from leaf to root
505        while level_size > 1 {
506            // Check if this is the last node at an odd-sized level (no real sibling)
507            let is_last_odd = position.is_multiple_of(2) && position as usize + 1 >= level_size;
508
509            let (left_node, right_node) = if is_last_odd {
510                // Node is duplicated - no sibling consumed from proof
511                (computed, computed)
512            } else if position.is_multiple_of(2) {
513                // Even position: sibling is to the right
514                let sibling = *sibling_iter.next().ok_or(Error::UnalignedProof)?;
515                (computed, sibling)
516            } else {
517                // Odd position: sibling is to the left
518                let sibling = *sibling_iter.next().ok_or(Error::UnalignedProof)?;
519                (sibling, computed)
520            };
521
522            // Compute the parent digest
523            computed = H::hash(&[left_node.as_ref(), right_node.as_ref()]);
524
525            // Move up the tree
526            position /= 2;
527            level_size = level_size.div_ceil(2);
528        }
529
530        // Ensure all siblings were consumed
531        if sibling_iter.next().is_some() {
532            return Err(Error::UnalignedProof);
533        }
534
535        // Finalize the root by incorporating the leaf count: H(leaf_count || tree_root)
536        // This binds the proof to the specific tree size, preventing malleability attacks.
537        let finalized = H::hash(&[&self.leaf_count.to_be_bytes(), computed.as_ref()]);
538
539        if finalized == *root {
540            Ok(())
541        } else {
542            Err(Error::InvalidProof(finalized.to_string(), root.to_string()))
543        }
544    }
545
546    /// Verifies that the given `elements` at their respective positions are included
547    /// in a Binary Merkle Tree with `root`.
548    ///
549    /// Elements can be provided in any order; positions are sorted internally.
550    /// Duplicate positions will cause verification to fail.
551    ///
552    /// The `leaf_count` stored in the proof is incorporated into the finalized root
553    /// computation, so any modification to it will cause verification to fail.
554    pub fn verify_multi_inclusion<H: Hasher<Digest = D>>(
555        &self,
556        elements: &[(D, u32)],
557        root: &D,
558    ) -> Result<(), Error> {
559        // Handle empty case
560        if elements.is_empty() {
561            if self.leaf_count == 0 && self.siblings.is_empty() {
562                // Compute finalized empty root: H(0 || empty_tree_root)
563                let empty_tree_root = H::hash(&[]);
564                let finalized = H::hash(&[&0u32.to_be_bytes(), empty_tree_root.as_ref()]);
565                if finalized == *root {
566                    return Ok(());
567                } else {
568                    return Err(Error::InvalidProof(finalized.to_string(), root.to_string()));
569                }
570            }
571            return Err(Error::NoLeaves);
572        }
573
574        // 1. Sort elements by position and check for duplicates/bounds
575        for (_, position) in elements {
576            if *position >= self.leaf_count {
577                return Err(Error::InvalidPosition(*position));
578            }
579        }
580        let mut sorted: Vec<(u32, D)> = Vec::with_capacity(elements.len());
581        let (leaf_chunks, leaf_remainder) = elements.as_chunks::<2>();
582        for chunk in leaf_chunks {
583            let (leaf_a, pos_a) = &chunk[0];
584            let (leaf_b, pos_b) = &chunk[1];
585            let (digest_a, digest_b) = H::hash_pair(
586                &[&pos_a.to_be_bytes(), leaf_a.as_ref()],
587                &[&pos_b.to_be_bytes(), leaf_b.as_ref()],
588            );
589            sorted.push((*pos_a, digest_a));
590            sorted.push((*pos_b, digest_b));
591        }
592        for (leaf, position) in leaf_remainder {
593            let digest = H::hash(&[&position.to_be_bytes(), leaf.as_ref()]);
594            sorted.push((*position, digest));
595        }
596        sorted.sort_unstable_by_key(|(pos, _)| *pos);
597
598        // Check for duplicates (adjacent elements with same position after sorting)
599        for i in 1..sorted.len() {
600            if sorted[i - 1].0 == sorted[i].0 {
601                return Err(Error::DuplicatePosition(sorted[i].0));
602            }
603        }
604
605        // 2. Iterate up the tree
606        // Since we process left-to-right and parent_pos = pos/2, next_level stays sorted.
607        let levels = levels_in_tree(self.leaf_count);
608        let mut level_size = self.leaf_count;
609        let mut sibling_iter = self.siblings.iter();
610        let mut current = sorted;
611        let mut next_level: Vec<(u32, D)> = Vec::with_capacity(current.len());
612        let mut parents: Vec<(u32, D, D)> = Vec::with_capacity(current.len());
613
614        for _ in 0..levels - 1 {
615            // First pass: determine each parent's (left, right) children without hashing, so
616            // independent parent digests can be batched together in the second pass.
617            let mut idx = 0;
618            while idx < current.len() {
619                let (pos, digest) = current[idx];
620                let parent_pos = pos / 2;
621
622                // Determine if we have the left or right child
623                let (left, right) = if pos.is_multiple_of(2) {
624                    // We are the LEFT child
625                    let left = digest;
626
627                    // Check if we have the right child in our current set
628                    let right = if idx + 1 < current.len() && current[idx + 1].0 == pos + 1 {
629                        idx += 1;
630                        current[idx].1
631                    } else if pos + 1 >= level_size {
632                        // If no right child exists in tree, duplicate left
633                        left
634                    } else {
635                        // Otherwise, must consume a sibling
636                        *sibling_iter.next().ok_or(Error::UnalignedProof)?
637                    };
638                    (left, right)
639                } else {
640                    // We are the RIGHT child
641                    // This implies the LEFT child was missing from 'current', so it must be a sibling.
642                    let right = digest;
643                    let left = *sibling_iter.next().ok_or(Error::UnalignedProof)?;
644                    (left, right)
645                };
646
647                parents.push((parent_pos, left, right));
648                idx += 1;
649            }
650
651            // Second pass: hash independent parent digests two at a time via `hash_pair`.
652            let (parent_chunks, parent_remainder) = parents.as_chunks::<2>();
653            for chunk in parent_chunks {
654                let (pos_a, left_a, right_a) = chunk[0];
655                let (pos_b, left_b, right_b) = chunk[1];
656                let (digest_a, digest_b) = H::hash_pair(
657                    &[left_a.as_ref(), right_a.as_ref()],
658                    &[left_b.as_ref(), right_b.as_ref()],
659                );
660                next_level.push((pos_a, digest_a));
661                next_level.push((pos_b, digest_b));
662            }
663            for &(pos, left, right) in parent_remainder {
664                next_level.push((pos, H::hash(&[left.as_ref(), right.as_ref()])));
665            }
666            parents.clear();
667
668            // Prepare for next level
669            core::mem::swap(&mut current, &mut next_level);
670            next_level.clear();
671            level_size = level_size.div_ceil(2);
672        }
673
674        // 3. Verify root
675        if sibling_iter.next().is_some() {
676            return Err(Error::UnalignedProof);
677        }
678
679        if current.len() != 1 {
680            return Err(Error::UnalignedProof);
681        }
682
683        // Finalize the root by incorporating the leaf count: H(leaf_count || tree_root)
684        // This binds the proof to the specific tree size, preventing malleability attacks.
685        let tree_root = current[0].1;
686        let finalized = H::hash(&[&self.leaf_count.to_be_bytes(), tree_root.as_ref()]);
687
688        if finalized == *root {
689            Ok(())
690        } else {
691            Err(Error::InvalidProof(finalized.to_string(), root.to_string()))
692        }
693    }
694
695    /// Verifies that a contiguous range of `leaves` starting at `position` are included
696    /// in a Binary Merkle Tree with `root`.
697    ///
698    /// This is a convenience method for verifying range proofs. The leaves must be
699    /// in order starting from `position`.
700    ///
701    /// The `leaf_count` stored in the proof is incorporated into the finalized root
702    /// computation, so any modification to it will cause verification to fail.
703    pub fn verify_range_inclusion<H: Hasher<Digest = D>>(
704        &self,
705        position: u32,
706        leaves: &[D],
707        root: &D,
708    ) -> Result<(), Error> {
709        // For empty trees, only position 0 with empty leaves is valid
710        if leaves.is_empty() && position != 0 {
711            return Err(Error::InvalidPosition(position));
712        }
713        if !leaves.is_empty() {
714            let leaves_len =
715                u32::try_from(leaves.len()).map_err(|_| Error::InvalidPosition(position))?;
716            let end = position
717                .checked_add(leaves_len - 1)
718                .ok_or(Error::InvalidPosition(position))?;
719            if end >= self.leaf_count {
720                return Err(Error::InvalidPosition(end));
721            }
722        }
723
724        // Convert to format expected by verify_multi_inclusion
725        let elements: Vec<(D, u32)> = leaves
726            .iter()
727            .enumerate()
728            .map(|(i, leaf)| (*leaf, position + i as u32))
729            .collect();
730        self.verify_multi_inclusion::<H>(&elements, root)
731    }
732}
733
734#[cfg(test)]
735mod tests {
736    use super::*;
737    use commonware_codec::{Decode, Encode};
738    use commonware_cryptography::sha256::{Digest, Sha256};
739    use rstest::rstest;
740
741    /// Regression test for https://github.com/commonwarexyz/monorepo/issues/2837
742    ///
743    /// Before the fix, two proofs with identical siblings but different leaf_count
744    /// values would both verify successfully against the same root, enabling
745    /// proof malleability attacks.
746    #[test]
747    fn issue_2837_regression() {
748        // Create a tree with 255 leaves (as in the issue report)
749        let digests: Vec<Digest> = (0..255u32)
750            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
751            .collect();
752
753        let mut builder = Builder::<Sha256>::new(255);
754        for digest in &digests {
755            builder.add(digest);
756        }
757        let tree = builder.build();
758        let root = tree.root();
759
760        // Get a valid proof for position 0
761        let original_proof = tree.proof(0).unwrap();
762        assert_eq!(original_proof.leaf_count, 255);
763
764        // Original proof should verify
765        assert!(
766            original_proof
767                .verify_element_inclusion::<Sha256>(&digests[0], 0, &root)
768                .is_ok(),
769            "Original proof should verify"
770        );
771
772        // Create a malleated proof with leaf_count=254 but same siblings
773        // (This is the exact attack from issue #2837)
774        let malleated_proof = Proof {
775            leaf_count: 254,
776            siblings: original_proof.siblings,
777        };
778
779        // Malleated proof should NOT verify because the root now incorporates
780        // the leaf_count: root = H(leaf_count || tree_root)
781        let result = malleated_proof.verify_element_inclusion::<Sha256>(&digests[0], 0, &root);
782        assert!(
783            result.is_err(),
784            "Malleated proof with wrong leaf_count must fail verification"
785        );
786    }
787
788    #[test]
789    fn test_tampered_proof_no_siblings() {
790        // Create transactions and digests
791        let txs = [b"tx1", b"tx2", b"tx3", b"tx4"];
792        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
793        let element = &digests[0];
794
795        // Build tree
796        let mut builder = Builder::<Sha256>::new(txs.len());
797        for digest in &digests {
798            builder.add(digest);
799        }
800        let tree = builder.build();
801        let root = tree.root();
802
803        // Build proof
804        let mut proof = tree.proof(0).unwrap();
805
806        // Tamper with proof
807        proof.siblings = Vec::new();
808
809        // Fail verification with an empty proof.
810        assert!(
811            proof
812                .verify_element_inclusion::<Sha256>(element, 0, &root)
813                .is_err()
814        );
815    }
816
817    #[test]
818    fn test_tampered_proof_extra_sibling() {
819        // Create transactions and digests
820        let txs = [b"tx1", b"tx2", b"tx3", b"tx4"];
821        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
822        let element = &digests[0];
823
824        // Build tree
825        let mut builder = Builder::<Sha256>::new(txs.len());
826        for digest in &digests {
827            builder.add(digest);
828        }
829        let tree = builder.build();
830        let root = tree.root();
831
832        // Build proof
833        let mut proof = tree.proof(0).unwrap();
834
835        // Tamper with proof
836        proof.siblings.push(*element);
837
838        // Fail verification with extra sibling
839        assert!(
840            proof
841                .verify_element_inclusion::<Sha256>(element, 0, &root)
842                .is_err()
843        );
844    }
845
846    #[test]
847    fn test_invalid_proof_wrong_element() {
848        // Create transactions and digests
849        let txs = [b"tx1", b"tx2", b"tx3", b"tx4"];
850        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
851
852        // Build tree
853        let mut builder = Builder::<Sha256>::new(txs.len());
854        for digest in &digests {
855            builder.add(digest);
856        }
857        let tree = builder.build();
858        let root = tree.root();
859
860        // Generate a valid proof for leaf at index 2.
861        let proof = tree.proof(2).unwrap();
862
863        // Use a wrong element (e.g. hash of a different transaction).
864        let wrong_leaf = Sha256::hash(&[b"wrong_tx"]);
865        assert!(
866            proof
867                .verify_element_inclusion::<Sha256>(&wrong_leaf, 2, &root)
868                .is_err()
869        );
870    }
871
872    #[test]
873    fn test_invalid_proof_wrong_index() {
874        // Create transactions and digests
875        let txs = [b"tx1", b"tx2", b"tx3", b"tx4"];
876        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
877
878        // Build tree
879        let mut builder = Builder::<Sha256>::new(txs.len());
880        for digest in &digests {
881            builder.add(digest);
882        }
883        let tree = builder.build();
884        let root = tree.root();
885
886        // Generate a valid proof for leaf at index 1.
887        let proof = tree.proof(1).unwrap();
888
889        // Use an incorrect index (e.g. 2 instead of 1).
890        assert!(
891            proof
892                .verify_element_inclusion::<Sha256>(&digests[1], 2, &root)
893                .is_err()
894        );
895    }
896
897    #[test]
898    fn test_invalid_proof_wrong_root() {
899        // Create transactions and digests
900        let txs = [b"tx1", b"tx2", b"tx3", b"tx4"];
901        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
902
903        // Build tree
904        let mut builder = Builder::<Sha256>::new(txs.len());
905        for digest in &digests {
906            builder.add(digest);
907        }
908        let tree = builder.build();
909
910        // Generate a valid proof for leaf at index 0.
911        let proof = tree.proof(0).unwrap();
912
913        // Use a wrong root (hash of a different input).
914        let wrong_root = Sha256::hash(&[b"wrong_root"]);
915        assert!(
916            proof
917                .verify_element_inclusion::<Sha256>(&digests[0], 0, &wrong_root)
918                .is_err()
919        );
920    }
921
922    #[test]
923    fn test_invalid_proof_serialization_truncated() {
924        // Create transactions and digests
925        let txs = [b"tx1", b"tx2", b"tx3"];
926        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
927
928        // Build tree
929        let mut builder = Builder::<Sha256>::new(txs.len());
930        for digest in &digests {
931            builder.add(digest);
932        }
933        let tree = builder.build();
934
935        // Generate a valid proof for leaf at index 1.
936        let proof = tree.proof(1).unwrap();
937        let mut serialized = proof.encode();
938
939        // Truncate one byte.
940        serialized.truncate(serialized.len() - 1);
941        assert!(Proof::<Digest>::decode_cfg(&mut serialized, &1).is_err());
942    }
943
944    #[test]
945    fn test_invalid_proof_serialization_extra() {
946        // Create transactions and digests
947        let txs = [b"tx1", b"tx2", b"tx3"];
948        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
949
950        // Build tree
951        let mut builder = Builder::<Sha256>::new(txs.len());
952        for digest in &digests {
953            builder.add(digest);
954        }
955        let tree = builder.build();
956
957        // Generate a valid proof for leaf at index 1.
958        let proof = tree.proof(1).unwrap();
959        let mut serialized = proof.encode_mut();
960
961        // Append an extra byte.
962        serialized.extend_from_slice(&[0u8]);
963        assert!(Proof::<Digest>::decode_cfg(&mut serialized, &1).is_err());
964    }
965
966    #[test]
967    fn test_invalid_proof_modified_hash() {
968        // Create transactions and digests
969        let txs = [b"tx1", b"tx2", b"tx3", b"tx4"];
970        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
971
972        // Build tree
973        let mut builder = Builder::<Sha256>::new(txs.len());
974        for digest in &digests {
975            builder.add(digest);
976        }
977        let tree = builder.build();
978        let root = tree.root();
979
980        // Generate a valid proof for leaf at index 2.
981        let mut proof = tree.proof(2).unwrap();
982
983        // Modify the first hash in the proof.
984        proof.siblings[0] = Sha256::hash(&[b"modified"]);
985        assert!(
986            proof
987                .verify_element_inclusion::<Sha256>(&digests[2], 2, &root)
988                .is_err()
989        );
990    }
991
992    #[test]
993    fn test_odd_tree_duplicate_index_proof() {
994        // Create transactions and digests
995        let txs = [b"tx1", b"tx2", b"tx3"];
996        let digests: Vec<Digest> = txs.iter().map(|tx| Sha256::hash(&[*tx])).collect();
997
998        // Build tree
999        let mut builder = Builder::<Sha256>::new(txs.len());
1000        for digest in &digests {
1001            builder.add(digest);
1002        }
1003        let tree = builder.build();
1004        let root = tree.root();
1005
1006        // The tree was built with 3 leaves; index 2 is the last valid index.
1007        let proof = tree.proof(2).unwrap();
1008
1009        // Verification should succeed for the proper index 2.
1010        assert!(
1011            proof
1012                .verify_element_inclusion::<Sha256>(&digests[2], 2, &root)
1013                .is_ok()
1014        );
1015
1016        // Should not be able to generate a proof for an out-of-range index (e.g. 3).
1017        assert!(tree.proof(3).is_err());
1018
1019        // Attempting to verify using an out-of-range index (e.g. 3, which would correspond
1020        // to a duplicate leaf that doesn't actually exist) should fail.
1021        assert!(
1022            proof
1023                .verify_element_inclusion::<Sha256>(&digests[2], 3, &root)
1024                .is_err()
1025        );
1026    }
1027
1028    #[test]
1029    fn test_range_proof_basic() {
1030        // Create test data
1031        let digests: Vec<Digest> = (0..8u32)
1032            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1033            .collect();
1034
1035        // Build tree
1036        let mut builder = Builder::<Sha256>::new(digests.len());
1037        for digest in &digests {
1038            builder.add(digest);
1039        }
1040        let tree = builder.build();
1041        let root = tree.root();
1042
1043        // Test range proof for elements 2-5
1044        let range_proof = tree.range_proof(2, 5).unwrap();
1045        let range_leaves = &digests[2..6];
1046
1047        assert!(
1048            range_proof
1049                .verify_range_inclusion::<Sha256>(2, range_leaves, &root)
1050                .is_ok()
1051        );
1052
1053        // Serialize and deserialize
1054        let mut serialized = range_proof.encode();
1055        let deserialized = Proof::<Digest>::decode_cfg(&mut serialized, &4).unwrap();
1056        assert!(
1057            deserialized
1058                .verify_range_inclusion::<Sha256>(2, range_leaves, &root)
1059                .is_ok()
1060        );
1061    }
1062
1063    #[test]
1064    fn test_range_proof_single_element() {
1065        // Create test data
1066        let digests: Vec<Digest> = (0..8u32)
1067            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1068            .collect();
1069
1070        // Build tree
1071        let mut builder = Builder::<Sha256>::new(digests.len());
1072        for digest in &digests {
1073            builder.add(digest);
1074        }
1075        let tree = builder.build();
1076        let root = tree.root();
1077
1078        // Test single element range proof
1079        for (i, digest) in digests.iter().enumerate() {
1080            let range_proof = tree.range_proof(i as u32, i as u32).unwrap();
1081
1082            let result = range_proof.verify_range_inclusion::<Sha256>(i as u32, &[*digest], &root);
1083            assert!(result.is_ok());
1084        }
1085    }
1086
1087    #[test]
1088    fn test_range_proof_full_tree() {
1089        // Create test data
1090        let digests: Vec<Digest> = (0..7u32)
1091            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1092            .collect();
1093
1094        // Build tree
1095        let mut builder = Builder::<Sha256>::new(digests.len());
1096        for digest in &digests {
1097            builder.add(digest);
1098        }
1099        let tree = builder.build();
1100        let root = tree.root();
1101
1102        // Test full tree range proof
1103        let range_proof = tree.range_proof(0, (digests.len() - 1) as u32).unwrap();
1104        assert!(
1105            range_proof
1106                .verify_range_inclusion::<Sha256>(0, &digests, &root)
1107                .is_ok()
1108        );
1109    }
1110
1111    #[test]
1112    fn test_range_proof_edge_cases() {
1113        // Create test data
1114        let digests: Vec<Digest> = (0..15u32)
1115            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1116            .collect();
1117
1118        // Build tree
1119        let mut builder = Builder::<Sha256>::new(digests.len());
1120        for digest in &digests {
1121            builder.add(digest);
1122        }
1123        let tree = builder.build();
1124        let root = tree.root();
1125
1126        // Test first half
1127        let range_proof = tree.range_proof(0, 7).unwrap();
1128        assert!(
1129            range_proof
1130                .verify_range_inclusion::<Sha256>(0, &digests[0..8], &root)
1131                .is_ok()
1132        );
1133
1134        // Test second half
1135        let range_proof = tree.range_proof(8, 14).unwrap();
1136        assert!(
1137            range_proof
1138                .verify_range_inclusion::<Sha256>(8, &digests[8..15], &root)
1139                .is_ok()
1140        );
1141
1142        // Test last elements
1143        let range_proof = tree.range_proof(13, 14).unwrap();
1144        assert!(
1145            range_proof
1146                .verify_range_inclusion::<Sha256>(13, &digests[13..15], &root)
1147                .is_ok()
1148        );
1149    }
1150
1151    #[test]
1152    fn test_range_proof_invalid_range() {
1153        // Create test data
1154        let digests: Vec<Digest> = (0..8u32)
1155            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1156            .collect();
1157
1158        // Build tree
1159        let mut builder = Builder::<Sha256>::new(digests.len());
1160        for digest in &digests {
1161            builder.add(digest);
1162        }
1163        let tree = builder.build();
1164
1165        // Test invalid ranges
1166        assert!(tree.range_proof(8, 8).is_err()); // Start out of bounds
1167        assert!(tree.range_proof(0, 8).is_err()); // End out of bounds
1168        assert!(tree.range_proof(5, 8).is_err()); // End out of bounds
1169        assert!(tree.range_proof(2, 1).is_err()); // Start > end
1170    }
1171
1172    #[test]
1173    fn test_range_proof_tampering() {
1174        // Create test data
1175        let digests: Vec<Digest> = (0..8u32)
1176            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1177            .collect();
1178
1179        // Build tree
1180        let mut builder = Builder::<Sha256>::new(digests.len());
1181        for digest in &digests {
1182            builder.add(digest);
1183        }
1184        let tree = builder.build();
1185        let root = tree.root();
1186
1187        // Get valid range proof
1188        let range_proof = tree.range_proof(2, 4).unwrap();
1189        let range_leaves = &digests[2..5];
1190
1191        // Test with wrong leaves
1192        let wrong_leaves = vec![
1193            Sha256::hash(&[b"wrong1"]),
1194            Sha256::hash(&[b"wrong2"]),
1195            Sha256::hash(&[b"wrong3"]),
1196        ];
1197        assert!(
1198            range_proof
1199                .verify_range_inclusion::<Sha256>(2, &wrong_leaves, &root)
1200                .is_err()
1201        );
1202
1203        // Test with wrong number of leaves
1204        assert!(
1205            range_proof
1206                .verify_range_inclusion::<Sha256>(2, &digests[2..4], &root)
1207                .is_err()
1208        );
1209
1210        // Test with tampered proof
1211        let mut tampered_proof = range_proof.clone();
1212        assert!(!tampered_proof.siblings.is_empty());
1213        // Tamper with the first sibling
1214        tampered_proof.siblings[0] = Sha256::hash(&[b"tampered"]);
1215        assert!(
1216            tampered_proof
1217                .verify_range_inclusion::<Sha256>(2, range_leaves, &root)
1218                .is_err()
1219        );
1220
1221        // Test with wrong root
1222        let wrong_root = Sha256::hash(&[b"wrong_root"]);
1223        assert!(
1224            range_proof
1225                .verify_range_inclusion::<Sha256>(2, range_leaves, &wrong_root)
1226                .is_err()
1227        );
1228    }
1229
1230    #[test]
1231    fn test_range_proof_various_sizes() {
1232        // Test range proofs for trees of various sizes
1233        for tree_size in [1, 2, 3, 4, 5, 7, 8, 15, 16, 31, 32, 63, 64] {
1234            let digests: Vec<Digest> = (0..tree_size as u32)
1235                .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1236                .collect();
1237
1238            // Build tree
1239            let mut builder = Builder::<Sha256>::new(digests.len());
1240            for digest in &digests {
1241                builder.add(digest);
1242            }
1243            let tree = builder.build();
1244            let root = tree.root();
1245
1246            // Test various range sizes
1247            for range_size in 1..=tree_size.min(8) {
1248                for start in 0..=(tree_size - range_size) {
1249                    let range_proof = tree
1250                        .range_proof(start as u32, (start + range_size - 1) as u32)
1251                        .unwrap();
1252                    let end = start + range_size;
1253                    assert!(
1254                        range_proof
1255                            .verify_range_inclusion::<Sha256>(
1256                                start as u32,
1257                                &digests[start..end],
1258                                &root
1259                            )
1260                            .is_ok(),
1261                        "Failed for tree_size={tree_size}, start={start}, range_size={range_size}"
1262                    );
1263                }
1264            }
1265        }
1266    }
1267
1268    #[test]
1269    fn test_range_proof_malicious_wrong_position() {
1270        // Create test data
1271        let digests: Vec<Digest> = (0..8u32)
1272            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1273            .collect();
1274
1275        // Build tree
1276        let mut builder = Builder::<Sha256>::new(digests.len());
1277        for digest in &digests {
1278            builder.add(digest);
1279        }
1280        let tree = builder.build();
1281        let root = tree.root();
1282
1283        // Get valid range proof for position 2 to 4
1284        let range_proof = tree.range_proof(2, 4).unwrap();
1285        let range_leaves = &digests[2..5];
1286
1287        // Try to verify with wrong position
1288        assert!(
1289            range_proof
1290                .verify_range_inclusion::<Sha256>(3, range_leaves, &root)
1291                .is_err()
1292        );
1293        assert!(
1294            range_proof
1295                .verify_range_inclusion::<Sha256>(1, range_leaves, &root)
1296                .is_err()
1297        );
1298    }
1299
1300    #[test]
1301    fn test_range_proof_malicious_reordered_leaves() {
1302        // Create test data
1303        let digests: Vec<Digest> = (0..8u32)
1304            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1305            .collect();
1306
1307        // Build tree
1308        let mut builder = Builder::<Sha256>::new(digests.len());
1309        for digest in &digests {
1310            builder.add(digest);
1311        }
1312        let tree = builder.build();
1313        let root = tree.root();
1314
1315        // Get valid range proof for position 2 to 4
1316        let range_proof = tree.range_proof(2, 4).unwrap();
1317
1318        // Try to verify with reordered leaves
1319        let reordered_leaves = vec![digests[3], digests[2], digests[4]];
1320        assert!(
1321            range_proof
1322                .verify_range_inclusion::<Sha256>(2, &reordered_leaves, &root)
1323                .is_err()
1324        );
1325    }
1326
1327    #[test]
1328    fn test_range_proof_malicious_extra_siblings() {
1329        // Create test data
1330        let digests: Vec<Digest> = (0..8u32)
1331            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1332            .collect();
1333
1334        // Build tree
1335        let mut builder = Builder::<Sha256>::new(digests.len());
1336        for digest in &digests {
1337            builder.add(digest);
1338        }
1339        let tree = builder.build();
1340        let root = tree.root();
1341
1342        // Get valid range proof
1343        let mut range_proof = tree.range_proof(2, 3).unwrap();
1344        let range_leaves = &digests[2..4];
1345
1346        // Tamper by adding extra siblings
1347        range_proof.siblings.push(Sha256::hash(&[b"extra"]));
1348        assert!(
1349            range_proof
1350                .verify_range_inclusion::<Sha256>(2, range_leaves, &root)
1351                .is_err()
1352        );
1353    }
1354
1355    #[test]
1356    fn test_range_proof_malicious_missing_siblings() {
1357        // Create test data
1358        let digests: Vec<Digest> = (0..8u32)
1359            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1360            .collect();
1361
1362        // Build tree
1363        let mut builder = Builder::<Sha256>::new(digests.len());
1364        for digest in &digests {
1365            builder.add(digest);
1366        }
1367        let tree = builder.build();
1368        let root = tree.root();
1369
1370        // Get valid range proof for a single element (which needs siblings)
1371        let mut range_proof = tree.range_proof(2, 2).unwrap();
1372        let range_leaves = &digests[2..3];
1373
1374        // The proof should have siblings
1375        assert!(!range_proof.siblings.is_empty());
1376
1377        // Remove a sibling
1378        range_proof.siblings.pop();
1379        assert!(
1380            range_proof
1381                .verify_range_inclusion::<Sha256>(2, range_leaves, &root)
1382                .is_err()
1383        );
1384    }
1385
1386    #[test]
1387    fn test_range_proof_integer_overflow_protection() {
1388        // Create test data
1389        let digests: Vec<Digest> = (0..8u32)
1390            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1391            .collect();
1392
1393        // Build tree
1394        let mut builder = Builder::<Sha256>::new(digests.len());
1395        for digest in &digests {
1396            builder.add(digest);
1397        }
1398        let tree = builder.build();
1399
1400        // Test overflow in range_proof generation
1401        assert!(tree.range_proof(u32::MAX, u32::MAX).is_err());
1402        assert!(tree.range_proof(u32::MAX - 1, u32::MAX).is_err());
1403        assert!(tree.range_proof(7, u32::MAX).is_err());
1404    }
1405
1406    #[test]
1407    fn test_range_proof_malicious_wrong_tree_structure() {
1408        // Create test data
1409        let digests: Vec<Digest> = (0..8u32)
1410            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1411            .collect();
1412
1413        // Build tree
1414        let mut builder = Builder::<Sha256>::new(digests.len());
1415        for digest in &digests {
1416            builder.add(digest);
1417        }
1418        let tree = builder.build();
1419        let root = tree.root();
1420
1421        // Get valid range proof
1422        let mut range_proof = tree.range_proof(2, 3).unwrap();
1423        let range_leaves = &digests[2..4];
1424
1425        // Add extra sibling (simulating proof from different tree structure)
1426        range_proof.siblings.push(Sha256::hash(&[b"fake_level"]));
1427        assert!(
1428            range_proof
1429                .verify_range_inclusion::<Sha256>(2, range_leaves, &root)
1430                .is_err()
1431        );
1432
1433        // Remove a sibling
1434        let mut range_proof = tree.range_proof(2, 2).unwrap();
1435        let range_leaves = &digests[2..3];
1436        assert!(!range_proof.siblings.is_empty());
1437        range_proof.siblings.pop();
1438        assert!(
1439            range_proof
1440                .verify_range_inclusion::<Sha256>(2, range_leaves, &root)
1441                .is_err()
1442        );
1443    }
1444
1445    #[test]
1446    fn test_range_proof_boundary_conditions() {
1447        // Test various power-of-2 boundary conditions
1448        for tree_size in [1, 2, 4, 8, 16, 32] {
1449            let digests: Vec<Digest> = (0..tree_size as u32)
1450                .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1451                .collect();
1452
1453            // Build tree
1454            let mut builder = Builder::<Sha256>::new(digests.len());
1455            for digest in &digests {
1456                builder.add(digest);
1457            }
1458            let tree = builder.build();
1459            let root = tree.root();
1460
1461            // Test edge cases
1462            // First element only
1463            let proof = tree.range_proof(0, 0).unwrap();
1464            assert!(
1465                proof
1466                    .verify_range_inclusion::<Sha256>(0, &digests[0..1], &root)
1467                    .is_ok()
1468            );
1469
1470            // Last element only
1471            let last_idx = tree_size - 1;
1472            let proof = tree.range_proof(last_idx as u32, last_idx as u32).unwrap();
1473            assert!(
1474                proof
1475                    .verify_range_inclusion::<Sha256>(
1476                        last_idx as u32,
1477                        &digests[last_idx..tree_size],
1478                        &root
1479                    )
1480                    .is_ok()
1481            );
1482
1483            // Full tree
1484            let proof = tree.range_proof(0, (tree_size - 1) as u32).unwrap();
1485            assert!(
1486                proof
1487                    .verify_range_inclusion::<Sha256>(0, &digests, &root)
1488                    .is_ok()
1489            );
1490        }
1491    }
1492
1493    #[test]
1494    fn test_empty_tree_proof() {
1495        // Build an empty tree
1496        let builder = Builder::<Sha256>::new(0);
1497        let tree = builder.build();
1498
1499        // Empty tree should fail for any position since there are no elements
1500        assert!(tree.proof(0).is_err());
1501        assert!(tree.proof(1).is_err());
1502        assert!(tree.proof(100).is_err());
1503    }
1504
1505    #[test]
1506    fn test_empty_tree_range_proof() {
1507        // Build an empty tree
1508        let builder = Builder::<Sha256>::new(0);
1509        let tree = builder.build();
1510        let root = tree.root();
1511
1512        // Empty tree should return default proof only for (0, 0)
1513        let range_proof = tree.range_proof(0, 0).unwrap();
1514        assert!(range_proof.siblings.is_empty());
1515        assert_eq!(range_proof, Proof::default());
1516
1517        // All other combinations should fail
1518        let invalid_ranges = vec![
1519            (0, 1),
1520            (0, 10),
1521            (1, 1),
1522            (1, 2),
1523            (5, 5),
1524            (10, 10),
1525            (0, u32::MAX),
1526            (u32::MAX, u32::MAX),
1527        ];
1528        for (start, end) in invalid_ranges {
1529            assert!(tree.range_proof(start, end).is_err());
1530        }
1531
1532        // Verify empty range proof against empty tree root
1533        let empty_leaves: &[Digest] = &[];
1534        assert!(
1535            range_proof
1536                .verify_range_inclusion::<Sha256>(0, empty_leaves, &root)
1537                .is_ok()
1538        );
1539
1540        // Should fail with non-empty leaves
1541        let non_empty_leaves = vec![Sha256::hash(&[b"leaf"])];
1542        assert!(
1543            range_proof
1544                .verify_range_inclusion::<Sha256>(0, &non_empty_leaves, &root)
1545                .is_err()
1546        );
1547
1548        // Should fail with wrong root
1549        let wrong_root = Sha256::hash(&[b"wrong"]);
1550        assert!(
1551            range_proof
1552                .verify_range_inclusion::<Sha256>(0, empty_leaves, &wrong_root)
1553                .is_err()
1554        );
1555
1556        // Should fail with wrong position
1557        assert!(
1558            range_proof
1559                .verify_range_inclusion::<Sha256>(1, empty_leaves, &root)
1560                .is_err()
1561        );
1562    }
1563
1564    #[test]
1565    fn test_empty_range_proof_serialization() {
1566        let proof = Proof::<Digest>::default();
1567        let mut serialized = proof.encode();
1568        let deserialized = Proof::<Digest>::decode_cfg(&mut serialized, &0).unwrap();
1569        assert_eq!(proof, deserialized);
1570    }
1571
1572    #[test]
1573    fn test_empty_tree_root_consistency() {
1574        // Create multiple empty trees and verify they have the same root
1575        let mut roots = Vec::new();
1576        for _ in 0..5 {
1577            let builder = Builder::<Sha256>::new(0);
1578            let tree = builder.build();
1579            roots.push(tree.root());
1580        }
1581
1582        // All empty trees should have the same root
1583        for i in 1..roots.len() {
1584            assert_eq!(roots[0], roots[i]);
1585        }
1586
1587        // The root should be the hash of empty data
1588        let empty = Sha256::hash(&[b""]);
1589        let expected_root = Sha256::hash(&[&0u32.to_be_bytes(), &empty]);
1590        assert_eq!(roots[0], expected_root);
1591    }
1592
1593    #[rstest]
1594    #[case::need_left_sibling(1, 2)] // Range starting at odd index (needs left sibling)
1595    #[case::need_right_sibling(4, 4)] // Range starting at even index
1596    #[case::full_tree(0, 16)] // Full tree (no siblings needed at leaf level)
1597    fn test_range_proof_siblings_usage(#[case] start: u32, #[case] count: u32) {
1598        // This test ensures that all siblings in a range proof are actually used during verification
1599        let digests: Vec<Digest> = (0..16u32)
1600            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1601            .collect();
1602
1603        // Build tree
1604        let mut builder = Builder::<Sha256>::new(digests.len());
1605        for digest in &digests {
1606            builder.add(digest);
1607        }
1608        let tree = builder.build();
1609        let root = tree.root();
1610
1611        let range_proof = tree.range_proof(start, start + count - 1).unwrap();
1612        let end = start as usize + count as usize;
1613
1614        // Verify the proof works
1615        assert!(
1616            range_proof
1617                .verify_range_inclusion::<Sha256>(start, &digests[start as usize..end], &root)
1618                .is_ok()
1619        );
1620
1621        // For each sibling, try tampering with it and verify the proof fails
1622        for sibling_idx in 0..range_proof.siblings.len() {
1623            let mut tampered_proof = range_proof.clone();
1624            tampered_proof.siblings[sibling_idx] = Sha256::hash(&[b"tampered"]);
1625            assert!(
1626                tampered_proof
1627                    .verify_range_inclusion::<Sha256>(start, &digests[start as usize..end], &root)
1628                    .is_err()
1629            );
1630        }
1631    }
1632
1633    // Test trees with odd sizes that require duplicate nodes
1634    #[rstest]
1635    fn test_range_proof_duplicate_node_edge_cases(
1636        #[values(3, 5, 7, 9, 11, 13, 15)] tree_size: usize,
1637    ) {
1638        let digests: Vec<Digest> = (0..tree_size as u32)
1639            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1640            .collect();
1641
1642        // Build tree
1643        let mut builder = Builder::<Sha256>::new(digests.len());
1644        for digest in &digests {
1645            builder.add(digest);
1646        }
1647        let tree = builder.build();
1648        let root = tree.root();
1649
1650        // Test range including the last element (which may require duplicate handling)
1651        let start = tree_size - 2;
1652        let proof = tree
1653            .range_proof(start as u32, (tree_size - 1) as u32)
1654            .unwrap();
1655        assert!(
1656            proof
1657                .verify_range_inclusion::<Sha256>(start as u32, &digests[start..tree_size], &root)
1658                .is_ok()
1659        );
1660    }
1661
1662    #[test]
1663    fn test_multi_proof_basic() {
1664        // Create test data
1665        let digests: Vec<Digest> = (0..8u32)
1666            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1667            .collect();
1668
1669        // Build tree
1670        let mut builder = Builder::<Sha256>::new(digests.len());
1671        for digest in &digests {
1672            builder.add(digest);
1673        }
1674        let tree = builder.build();
1675        let root = tree.root();
1676
1677        // Test multi-proof for non-contiguous positions [0, 3, 5]
1678        let positions = [0, 3, 5];
1679        let multi_proof = tree.multi_proof(positions).unwrap();
1680
1681        let elements: Vec<(Digest, u32)> = positions
1682            .iter()
1683            .map(|&p| (digests[p as usize], p))
1684            .collect();
1685        assert!(
1686            multi_proof
1687                .verify_multi_inclusion::<Sha256>(&elements, &root)
1688                .is_ok()
1689        );
1690    }
1691
1692    #[test]
1693    fn test_multi_proof_single_element() {
1694        // Create test data
1695        let digests: Vec<Digest> = (0..8u32)
1696            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1697            .collect();
1698
1699        // Build tree
1700        let mut builder = Builder::<Sha256>::new(digests.len());
1701        for digest in &digests {
1702            builder.add(digest);
1703        }
1704        let tree = builder.build();
1705        let root = tree.root();
1706
1707        // Test single element multi-proof for each position
1708        for (i, digest) in digests.iter().enumerate() {
1709            let multi_proof = tree.multi_proof([i as u32]).unwrap();
1710            let elements = [(*digest, i as u32)];
1711            assert!(
1712                multi_proof
1713                    .verify_multi_inclusion::<Sha256>(&elements, &root)
1714                    .is_ok(),
1715                "Failed for position {i}"
1716            );
1717        }
1718    }
1719
1720    #[test]
1721    fn test_multi_proof_all_elements() {
1722        // Create test data
1723        let digests: Vec<Digest> = (0..8u32)
1724            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1725            .collect();
1726
1727        // Build tree
1728        let mut builder = Builder::<Sha256>::new(digests.len());
1729        for digest in &digests {
1730            builder.add(digest);
1731        }
1732        let tree = builder.build();
1733        let root = tree.root();
1734
1735        // Test multi-proof for all elements
1736        let positions: Vec<u32> = (0..digests.len() as u32).collect();
1737        let multi_proof = tree.multi_proof(&positions).unwrap();
1738
1739        let elements: Vec<(Digest, u32)> = positions
1740            .iter()
1741            .map(|&p| (digests[p as usize], p))
1742            .collect();
1743        assert!(
1744            multi_proof
1745                .verify_multi_inclusion::<Sha256>(&elements, &root)
1746                .is_ok()
1747        );
1748
1749        // When proving all elements, we shouldn't need any siblings (all can be computed)
1750        assert!(multi_proof.siblings.is_empty());
1751    }
1752
1753    #[test]
1754    fn test_multi_proof_adjacent_elements() {
1755        // Create test data
1756        let digests: Vec<Digest> = (0..8u32)
1757            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1758            .collect();
1759
1760        // Build tree
1761        let mut builder = Builder::<Sha256>::new(digests.len());
1762        for digest in &digests {
1763            builder.add(digest);
1764        }
1765        let tree = builder.build();
1766        let root = tree.root();
1767
1768        // Test adjacent positions (should deduplicate shared siblings)
1769        let positions = [2, 3];
1770        let multi_proof = tree.multi_proof(positions).unwrap();
1771
1772        let elements: Vec<(Digest, u32)> = positions
1773            .iter()
1774            .map(|&p| (digests[p as usize], p))
1775            .collect();
1776        assert!(
1777            multi_proof
1778                .verify_multi_inclusion::<Sha256>(&elements, &root)
1779                .is_ok()
1780        );
1781    }
1782
1783    #[test]
1784    fn test_multi_proof_sparse_positions() {
1785        // Create test data
1786        let digests: Vec<Digest> = (0..16u32)
1787            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1788            .collect();
1789
1790        // Build tree
1791        let mut builder = Builder::<Sha256>::new(digests.len());
1792        for digest in &digests {
1793            builder.add(digest);
1794        }
1795        let tree = builder.build();
1796        let root = tree.root();
1797
1798        // Test widely separated positions
1799        let positions = [0, 7, 8, 15];
1800        let multi_proof = tree.multi_proof(positions).unwrap();
1801
1802        let elements: Vec<(Digest, u32)> = positions
1803            .iter()
1804            .map(|&p| (digests[p as usize], p))
1805            .collect();
1806        assert!(
1807            multi_proof
1808                .verify_multi_inclusion::<Sha256>(&elements, &root)
1809                .is_ok()
1810        );
1811    }
1812
1813    #[test]
1814    fn test_multi_proof_empty_tree() {
1815        // Build empty tree
1816        let builder = Builder::<Sha256>::new(0);
1817        let tree = builder.build();
1818
1819        // Empty tree with empty positions should return NoLeaves error
1820        // (we can't prove zero elements)
1821        assert!(matches!(
1822            tree.multi_proof(std::iter::empty::<u32>()),
1823            Err(Error::NoLeaves)
1824        ));
1825
1826        // Empty tree with any position should fail with InvalidPosition
1827        assert!(matches!(
1828            tree.multi_proof([0]),
1829            Err(Error::InvalidPosition(0))
1830        ));
1831    }
1832
1833    #[test]
1834    fn test_multi_proof_empty_positions() {
1835        // Create test data
1836        let digests: Vec<Digest> = (0..8u32)
1837            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1838            .collect();
1839
1840        // Build tree
1841        let mut builder = Builder::<Sha256>::new(digests.len());
1842        for digest in &digests {
1843            builder.add(digest);
1844        }
1845        let tree = builder.build();
1846
1847        // Empty positions should return error
1848        assert!(matches!(
1849            tree.multi_proof(std::iter::empty::<u32>()),
1850            Err(Error::NoLeaves)
1851        ));
1852    }
1853
1854    #[test]
1855    fn test_multi_proof_duplicate_positions_error() {
1856        // Create test data
1857        let digests: Vec<Digest> = (0..8u32)
1858            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1859            .collect();
1860
1861        // Build tree
1862        let mut builder = Builder::<Sha256>::new(digests.len());
1863        for digest in &digests {
1864            builder.add(digest);
1865        }
1866        let tree = builder.build();
1867
1868        // Duplicate positions should return error
1869        assert!(matches!(
1870            tree.multi_proof([1, 1]),
1871            Err(Error::DuplicatePosition(1))
1872        ));
1873        assert!(matches!(
1874            tree.multi_proof([0, 2, 2, 5]),
1875            Err(Error::DuplicatePosition(2))
1876        ));
1877    }
1878
1879    #[test]
1880    fn test_multi_proof_unsorted_input() {
1881        // Create test data
1882        let digests: Vec<Digest> = (0..8u32)
1883            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1884            .collect();
1885
1886        // Build tree
1887        let mut builder = Builder::<Sha256>::new(digests.len());
1888        for digest in &digests {
1889            builder.add(digest);
1890        }
1891        let tree = builder.build();
1892        let root = tree.root();
1893
1894        // Test with unsorted positions (should work - internal sorting)
1895        let positions = [5, 0, 3];
1896        let multi_proof = tree.multi_proof(positions).unwrap();
1897
1898        // Verify with unsorted elements (should work - internal sorting)
1899        let unsorted_elements = [(digests[5], 5), (digests[0], 0), (digests[3], 3)];
1900        assert!(
1901            multi_proof
1902                .verify_multi_inclusion::<Sha256>(&unsorted_elements, &root)
1903                .is_ok()
1904        );
1905    }
1906
1907    #[test]
1908    fn test_multi_proof_various_sizes() {
1909        // Test multi-proofs for trees of various sizes
1910        for tree_size in [1, 2, 3, 4, 5, 7, 8, 15, 16, 31, 32] {
1911            let digests: Vec<Digest> = (0..tree_size as u32)
1912                .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1913                .collect();
1914
1915            // Build tree
1916            let mut builder = Builder::<Sha256>::new(digests.len());
1917            for digest in &digests {
1918                builder.add(digest);
1919            }
1920            let tree = builder.build();
1921            let root = tree.root();
1922
1923            // Test various position combinations
1924            // First and last
1925            if tree_size >= 2 {
1926                let positions = [0, (tree_size - 1) as u32];
1927                let multi_proof = tree.multi_proof(positions).unwrap();
1928                let elements: Vec<(Digest, u32)> = positions
1929                    .iter()
1930                    .map(|&p| (digests[p as usize], p))
1931                    .collect();
1932                assert!(
1933                    multi_proof
1934                        .verify_multi_inclusion::<Sha256>(&elements, &root)
1935                        .is_ok(),
1936                    "Failed for tree_size={tree_size}, positions=[0, {}]",
1937                    tree_size - 1
1938                );
1939            }
1940
1941            // Every other element
1942            if tree_size >= 4 {
1943                let positions: Vec<u32> = (0..tree_size as u32).step_by(2).collect();
1944                let multi_proof = tree.multi_proof(&positions).unwrap();
1945                let elements: Vec<(Digest, u32)> = positions
1946                    .iter()
1947                    .map(|&p| (digests[p as usize], p))
1948                    .collect();
1949                assert!(
1950                    multi_proof
1951                        .verify_multi_inclusion::<Sha256>(&elements, &root)
1952                        .is_ok(),
1953                    "Failed for tree_size={tree_size}, every other element"
1954                );
1955            }
1956        }
1957    }
1958
1959    #[test]
1960    fn test_multi_proof_wrong_elements() {
1961        // Create test data
1962        let digests: Vec<Digest> = (0..8u32)
1963            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1964            .collect();
1965
1966        // Build tree
1967        let mut builder = Builder::<Sha256>::new(digests.len());
1968        for digest in &digests {
1969            builder.add(digest);
1970        }
1971        let tree = builder.build();
1972        let root = tree.root();
1973
1974        // Generate valid proof
1975        let positions = [0, 3, 5];
1976        let multi_proof = tree.multi_proof(positions).unwrap();
1977
1978        // Verify with wrong elements
1979        let wrong_elements = [
1980            (Sha256::hash(&[b"wrong1"]), 0),
1981            (digests[3], 3),
1982            (digests[5], 5),
1983        ];
1984        assert!(
1985            multi_proof
1986                .verify_multi_inclusion::<Sha256>(&wrong_elements, &root)
1987                .is_err()
1988        );
1989    }
1990
1991    #[test]
1992    fn test_multi_proof_wrong_positions() {
1993        // Create test data
1994        let digests: Vec<Digest> = (0..8u32)
1995            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
1996            .collect();
1997
1998        // Build tree
1999        let mut builder = Builder::<Sha256>::new(digests.len());
2000        for digest in &digests {
2001            builder.add(digest);
2002        }
2003        let tree = builder.build();
2004        let root = tree.root();
2005
2006        // Generate valid proof
2007        let positions = [0, 3, 5];
2008        let multi_proof = tree.multi_proof(positions).unwrap();
2009
2010        // Verify with wrong positions (same elements, different positions)
2011        let wrong_positions = [
2012            (digests[0], 1), // wrong position
2013            (digests[3], 3),
2014            (digests[5], 5),
2015        ];
2016        assert!(
2017            multi_proof
2018                .verify_multi_inclusion::<Sha256>(&wrong_positions, &root)
2019                .is_err()
2020        );
2021    }
2022
2023    #[test]
2024    fn test_multi_proof_wrong_root() {
2025        // Create test data
2026        let digests: Vec<Digest> = (0..8u32)
2027            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2028            .collect();
2029
2030        // Build tree
2031        let mut builder = Builder::<Sha256>::new(digests.len());
2032        for digest in &digests {
2033            builder.add(digest);
2034        }
2035        let tree = builder.build();
2036
2037        // Generate valid proof
2038        let positions = [0, 3, 5];
2039        let multi_proof = tree.multi_proof(positions).unwrap();
2040
2041        let elements: Vec<(Digest, u32)> = positions
2042            .iter()
2043            .map(|&p| (digests[p as usize], p))
2044            .collect();
2045
2046        // Verify with wrong root
2047        let wrong_root = Sha256::hash(&[b"wrong_root"]);
2048        assert!(
2049            multi_proof
2050                .verify_multi_inclusion::<Sha256>(&elements, &wrong_root)
2051                .is_err()
2052        );
2053    }
2054
2055    #[test]
2056    fn test_multi_proof_tampering() {
2057        // Create test data
2058        let digests: Vec<Digest> = (0..8u32)
2059            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2060            .collect();
2061
2062        // Build tree
2063        let mut builder = Builder::<Sha256>::new(digests.len());
2064        for digest in &digests {
2065            builder.add(digest);
2066        }
2067        let tree = builder.build();
2068        let root = tree.root();
2069
2070        // Generate valid proof
2071        let positions = [0, 5];
2072        let multi_proof = tree.multi_proof(positions).unwrap();
2073
2074        let elements: Vec<(Digest, u32)> = positions
2075            .iter()
2076            .map(|&p| (digests[p as usize], p))
2077            .collect();
2078
2079        // Tamper with sibling
2080        assert!(!multi_proof.siblings.is_empty());
2081        let mut modified = multi_proof.clone();
2082        modified.siblings[0] = Sha256::hash(&[b"tampered"]);
2083        assert!(
2084            modified
2085                .verify_multi_inclusion::<Sha256>(&elements, &root)
2086                .is_err()
2087        );
2088
2089        // Add extra sibling
2090        let mut extra = multi_proof.clone();
2091        extra.siblings.push(Sha256::hash(&[b"extra"]));
2092        assert!(
2093            extra
2094                .verify_multi_inclusion::<Sha256>(&elements, &root)
2095                .is_err()
2096        );
2097
2098        // Remove a sibling
2099        let mut missing = multi_proof;
2100        missing.siblings.pop();
2101        assert!(
2102            missing
2103                .verify_multi_inclusion::<Sha256>(&elements, &root)
2104                .is_err()
2105        );
2106    }
2107
2108    #[test]
2109    fn test_multi_proof_deduplication() {
2110        // Create test data
2111        let digests: Vec<Digest> = (0..16u32)
2112            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2113            .collect();
2114
2115        // Build tree
2116        let mut builder = Builder::<Sha256>::new(digests.len());
2117        for digest in &digests {
2118            builder.add(digest);
2119        }
2120        let tree = builder.build();
2121
2122        // Get individual proofs
2123        let individual_siblings: usize = [0u32, 1, 8, 9]
2124            .iter()
2125            .map(|&p| tree.proof(p).unwrap().siblings.len())
2126            .sum();
2127
2128        // Get multi-proof for same positions
2129        let multi_proof = tree.multi_proof([0, 1, 8, 9]).unwrap();
2130
2131        // Multi-proof should have fewer siblings due to deduplication
2132        assert!(
2133            multi_proof.siblings.len() < individual_siblings,
2134            "Multi-proof ({}) should have fewer siblings than sum of individual proofs ({})",
2135            multi_proof.siblings.len(),
2136            individual_siblings
2137        );
2138    }
2139
2140    #[test]
2141    fn test_multi_proof_serialization() {
2142        // Create test data
2143        let digests: Vec<Digest> = (0..8u32)
2144            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2145            .collect();
2146
2147        // Build tree
2148        let mut builder = Builder::<Sha256>::new(digests.len());
2149        for digest in &digests {
2150            builder.add(digest);
2151        }
2152        let tree = builder.build();
2153        let root = tree.root();
2154
2155        // Generate proof
2156        let positions = [0, 3, 5];
2157        let multi_proof = tree.multi_proof(positions).unwrap();
2158
2159        // Serialize and deserialize
2160        let serialized = multi_proof.encode();
2161        let deserialized = Proof::<Digest>::decode_cfg(serialized, &positions.len()).unwrap();
2162
2163        assert_eq!(multi_proof, deserialized);
2164
2165        // Verify deserialized proof works
2166        let elements: Vec<(Digest, u32)> = positions
2167            .iter()
2168            .map(|&p| (digests[p as usize], p))
2169            .collect();
2170        assert!(
2171            deserialized
2172                .verify_multi_inclusion::<Sha256>(&elements, &root)
2173                .is_ok()
2174        );
2175    }
2176
2177    #[test]
2178    fn test_multi_proof_serialization_truncated() {
2179        // Create test data
2180        let digests: Vec<Digest> = (0..8u32)
2181            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2182            .collect();
2183
2184        // Build tree
2185        let mut builder = Builder::<Sha256>::new(digests.len());
2186        for digest in &digests {
2187            builder.add(digest);
2188        }
2189        let tree = builder.build();
2190
2191        // Generate proof
2192        let positions = [0, 3, 5];
2193        let multi_proof = tree.multi_proof(positions).unwrap();
2194
2195        // Serialize and truncate
2196        let mut serialized = multi_proof.encode();
2197        serialized.truncate(serialized.len() - 1);
2198
2199        // Should fail to deserialize
2200        assert!(Proof::<Digest>::decode_cfg(&mut serialized, &positions.len()).is_err());
2201    }
2202
2203    #[test]
2204    fn test_multi_proof_serialization_extra() {
2205        // Create test data
2206        let digests: Vec<Digest> = (0..8u32)
2207            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2208            .collect();
2209
2210        // Build tree
2211        let mut builder = Builder::<Sha256>::new(digests.len());
2212        for digest in &digests {
2213            builder.add(digest);
2214        }
2215        let tree = builder.build();
2216
2217        // Generate proof
2218        let positions = [0, 3, 5];
2219        let multi_proof = tree.multi_proof(positions).unwrap();
2220
2221        // Serialize and add extra byte
2222        let mut serialized = multi_proof.encode_mut();
2223        serialized.extend_from_slice(&[0u8]);
2224
2225        // Should fail to deserialize
2226        assert!(Proof::<Digest>::decode_cfg(&mut serialized, &positions.len()).is_err());
2227    }
2228
2229    #[test]
2230    fn test_multi_proof_decode_insufficient_data() {
2231        let mut serialized = Vec::new();
2232        serialized.extend_from_slice(&8u32.encode()); // leaf_count
2233        serialized.extend_from_slice(&1usize.encode()); // claims 1 sibling but no data follows
2234
2235        // Should fail because the buffer claims 1 sibling but doesn't have the data
2236        let err = Proof::<Digest>::decode_cfg(serialized.as_slice(), &1).unwrap_err();
2237        assert!(matches!(err, commonware_codec::Error::EndOfBuffer));
2238    }
2239
2240    #[test]
2241    fn test_multi_proof_invalid_position() {
2242        // Create test data
2243        let digests: Vec<Digest> = (0..8u32)
2244            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2245            .collect();
2246
2247        // Build tree
2248        let mut builder = Builder::<Sha256>::new(digests.len());
2249        for digest in &digests {
2250            builder.add(digest);
2251        }
2252        let tree = builder.build();
2253
2254        // Test out of bounds position
2255        assert!(matches!(
2256            tree.multi_proof([0, 8]),
2257            Err(Error::InvalidPosition(8))
2258        ));
2259        assert!(matches!(
2260            tree.multi_proof([100]),
2261            Err(Error::InvalidPosition(100))
2262        ));
2263    }
2264
2265    #[test]
2266    fn test_multi_proof_verify_invalid_position() {
2267        // Create test data
2268        let digests: Vec<Digest> = (0..8u32)
2269            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2270            .collect();
2271
2272        // Build tree
2273        let mut builder = Builder::<Sha256>::new(digests.len());
2274        for digest in &digests {
2275            builder.add(digest);
2276        }
2277        let tree = builder.build();
2278        let root = tree.root();
2279
2280        // Generate valid proof
2281        let positions = [0, 3];
2282        let multi_proof = tree.multi_proof(positions).unwrap();
2283
2284        // Try to verify with out of bounds position
2285        let invalid_elements = [(digests[0], 0), (digests[3], 100)];
2286        assert!(
2287            multi_proof
2288                .verify_multi_inclusion::<Sha256>(&invalid_elements, &root)
2289                .is_err()
2290        );
2291    }
2292
2293    #[test]
2294    fn test_multi_proof_odd_tree_sizes() {
2295        // Test odd-sized trees that require node duplication
2296        for tree_size in [3, 5, 7, 9, 11, 13, 15] {
2297            let digests: Vec<Digest> = (0..tree_size as u32)
2298                .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2299                .collect();
2300
2301            // Build tree
2302            let mut builder = Builder::<Sha256>::new(digests.len());
2303            for digest in &digests {
2304                builder.add(digest);
2305            }
2306            let tree = builder.build();
2307            let root = tree.root();
2308
2309            // Test with positions including the last element
2310            let positions = [0, (tree_size - 1) as u32];
2311            let multi_proof = tree.multi_proof(positions).unwrap();
2312
2313            let elements: Vec<(Digest, u32)> = positions
2314                .iter()
2315                .map(|&p| (digests[p as usize], p))
2316                .collect();
2317            assert!(
2318                multi_proof
2319                    .verify_multi_inclusion::<Sha256>(&elements, &root)
2320                    .is_ok(),
2321                "Failed for tree_size={tree_size}"
2322            );
2323        }
2324    }
2325
2326    #[test]
2327    fn test_multi_proof_verify_empty_elements() {
2328        // Create a valid proof and try to verify with empty elements
2329        let digests: Vec<Digest> = (0..8u32)
2330            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2331            .collect();
2332
2333        let mut builder = Builder::<Sha256>::new(digests.len());
2334        for digest in &digests {
2335            builder.add(digest);
2336        }
2337        let tree = builder.build();
2338        let root = tree.root();
2339
2340        // Generate valid proof
2341        let positions = [0, 3];
2342        let multi_proof = tree.multi_proof(positions).unwrap();
2343
2344        // Try to verify with empty elements
2345        let empty_elements: &[(Digest, u32)] = &[];
2346        assert!(
2347            multi_proof
2348                .verify_multi_inclusion::<Sha256>(empty_elements, &root)
2349                .is_err()
2350        );
2351    }
2352
2353    #[test]
2354    fn test_multi_proof_default_verify() {
2355        // Default (empty) proof should only verify against empty tree
2356        let default_proof = Proof::<Digest>::default();
2357
2358        // Empty elements against default proof
2359        let empty_elements: &[(Digest, u32)] = &[];
2360
2361        // Build empty tree to get the empty root
2362        let builder = Builder::<Sha256>::new(0);
2363        let empty_tree = builder.build();
2364        let empty_root = empty_tree.root();
2365
2366        assert!(
2367            default_proof
2368                .verify_multi_inclusion::<Sha256>(empty_elements, &empty_root)
2369                .is_ok()
2370        );
2371
2372        // Should fail with wrong root
2373        let wrong_root = Sha256::hash(&[b"not_empty"]);
2374        assert!(
2375            default_proof
2376                .verify_multi_inclusion::<Sha256>(empty_elements, &wrong_root)
2377                .is_err()
2378        );
2379    }
2380
2381    #[test]
2382    fn test_multi_proof_single_leaf_tree() {
2383        // Edge case: tree with exactly one leaf
2384        let digest = Sha256::hash(&[b"only_leaf"]);
2385
2386        // Build single-leaf tree
2387        let mut builder = Builder::<Sha256>::new(1);
2388        builder.add(&digest);
2389        let tree = builder.build();
2390        let root = tree.root();
2391
2392        // Generate multi-proof for the only leaf
2393        let multi_proof = tree.multi_proof([0]).unwrap();
2394
2395        // Single leaf tree: leaf_count should be 1
2396        assert_eq!(multi_proof.leaf_count, 1);
2397
2398        // Single leaf tree: no siblings needed (leaf is the root after position hashing)
2399        assert!(
2400            multi_proof.siblings.is_empty(),
2401            "Single leaf tree should have no siblings"
2402        );
2403
2404        // Verify the proof
2405        let elements = [(digest, 0u32)];
2406        assert!(
2407            multi_proof
2408                .verify_multi_inclusion::<Sha256>(&elements, &root)
2409                .is_ok(),
2410            "Single leaf multi-proof verification failed"
2411        );
2412
2413        // Verify with wrong digest fails
2414        let wrong_digest = Sha256::hash(&[b"wrong"]);
2415        let wrong_elements = [(wrong_digest, 0u32)];
2416        assert!(
2417            multi_proof
2418                .verify_multi_inclusion::<Sha256>(&wrong_elements, &root)
2419                .is_err(),
2420            "Should fail with wrong digest"
2421        );
2422
2423        // Verify with wrong position fails
2424        let wrong_position_elements = [(digest, 1u32)];
2425        assert!(
2426            multi_proof
2427                .verify_multi_inclusion::<Sha256>(&wrong_position_elements, &root)
2428                .is_err(),
2429            "Should fail with invalid position"
2430        );
2431    }
2432
2433    #[test]
2434    fn test_multi_proof_malicious_leaf_count_zero() {
2435        // Attacker sets leaf_count = 0 but provides siblings
2436        let digests: Vec<Digest> = (0..8u32)
2437            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2438            .collect();
2439
2440        let mut builder = Builder::<Sha256>::new(digests.len());
2441        for digest in &digests {
2442            builder.add(digest);
2443        }
2444        let tree = builder.build();
2445        let root = tree.root();
2446
2447        // Generate valid proof and tamper with leaf_count
2448        let positions = [0, 3];
2449        let mut multi_proof = tree.multi_proof(positions).unwrap();
2450        multi_proof.leaf_count = 0;
2451
2452        let elements: Vec<(Digest, u32)> = positions
2453            .iter()
2454            .map(|&p| (digests[p as usize], p))
2455            .collect();
2456
2457        // Should fail - leaf_count=0 but we have elements
2458        assert!(
2459            multi_proof
2460                .verify_multi_inclusion::<Sha256>(&elements, &root)
2461                .is_err()
2462        );
2463    }
2464
2465    #[test]
2466    fn test_multi_proof_malicious_leaf_count_larger() {
2467        // Attacker inflates leaf_count to claim proof is for larger tree
2468        let digests: Vec<Digest> = (0..8u32)
2469            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2470            .collect();
2471
2472        let mut builder = Builder::<Sha256>::new(digests.len());
2473        for digest in &digests {
2474            builder.add(digest);
2475        }
2476        let tree = builder.build();
2477        let root = tree.root();
2478
2479        // Generate valid proof and inflate leaf_count
2480        let positions = [0, 3];
2481        let mut multi_proof = tree.multi_proof(positions).unwrap();
2482        let original_leaf_count = multi_proof.leaf_count;
2483        multi_proof.leaf_count = 1000;
2484
2485        let elements: Vec<(Digest, u32)> = positions
2486            .iter()
2487            .map(|&p| (digests[p as usize], p))
2488            .collect();
2489
2490        // Should fail - inflated leaf_count changes required siblings
2491        assert!(
2492            multi_proof
2493                .verify_multi_inclusion::<Sha256>(&elements, &root)
2494                .is_err(),
2495            "Should reject proof with inflated leaf_count ({} -> {})",
2496            original_leaf_count,
2497            multi_proof.leaf_count
2498        );
2499    }
2500
2501    #[test]
2502    fn test_multi_proof_malicious_leaf_count_smaller() {
2503        // Attacker deflates leaf_count to claim proof is for smaller tree
2504        let digests: Vec<Digest> = (0..8u32)
2505            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2506            .collect();
2507
2508        let mut builder = Builder::<Sha256>::new(digests.len());
2509        for digest in &digests {
2510            builder.add(digest);
2511        }
2512        let tree = builder.build();
2513        let root = tree.root();
2514
2515        // Generate valid proof and deflate leaf_count
2516        let positions = [0, 3];
2517        let mut multi_proof = tree.multi_proof(positions).unwrap();
2518        multi_proof.leaf_count = 4; // Smaller than actual tree
2519
2520        let elements: Vec<(Digest, u32)> = positions
2521            .iter()
2522            .map(|&p| (digests[p as usize], p))
2523            .collect();
2524
2525        // Should fail - deflated leaf_count changes tree structure
2526        assert!(
2527            multi_proof
2528                .verify_multi_inclusion::<Sha256>(&elements, &root)
2529                .is_err(),
2530            "Should reject proof with deflated leaf_count"
2531        );
2532    }
2533
2534    #[test]
2535    fn test_multi_proof_mismatched_element_count() {
2536        // Provide more or fewer elements than the proof was generated for
2537        let digests: Vec<Digest> = (0..8u32)
2538            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2539            .collect();
2540
2541        let mut builder = Builder::<Sha256>::new(digests.len());
2542        for digest in &digests {
2543            builder.add(digest);
2544        }
2545        let tree = builder.build();
2546        let root = tree.root();
2547
2548        // Generate proof for 2 positions
2549        let positions = [0, 3];
2550        let multi_proof = tree.multi_proof(positions).unwrap();
2551
2552        // Try to verify with only 1 element (too few)
2553        let too_few = [(digests[0], 0u32)];
2554        assert!(
2555            multi_proof
2556                .verify_multi_inclusion::<Sha256>(&too_few, &root)
2557                .is_err(),
2558            "Should reject when fewer elements provided than proof was generated for"
2559        );
2560
2561        // Try to verify with 3 elements (too many)
2562        let too_many = [(digests[0], 0u32), (digests[3], 3), (digests[5], 5)];
2563        assert!(
2564            multi_proof
2565                .verify_multi_inclusion::<Sha256>(&too_many, &root)
2566                .is_err(),
2567            "Should reject when more elements provided than proof was generated for"
2568        );
2569    }
2570
2571    #[test]
2572    fn test_multi_proof_swapped_siblings() {
2573        // Swap the order of siblings in the proof
2574        let digests: Vec<Digest> = (0..8u32)
2575            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2576            .collect();
2577
2578        let mut builder = Builder::<Sha256>::new(digests.len());
2579        for digest in &digests {
2580            builder.add(digest);
2581        }
2582        let tree = builder.build();
2583        let root = tree.root();
2584
2585        // Generate valid proof with multiple siblings
2586        let positions = [0, 5];
2587        let mut multi_proof = tree.multi_proof(positions).unwrap();
2588
2589        // Ensure we have at least 2 siblings to swap
2590        if multi_proof.siblings.len() >= 2 {
2591            // Swap first two siblings
2592            multi_proof.siblings.swap(0, 1);
2593
2594            let elements: Vec<(Digest, u32)> = positions
2595                .iter()
2596                .map(|&p| (digests[p as usize], p))
2597                .collect();
2598
2599            assert!(
2600                multi_proof
2601                    .verify_multi_inclusion::<Sha256>(&elements, &root)
2602                    .is_err(),
2603                "Should reject proof with swapped siblings"
2604            );
2605        }
2606    }
2607
2608    #[test]
2609    fn test_multi_proof_dos_large_leaf_count() {
2610        // Attacker sets massive leaf_count trying to cause DoS via memory allocation
2611        // The verify function should NOT allocate proportional to leaf_count
2612        let digests: Vec<Digest> = (0..4u32)
2613            .map(|i| Sha256::hash(&[&i.to_be_bytes()]))
2614            .collect();
2615
2616        let mut builder = Builder::<Sha256>::new(digests.len());
2617        for digest in &digests {
2618            builder.add(digest);
2619        }
2620        let tree = builder.build();
2621        let root = tree.root();
2622
2623        // Generate valid proof
2624        let positions = [0, 2];
2625        let mut multi_proof = tree.multi_proof(positions).unwrap();
2626
2627        // Set massive leaf_count (attacker trying to exhaust memory)
2628        multi_proof.leaf_count = u32::MAX;
2629
2630        let elements: Vec<(Digest, u32)> = positions
2631            .iter()
2632            .map(|&p| (digests[p as usize], p))
2633            .collect();
2634
2635        // This should fail quickly without allocating massive memory
2636        // The function is O(elements * levels), not O(leaf_count)
2637        let result = multi_proof.verify_multi_inclusion::<Sha256>(&elements, &root);
2638        assert!(result.is_err(), "Should reject malicious large leaf_count");
2639    }
2640
2641    #[cfg(feature = "arbitrary")]
2642    mod conformance {
2643        use super::*;
2644        use commonware_codec::conformance::CodecConformance;
2645        use commonware_conformance::Conformance;
2646        use commonware_cryptography::sha256::Digest as Sha256Digest;
2647
2648        fn test_merkle_tree(n: usize) -> Digest {
2649            // Build tree
2650            let mut digests = Vec::with_capacity(n);
2651            let mut builder = Builder::<Sha256>::new(n);
2652            for i in 0..n {
2653                let digest = Sha256::hash(&[&i.to_be_bytes()]);
2654                builder.add(&digest);
2655                digests.push(digest);
2656            }
2657            let tree = builder.build();
2658            let root = tree.root();
2659
2660            // For each leaf, generate and verify its proof
2661            for (i, leaf) in digests.iter().enumerate() {
2662                // Generate proof
2663                let proof = tree.proof(i as u32).unwrap();
2664                assert!(
2665                    proof
2666                        .verify_element_inclusion::<Sha256>(leaf, i as u32, &root)
2667                        .is_ok(),
2668                    "correct fail for size={n} leaf={i}"
2669                );
2670
2671                // Serialize and deserialize the proof
2672                let serialized = proof.encode();
2673                let deserialized = Proof::<Digest>::decode_cfg(serialized, &1).unwrap();
2674                assert!(
2675                    deserialized
2676                        .verify_element_inclusion::<Sha256>(leaf, i as u32, &root)
2677                        .is_ok(),
2678                    "deserialize fail for size={n} leaf={i}"
2679                );
2680
2681                // Modify a sibling hash and ensure the proof fails
2682                if !proof.siblings.is_empty() {
2683                    let mut update_tamper = proof.clone();
2684                    update_tamper.siblings[0] = Sha256::hash(&[b"tampered"]);
2685                    assert!(
2686                        update_tamper
2687                            .verify_element_inclusion::<Sha256>(leaf, i as u32, &root)
2688                            .is_err(),
2689                        "modify fail for size={n} leaf={i}"
2690                    );
2691                }
2692
2693                // Add a sibling hash and ensure the proof fails
2694                let mut add_tamper = proof.clone();
2695                add_tamper.siblings.push(Sha256::hash(&[b"tampered"]));
2696                assert!(
2697                    add_tamper
2698                        .verify_element_inclusion::<Sha256>(leaf, i as u32, &root)
2699                        .is_err(),
2700                    "add fail for size={n} leaf={i}"
2701                );
2702
2703                // Remove a sibling hash and ensure the proof fails
2704                if !proof.siblings.is_empty() {
2705                    let mut remove_tamper = proof.clone();
2706                    remove_tamper.siblings.pop();
2707                    assert!(
2708                        remove_tamper
2709                            .verify_element_inclusion::<Sha256>(leaf, i as u32, &root)
2710                            .is_err(),
2711                        "remove fail for size={n} leaf={i}"
2712                    );
2713                }
2714            }
2715
2716            // Test proof for larger than size
2717            assert!(tree.proof(n as u32).is_err());
2718
2719            // Return the root so we can ensure we don't silently change.
2720            root
2721        }
2722
2723        struct RootConformance;
2724
2725        impl Conformance for RootConformance {
2726            async fn commit(seed: u64) -> Vec<u8> {
2727                let root = test_merkle_tree(seed as usize);
2728                root.to_vec()
2729            }
2730        }
2731
2732        commonware_conformance::conformance_tests! {
2733            CodecConformance<Proof<Sha256Digest>>,
2734            RootConformance => 200
2735        }
2736    }
2737}