From 7638c43a7cb46bc338a3558c4ac66d7357ff4ba2 Mon Sep 17 00:00:00 2001
From: Alfredo Garcia <oxarbitrage@gmail.com>
Date: Mon, 21 Jun 2021 23:06:52 -0300
Subject: [PATCH] Move network_upgrade check into zebra-chain (#2354)

* move network_upgrade check into zebra-chain

* fix the errors

* rename function

* typo fix

* rename the check function

* make changes from last code review
---
 zebra-chain/src/block.rs                  | 32 ++++++++++++++++++++++-
 zebra-chain/src/block/error.rs            | 10 +++++++
 zebra-chain/src/transaction.rs            | 16 ++++++++++++
 zebra-chain/src/transaction/tests/prop.rs |  8 +-----
 zebra-consensus/src/block/check.rs        | 18 ++-----------
 5 files changed, 60 insertions(+), 24 deletions(-)
 create mode 100644 zebra-chain/src/block/error.rs

diff --git a/zebra-chain/src/block.rs b/zebra-chain/src/block.rs
index 263646cd9..c96032ecb 100644
--- a/zebra-chain/src/block.rs
+++ b/zebra-chain/src/block.rs
@@ -1,6 +1,7 @@
 //! Blocks and block-related structures (heights, headers, etc.)
 
 mod commitment;
+mod error;
 mod hash;
 mod header;
 mod height;
@@ -28,7 +29,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::{
     fmt::DisplayToDebug,
-    parameters::Network,
+    parameters::{Network, NetworkUpgrade},
     serialization::{TrustedPreallocate, MAX_PROTOCOL_MESSAGE_LEN},
     transaction::Transaction,
     transparent,
@@ -86,6 +87,35 @@ impl Block {
             Some(height) => Commitment::from_bytes(self.header.commitment_bytes, network, height),
         }
     }
+
+    /// Check if the `network_upgrade` fields from each transaction in the block matches
+    /// the network upgrade calculated from the `network` and block height.
+    ///
+    /// # Consensus rule:
+    ///
+    ///  The nConsensusBranchId field MUST match the consensus branch ID used for
+    ///  SIGHASH transaction hashes, as specifed in [ZIP-244] ([7.1]).
+    ///
+    /// [ZIP-244]: https://zips.z.cash/zip-0244
+    /// [7.1]: https://zips.z.cash/protocol/nu5.pdf#txnencodingandconsensus
+    pub fn check_transaction_network_upgrade_consistency(
+        &self,
+        network: Network,
+    ) -> Result<(), error::BlockError> {
+        let block_nu =
+            NetworkUpgrade::current(network, self.coinbase_height().expect("a valid height"));
+
+        if self
+            .transactions
+            .iter()
+            .filter_map(|trans| trans.as_ref().network_upgrade())
+            .any(|trans_nu| trans_nu != block_nu)
+        {
+            return Err(error::BlockError::WrongTransactionConsensusBranchId);
+        }
+
+        Ok(())
+    }
 }
 
 impl<'a> From<&'a Block> for Hash {
diff --git a/zebra-chain/src/block/error.rs b/zebra-chain/src/block/error.rs
new file mode 100644
index 000000000..882f38a65
--- /dev/null
+++ b/zebra-chain/src/block/error.rs
@@ -0,0 +1,10 @@
+//! Errors that can occur when checking Block consensus rules.
+
+use thiserror::Error;
+
+#[allow(dead_code, missing_docs)]
+#[derive(Error, Debug, PartialEq)]
+pub enum BlockError {
+    #[error("transaction has wrong consensus branch id for block network upgrade")]
+    WrongTransactionConsensusBranchId,
+}
diff --git a/zebra-chain/src/transaction.rs b/zebra-chain/src/transaction.rs
index f4b1d87e6..d4b086631 100644
--- a/zebra-chain/src/transaction.rs
+++ b/zebra-chain/src/transaction.rs
@@ -192,6 +192,22 @@ impl Transaction {
         }
     }
 
+    /// Get this transaction's network upgrade field, if any.
+    /// This field is serialized as `nConsensusBranchId` ([7.1]).
+    ///
+    /// [7.1]: https://zips.z.cash/protocol/nu5.pdf#txnencodingandconsensus
+    pub fn network_upgrade(&self) -> Option<NetworkUpgrade> {
+        match self {
+            Transaction::V1 { .. }
+            | Transaction::V2 { .. }
+            | Transaction::V3 { .. }
+            | Transaction::V4 { .. } => None,
+            Transaction::V5 {
+                network_upgrade, ..
+            } => Some(*network_upgrade),
+        }
+    }
+
     // transparent
 
     /// Access the transparent inputs of this transaction, regardless of version.
diff --git a/zebra-chain/src/transaction/tests/prop.rs b/zebra-chain/src/transaction/tests/prop.rs
index e804da026..f1cebf535 100644
--- a/zebra-chain/src/transaction/tests/prop.rs
+++ b/zebra-chain/src/transaction/tests/prop.rs
@@ -97,13 +97,7 @@ fn transaction_valid_network_upgrade_strategy() -> Result<()> {
     });
 
     proptest!(|((network, block) in strategy)| {
-        // TODO: replace with check_transaction_network_upgrade from #2343
-        let block_network_upgrade = NetworkUpgrade::current(network, block.coinbase_height().unwrap());
-        for transaction in block.transactions {
-            if let Transaction::V5 { network_upgrade, .. } = transaction.as_ref() {
-                prop_assert_eq!(network_upgrade, &block_network_upgrade);
-            }
-        }
+        block.check_transaction_network_upgrade_consistency(network)?;
     });
 
     Ok(())
diff --git a/zebra-consensus/src/block/check.rs b/zebra-consensus/src/block/check.rs
index b56dd4577..e7620067b 100644
--- a/zebra-consensus/src/block/check.rs
+++ b/zebra-consensus/src/block/check.rs
@@ -189,23 +189,9 @@ pub fn merkle_root_validity(
     block: &Block,
     transaction_hashes: &[transaction::Hash],
 ) -> Result<(), BlockError> {
-    use transaction::Transaction;
-    let block_nu =
-        NetworkUpgrade::current(network, block.coinbase_height().expect("a valid height"));
-
     if block
-        .transactions
-        .iter()
-        .filter_map(|trans| match *trans.as_ref() {
-            Transaction::V5 {
-                network_upgrade, ..
-            } => Some(network_upgrade),
-            Transaction::V1 { .. }
-            | Transaction::V2 { .. }
-            | Transaction::V3 { .. }
-            | Transaction::V4 { .. } => None,
-        })
-        .any(|trans_nu| trans_nu != block_nu)
+        .check_transaction_network_upgrade_consistency(network)
+        .is_err()
     {
         return Err(BlockError::WrongTransactionConsensusBranchId);
     }