|
| 1 | +// Copyright 2020-2021 IOTA Stiftung |
| 2 | +// SPDX-License-Identifier: Apache-2.0 |
| 3 | + |
| 4 | +use bee_message::MessageId; |
| 5 | + |
| 6 | +use crypto::hashes::{Digest, Output}; |
| 7 | + |
| 8 | +use std::marker::PhantomData; |
| 9 | + |
| 10 | +/// Leaf domain separation prefix. |
| 11 | +const LEAF_HASH_PREFIX: u8 = 0x00; |
| 12 | +/// Node domain separation prefix. |
| 13 | +const NODE_HASH_PREFIX: u8 = 0x01; |
| 14 | + |
| 15 | +/// Computes the largest power of two inferior to `n`. |
| 16 | +fn largest_power_of_two(n: u32) -> usize { |
| 17 | + 1 << (32 - n.leading_zeros() - 1) |
| 18 | +} |
| 19 | + |
| 20 | +/// A Merkle hasher based on a digest function. |
| 21 | +pub(crate) struct MerkleHasher<D> { |
| 22 | + marker: PhantomData<D>, |
| 23 | +} |
| 24 | + |
| 25 | +impl<D: Default + Digest> MerkleHasher<D> { |
| 26 | + /// Creates a new Merkle hasher. |
| 27 | + pub(crate) fn new() -> Self { |
| 28 | + Self { marker: PhantomData } |
| 29 | + } |
| 30 | + |
| 31 | + /// Returns the digest of the empty hash. |
| 32 | + fn empty(&mut self) -> Output<D> { |
| 33 | + D::digest(&[]) |
| 34 | + } |
| 35 | + |
| 36 | + /// Returns the digest of a Merkle leaf. |
| 37 | + fn leaf(&mut self, message_id: MessageId) -> Output<D> { |
| 38 | + let mut hasher = D::default(); |
| 39 | + |
| 40 | + hasher.update([LEAF_HASH_PREFIX]); |
| 41 | + hasher.update(message_id); |
| 42 | + hasher.finalize() |
| 43 | + } |
| 44 | + |
| 45 | + /// Returns the digest of a Merkle node. |
| 46 | + fn node(&mut self, message_ids: &[MessageId]) -> Output<D> { |
| 47 | + let mut hasher = D::default(); |
| 48 | + let (left, right) = message_ids.split_at(largest_power_of_two(message_ids.len() as u32 - 1)); |
| 49 | + |
| 50 | + hasher.update([NODE_HASH_PREFIX]); |
| 51 | + hasher.update(self.digest_inner(left)); |
| 52 | + hasher.update(self.digest_inner(right)); |
| 53 | + hasher.finalize() |
| 54 | + } |
| 55 | + |
| 56 | + /// Returns the digest of a list of hashes as an `Output<D>`. |
| 57 | + fn digest_inner(&mut self, message_ids: &[MessageId]) -> Output<D> { |
| 58 | + match message_ids.len() { |
| 59 | + 0 => self.empty(), |
| 60 | + 1 => self.leaf(message_ids[0]), |
| 61 | + _ => self.node(message_ids), |
| 62 | + } |
| 63 | + } |
| 64 | + |
| 65 | + /// Returns the digest of a list of hashes as a `Vec<u8>`. |
| 66 | + pub(crate) fn digest(&mut self, message_ids: &[MessageId]) -> Vec<u8> { |
| 67 | + self.digest_inner(message_ids).to_vec() |
| 68 | + } |
| 69 | +} |
| 70 | + |
| 71 | +#[cfg(test)] |
| 72 | +mod tests { |
| 73 | + |
| 74 | + use super::*; |
| 75 | + |
| 76 | + use crypto::hashes::blake2b::Blake2b256; |
| 77 | + |
| 78 | + use std::str::FromStr; |
| 79 | + |
| 80 | + #[test] |
| 81 | + fn tree() { |
| 82 | + let mut hashes = Vec::new(); |
| 83 | + |
| 84 | + for hash in [ |
| 85 | + "52fdfc072182654f163f5f0f9a621d729566c74d10037c4d7bbb0407d1e2c649", |
| 86 | + "81855ad8681d0d86d1e91e00167939cb6694d2c422acd208a0072939487f6999", |
| 87 | + "eb9d18a44784045d87f3c67cf22746e995af5a25367951baa2ff6cd471c483f1", |
| 88 | + "5fb90badb37c5821b6d95526a41a9504680b4e7c8b763a1b1d49d4955c848621", |
| 89 | + "6325253fec738dd7a9e28bf921119c160f0702448615bbda08313f6a8eb668d2", |
| 90 | + "0bf5059875921e668a5bdf2c7fc4844592d2572bcd0668d2d6c52f5054e2d083", |
| 91 | + "6bf84c7174cb7476364cc3dbd968b0f7172ed85794bb358b0c3b525da1786f9f", |
| 92 | + ] |
| 93 | + .iter() |
| 94 | + { |
| 95 | + hashes.push(MessageId::from_str(hash).unwrap()); |
| 96 | + } |
| 97 | + |
| 98 | + let hash = MerkleHasher::<Blake2b256>::new().digest(&hashes); |
| 99 | + |
| 100 | + assert_eq!( |
| 101 | + hex::encode(hash), |
| 102 | + "bf67ce7ba23e8c0951b5abaec4f5524360d2c26d971ff226d3359fa70cdb0beb" |
| 103 | + ) |
| 104 | + } |
| 105 | +} |
0 commit comments