Merge pull request #2167 from AleoHQ/fix/from-bytes-bounds

howardwu · web-flow · commit 54447585f747 · 2023-11-09T18:13:23.000-08:00
[TOB] Adds check for possible panics in FromBytes impls
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/console/network/src/lib.rs b/console/network/src/lib.rs
@@ -109,8 +109,8 @@ pub trait Network:
     const BLOCK_TIME: u16 = 10;
     /// The coinbase puzzle degree.
     const COINBASE_PUZZLE_DEGREE: u32 = (1 << 13) - 1; // 8,191
-    /// The maximum number of prover solutions that can be included per block.
-    const MAX_PROVER_SOLUTIONS: usize = 1 << 8; // 256 prover solutions
+    /// The maximum number of solutions that can be included per block.
+    const MAX_SOLUTIONS: usize = 1 << 8; // 256 solutions
     /// The number of blocks per epoch.
     const NUM_BLOCKS_PER_EPOCH: u32 = 3600 / Self::BLOCK_TIME as u32; // 360 blocks == ~1 hour
 
diff --git a/ledger/block/Cargo.toml b/ledger/block/Cargo.toml
@@ -100,6 +100,10 @@ package = "snarkvm-ledger-committee"
 path = "../../ledger/committee"
 features = [ "test-helpers" ]
 
+[dev-dependencies.ledger-narwhal-batch-header]
+package = "snarkvm-ledger-narwhal-batch-header"
+path = "../narwhal/batch-header"
+
 [dev-dependencies.ledger-query]
 package = "snarkvm-ledger-query"
 path = "../query"
diff --git a/ledger/block/src/transactions/mod.rs b/ledger/block/src/transactions/mod.rs
@@ -339,3 +339,18 @@ pub mod test_helpers {
         crate::test_helpers::sample_genesis_block(rng).transactions().clone()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    type CurrentNetwork = console::network::Testnet3;
+
+    #[test]
+    fn test_max_transactions() {
+        assert_eq!(
+            Transactions::<CurrentNetwork>::MAX_TRANSACTIONS,
+            ledger_narwhal_batch_header::BatchHeader::<CurrentNetwork>::MAX_TRANSACTIONS
+        );
+    }
+}
diff --git a/ledger/block/src/verify.rs b/ledger/block/src/verify.rs
@@ -266,10 +266,10 @@ impl<N: Network> Block<N> {
             Some(coinbase) => {
                 // Ensure the number of solutions is within the allowed range.
                 ensure!(
-                    coinbase.len() <= N::MAX_PROVER_SOLUTIONS,
+                    coinbase.len() <= N::MAX_SOLUTIONS,
                     "Block {height} contains too many prover solutions (found '{}', expected '{}')",
                     coinbase.len(),
-                    N::MAX_PROVER_SOLUTIONS
+                    N::MAX_SOLUTIONS
                 );
                 // Ensure the solutions are not accepted after the block height at year 10.
                 if height > block_height_at_year(N::BLOCK_TIME, 10) {
diff --git a/ledger/coinbase/benches/coinbase_puzzle.rs b/ledger/coinbase/benches/coinbase_puzzle.rs
@@ -95,7 +95,7 @@ fn coinbase_puzzle_verify(c: &mut Criterion) {
         let puzzle = CoinbasePuzzleInst::trim(&universal_srs, config).unwrap();
         let epoch_challenge = sample_epoch_challenge(degree, rng);
 
-        for batch_size in [10, 100, <Testnet3 as Network>::MAX_PROVER_SOLUTIONS] {
+        for batch_size in [10, 100, <Testnet3 as Network>::MAX_SOLUTIONS] {
             let solutions = (0..batch_size)
                 .map(|_| {
                     let (address, nonce) = sample_address_and_nonce(rng);
diff --git a/ledger/coinbase/src/helpers/coinbase_solution/mod.rs b/ledger/coinbase/src/helpers/coinbase_solution/mod.rs
@@ -33,11 +33,11 @@ impl<N: Network> CoinbaseSolution<N> {
         // Ensure the solutions are not empty.
         ensure!(!solutions.is_empty(), "There are no solutions to verify for the coinbase puzzle");
         // Ensure the number of partial solutions does not exceed `MAX_PROVER_SOLUTIONS`.
-        if solutions.len() > N::MAX_PROVER_SOLUTIONS {
+        if solutions.len() > N::MAX_SOLUTIONS {
             bail!(
                 "The solutions exceed the allowed number of partial solutions. ({} > {})",
                 solutions.len(),
-                N::MAX_PROVER_SOLUTIONS
+                N::MAX_SOLUTIONS
             );
         }
         // Ensure the puzzle commitments are unique.
diff --git a/ledger/coinbase/src/lib.rs b/ledger/coinbase/src/lib.rs
@@ -171,11 +171,11 @@ impl<N: Network> CoinbasePuzzle<N> {
         ensure!(!solutions.is_empty(), "There are no solutions to verify for the coinbase puzzle");
 
         // Ensure the number of partial solutions does not exceed `MAX_PROVER_SOLUTIONS`.
-        if solutions.len() > N::MAX_PROVER_SOLUTIONS {
+        if solutions.len() > N::MAX_SOLUTIONS {
             bail!(
                 "The solutions exceed the allowed number of partial solutions. ({} > {})",
                 solutions.len(),
-                N::MAX_PROVER_SOLUTIONS
+                N::MAX_SOLUTIONS
             );
         }
 
diff --git a/ledger/coinbase/src/tests.rs b/ledger/coinbase/src/tests.rs
@@ -133,7 +133,7 @@ fn test_profiler() -> Result<()> {
     // Generate proof inputs
     let epoch_challenge = EpochChallenge::new(rng.next_u32(), Default::default(), degree).unwrap();
 
-    for batch_size in [10, 100, <Testnet3 as Network>::MAX_PROVER_SOLUTIONS] {
+    for batch_size in [10, 100, <Testnet3 as Network>::MAX_SOLUTIONS] {
         // Generate the solutions.
         let solutions = (0..batch_size)
             .map(|_| {
diff --git a/ledger/committee/Cargo.toml b/ledger/committee/Cargo.toml
@@ -82,3 +82,7 @@ version = "1"
 [dev-dependencies.snarkvm-ledger-committee]
 path = "."
 features = [ "prop-tests" ]
+
+[dev-dependencies.ledger-narwhal-batch-header]
+package = "snarkvm-ledger-narwhal-batch-header"
+path = "../narwhal/batch-header"
diff --git a/ledger/committee/src/bytes.rs b/ledger/committee/src/bytes.rs
@@ -28,6 +28,13 @@ impl<N: Network> FromBytes for Committee<N> {
         let starting_round = u64::read_le(&mut reader)?;
         // Read the number of members.
         let num_members = u16::read_le(&mut reader)?;
+        // Ensure the number of members is within the allowed limit.
+        if num_members > Self::MAX_COMMITTEE_SIZE {
+            return Err(error(format!(
+                "Committee cannot exceed {} members, found {num_members}",
+                Self::MAX_COMMITTEE_SIZE,
+            )));
+        }
         // Read the members.
         let mut members = IndexMap::with_capacity(num_members as usize);
         for _ in 0..num_members {
diff --git a/ledger/committee/src/lib.rs b/ledger/committee/src/lib.rs
@@ -405,4 +405,12 @@ mod tests {
             }
         }
     }
+
+    #[test]
+    fn test_maximum_committee_size() {
+        assert_eq!(
+            Committee::<CurrentNetwork>::MAX_COMMITTEE_SIZE as usize,
+            ledger_narwhal_batch_header::BatchHeader::<CurrentNetwork>::MAX_CERTIFICATES
+        );
+    }
 }
diff --git a/ledger/narwhal/batch-header/src/bytes.rs b/ledger/narwhal/batch-header/src/bytes.rs
@@ -34,18 +34,32 @@ impl<N: Network> FromBytes for BatchHeader<N> {
         let timestamp = i64::read_le(&mut reader)?;
 
         // Read the number of transmission IDs.
-        let num_transmissions = u32::read_le(&mut reader)?;
+        let num_transmission_ids = u32::read_le(&mut reader)?;
+        // Ensure the number of transmission IDs is within bounds.
+        if num_transmission_ids as usize > Self::MAX_TRANSMISSIONS {
+            return Err(error(format!(
+                "Number of transmission IDs ({num_transmission_ids}) exceeds the maximum ({})",
+                Self::MAX_TRANSMISSIONS,
+            )));
+        }
         // Read the transmission IDs.
         let mut transmission_ids = IndexSet::new();
-        for _ in 0..num_transmissions {
+        for _ in 0..num_transmission_ids {
             // Insert the transmission ID.
             transmission_ids.insert(TransmissionID::read_le(&mut reader)?);
         }
 
         // Read the number of previous certificate IDs.
         let num_previous_certificate_ids = u32::read_le(&mut reader)?;
+        // Ensure the number of previous certificate IDs is within bounds.
+        if num_previous_certificate_ids as usize > Self::MAX_CERTIFICATES {
+            return Err(error(format!(
+                "Number of previous certificate IDs ({num_previous_certificate_ids}) exceeds the maximum ({})",
+                Self::MAX_CERTIFICATES
+            )));
+        }
         // Read the previous certificate IDs.
-        let mut previous_certificate_ids = IndexSet::with_capacity(num_previous_certificate_ids as usize);
+        let mut previous_certificate_ids = IndexSet::new();
         for _ in 0..num_previous_certificate_ids {
             // Read the certificate ID.
             previous_certificate_ids.insert(Field::read_le(&mut reader)?);
diff --git a/ledger/narwhal/batch-header/src/lib.rs b/ledger/narwhal/batch-header/src/lib.rs
@@ -46,6 +46,17 @@ pub struct BatchHeader<N: Network> {
     signature: Signature<N>,
 }
 
+impl<N: Network> BatchHeader<N> {
+    /// The maximum number of certificates in a batch.
+    pub const MAX_CERTIFICATES: usize = 200;
+    /// The maximum number of solutions in a batch.
+    pub const MAX_SOLUTIONS: usize = N::MAX_SOLUTIONS;
+    /// The maximum number of transactions in a batch.
+    pub const MAX_TRANSACTIONS: usize = usize::pow(2, console::program::TRANSACTIONS_DEPTH as u32);
+    /// The maximum number of transmissions in a batch.
+    pub const MAX_TRANSMISSIONS: usize = Self::MAX_SOLUTIONS + Self::MAX_TRANSACTIONS;
+}
+
 impl<N: Network> BatchHeader<N> {
     /// Initializes a new batch header.
     pub fn new<R: Rng + CryptoRng>(
diff --git a/ledger/narwhal/subdag/Cargo.toml b/ledger/narwhal/subdag/Cargo.toml
@@ -39,6 +39,11 @@ package = "snarkvm-ledger-narwhal-batch-certificate"
 path = "../batch-certificate"
 version = "=0.16.8"
 
+[dependencies.narwhal-batch-header]
+package = "snarkvm-ledger-narwhal-batch-header"
+path = "../batch-header"
+version = "=0.16.8"
+
 [dependencies.narwhal-transmission-id]
 package = "snarkvm-ledger-narwhal-transmission-id"
 path = "../transmission-id"
diff --git a/ledger/narwhal/subdag/src/bytes.rs b/ledger/narwhal/subdag/src/bytes.rs
@@ -21,20 +21,31 @@ impl<N: Network> FromBytes for Subdag<N> {
         let version = u8::read_le(&mut reader)?;
         // Ensure the version is valid.
         if version != 1 {
-            return Err(error("Invalid batch version"));
+            return Err(error(format!("Invalid subdag version ({version})")));
         }
 
         // Read the number of rounds.
         let num_rounds = u32::read_le(&mut reader)?;
+        // Ensure the number of rounds is within bounds.
+        if num_rounds as usize > Self::MAX_ROUNDS {
+            return Err(error(format!("Number of rounds ({num_rounds}) exceeds the maximum ({})", Self::MAX_ROUNDS)));
+        }
         // Read the round certificates.
         let mut subdag = BTreeMap::new();
         for _ in 0..num_rounds {
             // Read the round.
             let round = u64::read_le(&mut reader)?;
             // Read the number of certificates.
             let num_certificates = u32::read_le(&mut reader)?;
+            // Ensure the number of certificates is within bounds.
+            if num_certificates as usize > BatchHeader::<N>::MAX_CERTIFICATES {
+                return Err(error(format!(
+                    "Number of certificates ({num_certificates}) exceeds the maximum ({})",
+                    BatchHeader::<N>::MAX_CERTIFICATES
+                )));
+            }
             // Read the certificates.
-            let mut certificates = IndexSet::with_capacity(num_certificates as usize);
+            let mut certificates = IndexSet::new();
             for _ in 0..num_certificates {
                 // Read the certificate.
                 certificates.insert(BatchCertificate::read_le(&mut reader)?);
@@ -43,7 +54,7 @@ impl<N: Network> FromBytes for Subdag<N> {
             subdag.insert(round, certificates);
         }
         // Return the subdag.
-        Self::from(subdag).map_err(|e| error(e.to_string()))
+        Self::from(subdag).map_err(error)
     }
 }
 
@@ -53,13 +64,13 @@ impl<N: Network> ToBytes for Subdag<N> {
         // Write the version.
         1u8.write_le(&mut writer)?;
         // Write the number of rounds.
-        u32::try_from(self.subdag.len()).map_err(|e| error(e.to_string()))?.write_le(&mut writer)?;
+        u32::try_from(self.subdag.len()).map_err(error)?.write_le(&mut writer)?;
         // Write the round certificates.
         for (round, certificates) in &self.subdag {
             // Write the round.
             round.write_le(&mut writer)?;
             // Write the number of certificates.
-            u32::try_from(certificates.len()).map_err(|e| error(e.to_string()))?.write_le(&mut writer)?;
+            u32::try_from(certificates.len()).map_err(error)?.write_le(&mut writer)?;
             // Write the certificates.
             for certificate in certificates {
                 // Write the certificate.
diff --git a/ledger/narwhal/subdag/src/lib.rs b/ledger/narwhal/subdag/src/lib.rs
@@ -21,6 +21,7 @@ mod string;
 
 use console::{account::Address, prelude::*, program::SUBDAG_CERTIFICATES_DEPTH, types::Field};
 use narwhal_batch_certificate::BatchCertificate;
+use narwhal_batch_header::BatchHeader;
 use narwhal_transmission_id::TransmissionID;
 
 use indexmap::IndexSet;
@@ -101,6 +102,11 @@ impl<N: Network> Subdag<N> {
     }
 }
 
+impl<N: Network> Subdag<N> {
+    /// The maximum number of rounds in a subdag (bounded up to GC depth).
+    pub const MAX_ROUNDS: usize = 50;
+}
+
 impl<N: Network> Subdag<N> {
     /// Returns the anchor round.
     pub fn anchor_round(&self) -> u64 {
diff --git a/utilities/src/bytes.rs b/utilities/src/bytes.rs
@@ -145,6 +145,11 @@ impl<'de, T: FromBytes> FromBytesDeserializer<T> {
             false => (size_b, size_a),
         };
 
+        // Ensure 'size_b' is within bounds.
+        if size_b > i32::MAX as usize {
+            return Err(D::Error::custom(format!("size_b ({size_b}) exceeds maximum")));
+        }
+
         // Reserve a new `Vec` with the larger size capacity.
         let mut buffer = Vec::with_capacity(size_b);
 
