1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
//! Merkle tree test implementation
//!
//! # Usage
//!
//! Add this to your `Cargo.toml`
//!
//! ```toml
//! [dependencies]
//! merkle_test = "0.1"
//! ```
//!
//! ```rust
//! extern crate merkle_test;
//! use merkle_test::MerkleTree;
//! use merkle_test::crypto_reexport::sha2::Sha256;
//!
//! fn main() {
//!     let hashes: Vec<Vec<u8>> = vec![];
//!     let mut tree = MerkleTree::new(Sha256::new()); // use Sha256
//!     hashes.iter().for_each(|hash| tree.add_hash(hash.to_vec())); // add hashes to tree
//!
//!     let root = tree.root(); // calculate merkle root
//!     println!("Merkle root: {:?}", root);
//! }
//! ```
//!

extern crate crypto;

use crypto::digest::Digest;

mod error;
pub use error::MerkleTreeError;

type Hash = Vec<u8>;
type Round = Vec<Hash>;

/// Mercle tree implementation
pub struct MerkleTree<D: Digest> {
    digest: D,
    state: Vec<Hash>,
}
impl<D: Digest> MerkleTree<D> {
    /// Consructs new MerkleTree with specific [Digest](https://docs.rs/rust-crypto/0.2.36/crypto/digest/trait.Digest.html)
    pub fn new(digest: D) -> MerkleTree<D> {
        let state = Vec::new();
        MerkleTree { digest, state }
    }
    /// Add hash to tree (hash must be impl `Into<Vec<u8>>`)
    pub fn add_hash<H: Into<Hash>>(&mut self, hash: H) {
        self.state.push(hash.into());
    }
    /// Calulates root hash of tree state
    pub fn root(&mut self) -> Result<Hash, MerkleTreeError> {
        if self.state.len() < 2 {
            return Err(MerkleTreeError::NotEnoughHashes)
        }

        let mut current_round = self.state.clone() as Round;
        loop {
            current_round = self.round(current_round);
            if current_round.len() == 1 {
                return Ok(current_round[0].clone());
            }
        }
    }
    /// Calculate tree round
    fn round(&mut self, prev_round: Round) -> Round {
        prev_round
            .chunks(2)
            .map(|chunk| if chunk.len() == 2 {
                self.digest.input(&chunk[0]);
                self.digest.input(&chunk[1]);
                let mut pair_hash = vec![0; self.digest.output_bytes()];
                self.digest.result(&mut pair_hash);
                self.digest.reset();
                pair_hash
            } else {
                chunk[0].clone()
            })
            .collect()
    }
}

pub mod crypto_reexport {
    pub use crypto::*;
}