Add ZIP-221 history tree to non-finalized state (#2583)

* Refactor HistoryTree into NonEmptyHistoryTree and HistoryTree * HistoryTree: use Deref instead of AsRef; remove unneeded PartialEq * 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> * Apply suggestions from code review Co-authored-by: teor <teor@riseup.net> * Improvements from code review * Restore blocks returned by PreparedChain since other tests broken; adjust tests with history trees Co-authored-by: teor <teor@riseup.net> Co-authored-by: Deirdre Connolly <deirdre@zfnd.org>
ZcashFoundation · Aug 11, 2021 · 94175c6 · 94175c6
1 parent 298ecec
commit 94175c6
Show file tree

Hide file tree

Showing 9 changed files with 362 additions and 57 deletions.
diff --git a/zebra-chain/src/block/commitment.rs b/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 {

diff --git a/zebra-chain/src/history_tree.rs b/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 {
@@ -403,6 +413,30 @@ impl Clone for NonEmptyHistoryTree {
 pub struct HistoryTree(Option<NonEmptyHistoryTree>);
 
 impl HistoryTree {
+    /// Create a HistoryTree from a block.
+    ///
+    /// If the block is pre-Heartwood, it returns an empty history tree.
+    pub fn from_block(
+        network: Network,
+        block: Arc<Block>,
+        sapling_root: &sapling::tree::Root,
+        orchard_root: &orchard::tree::Root,
+    ) -> Result<Self, HistoryTreeError> {
+        let heartwood_height = NetworkUpgrade::Heartwood
+            .activation_height(network)
+            .expect("Heartwood height is known");
+        match block
+            .coinbase_height()
+            .expect("must have height")
+            .cmp(&heartwood_height)
+        {
+            std::cmp::Ordering::Less => Ok(HistoryTree(None)),
+            _ => Ok(
+                NonEmptyHistoryTree::from_block(network, block, sapling_root, orchard_root)?.into(),
+            ),
+        }
+    }
+
     /// Push a block to a maybe-existing HistoryTree, handling network upgrades.
     ///
     /// The tree is updated in-place. It is created when pushing the Heartwood

diff --git a/zebra-state/src/error.rs b/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.

diff --git a/zebra-state/src/service/arbitrary.rs b/zebra-state/src/service/arbitrary.rs
@@ -10,6 +10,7 @@ use proptest::{
 use zebra_chain::{
     block::{self, Block},
     fmt::SummaryDebug,
+    history_tree::HistoryTree,
     parameters::NetworkUpgrade,
     LedgerState,
 };
@@ -34,17 +35,24 @@ pub struct PreparedChainTree {
     chain: Arc<SummaryDebug<Vec<PreparedBlock>>>,
     count: BinarySearch,
     network: Network,
+    history_tree: HistoryTree,
 }
 
 impl ValueTree for PreparedChainTree {
     type Value = (
         Arc<SummaryDebug<Vec<PreparedBlock>>>,
         <BinarySearch as ValueTree>::Value,
         Network,
+        HistoryTree,
     );
 
     fn current(&self) -> Self::Value {
-        (self.chain.clone(), self.count.current(), self.network)
+        (
+            self.chain.clone(),
+            self.count.current(),
+            self.network,
+            self.history_tree.clone(),
+        )
     }
 
     fn simplify(&mut self) -> bool {
@@ -59,7 +67,36 @@ impl ValueTree for PreparedChainTree {
 #[derive(Debug, Default)]
 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>>>)>>,
+    chain: std::sync::Mutex<Option<(Network, Arc<SummaryDebug<Vec<PreparedBlock>>>, HistoryTree)>>,
+    // the strategy for generating LedgerStates. If None, it calls [`LedgerState::genesis_strategy`].
+    ledger_strategy: Option<BoxedStrategy<LedgerState>>,
+}
+
+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);
+
+        PreparedChain {
+            ledger_strategy: Some(LedgerState::height_strategy(
+                height,
+                NetworkUpgrade::Nu5,
+                None,
+                false,
+            )),
+            ..Default::default()
+        }
+    }
 }
 
 impl Strategy for PreparedChain {
@@ -70,7 +107,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 default_ledger_strategy =
+                LedgerState::genesis_strategy(NetworkUpgrade::Nu5, None, false);
+            let ledger_strategy = self
+                .ledger_strategy
+                .as_ref()
+                .unwrap_or(&default_ledger_strategy);
 
             let (network, blocks) = ledger_strategy
                 .prop_flat_map(|ledger| {
@@ -93,7 +135,16 @@ impl Strategy for PreparedChain {
                 })
                 .new_tree(runner)?
                 .current();
-            *chain = Some((network, Arc::new(SummaryDebug(blocks))));
+            // Generate a history tree from the first block
+            let history_tree = HistoryTree::from_block(
+                network,
+                blocks[0].block.clone(),
+                // Dummy roots since this is only used for tests
+                &Default::default(),
+                &Default::default(),
+            )
+            .expect("history tree should be created");
+            *chain = Some((network, Arc::new(SummaryDebug(blocks)), history_tree));
         }
 
         let chain = chain.clone().expect("should be generated");
@@ -102,6 +153,7 @@ impl Strategy for PreparedChain {
             chain: chain.1,
             count,
             network: chain.0,
+            history_tree: chain.2,
         })
     }
 }
diff --git a/zebra-state/src/service/finalized_state/tests/prop.rs b/zebra-state/src/service/finalized_state/tests/prop.rs
@@ -21,7 +21,7 @@ fn blocks_with_v5_transactions() -> Result<()> {
         .ok()
         .and_then(|v| v.parse().ok())
         .unwrap_or(DEFAULT_PARTIAL_CHAIN_PROPTEST_CASES)),
-        |((chain, count, network) in PreparedChain::default())| {
+        |((chain, count, network, _history_tree) in PreparedChain::default())| {
             let mut state = FinalizedState::new(&Config::ephemeral(), network);
             let mut height = Height(0);
             // use `count` to minimize test failures, so they are easier to diagnose

diff --git a/zebra-state/src/service/non_finalized_state.rs b/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()
                     })