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