Restore the previous non-finalized chain if a block is invalid (#2478)

Previously, we would implicitly drop the full non-finalized chain,
and reset the state to the finalized tip.

zebra-state/src/service.rs

@@ -174,6 +174,13 @@ impl StateService {
     /// Run contextual validation on the prepared block and add it to the
     /// non-finalized state if it is contextually valid.
     fn validate_and_commit(&mut self, prepared: PreparedBlock) -> Result<(), CommitBlockError> {
+        let mandatory_checkpoint = self.network.mandatory_checkpoint_height();
+        if prepared.height <= mandatory_checkpoint {
+            panic!(
+                "invalid non-finalized block height: the canopy checkpoint is mandatory, pre-canopy blocks, and the canopy activation block, must be committed to the state as finalized blocks"
+            );
+        }
+
         self.check_contextual_validity(&prepared)?;
         let parent_hash = prepared.block.header.previous_block_hash;
 

zebra-state/src/service/non_finalized_state.rs

@@ -78,19 +78,22 @@ impl NonFinalizedState {
         let parent_hash = prepared.block.header.previous_block_hash;
         let (height, hash) = (prepared.height, prepared.hash);
 
-        let mandatory_checkpoint = self.network.mandatory_checkpoint_height();
+        let parent_chain = self.parent_chain(parent_hash)?;
-        if height <= mandatory_checkpoint {
+
-            panic!(
+        match parent_chain.clone().push(prepared) {
-                "invalid non-finalized block height: the canopy checkpoint is mandatory, pre-canopy blocks, and the canopy activation block, must be committed to the state as finalized blocks"
+            Ok(child_chain) => {
-            );
+                // if the block is valid, keep the child chain, and drop the parent chain
+                self.chain_set.insert(Box::new(child_chain));
+                self.update_metrics_for_committed_block(height, hash);
+                Ok(())
+            }
+            Err(err) => {
+                // if the block is invalid, restore the unmodified parent chain
+                // (the child chain might have been modified before the error)
+                self.chain_set.insert(parent_chain);
+                Err(err)
+            }
         }
-
-        let mut parent_chain = self.parent_chain(parent_hash)?;
-
-        parent_chain.push(prepared)?;
-        self.chain_set.insert(parent_chain);
-        self.update_metrics_for_committed_block(height, hash);
-        Ok(())
     }
 
     /// Commit block to the non-finalized state as a new chain where its parent
@@ -99,9 +102,11 @@ impl NonFinalizedState {
         &mut self,
         prepared: PreparedBlock,
     ) -> Result<(), ValidateContextError> {
-        let mut chain = Chain::default();
+        let chain = Chain::default();
         let (height, hash) = (prepared.height, prepared.hash);
-        chain.push(prepared)?;
+
+        // if the block is invalid, drop the newly created chain fork
+        let chain = chain.push(prepared)?;
         self.chain_set.insert(Box::new(chain));
         self.update_metrics_for_committed_block(height, hash);
         Ok(())

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

@@ -30,14 +30,16 @@ pub struct Chain {
 }
 
 impl Chain {
-    /// Push a contextually valid non-finalized block into a chain as the new tip.
+    /// Push a contextually valid non-finalized block into this chain as the new tip.
+    ///
+    /// If the block is invalid, drop this chain and return an error.
     #[instrument(level = "debug", skip(self, block), fields(block = %block.block))]
-    pub fn push(&mut self, block: PreparedBlock) -> Result<(), ValidateContextError> {
+    pub fn push(mut self, block: PreparedBlock) -> Result<Chain, ValidateContextError> {
         // update cumulative data members
         self.update_chain_state_with(&block)?;
         tracing::debug!(block = %block.block, "adding block to chain");
         self.blocks.insert(block.height, block);
-        Ok(())
+        Ok(self)
     }
 
     /// Remove the lowest height block of the non-finalized portion of a chain.

zebra-state/src/service/non_finalized_state/tests/prop.rs

@@ -20,10 +20,10 @@ fn forked_equals_pushed() -> Result<()> {
             let mut partial_chain = Chain::default();
 
             for block in chain.iter().take(count) {
-                partial_chain.push(block.clone())?;
+                partial_chain = partial_chain.push(block.clone())?;
             }
             for block in chain.iter() {
-                full_chain.push(block.clone())?;
+                full_chain = full_chain.push(block.clone())?;
             }
 
             let forked = full_chain.fork(fork_tip_hash).expect("fork works").expect("hash is present");
@@ -48,10 +48,10 @@ fn finalized_equals_pushed() -> Result<()> {
         let mut partial_chain = Chain::default();
 
         for block in chain.iter().skip(finalized_count) {
-            partial_chain.push(block.clone())?;
+            partial_chain = partial_chain.push(block.clone())?;
         }
         for block in chain.iter() {
-            full_chain.push(block.clone())?;
+            full_chain = full_chain.push(block.clone())?;
         }
 
         for _ in 0..finalized_count {

zebra-state/src/service/non_finalized_state/tests/vectors.rs

@@ -23,7 +23,7 @@ fn construct_single() -> Result<()> {
         zebra_test::vectors::BLOCK_MAINNET_434873_BYTES.zcash_deserialize_into()?;
 
     let mut chain = Chain::default();
-    chain.push(block.prepare())?;
+    chain = chain.push(block.prepare())?;
 
     assert_eq!(1, chain.blocks.len());
 
@@ -47,7 +47,7 @@ fn construct_many() -> Result<()> {
     let mut chain = Chain::default();
 
     for block in blocks {
-        chain.push(block.prepare())?;
+        chain = chain.push(block.prepare())?;
     }
 
     assert_eq!(100, chain.blocks.len());
@@ -64,10 +64,10 @@ fn ord_matches_work() -> Result<()> {
     let more_block = less_block.clone().set_work(10);
 
     let mut lesser_chain = Chain::default();
-    lesser_chain.push(less_block.prepare())?;
+    lesser_chain = lesser_chain.push(less_block.prepare())?;
 
     let mut bigger_chain = Chain::default();
-    bigger_chain.push(more_block.prepare())?;
+    bigger_chain = bigger_chain.push(more_block.prepare())?;
 
     assert!(bigger_chain > lesser_chain);