2018-08-08 07:41:11 -07:00
|
|
|
use std::mem;
|
|
|
|
|
2018-10-29 07:36:56 -07:00
|
|
|
use serde_derive::{Deserialize, Serialize};
|
2018-08-09 02:51:31 -07:00
|
|
|
use tiny_keccak::sha3_256;
|
2018-08-08 07:41:11 -07:00
|
|
|
|
2018-08-09 02:51:31 -07:00
|
|
|
pub type Digest = [u8; 32];
|
2018-08-08 07:41:11 -07:00
|
|
|
|
|
|
|
/// A Merkle tree: The leaves are values and their hashes. Each level consists of the hashes of
|
|
|
|
/// pairs of values on the previous level. The root is the value in the first level with only one
|
|
|
|
/// entry.
|
|
|
|
#[derive(Debug)]
|
|
|
|
pub struct MerkleTree<T> {
|
|
|
|
levels: Vec<Vec<Digest>>,
|
|
|
|
values: Vec<T>,
|
|
|
|
root_hash: Digest,
|
|
|
|
}
|
|
|
|
|
|
|
|
impl<T: AsRef<[u8]> + Clone> MerkleTree<T> {
|
|
|
|
/// Creates a new Merkle tree with the given values.
|
|
|
|
pub fn from_vec(values: Vec<T>) -> Self {
|
|
|
|
let mut levels = Vec::new();
|
|
|
|
let mut cur_lvl: Vec<Digest> = values.iter().map(hash).collect();
|
|
|
|
while cur_lvl.len() > 1 {
|
|
|
|
let next_lvl = cur_lvl.chunks(2).map(hash_chunk).collect();
|
|
|
|
levels.push(mem::replace(&mut cur_lvl, next_lvl));
|
|
|
|
}
|
|
|
|
let root_hash = cur_lvl[0];
|
|
|
|
MerkleTree {
|
|
|
|
levels,
|
|
|
|
values,
|
|
|
|
root_hash,
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the proof for entry `index`, if that is a valid index.
|
|
|
|
pub fn proof(&self, index: usize) -> Option<Proof<T>> {
|
|
|
|
let value = self.values.get(index)?.clone();
|
|
|
|
let mut lvl_i = index;
|
|
|
|
let mut digests = Vec::new();
|
|
|
|
for level in &self.levels {
|
|
|
|
// Insert the sibling hash if there is one.
|
|
|
|
if let Some(digest) = level.get(lvl_i ^ 1) {
|
2018-08-09 02:51:31 -07:00
|
|
|
digests.push(*digest);
|
2018-08-08 07:41:11 -07:00
|
|
|
}
|
|
|
|
lvl_i /= 2;
|
|
|
|
}
|
|
|
|
Some(Proof {
|
|
|
|
index,
|
|
|
|
digests,
|
|
|
|
value,
|
2018-08-09 02:51:31 -07:00
|
|
|
root_hash: self.root_hash,
|
2018-08-08 07:41:11 -07:00
|
|
|
})
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the root hash of the tree.
|
2018-08-09 02:51:31 -07:00
|
|
|
pub fn root_hash(&self) -> &Digest {
|
|
|
|
&self.root_hash
|
2018-08-08 07:41:11 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns a the slice containing all leaf values.
|
|
|
|
pub fn values(&self) -> &[T] {
|
|
|
|
&self.values
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Consumes the tree, and returns the vector of leaf values.
|
|
|
|
pub fn into_values(self) -> Vec<T> {
|
|
|
|
self.values
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
/// A proof that a value is at a particular index in the Merkle tree specified by its root hash.
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
|
|
|
|
pub struct Proof<T> {
|
|
|
|
value: T,
|
|
|
|
index: usize,
|
2018-08-09 02:51:31 -07:00
|
|
|
digests: Vec<Digest>,
|
|
|
|
root_hash: Digest,
|
2018-08-08 07:41:11 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
impl<T: AsRef<[u8]>> Proof<T> {
|
|
|
|
/// Returns `true` if the digests in this proof constitute a valid branch in a Merkle tree with
|
|
|
|
/// the root hash.
|
|
|
|
pub fn validate(&self, n: usize) -> bool {
|
|
|
|
let mut digest = hash(&self.value);
|
|
|
|
let mut lvl_i = self.index;
|
|
|
|
let mut lvl_n = n;
|
|
|
|
let mut digest_itr = self.digests.iter();
|
|
|
|
while lvl_n > 1 {
|
|
|
|
if lvl_i ^ 1 < lvl_n {
|
|
|
|
digest = match digest_itr.next() {
|
|
|
|
None => return false, // Not enough levels in the proof.
|
|
|
|
Some(sibling) if lvl_i & 1 == 1 => hash_pair(&sibling, &digest),
|
|
|
|
Some(sibling) => hash_pair(&digest, &sibling),
|
|
|
|
};
|
|
|
|
}
|
|
|
|
lvl_i /= 2; // Our index on the next level.
|
|
|
|
lvl_n = (lvl_n + 1) / 2; // The next level's size.
|
|
|
|
}
|
|
|
|
if digest_itr.next().is_some() {
|
|
|
|
return false; // Too many levels in the proof.
|
|
|
|
}
|
2018-08-09 02:51:31 -07:00
|
|
|
digest == self.root_hash
|
2018-08-08 07:41:11 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the index of this proof's value in the tree.
|
|
|
|
pub fn index(&self) -> usize {
|
|
|
|
self.index
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the tree's root hash.
|
2018-08-09 02:51:31 -07:00
|
|
|
pub fn root_hash(&self) -> &Digest {
|
|
|
|
&self.root_hash
|
2018-08-08 07:41:11 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the leaf value.
|
|
|
|
pub fn value(&self) -> &T {
|
|
|
|
&self.value
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Consumes the proof and returns the leaf value.
|
|
|
|
pub fn into_value(self) -> T {
|
|
|
|
self.value
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Takes a chunk of one or two digests. In the former case, returns the digest itself, in the
|
|
|
|
/// latter, it returns the hash of the two digests.
|
|
|
|
fn hash_chunk(chunk: &[Digest]) -> Digest {
|
|
|
|
if chunk.len() == 1 {
|
|
|
|
chunk[0]
|
|
|
|
} else {
|
|
|
|
hash_pair(&chunk[0], &chunk[1])
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the hash of the concatenated bytes of `d0` and `d1`.
|
|
|
|
fn hash_pair<T0: AsRef<[u8]>, T1: AsRef<[u8]>>(v0: &T0, v1: &T1) -> Digest {
|
|
|
|
let bytes: Vec<u8> = v0.as_ref().iter().chain(v1.as_ref()).cloned().collect();
|
|
|
|
hash(&bytes)
|
|
|
|
}
|
|
|
|
|
|
|
|
/// Returns the SHA-256 hash of the value's `[u8]` representation.
|
|
|
|
fn hash<T: AsRef<[u8]>>(value: T) -> Digest {
|
2018-08-09 02:51:31 -07:00
|
|
|
sha3_256(value.as_ref())
|
2018-08-08 07:41:11 -07:00
|
|
|
}
|
|
|
|
|
|
|
|
#[cfg(test)]
|
|
|
|
mod tests {
|
|
|
|
use super::MerkleTree;
|
|
|
|
|
|
|
|
#[test]
|
|
|
|
fn test_merkle() {
|
|
|
|
for &n in &[4, 7, 8, 9, 17] {
|
|
|
|
let tree = MerkleTree::from_vec((0..n).map(|i| vec![i as u8]).collect());
|
|
|
|
for i in 0..n {
|
|
|
|
let proof = tree.proof(i).expect("couldn't get proof");
|
|
|
|
assert!(proof.validate(n));
|
|
|
|
}
|
|
|
|
assert!(tree.proof(n).is_none());
|
|
|
|
}
|
|
|
|
}
|
|
|
|
}
|