Merge pull request #2537 from ProvableHQ/fix/subdag-verification

[Fix] Add subdag verification checks

ledger/block/src/verify.rs

@@ -173,6 +173,15 @@ impl<N: Network> Block<N> {
                     subdag.anchor_round(),
                     previous_round
                 );
+                // Ensure that the rounds in the subdag are sequential.
+                if previous_round != 0 {
+                    for round in previous_round..=subdag.anchor_round() {
+                        ensure!(
+                            subdag.contains_key(&round),
+                            "Subdag is missing round {round} in block {expected_height}",
+                        );
+                    }
+                }
                 // Output the subdag anchor round.
                 subdag.anchor_round()
             }

ledger/src/check_next_block.rs

@@ -117,6 +117,9 @@ impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
             ratified_finalize_operations,
         )?;
 
+        // Determine if the block subdag is correctly constructed and is not a combination of multiple subdags.
+        self.check_block_subdag_atomicity(block)?;
+
         // Ensure that each existing solution ID from the block exists in the ledger.
         for existing_solution_id in expected_existing_solution_ids {
             if !self.contains_solution_id(&existing_solution_id)? {
@@ -133,4 +136,66 @@ impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
 
         Ok(())
     }
+
+    /// Checks that the block subdag can not be split into multiple valid subdags.
+    fn check_block_subdag_atomicity(&self, block: &Block<N>) -> Result<()> {
+        // Returns `true` if there is a path from the previous certificate to the current certificate.
+        fn is_linked<N: Network>(
+            subdag: &Subdag<N>,
+            previous_certificate: &BatchCertificate<N>,
+            current_certificate: &BatchCertificate<N>,
+        ) -> Result<bool> {
+            // Initialize the list containing the traversal.
+            let mut traversal = vec![current_certificate];
+            // Iterate over the rounds from the current certificate to the previous certificate.
+            for round in (previous_certificate.round()..current_certificate.round()).rev() {
+                // Retrieve all of the certificates for this past round.
+                let certificates = subdag.get(&round).ok_or(anyhow!("No certificates found for round {round}"))?;
+                // Filter the certificates to only include those that are in the traversal.
+                traversal = certificates
+                    .into_iter()
+                    .filter(|p| traversal.iter().any(|c| c.previous_certificate_ids().contains(&p.id())))
+                    .collect();
+            }
+            Ok(traversal.contains(&previous_certificate))
+        }
+
+        // Check if the block has a subdag.
+        let subdag = match block.authority() {
+            Authority::Quorum(subdag) => subdag,
+            _ => return Ok(()),
+        };
+
+        // Iterate over the rounds to find possible leader certificates.
+        for round in (self.latest_round().saturating_add(2)..=subdag.anchor_round().saturating_sub(2)).rev().step_by(2)
+        {
+            // Retrieve the previous committee lookback.
+            let previous_committee_lookback = self
+                .get_committee_lookback_for_round(round)?
+                .ok_or_else(|| anyhow!("No committee lookback found for round {round}"))?;
+
+            // Compute the leader for the commit round.
+            let computed_leader = previous_committee_lookback
+                .get_leader(round)
+                .map_err(|e| anyhow!("Failed to compute leader for round {round}: {e}"))?;
+
+            // Retrieve the previous leader certificates.
+            let previous_certificate = match subdag.get(&round).and_then(|certificates| {
+                certificates.iter().find(|certificate| certificate.author() == computed_leader)
+            }) {
+                Some(cert) => cert,
+                None => continue,
+            };
+
+            // Determine if there is a path between the previous certificate and the subdag's leader certificate.
+            if is_linked(subdag, previous_certificate, subdag.leader_certificate())? {
+                bail!(
+                    "The previous certificate should not be linked to the current certificate in block {}",
+                    block.height()
+                );
+            }
+        }
+
+        Ok(())
+    }
 }

ledger/src/get.rs

@@ -25,6 +25,22 @@ impl<N: Network, C: ConsensusStorage<N>> Ledger<N, C> {
         self.vm.finalize_store().committee_store().get_committee_for_round(round)
     }
 
+    /// Returns the committee lookback for the given round.
+    pub fn get_committee_lookback_for_round(&self, round: u64) -> Result<Option<Committee<N>>> {
+        // Get the round number for the previous committee. Note, we subtract 2 from odd rounds,
+        // because committees are updated in even rounds.
+        let previous_round = match round % 2 == 0 {
+            true => round.saturating_sub(1),
+            false => round.saturating_sub(2),
+        };
+
+        // Get the committee lookback round.
+        let committee_lookback_round = previous_round.saturating_sub(Committee::<N>::COMMITTEE_LOOKBACK_RANGE);
+
+        // Retrieve the committee for the committee lookback round.
+        self.get_committee_for_round(committee_lookback_round)
+    }
+
     /// Returns the state root that contains the given `block height`.
     pub fn get_state_root(&self, block_height: u32) -> Result<Option<N::StateRoot>> {
         self.vm.block_store().get_state_root(block_height)

ledger/src/tests.rs

@@ -25,20 +25,155 @@ use console::{
     program::{Entry, Identifier, Literal, Plaintext, ProgramID, Value},
     types::U16,
 };
-use ledger_block::{ConfirmedTransaction, Execution, Ratify, Rejected, Transaction};
+use ledger_authority::Authority;
+use ledger_block::{Block, ConfirmedTransaction, Execution, Ratify, Rejected, Transaction};
 use ledger_committee::{Committee, MIN_VALIDATOR_STAKE};
+use ledger_narwhal::{BatchCertificate, BatchHeader, Data, Subdag, Transmission, TransmissionID};
 use ledger_store::{helpers::memory::ConsensusMemory, ConsensusStore};
 use snarkvm_utilities::try_vm_runtime;
 use synthesizer::{program::Program, vm::VM, Stack};
 
-use indexmap::IndexMap;
+use indexmap::{IndexMap, IndexSet};
 use rand::seq::SliceRandom;
+use std::collections::{BTreeMap, HashMap};
+use time::OffsetDateTime;
 
 /// Initializes a sample VM.
 fn sample_vm() -> VM<CurrentNetwork, ConsensusMemory<CurrentNetwork>> {
     VM::from(ConsensusStore::<CurrentNetwork, ConsensusMemory<CurrentNetwork>>::open(None).unwrap()).unwrap()
 }
 
+/// Extract the transmissions from a block.
+fn extract_transmissions(
+    block: &Block<CurrentNetwork>,
+) -> IndexMap<TransmissionID<CurrentNetwork>, Transmission<CurrentNetwork>> {
+    let mut transmissions = IndexMap::new();
+    for tx in block.transactions().iter() {
+        let checksum = Data::Object(tx.transaction().clone()).to_checksum::<CurrentNetwork>().unwrap();
+        transmissions.insert(TransmissionID::from((&tx.id(), &checksum)), tx.transaction().clone().into());
+    }
+    if let Some(coinbase_solution) = block.solutions().as_ref() {
+        for (_, solution) in coinbase_solution.iter() {
+            let checksum = Data::Object(*solution).to_checksum::<CurrentNetwork>().unwrap();
+            transmissions.insert(TransmissionID::from((solution.id(), checksum)), (*solution).into());
+        }
+    }
+    transmissions
+}
+
+/// Construct `num_blocks` quorum blocks given a set of validator private keys and the genesis block.
+fn construct_quorum_blocks(
+    private_keys: Vec<PrivateKey<CurrentNetwork>>,
+    genesis: Block<CurrentNetwork>,
+    num_blocks: u64,
+    rng: &mut TestRng,
+) -> Vec<Block<CurrentNetwork>> {
+    // Initialize the ledger with the genesis block.
+    let ledger =
+        Ledger::<CurrentNetwork, ConsensusMemory<CurrentNetwork>>::load(genesis.clone(), StorageMode::Production)
+            .unwrap();
+
+    // Initialize the round parameters.
+    assert!(num_blocks > 0);
+    assert!(num_blocks < 25);
+    let rounds_per_commit = 2;
+    let final_round = num_blocks.saturating_mul(rounds_per_commit);
+
+    // Sample rounds of batch certificates starting at the genesis round from a static set of 4 authors.
+    let (round_to_certificates_map, committee) = {
+        let committee = ledger.latest_committee().unwrap();
+        let mut round_to_certificates_map: HashMap<u64, IndexSet<BatchCertificate<CurrentNetwork>>> = HashMap::new();
+        let mut previous_certificates: IndexSet<BatchCertificate<CurrentNetwork>> = IndexSet::with_capacity(4);
+
+        // Create certificates for each round.
+        for round in 1..=final_round {
+            let mut current_certificates = IndexSet::new();
+            let previous_certificate_ids =
+                if round <= 1 { IndexSet::new() } else { previous_certificates.iter().map(|c| c.id()).collect() };
+
+            for (i, private_key_1) in private_keys.iter().enumerate() {
+                let batch_header = BatchHeader::new(
+                    private_key_1,
+                    round,
+                    OffsetDateTime::now_utc().unix_timestamp(),
+                    committee.id(),
+                    Default::default(),
+                    previous_certificate_ids.clone(),
+                    rng,
+                )
+                .unwrap();
+                // Add signatures for the batch headers. This creates a fully connected DAG.
+                let signatures = private_keys
+                    .iter()
+                    .enumerate()
+                    .filter(|&(j, _)| i != j)
+                    .map(|(_, private_key_2)| private_key_2.sign(&[batch_header.batch_id()], rng).unwrap())
+                    .collect();
+                current_certificates.insert(BatchCertificate::from(batch_header, signatures).unwrap());
+            }
+
+            round_to_certificates_map.insert(round, current_certificates.clone());
+            previous_certificates = current_certificates;
+        }
+        (round_to_certificates_map, committee)
+    };
+
+    // Helper function to create a quorum block.
+    fn create_next_quorum_block(
+        ledger: &Ledger<CurrentNetwork, ConsensusMemory<CurrentNetwork>>,
+        round: u64,
+        leader_certificate: &BatchCertificate<CurrentNetwork>,
+        previous_leader_certificate: Option<&BatchCertificate<CurrentNetwork>>,
+        round_to_certificates_map: &HashMap<u64, IndexSet<BatchCertificate<CurrentNetwork>>>,
+        rng: &mut TestRng,
+    ) -> Block<CurrentNetwork> {
+        // Construct the subdag for the block.
+        let mut subdag_map = BTreeMap::new();
+        // Add the leader certificate.
+        subdag_map.insert(round, [leader_certificate.clone()].into());
+        // Add the certificates of the previous round.
+        subdag_map.insert(round - 1, round_to_certificates_map.get(&(round - 1)).unwrap().clone());
+        // Add the certificates from the previous leader round, excluding the previous leader certificate.
+        // This assumes the number of rounds per commit is 2.
+        if let Some(prev_leader_cert) = previous_leader_certificate {
+            let mut previous_leader_round_certificates =
+                round_to_certificates_map.get(&(round - 2)).cloned().unwrap_or_default();
+            previous_leader_round_certificates.shift_remove(prev_leader_cert);
+            subdag_map.insert(round - 2, previous_leader_round_certificates);
+        }
+        // Construct the block.
+        let subdag = Subdag::from(subdag_map).unwrap();
+        let block = ledger.prepare_advance_to_next_quorum_block(subdag, Default::default(), rng).unwrap();
+        ledger.check_next_block(&block, rng).unwrap();
+        block
+    }
+
+    // Track the blocks that are created.
+    let mut blocks = Vec::new();
+    let mut previous_leader_certificate: Option<&BatchCertificate<CurrentNetwork>> = None;
+
+    // Construct the blocks.
+    for block_height in 1..=num_blocks {
+        let round = block_height.saturating_mul(rounds_per_commit);
+        let leader = committee.get_leader(round).unwrap();
+        let leader_certificate =
+            round_to_certificates_map.get(&round).unwrap().iter().find(|c| c.author() == leader).unwrap();
+        let block = create_next_quorum_block(
+            &ledger,
+            round,
+            leader_certificate,
+            previous_leader_certificate,
+            &round_to_certificates_map,
+            rng,
+        );
+        ledger.advance_to_next_block(&block).unwrap();
+        previous_leader_certificate = Some(leader_certificate);
+        blocks.push(block);
+    }
+
+    blocks
+}
+
 #[test]
 fn test_load() {
     let rng = &mut TestRng::default();
@@ -2906,3 +3041,103 @@ mod valid_solutions {
         assert_eq!(*block_aborted_solution_id, invalid_solution.id(), "Aborted solutions do not match");
     }
 }
+
+#[test]
+fn test_forged_block_subdags() {
+    let rng = &mut TestRng::default();
+
+    // Sample the genesis private key.
+    let private_key = PrivateKey::<CurrentNetwork>::new(rng).unwrap();
+    // Initialize the store.
+    let store = ConsensusStore::<_, ConsensusMemory<_>>::open(None).unwrap();
+    // Create a genesis block with a seeded RNG to reproduce the same genesis private keys.
+    let seed: u64 = rng.gen();
+    let genesis_rng = &mut TestRng::from_seed(seed);
+    let genesis = VM::from(store).unwrap().genesis_beacon(&private_key, genesis_rng).unwrap();
+
+    // Extract the private keys from the genesis committee by using the same RNG to sample private keys.
+    let genesis_rng = &mut TestRng::from_seed(seed);
+    let private_keys = [
+        private_key,
+        PrivateKey::new(genesis_rng).unwrap(),
+        PrivateKey::new(genesis_rng).unwrap(),
+        PrivateKey::new(genesis_rng).unwrap(),
+    ];
+
+    // Construct 3 quorum blocks.
+    let mut quorum_blocks = construct_quorum_blocks(private_keys.to_vec(), genesis.clone(), 3, rng);
+
+    // Extract the individual blocks.
+    let block_1 = quorum_blocks.remove(0);
+    let block_2 = quorum_blocks.remove(0);
+    let block_3 = quorum_blocks.remove(0);
+
+    // Construct the ledger.
+    let ledger =
+        Ledger::<CurrentNetwork, ConsensusMemory<CurrentNetwork>>::load(genesis, StorageMode::Production).unwrap();
+    ledger.advance_to_next_block(&block_1).unwrap();
+    ledger.check_next_block(&block_2, rng).unwrap();
+
+    ////////////////////////////////////////////////////////////////////////////
+    // Attack 1: Forge block 2' with the subdag of block 3.
+    ////////////////////////////////////////////////////////////////////////////
+    {
+        let block_3_subdag =
+            if let Authority::Quorum(subdag) = block_3.authority() { subdag } else { unreachable!("") };
+
+        // Fetch the transmissions.
+        let transmissions = extract_transmissions(&block_3);
+
+        // Forge the block.
+        let forged_block_2 = ledger
+            .prepare_advance_to_next_quorum_block(block_3_subdag.clone(), transmissions, &mut rand::thread_rng())
+            .unwrap();
+
+        assert_ne!(forged_block_2, block_2);
+
+        // Attempt to verify the forged block.
+        assert!(ledger.check_next_block(&forged_block_2, &mut rand::thread_rng()).is_err());
+    }
+
+    ////////////////////////////////////////////////////////////////////////////
+    // Attack 2: Forge block 2' with the combined subdag of block 2 and 3.
+    ////////////////////////////////////////////////////////////////////////////
+    {
+        // Fetch the subdags.
+        let block_2_subdag =
+            if let Authority::Quorum(subdag) = block_2.authority() { subdag } else { unreachable!("") };
+        let block_3_subdag =
+            if let Authority::Quorum(subdag) = block_3.authority() { subdag } else { unreachable!("") };
+
+        // Combined the subdags.
+        let mut combined_subdag = block_2_subdag.deref().clone();
+        for (round, certificates) in block_3_subdag.iter() {
+            combined_subdag
+                .entry(*round)
+                .and_modify(|c| c.extend(certificates.clone()))
+                .or_insert(certificates.clone());
+        }
+
+        // Fetch the transmissions.
+        let block_2_transmissions = extract_transmissions(&block_2);
+        let block_3_transmissions = extract_transmissions(&block_3);
+
+        // Combine the transmissions.
+        let mut combined_transmissions = block_2_transmissions;
+        combined_transmissions.extend(block_3_transmissions);
+
+        // Forge the block.
+        let forged_block_2_from_both_subdags = ledger
+            .prepare_advance_to_next_quorum_block(
+                Subdag::from(combined_subdag).unwrap(),
+                combined_transmissions,
+                &mut rand::thread_rng(),
+            )
+            .unwrap();
+
+        assert_ne!(forged_block_2_from_both_subdags, block_1);
+
+        // Attempt to verify the forged block.
+        assert!(ledger.check_next_block(&forged_block_2_from_both_subdags, &mut rand::thread_rng()).is_err());
+    }
+}