A simple merkle tree

hash.rs

use std::fmt;

use ring::digest;
use ring::constant_time::verify_slices_are_equal;
use rustc_serialize::hex::{ ToHex, FromHex, FromHexError };
use rustc_serialize::base64;
use rustc_serialize::base64::{ ToBase64, FromBase64, FromBase64Error };

/// Errors that occur when parsing a hash value
pub enum ParseError {
    InvalidCharacter,
    InvalidLength,
}

#[derive(Copy, Clone)]
pub struct Hash256 {
    hash: [u8; 32],
}

impl Hash256 {
    pub fn from_bytes<'a, B>(buf: B) -> Result<Hash256, ParseError>
    where B : Into<&'a [u8]> {
        let buf = buf.into();
        if buf.len() == 32 {
            let mut result = Hash256 { hash: [0; 32] };
            result.hash.copy_from_slice(buf);
            Ok(result)
        } else {
            Err(ParseError::InvalidLength)
        }
    }

    pub fn from_hex<'a, S>(buf: S) -> Result<Hash256, ParseError>
    where S : Into<&'a str>{
        buf.into().from_hex()
            .map_err(|err| match err {
                FromHexError::InvalidHexCharacter(_,_) => ParseError::InvalidCharacter,
                FromHexError::InvalidHexLength => ParseError::InvalidLength,
            })
            .and_then(|buf| Self::from_bytes(buf.as_slice()))
    }

    pub fn from_base64<'a, S>(buf: S) -> Result<Hash256, ParseError>
    where S : Into<&'a str>{
        buf.into().from_base64()
            .map_err(|err| match err {
                FromBase64Error::InvalidBase64Byte(_,_) => ParseError::InvalidCharacter,
                FromBase64Error::InvalidBase64Length => ParseError::InvalidLength,
            })
            .and_then(|buf| Self::from_bytes(buf.as_slice()))
    }

    fn from_digest(digest: &digest::Digest) -> Hash256 {
        let mut result = Hash256 { hash: [0; 32] };
        result.hash.copy_from_slice(digest.as_ref());
        result
    }

    pub fn double_sha_str<'a, S>(buf: S) -> Hash256
    where S : Into<&'a str> {
        Self::double_sha_bytes(buf.into().as_bytes())
    }

    pub fn double_sha_bytes<'a, B>(buf: B) -> Hash256
    where B : Into<&'a [u8]> {
        let mut sha = digest::Context::new(&digest::SHA256);
        sha.update(buf.into());
        let first = sha.finish();
        let second = digest::digest(&digest::SHA256, first.as_ref());

        Self::from_digest(&second)
    }

    pub fn double_sha_concat(a: &Hash256, b: &Hash256) -> Hash256 {
        let mut sha = digest::Context::new(&digest::SHA256);
        sha.update(&a.hash);
        sha.update(&b.hash);

        let first = sha.finish();
        let second = digest::digest(&digest::SHA256, first.as_ref());

        Self::from_digest(&second)
    }

    pub fn to_hex_prefix(&self) -> String {
        self.hash[..4].to_hex()
    }

    // done in "constant_time" to avoid timing attacks
    pub fn verify_equal(&self, other: &Hash256) -> Result<(), ()> {
        verify_slices_are_equal(&self.hash, &other.hash)
            .map_err(|_| ())
    }
}

impl AsRef<[u8]> for Hash256 {
    fn as_ref<'a>(&'a self) -> &'a [u8] {
        &self.hash
    }
}

impl ToHex for Hash256 {
    fn to_hex(&self) -> String {
        self.hash.to_hex()
    }
}

impl ToBase64 for Hash256 {
    fn to_base64(&self, config: base64::Config) -> String {
        self.hash.to_base64(config)
    }
}

impl fmt::Debug for Hash256 {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "{}", &self.to_hex_prefix() )
    }
}

impl fmt::Display for Hash256 {
    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        f.write_str(&self.to_hex_prefix())
    }
}

impl PartialEq for Hash256 {
    fn eq(&self, other: &Hash256) -> bool {
        self.verify_equal(other).is_ok()
    }
}

merkle.rs

use std::mem;
use std::ops;
use hash::Hash256;

/// A representation of a large list of hashes. Typically only the root hash is
/// passed around; it contains enough information to verify the entire tree.
///
/// A merkle tree is supierior to a linear hashing when either a few updates
/// need to happen or to prove a single hash is in the tree. For the later a
/// Merkle path is constructed.
pub struct MerkleTree {
    // root of the tree is levels[0]
    levels: Vec<MerkleLevel>
}

impl MerkleTree {
    /// Build a merkle tree with the provided hashes at the leaves
    pub fn build_tree(leaves: Vec<Hash256>) -> MerkleTree {
        assert!(leaves.len() > 0);
        let mut levels: Vec<_> =  ByLevelBuilder::new(leaves).collect();
        levels.reverse();
        MerkleTree { levels: levels }
    }

    /// Calculate the root of the merkle tree, but without putting
    /// the whole tree in memory.
    pub fn build_root(leaves: Vec<Hash256>) -> Hash256 {
        ByLevelBuilder::new(leaves)
            .last().unwrap()
            .row[0]
    }

    pub fn to_vec(&self) -> Vec<Hash256> {
        self.levels.iter()
            .rev()
            .flat_map( |ref lvl| lvl.row.iter() )
            .cloned()
            .collect()
    }

    pub fn root_hash(&self) -> &Hash256 {
        &self.levels[0].row[0]
    }

    pub fn leaves(&self) -> &[Hash256] {
        &self.levels[self.levels.len() - 1].row
    }

    /// Update the hash of one of the leaves. This updates lg(2) of
    /// the hashes, all the way up to and including the root.
    pub fn update_leaf(&mut self, idx: usize, hash: Hash256) {
        let mut idx = idx;
        let mut hash = hash;
        let mut level_idx = self.levels.len() - 1;
        while level_idx > 0 {
            let row = &mut self.levels[level_idx].row;
            row[idx] = hash;

            let parent_idx = idx >> 1;
            let a = &row[(parent_idx << 1)];
            let b = &row.get((parent_idx << 1) + 1)
                .unwrap_or(a);

            hash = Hash256::double_sha_concat(a, b);
            idx = parent_idx;
            level_idx -= 1;
        }

        self.levels[0].row[0] = hash;
    }

    /// Create a path that verifies the existence of a leaf
    /// in the tree. Note: the returned path also verifies one
    /// of the leaf siblings (if there is one).
    pub fn merkle_path(&self, idx: usize) -> MerklePath {
        if self.levels.len() == 1 {
            assert!(idx == 0);
            return MerklePath { path: self.levels[0].row.clone() };
        }

        let mut path = Vec::new();

        let mut idx = idx;
        let mut level = self.levels.len() - 1;
        while level > 0 {
            let row = &self.levels[level].row;
            let parent_idx = idx >> 1;
            let a = &row[(parent_idx << 1)];
            let b = row.get((parent_idx << 1) + 1)
                .unwrap_or(a);
            path.push(*a);
            path.push(*b);
            level -= 1;
            idx = parent_idx;
        }

        path.push(self.levels[0].row[0]);

        MerklePath { path: path }
    }
}

struct MerkleLevel {
    row: Vec<Hash256>
}

impl MerkleLevel {
    fn build_next_level(&self) -> Self {
        let mut next_row = Vec::new();
        let len = self.row.len();

        let mut i = 0;
        while i < len {
            let ref a = self.row[i];
            let ref b = self.row.get(i+1).unwrap_or(a);

            let parent = Hash256::double_sha_concat(&a, &b);

            next_row.push(parent);

            i += 2
        }

        MerkleLevel { row: next_row }
    }
}

pub struct MerklePath {
    path: Vec<Hash256>,
}

impl MerklePath {
    pub fn verify_root(&self, expected: &Hash256) -> Result<(), ()> {
        self.path[self.path.len() - 1].verify_equal(expected)
    }

    pub fn verify_leaf(&self, expected: &Hash256) -> Result<(), ()> {
        if self.path.len() == 1 {
            self.verify_root(expected)
        } else {
            // Don't know if the leaf is at position 0 or 1
            self.path[0].verify_equal(expected)
                .or_else(|_| self.path[1].verify_equal(expected))
        }
    }
}

impl ops::Deref for MerklePath {
    type Target = [Hash256];
    fn deref(&self) -> &Self::Target {
        &self.path
    }
}

struct ByLevelBuilder {
    cur : MerkleLevel,
}

impl ByLevelBuilder {
    fn new(row: Vec<Hash256>) -> Self {
        ByLevelBuilder { cur : MerkleLevel { row : row } }
    }
}

impl Iterator for ByLevelBuilder {
    type Item = MerkleLevel;

    fn next(&mut self) -> Option<Self::Item> {
        match self.cur.row.len() {
            0 => None,
            1  => {
                let empty = MerkleLevel { row : Vec::new() };
                Some(mem::replace(&mut self.cur, empty))
            }
            _ => {
                let next_row = self.cur.build_next_level();
                Some(mem::replace(&mut self.cur, next_row))
            }
        }
    }
}


#[cfg(test)]
mod test {
    use super::*;
    use hash::Hash256;

    fn jon() -> Hash256 { Hash256::double_sha_str("jon") }
    fn paul() -> Hash256 { Hash256::double_sha_str("Paul") }
    fn george() -> Hash256 { Hash256::double_sha_str("George") }
    fn ringo() -> Hash256 { Hash256::double_sha_str("Ringo") }
    fn pete() -> Hash256 { Hash256::double_sha_str("Pete") }
    fn stuart() -> Hash256 { Hash256::double_sha_str("Stuart") }

    // the tree made herein is full
    #[test]
    fn build_tree_filled() {
        let (jon, paul, george, ringo) = (jon(), paul(), george(), ringo());
        let beatles_hashes = vec![jon, paul, george, ringo];

        let jon_paul = Hash256::double_sha_concat(&jon, &paul);
        let george_ringo = Hash256::double_sha_concat(&george, &ringo);
        let jon_paul_george_ringo = Hash256::double_sha_concat(&jon_paul, &george_ringo);

        assert_eq!(MerkleTree::build_root(beatles_hashes.clone()), jon_paul_george_ringo);

        let mut tree = MerkleTree::build_tree(beatles_hashes.clone());
        let all = tree.to_vec();

        assert_eq!(tree.root_hash(), &jon_paul_george_ringo);

        assert_eq!(all.len(), 7);
        assert_eq!(&all[..4], beatles_hashes.as_slice());
        assert_eq!(all[4], jon_paul);
        assert_eq!(all[5], george_ringo);
        assert_eq!(all[6], jon_paul_george_ringo);

        let path = tree.merkle_path(1);
        assert!(path.verify_root(&jon_paul_george_ringo).is_ok());
        assert!(path.verify_root(&jon).is_err());

        assert!(path.verify_leaf(&paul).is_ok());
        assert!(path.verify_leaf(&ringo).is_err());

        // For you historians let's go back further
        let pete = pete();
        let george_pete = Hash256::double_sha_concat(&george, &pete);
        let jon_paul_george_pete = Hash256::double_sha_concat(&jon_paul, &george_pete);

        tree.update_leaf(3, pete);
        let all = tree.to_vec();

        assert_eq!(tree.root_hash(), &jon_paul_george_pete);

        assert_eq!(all.len(), 7);
        assert_eq!(all[..3], beatles_hashes[..3]);
        assert_eq!(all[3], pete);
        assert_eq!(all[4], jon_paul);
        assert_eq!(all[5], george_pete);
        assert_eq!(all[6], jon_paul_george_pete);
        assert_eq!(path.len(), 5);
    }

    #[test]
    fn single_hash() {
        let george = george();

        assert_eq!(MerkleTree::build_root(vec![george]), george);

        let tree = MerkleTree::build_tree(vec![george]);
        let all = tree.to_vec();

        assert_eq!(tree.root_hash(), &george);

        assert_eq!(all.len(), 1);
        assert_eq!(all, vec![george]);

        let path = tree.merkle_path(0);
        assert!(path.verify_leaf(&george).is_ok());
        assert!(path.verify_root(&george).is_ok());
        assert_eq!(path.len(), 1);
    }

    // the tree made herein does not fill up the leaf level
    #[test]
    fn build_tree_uneven() {
        let (jon, paul, george, ringo, pete) = (jon(), paul(), george(), ringo(), pete());
        let beatles_hashes = vec![jon, paul, george, ringo, pete];

        let jon_paul = Hash256::double_sha_concat(&jon, &paul);
        let george_ringo = Hash256::double_sha_concat(&george, &ringo);
        let pete_pete = Hash256::double_sha_concat(&pete, &pete);
        let jon_paul_george_ringo = Hash256::double_sha_concat(&jon_paul, &george_ringo);
        let pete_pete_pete = Hash256::double_sha_concat(&pete_pete, &pete_pete);
        let jon_paul_george_ringo_pete = Hash256::double_sha_concat(&jon_paul_george_ringo, &pete_pete_pete);

        assert_eq!(MerkleTree::build_root(beatles_hashes.clone()), jon_paul_george_ringo_pete);

        let mut tree = MerkleTree::build_tree(beatles_hashes.clone());
        let all = tree.to_vec();

        assert_eq!(tree.root_hash(), &jon_paul_george_ringo_pete);

        assert_eq!(all.len(), 11);
        assert_eq!(&all[..5], beatles_hashes.as_slice());
        assert_eq!(all[5], jon_paul);
        assert_eq!(all[6], george_ringo);
        assert_eq!(all[7], pete_pete);
        assert_eq!(all[8], jon_paul_george_ringo);
        assert_eq!(all[9], pete_pete_pete);
        assert_eq!(all[10], jon_paul_george_ringo_pete);

        let path = tree.merkle_path(1);
        assert!(path.verify_root(&jon_paul_george_ringo_pete).is_ok());
        assert!(path.verify_root(&jon).is_err());

        assert!(path.verify_leaf(&paul).is_ok());
        assert!(path.verify_leaf(&ringo).is_err());
        assert_eq!(path.len(), 7);

        // ...and the path update
        let stuart = stuart();
        let stuart_stuart = Hash256::double_sha_concat(&stuart, &stuart);
        let stuart_stuart_stuart = Hash256::double_sha_concat(&stuart_stuart, &stuart_stuart);
        let jon_paul_george_ringo_stuart = Hash256::double_sha_concat(&jon_paul_george_ringo, &stuart_stuart_stuart);

        tree.update_leaf(4, stuart);
        let all = tree.to_vec();

        assert_eq!(tree.root_hash(), &jon_paul_george_ringo_stuart);

        assert_eq!(all.len(), 11);
        assert_eq!(all[..4], beatles_hashes[..4]);
        assert_eq!(all[4], stuart);
        assert_eq!(all[5], jon_paul);
        assert_eq!(all[6], george_ringo);
        assert_eq!(all[7], stuart_stuart);
        assert_eq!(all[8], jon_paul_george_ringo);
        assert_eq!(all[9], stuart_stuart_stuart);
        assert_eq!(all[10], jon_paul_george_ringo_stuart);
    }
}