ZIP-221: Validate chain history commitments in the non-finalized state (

#2301)

* sketch of implementation

* refined implementation; still incomplete

* update librustzcash, change zcash_history to work with it

* simplified code per review; renamed MMR to HistoryTree

* expand HistoryTree implementation

* handle and propagate errors

* simplify check.rs tracing

* add suggested TODO

* add HistoryTree::prune

* fix bug in pruning

* fix compilation of tests; still need to make them pass

* Apply suggestions from code review

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

* Apply suggestions from code review

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

* improvements from code review

* improve check.rs comments and variable names

* fix HistoryTree which should use BTreeMap and not HashMap; fix non_finalized_state prop tests

* fix finalized_state proptest

* fix non_finalized_state tests by setting the correct commitments

* renamed mmr.rs to history_tree.rs

* Add HistoryTree struct

* expand non_finalized_state protest

* fix typo

* Add HistoryTree struct

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

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

* fix formatting

* Apply suggestions from code review

Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

* history_tree.rs: fixes from code review

* fixes to work with updated HistoryTree

* Improvements from code review

* Add Debug implementations to allow comparing Chains with proptest_assert_eq

Co-authored-by: teor <teor@riseup.net>
Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>

zebra-chain/src/block/commitment.rs

@@ -149,6 +149,12 @@ impl From<[u8; 32]> for ChainHistoryMmrRootHash {
     }
 }
 
+impl From<ChainHistoryMmrRootHash> for [u8; 32] {
+    fn from(hash: ChainHistoryMmrRootHash) -> Self {
+        hash.0
+    }
+}
+
 /// A block commitment to chain history and transaction auth.
 /// - the chain history tree for all ancestors in the current network upgrade,
 ///   and
@@ -170,7 +176,7 @@ pub struct ChainHistoryBlockTxAuthCommitmentHash([u8; 32]);
 /// implement, and ensures that we don't reject blocks or transactions
 /// for a non-enumerated reason.
 #[allow(dead_code, missing_docs)]
-#[derive(Error, Debug, PartialEq)]
+#[derive(Error, Debug, PartialEq, Eq)]
 pub enum CommitmentError {
     #[error("invalid final sapling root: expected {expected:?}, actual: {actual:?}")]
     InvalidFinalSaplingRoot {

zebra-chain/src/history_tree.rs

@@ -33,6 +33,16 @@ pub enum HistoryTreeError {
     IOError(#[from] io::Error),
 }
 
+impl PartialEq for HistoryTreeError {
+    fn eq(&self, other: &Self) -> bool {
+        // Workaround since subtypes do not implement Eq.
+        // This is only used for tests anyway.
+        format!("{:?}", self) == format!("{:?}", other)
+    }
+}
+
+impl Eq for HistoryTreeError {}
+
 /// The inner [Tree] in one of its supported versions.
 #[derive(Debug)]
 enum InnerHistoryTree {

zebra-state/src/error.rs

@@ -4,8 +4,8 @@ use chrono::{DateTime, Utc};
 use thiserror::Error;
 
 use zebra_chain::{
-    amount, block, orchard, sapling, sprout, transparent, value_balance::ValueBalanceError,
-    work::difficulty::CompactDifficulty,
+    amount, block, history_tree::HistoryTreeError, orchard, sapling, sprout, transparent,
+    value_balance::ValueBalanceError, work::difficulty::CompactDifficulty,
 };
 
 use crate::constants::MIN_TRANSPARENT_COINBASE_MATURITY;
@@ -36,12 +36,12 @@ impl From<BoxError> for CloneError {
 pub type BoxError = Box<dyn std::error::Error + Send + Sync + 'static>;
 
 /// An error describing the reason a block could not be committed to the state.
-#[derive(Debug, Error, Clone, PartialEq, Eq)]
+#[derive(Debug, Error, PartialEq, Eq)]
 #[error("block is not contextually valid")]
 pub struct CommitBlockError(#[from] ValidateContextError);
 
 /// An error describing why a block failed contextual validation.
-#[derive(Debug, Error, Clone, PartialEq, Eq)]
+#[derive(Debug, Error, PartialEq, Eq)]
 #[non_exhaustive]
 #[allow(missing_docs)]
 pub enum ValidateContextError {
@@ -185,6 +185,9 @@ pub enum ValidateContextError {
 
     #[error("error in Orchard note commitment tree")]
     OrchardNoteCommitmentTreeError(#[from] zebra_chain::orchard::tree::NoteCommitmentTreeError),
+
+    #[error("error building the history tree")]
+    HistoryTreeError(#[from] HistoryTreeError),
 }
 
 /// Trait for creating the corresponding duplicate nullifier error from a nullifier.

zebra-state/src/service/arbitrary.rs

@@ -8,7 +8,7 @@ use proptest::{
 };
 
 use zebra_chain::{
-    block::{self, Block},
+    block::{self, Block, Height},
     fmt::SummaryDebug,
     parameters::NetworkUpgrade,
     LedgerState,
@@ -60,6 +60,30 @@ impl ValueTree for PreparedChainTree {
 pub struct PreparedChain {
     // the proptests are threaded (not async), so we want to use a threaded mutex here
     chain: std::sync::Mutex<Option<(Network, Arc<SummaryDebug<Vec<PreparedBlock>>>)>>,
+    // the height from which to start the chain. If None, starts at the genesis block
+    start_height: Option<Height>,
+}
+
+impl PreparedChain {
+    /// Create a PreparedChain strategy with Heartwood-onward blocks.
+    #[cfg(test)]
+    pub(crate) fn new_heartwood() -> Self {
+        // The history tree only works with Heartwood onward.
+        // Since the network will be chosen later, we pick the larger
+        // between the mainnet and testnet Heartwood activation heights.
+        let main_height = NetworkUpgrade::Heartwood
+            .activation_height(Network::Mainnet)
+            .expect("must have height");
+        let test_height = NetworkUpgrade::Heartwood
+            .activation_height(Network::Testnet)
+            .expect("must have height");
+        let height = (std::cmp::max(main_height, test_height) + 1).expect("must be valid");
+
+        PreparedChain {
+            start_height: Some(height),
+            ..Default::default()
+        }
+    }
 }
 
 impl Strategy for PreparedChain {
@@ -70,7 +94,12 @@ impl Strategy for PreparedChain {
         let mut chain = self.chain.lock().unwrap();
         if chain.is_none() {
             // TODO: use the latest network upgrade (#1974)
-            let ledger_strategy = LedgerState::genesis_strategy(NetworkUpgrade::Nu5, None, false);
+            let ledger_strategy = match self.start_height {
+                Some(start_height) => {
+                    LedgerState::height_strategy(start_height, NetworkUpgrade::Nu5, None, false)
+                }
+                None => LedgerState::genesis_strategy(NetworkUpgrade::Nu5, None, false),
+            };
 
             let (network, blocks) = ledger_strategy
                 .prop_flat_map(|ledger| {
@@ -97,7 +126,9 @@ impl Strategy for PreparedChain {
         }
 
         let chain = chain.clone().expect("should be generated");
-        let count = (1..chain.1.len()).new_tree(runner)?;
+        // `count` must be 1 less since the first block is used to build the
+        // history tree.
+        let count = (1..chain.1.len() - 1).new_tree(runner)?;
         Ok(PreparedChainTree {
             chain: chain.1,
             count,

zebra-state/src/service/non_finalized_state.rs

@@ -14,6 +14,7 @@ use std::{collections::BTreeSet, mem, ops::Deref, sync::Arc};
 
 use zebra_chain::{
     block::{self, Block},
+    history_tree::HistoryTree,
     orchard,
     parameters::Network,
     sapling,
@@ -129,6 +130,7 @@ impl NonFinalizedState {
             parent_hash,
             finalized_state.sapling_note_commitment_tree(),
             finalized_state.orchard_note_commitment_tree(),
+            finalized_state.history_tree(),
         )?;
 
         // We might have taken a chain, so all validation must happen within
@@ -161,8 +163,10 @@ impl NonFinalizedState {
         finalized_state: &FinalizedState,
     ) -> Result<(), ValidateContextError> {
         let chain = Chain::new(
+            self.network,
             finalized_state.sapling_note_commitment_tree(),
             finalized_state.orchard_note_commitment_tree(),
+            finalized_state.history_tree(),
         );
         let (height, hash) = (prepared.height, prepared.hash);
 
@@ -355,13 +359,14 @@ impl NonFinalizedState {
     /// The chain can be an existing chain in the non-finalized state or a freshly
     /// created fork, if needed.
     ///
-    /// The note commitment trees must be the trees of the finalized tip.
+    /// The trees must be the trees of the finalized tip.
     /// They are used to recreate the trees if a fork is needed.
     fn parent_chain(
         &mut self,
         parent_hash: block::Hash,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Result<Box<Chain>, ValidateContextError> {
         match self.take_chain_if(|chain| chain.non_finalized_tip_hash() == parent_hash) {
             // An existing chain in the non-finalized state
@@ -376,6 +381,7 @@ impl NonFinalizedState {
                                 parent_hash,
                                 sapling_note_commitment_tree.clone(),
                                 orchard_note_commitment_tree.clone(),
+                                history_tree.clone(),
                             )
                             .transpose()
                     })

zebra-state/src/service/non_finalized_state/chain.rs

@@ -8,14 +8,16 @@ use multiset::HashMultiSet;
 use tracing::instrument;
 
 use zebra_chain::{
-    block, orchard, primitives::Groth16Proof, sapling, sprout, transaction,
-    transaction::Transaction::*, transparent, work::difficulty::PartialCumulativeWork,
+    block, history_tree::HistoryTree, orchard, parameters::Network, primitives::Groth16Proof,
+    sapling, sprout, transaction, transaction::Transaction::*, transparent,
+    work::difficulty::PartialCumulativeWork,
 };
 
 use crate::{service::check, ContextuallyValidBlock, PreparedBlock, ValidateContextError};
 
 #[derive(Debug, Clone)]
 pub struct Chain {
+    network: Network,
     /// The contextually valid blocks which form this non-finalized partial chain, in height order.
     pub(crate) blocks: BTreeMap<block::Height, ContextuallyValidBlock>,
 
@@ -37,6 +39,8 @@ pub struct Chain {
     pub(super) sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
     /// The Orchard note commitment tree of the tip of this Chain.
     pub(super) orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+    /// The ZIP-221 history tree of the tip of this Chain.
+    pub(crate) history_tree: HistoryTree,
 
     /// The Sapling anchors created by `blocks`.
     pub(super) sapling_anchors: HashMultiSet<sapling::tree::Root>,
@@ -59,12 +63,15 @@ pub struct Chain {
 }
 
 impl Chain {
-    // Create a new Chain with the given note commitment trees.
+    /// Create a new Chain with the given note commitment trees.
     pub(crate) fn new(
+        network: Network,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Self {
         Self {
+            network,
             blocks: Default::default(),
             height_by_hash: Default::default(),
             tx_by_hash: Default::default(),
@@ -80,6 +87,7 @@ impl Chain {
             sapling_nullifiers: Default::default(),
             orchard_nullifiers: Default::default(),
             partial_cumulative_work: Default::default(),
+            history_tree,
         }
     }
 
@@ -95,6 +103,8 @@ impl Chain {
     /// even if the blocks in the two chains are equal.
     #[cfg(test)]
     pub(crate) fn eq_internal_state(&self, other: &Chain) -> bool {
+        use zebra_chain::history_tree::NonEmptyHistoryTree;
+
         // this method must be updated every time a field is added to Chain
 
         // blocks, heights, hashes
@@ -110,6 +120,9 @@ impl Chain {
             self.sapling_note_commitment_tree.root() == other.sapling_note_commitment_tree.root() &&
             self.orchard_note_commitment_tree.root() == other.orchard_note_commitment_tree.root() &&
 
+            // history tree
+            self.history_tree.as_ref().map(NonEmptyHistoryTree::hash) == other.history_tree.as_ref().map(NonEmptyHistoryTree::hash) &&
+
             // anchors
             self.sapling_anchors == other.sapling_anchors &&
             self.sapling_anchors_by_height == other.sapling_anchors_by_height &&
@@ -169,19 +182,24 @@ impl Chain {
     /// Fork a chain at the block with the given hash, if it is part of this
     /// chain.
     ///
-    /// The note commitment trees must be the trees of the finalized tip.
+    /// The trees must match the trees of the finalized tip and are used
+    /// to rebuild them after the fork.
     pub fn fork(
         &self,
         fork_tip: block::Hash,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Result<Option<Self>, ValidateContextError> {
         if !self.height_by_hash.contains_key(&fork_tip) {
             return Ok(None);
         }
 
-        let mut forked =
-            self.with_trees(sapling_note_commitment_tree, orchard_note_commitment_tree);
+        let mut forked = self.with_trees(
+            sapling_note_commitment_tree,
+            orchard_note_commitment_tree,
+            history_tree,
+        );
 
         while forked.non_finalized_tip_hash() != fork_tip {
             forked.pop_tip();
@@ -206,6 +224,24 @@ impl Chain {
                         .expect("must work since it was already appended before the fork");
                 }
             }
+
+            // Note that anchors don't need to be recreated since they are already
+            // handled in revert_chain_state_with.
+
+            let sapling_root = forked
+                .sapling_anchors_by_height
+                .get(&block.height)
+                .expect("Sapling anchors must exist for pre-fork blocks");
+            let orchard_root = forked
+                .orchard_anchors_by_height
+                .get(&block.height)
+                .expect("Orchard anchors must exist for pre-fork blocks");
+            forked.history_tree.push(
+                self.network,
+                block.block.clone(),
+                *sapling_root,
+                *orchard_root,
+            )?;
         }
 
         Ok(Some(forked))
@@ -267,8 +303,10 @@ impl Chain {
         &self,
         sapling_note_commitment_tree: sapling::tree::NoteCommitmentTree,
         orchard_note_commitment_tree: orchard::tree::NoteCommitmentTree,
+        history_tree: HistoryTree,
     ) -> Self {
         Chain {
+            network: self.network,
             blocks: self.blocks.clone(),
             height_by_hash: self.height_by_hash.clone(),
             tx_by_hash: self.tx_by_hash.clone(),
@@ -284,6 +322,7 @@ impl Chain {
             sapling_nullifiers: self.sapling_nullifiers.clone(),
             orchard_nullifiers: self.orchard_nullifiers.clone(),
             partial_cumulative_work: self.partial_cumulative_work,
+            history_tree,
         }
     }
 }
@@ -395,12 +434,19 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
         // Having updated all the note commitment trees and nullifier sets in
         // this block, the roots of the note commitment trees as of the last
         // transaction are the treestates of this block.
-        let root = self.sapling_note_commitment_tree.root();
-        self.sapling_anchors.insert(root);
-        self.sapling_anchors_by_height.insert(height, root);
-        let root = self.orchard_note_commitment_tree.root();
-        self.orchard_anchors.insert(root);
-        self.orchard_anchors_by_height.insert(height, root);
+        let sapling_root = self.sapling_note_commitment_tree.root();
+        self.sapling_anchors.insert(sapling_root);
+        self.sapling_anchors_by_height.insert(height, sapling_root);
+        let orchard_root = self.orchard_note_commitment_tree.root();
+        self.orchard_anchors.insert(orchard_root);
+        self.orchard_anchors_by_height.insert(height, orchard_root);
+
+        self.history_tree.push(
+            self.network,
+            contextually_valid.block.clone(),
+            sapling_root,
+            orchard_root,
+        )?;
 
         Ok(())
     }
@@ -429,6 +475,11 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             .expect("work has already been validated");
         self.partial_cumulative_work -= block_work;
 
+        // Note: the history tree is not modified in this method.
+        // This method is called on two scenarios:
+        // - When popping the root: the history tree does not change.
+        // - When popping the tip: the history tree is rebuilt in fork().
+
         // for each transaction in block
         for (transaction, transaction_hash) in
             block.transactions.iter().zip(transaction_hashes.iter())