ZIP-221: integrate MMR tree from librustcash (without Orchard) (#2227)

* add zcash_history.rs with librustzcash Tree wrapper

* Apply suggestions from code review

Co-authored-by: teor <teor@riseup.net>

* Apply changes from code review

* Update zebra-chain/src/primitives/zcash_history.rs

Co-authored-by: teor <teor@riseup.net>

* Apply changes from code review

* Add Entry struct; return Result where needed; add test

* Apply suggestions from code review

Co-authored-by: teor <teor@riseup.net>

* zcash_history: improve naming style with `inner`

* zcash_history: check if block has the correct network upgrade when adding to tree

* zcash_history: test improvements

* zcash_history: split Tree::new into new_from_block and new_from_cache

* zcash_history: move tests to their own file

* remove unneeded empty line in Cargo.toml

Co-authored-by: teor <teor@riseup.net>

Cargo.lock

@@ -283,6 +283,16 @@ dependencies = [
  "subtle",
 ]
 
+[[package]]
+name = "bigint"
+version = "4.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c0e8c8a600052b52482eff2cf4d810e462fdff1f656ac1ecb6232132a1ed7def"
+dependencies = [
+ "byteorder",
+ "crunchy 0.1.6",
+]
+
 [[package]]
 name = "bincode"
 version = "1.3.3"
@@ -885,6 +895,12 @@ dependencies = [
  "lazy_static",
 ]
 
+[[package]]
+name = "crunchy"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2f4a431c5c9f662e1200b7c7f02c34e91361150e382089a8f2dec3ba680cbda"
+
 [[package]]
 name = "crunchy"
 version = "0.2.2"
@@ -3944,7 +3960,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e11fe9a9348741cf134085ad57c249508345fe16411b3d7fb4ff2da2f1d6382e"
 dependencies = [
  "byteorder",
- "crunchy",
+ "crunchy 0.2.2",
  "hex",
  "static_assertions",
 ]
@@ -4395,6 +4411,16 @@ dependencies = [
  "zeroize",
 ]
 
+[[package]]
+name = "zcash_history"
+version = "0.2.0"
+source = "git+https://github.com/zcash/librustzcash.git#d50bb12a97da768dc8f3ee39b81f84262103e6eb"
+dependencies = [
+ "bigint",
+ "blake2b_simd",
+ "byteorder",
+]
+
 [[package]]
 name = "zcash_script"
 version = "0.1.6-alpha.0"
@@ -4413,6 +4439,7 @@ version = "1.0.0-alpha.10"
 dependencies = [
  "aes",
  "bech32",
+ "bigint",
  "bincode",
  "bitflags",
  "bitvec 0.17.4",
@@ -4452,6 +4479,7 @@ dependencies = [
  "thiserror",
  "tracing",
  "x25519-dalek",
+ "zcash_history",
  "zebra-test",
 ]
 

zebra-chain/Cargo.toml

@@ -44,6 +44,8 @@ sha2 = { version = "0.9.5", features=["compress"] }
 subtle = "2.4"
 thiserror = "1"
 x25519-dalek = { version = "1.1", features = ["serde"] }
+zcash_history = { git = "https://github.com/zcash/librustzcash.git" }
+bigint = "4"
 
 proptest = { version = "0.10", optional = true }
 proptest-derive = { version = "0.3.0", optional = true }

zebra-chain/src/block.rs

@@ -15,7 +15,7 @@ pub mod tests;
 
 use std::fmt;
 
-pub use commitment::{Commitment, CommitmentError};
+pub use commitment::{ChainHistoryMmrRootHash, Commitment, CommitmentError};
 pub use hash::Hash;
 pub use header::{BlockTimeError, CountedHeader, Header};
 pub use height::Height;

zebra-chain/src/block/commitment.rs

@@ -143,6 +143,12 @@ impl Commitment {
 #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
 pub struct ChainHistoryMmrRootHash([u8; 32]);
 
+impl From<[u8; 32]> for ChainHistoryMmrRootHash {
+    fn from(hash: [u8; 32]) -> Self {
+        ChainHistoryMmrRootHash(hash)
+    }
+}
+
 /// A block commitment to chain history and transaction auth.
 /// - the chain history tree for all ancestors in the current network upgrade,
 ///   and

zebra-chain/src/primitives.rs

@@ -13,3 +13,5 @@ pub use redjubjub;
 pub use x25519_dalek as x25519;
 
 pub use proofs::{Bctv14Proof, Groth16Proof, Halo2Proof, ZkSnarkProof};
+
+pub mod zcash_history;

zebra-chain/src/primitives/zcash_history.rs

@@ -0,0 +1,286 @@
+//! Contains code that interfaces with the zcash_history crate from
+//! librustzcash.
+
+// TODO: remove after this module gets to be used
+#![allow(dead_code)]
+
+mod tests;
+
+use std::{collections::HashMap, convert::TryInto, io, sync::Arc};
+
+use crate::{
+    block::{Block, ChainHistoryMmrRootHash},
+    parameters::{ConsensusBranchId, Network, NetworkUpgrade},
+    sapling,
+};
+
+/// A MMR Tree using zcash_history::Tree.
+///
+/// Currently it should not be used as a long-term data structure because it
+/// may grow without limits.
+pub struct Tree {
+    network: Network,
+    network_upgrade: NetworkUpgrade,
+    inner: zcash_history::Tree,
+}
+
+/// An encoded tree node data.
+pub struct NodeData {
+    inner: [u8; zcash_history::MAX_NODE_DATA_SIZE],
+}
+
+impl From<&zcash_history::NodeData> for NodeData {
+    /// Convert from librustzcash.
+    fn from(inner_node: &zcash_history::NodeData) -> Self {
+        let mut node = NodeData {
+            inner: [0; zcash_history::MAX_NODE_DATA_SIZE],
+        };
+        inner_node
+            .write(&mut &mut node.inner[..])
+            .expect("buffer has the proper size");
+        node
+    }
+}
+
+/// An encoded entry in the tree.
+/// Contains the node data and information about its position in the tree.
+pub struct Entry {
+    inner: [u8; zcash_history::MAX_ENTRY_SIZE],
+}
+
+impl From<zcash_history::Entry> for Entry {
+    /// Convert from librustzcash.
+    fn from(inner_entry: zcash_history::Entry) -> Self {
+        let mut entry = Entry {
+            inner: [0; zcash_history::MAX_ENTRY_SIZE],
+        };
+        inner_entry
+            .write(&mut &mut entry.inner[..])
+            .expect("buffer has the proper size");
+        entry
+    }
+}
+
+impl Entry {
+    /// Create a leaf Entry for the given block, its network, and the root of its
+    /// Sapling note commitment tree.
+    fn new_leaf(block: Arc<Block>, network: Network, sapling_root: &sapling::tree::Root) -> Self {
+        let node_data = block_to_history_node(block, network, sapling_root);
+        let inner_entry: zcash_history::Entry = node_data.into();
+        inner_entry.into()
+    }
+
+    /// Create a node (non-leaf) Entry from the encoded node data and the indices of
+    /// its children (in the array representation of the MMR tree).
+    fn new_node(
+        branch_id: ConsensusBranchId,
+        data: NodeData,
+        left_idx: u32,
+        right_idx: u32,
+    ) -> Result<Self, io::Error> {
+        let node_data = zcash_history::NodeData::from_bytes(branch_id.into(), data.inner)?;
+        let inner_entry = zcash_history::Entry::new(
+            node_data,
+            zcash_history::EntryLink::Stored(left_idx),
+            zcash_history::EntryLink::Stored(right_idx),
+        );
+        Ok(inner_entry.into())
+    }
+}
+
+impl Tree {
+    /// Create a MMR tree with the given length from the given cache of nodes.
+    ///
+    /// The `peaks` are the peaks of the MMR tree to build and their position in the
+    /// array representation of the tree.
+    /// The `extra` are extra nodes that enable removing nodes from the tree, and their position.
+    ///
+    /// Note that the length is usually larger than the length of `peaks` and `extra`, since
+    /// you don't need to pass every node, just the peaks of the tree (plus extra).
+    ///
+    /// # Panics
+    ///
+    /// Will panic if `peaks` is empty.
+    fn new_from_cache(
+        network: Network,
+        network_upgrade: NetworkUpgrade,
+        length: u32,
+        peaks: &HashMap<u32, &Entry>,
+        extra: &HashMap<u32, &Entry>,
+    ) -> Result<Self, io::Error> {
+        let branch_id = network_upgrade
+            .branch_id()
+            .expect("unexpected pre-Overwinter MMR history tree");
+        let mut peaks_vec = Vec::new();
+        for (idx, entry) in peaks {
+            let inner_entry = zcash_history::Entry::from_bytes(branch_id.into(), entry.inner)?;
+            peaks_vec.push((*idx, inner_entry));
+        }
+        let mut extra_vec = Vec::new();
+        for (idx, entry) in extra {
+            let inner_entry = zcash_history::Entry::from_bytes(branch_id.into(), entry.inner)?;
+            extra_vec.push((*idx, inner_entry));
+        }
+        let inner = zcash_history::Tree::new(length, peaks_vec, extra_vec);
+        Ok(Tree {
+            network,
+            network_upgrade,
+            inner,
+        })
+    }
+
+    /// Create a single-node MMR tree for the given block.
+    ///
+    /// `sapling_root` is the root of the Sapling note commitment tree of the block.
+    fn new_from_block(
+        network: Network,
+        block: Arc<Block>,
+        sapling_root: &sapling::tree::Root,
+    ) -> Result<Self, io::Error> {
+        let height = block
+            .coinbase_height()
+            .expect("block must have coinbase height during contextual verification");
+        let network_upgrade = NetworkUpgrade::current(network, height);
+        let entry0 = Entry::new_leaf(block, network, sapling_root);
+        let mut peaks = HashMap::new();
+        peaks.insert(0u32, &entry0);
+        Tree::new_from_cache(network, network_upgrade, 1, &peaks, &HashMap::new())
+    }
+
+    /// Append a new block to the tree, as a new leaf.
+    ///
+    /// `sapling_root` is the root of the Sapling note commitment tree of the block.
+    ///
+    /// Returns a vector of nodes added to the tree (leaf + internal nodes).
+    ///
+    /// # Panics
+    ///
+    /// Panics if the network upgrade of the given block is different from
+    /// the network upgrade of the other blocks in the tree.
+    fn append_leaf(
+        &mut self,
+        block: Arc<Block>,
+        sapling_root: &sapling::tree::Root,
+    ) -> Result<Vec<NodeData>, zcash_history::Error> {
+        let height = block
+            .coinbase_height()
+            .expect("block must have coinbase height during contextual verification");
+        let network_upgrade = NetworkUpgrade::current(self.network, height);
+        if self.network_upgrade != network_upgrade {
+            panic!(
+                "added block from network upgrade {:?} but MMR tree is restricted to {:?}",
+                network_upgrade, self.network_upgrade
+            );
+        }
+
+        let node_data = block_to_history_node(block, self.network, sapling_root);
+        let appended = self.inner.append_leaf(node_data)?;
+
+        let mut new_nodes = Vec::new();
+        for entry in appended {
+            let mut node = NodeData {
+                inner: [0; zcash_history::MAX_NODE_DATA_SIZE],
+            };
+            self.inner
+                .resolve_link(entry)
+                .expect("entry was just generated so it must be valid")
+                .data()
+                .write(&mut &mut node.inner[..])
+                .expect("buffer was created with enough capacity");
+            new_nodes.push(node);
+        }
+        Ok(new_nodes)
+    }
+
+    /// Append multiple blocks to the tree.
+    fn append_leaf_iter(
+        &mut self,
+        vals: impl Iterator<Item = (Arc<Block>, sapling::tree::Root)>,
+    ) -> Result<Vec<NodeData>, zcash_history::Error> {
+        let mut new_nodes = Vec::new();
+        for (block, root) in vals {
+            new_nodes.append(&mut self.append_leaf(block, &root)?);
+        }
+        Ok(new_nodes)
+    }
+
+    /// Remove the last leaf (block) from the tree.
+    ///
+    /// Returns the number of nodes removed from the tree after the operation.
+    fn truncate_leaf(&mut self) -> Result<u32, zcash_history::Error> {
+        self.inner.truncate_leaf()
+    }
+
+    /// Return the root hash of the tree, i.e. `hashChainHistoryRoot`.
+    fn hash(&self) -> ChainHistoryMmrRootHash {
+        // Both append_leaf() and truncate_leaf() leave a root node, so it should
+        // always exist.
+        self.inner
+            .root_node()
+            .expect("must have root node")
+            .data()
+            .hash()
+            .into()
+    }
+}
+
+/// Convert a Block into a zcash_history::NodeData used in the MMR tree.
+///
+/// `sapling_root` is the root of the Sapling note commitment tree of the block.
+fn block_to_history_node(
+    block: Arc<Block>,
+    network: Network,
+    sapling_root: &sapling::tree::Root,
+) -> zcash_history::NodeData {
+    let height = block
+        .coinbase_height()
+        .expect("block must have coinbase height during contextual verification");
+    let branch_id = ConsensusBranchId::current(network, height)
+        .expect("must have branch ID for chain history network upgrades");
+    let block_hash = block.hash().0;
+    let time: u32 = block
+        .header
+        .time
+        .timestamp()
+        .try_into()
+        .expect("deserialized and generated timestamps are u32 values");
+    let target = block.header.difficulty_threshold.0;
+    let sapling_root: [u8; 32] = sapling_root.into();
+    let work = block
+        .header
+        .difficulty_threshold
+        .to_work()
+        .expect("work must be valid during contextual verification");
+    // There is no direct `std::primitive::u128` to `bigint::U256` conversion
+    let work = bigint::U256::from_big_endian(&work.as_u128().to_be_bytes());
+
+    let sapling_tx_count = count_sapling_transactions(block);
+
+    zcash_history::NodeData {
+        consensus_branch_id: branch_id.into(),
+        subtree_commitment: block_hash,
+        start_time: time,
+        end_time: time,
+        start_target: target,
+        end_target: target,
+        start_sapling_root: sapling_root,
+        end_sapling_root: sapling_root,
+        subtree_total_work: work,
+        start_height: height.0 as u64,
+        end_height: height.0 as u64,
+        sapling_tx: sapling_tx_count,
+    }
+}
+
+/// Count how many Sapling transactions exist in a block,
+/// i.e. transactions "where either of vSpendsSapling or vOutputsSapling is non-empty"
+/// (https://zips.z.cash/zip-0221#tree-node-specification).
+fn count_sapling_transactions(block: Arc<Block>) -> u64 {
+    block
+        .transactions
+        .iter()
+        .filter(|tx| tx.has_sapling_shielded_data())
+        .count()
+        .try_into()
+        .expect("number of transactions must fit u64")
+}

zebra-chain/src/primitives/zcash_history/tests.rs

@@ -0,0 +1,4 @@
+//! Tests for Zebra history trees
+
+#[cfg(test)]
+mod vectors;

zebra-chain/src/primitives/zcash_history/tests/vectors.rs

@@ -0,0 +1,85 @@
+use crate::{
+    block::Commitment::{self, ChainHistoryActivationReserved},
+    serialization::ZcashDeserializeInto,
+};
+
+use crate::primitives::zcash_history::*;
+use color_eyre::eyre;
+use eyre::Result;
+use zebra_test::vectors::{
+    MAINNET_BLOCKS, MAINNET_FINAL_SAPLING_ROOTS, TESTNET_BLOCKS, TESTNET_FINAL_SAPLING_ROOTS,
+};
+
+/// Test the MMR tree using the activation block of a network upgrade
+/// and its next block.
+#[test]
+fn tree() -> Result<()> {
+    tree_for_network_upgrade(Network::Mainnet, NetworkUpgrade::Heartwood)?;
+    tree_for_network_upgrade(Network::Testnet, NetworkUpgrade::Heartwood)?;
+    tree_for_network_upgrade(Network::Mainnet, NetworkUpgrade::Canopy)?;
+    tree_for_network_upgrade(Network::Testnet, NetworkUpgrade::Canopy)?;
+    Ok(())
+}
+
+fn tree_for_network_upgrade(network: Network, network_upgrade: NetworkUpgrade) -> Result<()> {
+    let (blocks, sapling_roots) = match network {
+        Network::Mainnet => (&*MAINNET_BLOCKS, &*MAINNET_FINAL_SAPLING_ROOTS),
+        Network::Testnet => (&*TESTNET_BLOCKS, &*TESTNET_FINAL_SAPLING_ROOTS),
+    };
+    let height = network_upgrade.activation_height(network).unwrap().0;
+
+    // Load Block 0 (activation block of the given network upgrade)
+    let block0 = Arc::new(
+        blocks
+            .get(&height)
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+
+    // Check its commitment
+    let commitment0 = block0.commitment(network)?;
+    if network_upgrade == NetworkUpgrade::Heartwood {
+        // Heartwood is the only upgrade that has a reserved value.
+        // (For other upgrades we could compare with the expected commitment,
+        // but we haven't calculated them.)
+        assert_eq!(commitment0, ChainHistoryActivationReserved);
+    }
+
+    // Build initial MMR tree with only Block 0
+    let sapling_root0 =
+        sapling::tree::Root(**sapling_roots.get(&height).expect("test vector exists"));
+    let mut tree = Tree::new_from_block(network, block0, &sapling_root0)?;
+
+    // Compute root hash of the MMR tree, which will be included in the next block
+    let hash0 = tree.hash();
+
+    // Load Block 1 (activation + 1)
+    let block1 = Arc::new(
+        blocks
+            .get(&(height + 1))
+            .expect("test vector exists")
+            .zcash_deserialize_into::<Block>()
+            .expect("block is structurally valid"),
+    );
+
+    // Check its commitment
+    let commitment1 = block1.commitment(network)?;
+    assert_eq!(commitment1, Commitment::ChainHistoryRoot(hash0));
+
+    // Append Block to MMR tree
+    let sapling_root1 = sapling::tree::Root(
+        **sapling_roots
+            .get(&(height + 1))
+            .expect("test vector exists"),
+    );
+    let append = tree.append_leaf(block1, &sapling_root1).unwrap();
+
+    // Tree how has 3 nodes: two leafs for each block, and one parent node
+    // which is the new root
+    assert_eq!(tree.inner.len(), 3);
+    // Two nodes were appended: the new leaf and the parent node
+    assert_eq!(append.len(), 2);
+
+    Ok(())
+}

zebra-chain/src/transaction.rs

@@ -393,6 +393,21 @@ impl Transaction {
         }
     }
 
+    /// Return if the transaction has any Sapling shielded data.
+    pub fn has_sapling_shielded_data(&self) -> bool {
+        match self {
+            Transaction::V1 { .. } | Transaction::V2 { .. } | Transaction::V3 { .. } => false,
+            Transaction::V4 {
+                sapling_shielded_data,
+                ..
+            } => sapling_shielded_data.is_some(),
+            Transaction::V5 {
+                sapling_shielded_data,
+                ..
+            } => sapling_shielded_data.is_some(),
+        }
+    }
+
     // orchard
 
     /// Access the [`orchard::ShieldedData`] in this transaction, if there are any,