Parse ConsensusBranchId into NetworkUpgrade for transaction v5 (#2075)

* add consensus_branch_id field to transaction v5

* clippy

* rustfmt

* replace consensus_branch_id with network_upgrade

* remove unintended test files

* change method name

* some clanups

* add network_upgrade as a constant in tests

* use std in created function

* add comment to manual arbitrary impl

* create custom strategy to deal with NetworkUpgrade

* Add a missing TODO comment

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

zebra-chain/src/parameters/network_upgrade.rs

@@ -14,7 +14,7 @@ use chrono::{DateTime, Duration, Utc};
 ///
 /// Network upgrades can change the Zcash network protocol or consensus rules in
 /// incompatible ways.
-#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq)]
+#[derive(Copy, Clone, Debug, Eq, Hash, PartialEq, Serialize, Deserialize)]
 pub enum NetworkUpgrade {
     /// The Zcash protocol for a Genesis block.
     ///
@@ -301,6 +301,13 @@ impl NetworkUpgrade {
             Network::Testnet => height >= TESTNET_MAX_TIME_START_HEIGHT,
         }
     }
+    /// Returns the NetworkUpgrade given an u32 as ConsensusBranchId
+    pub fn from_branch_id(branch_id: u32) -> Option<NetworkUpgrade> {
+        CONSENSUS_BRANCH_IDS
+            .iter()
+            .find(|id| id.1 == ConsensusBranchId(branch_id))
+            .map(|nu| nu.0)
+    }
 }
 
 impl ConsensusBranchId {

zebra-chain/src/transaction.rs

@@ -97,6 +97,10 @@ pub enum Transaction {
     },
     /// A `version = 5` transaction, which supports `Sapling` and `Orchard`.
     V5 {
+        /// The Network Upgrade for this transaction.
+        ///
+        /// Derived from the ConsensusBranchId field.
+        network_upgrade: NetworkUpgrade,
         /// The earliest time or block height that this transaction can be added to the
         /// chain.
         lock_time: LockTime,

zebra-chain/src/transaction/arbitrary.rs

@@ -105,6 +105,7 @@ impl Transaction {
     /// Generate a proptest strategy for V5 Transactions
     pub fn v5_strategy(ledger_state: LedgerState) -> BoxedStrategy<Self> {
         (
+            Self::branch_id_strategy(),
             any::<LockTime>(),
             any::<block::Height>(),
             transparent::Input::vec_strategy(ledger_state, 10),
@@ -112,8 +113,16 @@ impl Transaction {
             option::of(any::<sapling::ShieldedData<sapling::SharedAnchor>>()),
         )
             .prop_map(
-                |(lock_time, expiry_height, inputs, outputs, sapling_shielded_data)| {
+                |(
+                    network_upgrade,
+                    lock_time,
+                    expiry_height,
+                    inputs,
+                    outputs,
+                    sapling_shielded_data,
+                )| {
                     Transaction::V5 {
+                        network_upgrade,
                         lock_time,
                         expiry_height,
                         inputs,
@@ -125,6 +134,20 @@ impl Transaction {
             .boxed()
     }
 
+    // A custom strategy to use only some of the NetworkUpgrade values
+    fn branch_id_strategy() -> BoxedStrategy<NetworkUpgrade> {
+        prop_oneof![
+            Just(NetworkUpgrade::Overwinter),
+            Just(NetworkUpgrade::Sapling),
+            Just(NetworkUpgrade::Blossom),
+            Just(NetworkUpgrade::Heartwood),
+            Just(NetworkUpgrade::Canopy),
+            Just(NetworkUpgrade::Nu5),
+        // TODO: add future network upgrades
+        ]
+        .boxed()
+    }
+
     /// Proptest Strategy for creating a Vector of transactions where the first
     /// transaction is always the only coinbase transaction
     pub fn vec_strategy(

zebra-chain/src/transaction/serialize.rs

@@ -343,6 +343,7 @@ impl ZcashSerialize for Transaction {
             }
 
             Transaction::V5 {
+                network_upgrade,
                 lock_time,
                 expiry_height,
                 inputs,
@@ -356,6 +357,13 @@ impl ZcashSerialize for Transaction {
                 writer.write_u32::<LittleEndian>(5 | (1 << 31))?;
                 writer.write_u32::<LittleEndian>(TX_V5_VERSION_GROUP_ID)?;
 
+                // header: Write the nConsensusBranchId
+                writer.write_u32::<LittleEndian>(u32::from(
+                    network_upgrade
+                        .branch_id()
+                        .expect("valid transactions must have a network upgrade with a branch id"),
+                ))?;
+
                 // transaction validity time and height limits
                 lock_time.zcash_serialize(&mut writer)?;
                 writer.write_u32::<LittleEndian>(expiry_height.0)?;
@@ -491,6 +499,13 @@ impl ZcashDeserialize for Transaction {
                 if id != TX_V5_VERSION_GROUP_ID {
                     return Err(SerializationError::Parse("expected TX_V5_VERSION_GROUP_ID"));
                 }
+                // convert the nConsensusBranchId to a NetworkUpgrade
+                let network_upgrade = NetworkUpgrade::from_branch_id(
+                    reader.read_u32::<LittleEndian>()?,
+                )
+                .ok_or(SerializationError::Parse(
+                    "expected a valid network upgrade from the consensus branch id",
+                ))?;
 
                 // transaction validity time and height limits
                 let lock_time = LockTime::zcash_deserialize(&mut reader)?;
@@ -514,6 +529,7 @@ impl ZcashDeserialize for Transaction {
                 }
 
                 Ok(Transaction::V5 {
+                    network_upgrade,
                     lock_time,
                     expiry_height,
                     inputs,

zebra-chain/src/transaction/tests/vectors.rs

@@ -2,6 +2,8 @@ use super::super::*;
 
 use crate::{
     block::{Block, MAX_BLOCK_BYTES},
+    parameters::NetworkUpgrade::Nu5,
+    sapling::{PerSpendAnchor, SharedAnchor},
     serialization::{ZcashDeserialize, ZcashDeserializeInto, ZcashSerialize},
 };
 
@@ -102,6 +104,7 @@ fn empty_v5_round_trip() {
     zebra_test::init();
 
     let tx = Transaction::V5 {
+        network_upgrade: NETWORK_UPGRADE,
         lock_time: LockTime::min_lock_time(),
         expiry_height: block::Height(0),
         inputs: Vec::new(),
@@ -261,3 +264,138 @@ fn fake_v5_round_trip() {
         );
     }
 }
+
+// Utility functions
+
+/// The network upgrade for any fake transactions we will create.
+const NETWORK_UPGRADE: NetworkUpgrade = Nu5;
+
+/// Convert `trans` into a fake v5 transaction,
+/// converting sapling shielded data from v4 to v5 if possible.
+fn transaction_to_fake_v5(trans: &Transaction) -> Transaction {
+    use Transaction::*;
+
+    match trans {
+        V1 {
+            inputs,
+            outputs,
+            lock_time,
+        } => V5 {
+            network_upgrade: NETWORK_UPGRADE,
+            inputs: inputs.to_vec(),
+            outputs: outputs.to_vec(),
+            lock_time: *lock_time,
+            expiry_height: block::Height(0),
+            sapling_shielded_data: None,
+        },
+        V2 {
+            inputs,
+            outputs,
+            lock_time,
+            ..
+        } => V5 {
+            network_upgrade: NETWORK_UPGRADE,
+            inputs: inputs.to_vec(),
+            outputs: outputs.to_vec(),
+            lock_time: *lock_time,
+            expiry_height: block::Height(0),
+            sapling_shielded_data: None,
+        },
+        V3 {
+            inputs,
+            outputs,
+            lock_time,
+            expiry_height,
+            ..
+        } => V5 {
+            network_upgrade: NETWORK_UPGRADE,
+            inputs: inputs.to_vec(),
+            outputs: outputs.to_vec(),
+            lock_time: *lock_time,
+            expiry_height: *expiry_height,
+            sapling_shielded_data: None,
+        },
+        V4 {
+            inputs,
+            outputs,
+            lock_time,
+            expiry_height,
+            sapling_shielded_data,
+            ..
+        } => V5 {
+            network_upgrade: NETWORK_UPGRADE,
+            inputs: inputs.to_vec(),
+            outputs: outputs.to_vec(),
+            lock_time: *lock_time,
+            expiry_height: *expiry_height,
+            sapling_shielded_data: sapling_shielded_data
+                .clone()
+                .map(sapling_shielded_v4_to_fake_v5)
+                .flatten(),
+        },
+        v5 @ V5 { .. } => v5.clone(),
+    }
+}
+
+/// Convert a v4 sapling shielded data into a fake v5 sapling shielded data,
+/// if possible.
+fn sapling_shielded_v4_to_fake_v5(
+    v4_shielded: sapling::ShieldedData<PerSpendAnchor>,
+) -> Option<sapling::ShieldedData<SharedAnchor>> {
+    use sapling::ShieldedData;
+    use sapling::TransferData::*;
+
+    let unique_anchors: Vec<_> = v4_shielded
+        .spends()
+        .map(|spend| spend.per_spend_anchor)
+        .unique()
+        .collect();
+
+    let fake_spends: Vec<_> = v4_shielded
+        .spends()
+        .cloned()
+        .map(sapling_spend_v4_to_fake_v5)
+        .collect();
+
+    let transfers = match v4_shielded.transfers {
+        SpendsAndMaybeOutputs { maybe_outputs, .. } => {
+            let shared_anchor = match unique_anchors.as_slice() {
+                [unique_anchor] => *unique_anchor,
+                // Multiple different anchors, can't convert to v5
+                _ => return None,
+            };
+
+            SpendsAndMaybeOutputs {
+                shared_anchor,
+                spends: fake_spends.try_into().unwrap(),
+                maybe_outputs,
+            }
+        }
+        JustOutputs { outputs } => JustOutputs { outputs },
+    };
+
+    let fake_shielded_v5 = ShieldedData::<SharedAnchor> {
+        value_balance: v4_shielded.value_balance,
+        transfers,
+        binding_sig: v4_shielded.binding_sig,
+    };
+
+    Some(fake_shielded_v5)
+}
+
+/// Convert a v4 sapling spend into a fake v5 sapling spend.
+fn sapling_spend_v4_to_fake_v5(
+    v4_spend: sapling::Spend<PerSpendAnchor>,
+) -> sapling::Spend<SharedAnchor> {
+    use sapling::Spend;
+
+    Spend::<SharedAnchor> {
+        cv: v4_spend.cv,
+        per_spend_anchor: FieldNotPresent,
+        nullifier: v4_spend.nullifier,
+        rk: v4_spend.rk,
+        zkproof: v4_spend.zkproof,
+        spend_auth_sig: v4_spend.spend_auth_sig,
+    }
+}
+