Consensus: fix panic in batch verifier (#152)

* fix panic in batch verifier * docs * review comments * Update consensus/rules/src/batch_verifier.rs Co-authored-by: hinto-janai <[email protected]> --------- Co-authored-by: hinto-janai <[email protected]>
Cuprate · Jun 7, 2024 · 07f61bd · 07f61bd
1 parent 6df67bb
commit 07f61bd
Show file tree

Hide file tree

Showing 7 changed files with 125 additions and 39 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/consensus/rules/src/batch_verifier.rs b/consensus/rules/src/batch_verifier.rs
@@ -0,0 +1,29 @@
+use multiexp::BatchVerifier as InternalBatchVerifier;
+
+/// This trait represents a batch verifier.
+///
+/// A batch verifier is used to speed up verification by verifying multiple transactions together.
+///
+/// Not all proofs can be batched and at its core it's intended to verify a series of statements are
+/// each equivalent to zero.
+pub trait BatchVerifier {
+    /// Queue a statement for batch verification.
+    ///
+    /// # Panics
+    /// This function may panic if `stmt` contains calls to `rayon`'s parallel iterators, e.g. `par_iter()`.
+    // TODO: remove the panics by adding a generic API upstream.
+    fn queue_statement<R>(
+        &mut self,
+        stmt: impl FnOnce(&mut InternalBatchVerifier<(), dalek_ff_group::EdwardsPoint>) -> R,
+    ) -> R;
+}
+
+// impl this for a single threaded batch verifier.
+impl BatchVerifier for &'_ mut InternalBatchVerifier<(), dalek_ff_group::EdwardsPoint> {
+    fn queue_statement<R>(
+        &mut self,
+        stmt: impl FnOnce(&mut InternalBatchVerifier<(), dalek_ff_group::EdwardsPoint>) -> R,
+    ) -> R {
+        stmt(self)
+    }
+}
diff --git a/consensus/rules/src/lib.rs b/consensus/rules/src/lib.rs
@@ -1,5 +1,6 @@
 use std::time::{SystemTime, UNIX_EPOCH};
 
+pub mod batch_verifier;
 pub mod blocks;
 mod decomposed_amount;
 pub mod genesis;

diff --git a/consensus/rules/src/transactions.rs b/consensus/rules/src/transactions.rs
@@ -3,10 +3,10 @@ use std::cmp::Ordering;
 use monero_serai::ringct::RctType;
 
 use monero_serai::transaction::{Input, Output, Timelock, Transaction};
-use multiexp::BatchVerifier;
 
 use crate::{
-    blocks::penalty_free_zone, check_point_canonically_encoded, is_decomposed_amount, HardFork,
+    batch_verifier::BatchVerifier, blocks::penalty_free_zone, check_point_canonically_encoded,
+    is_decomposed_amount, HardFork,
 };
 
 mod contextual_data;
@@ -606,7 +606,7 @@ pub fn check_transaction_semantic(
     tx_weight: usize,
     tx_hash: &[u8; 32],
     hf: &HardFork,
-    verifier: &mut BatchVerifier<(), dalek_ff_group::EdwardsPoint>,
+    verifier: impl BatchVerifier,
 ) -> Result<u64, TransactionError> {
     // <https://monero-book.cuprate.org/consensus_rules/transactions.html#transaction-size>
     if tx_blob_size > MAX_TX_BLOB_SIZE

diff --git a/consensus/rules/src/transactions/ring_ct.rs b/consensus/rules/src/transactions/ring_ct.rs
@@ -9,12 +9,11 @@ use monero_serai::{
     transaction::{Input, Transaction},
     H,
 };
-use multiexp::BatchVerifier;
 use rand::thread_rng;
 #[cfg(feature = "rayon")]
 use rayon::prelude::*;
 
-use crate::{transactions::Rings, try_par_iter, HardFork};
+use crate::{batch_verifier::BatchVerifier, transactions::Rings, try_par_iter, HardFork};
 
 /// This constant contains the IDs of 2 transactions that should be allowed after the fork the ringCT
 /// type they used should be banned.
@@ -91,7 +90,7 @@ fn simple_type_balances(rct_sig: &RctSignatures) -> Result<(), RingCTError> {
 /// <https://monero-book.cuprate.org/consensus_rules/ring_ct/bulletproofs+.html>
 fn check_output_range_proofs(
     rct_sig: &RctSignatures,
-    verifier: &mut BatchVerifier<(), dalek_ff_group::EdwardsPoint>,
+    mut verifier: impl BatchVerifier,
 ) -> Result<(), RingCTError> {
     let commitments = &rct_sig.base.commitments;
 
@@ -109,7 +108,9 @@ fn check_output_range_proofs(
             }),
         RctPrunable::MlsagBulletproofs { bulletproofs, .. }
         | RctPrunable::Clsag { bulletproofs, .. } => {
-            if bulletproofs.batch_verify(&mut thread_rng(), verifier, (), commitments) {
+            if verifier.queue_statement(|verifier| {
+                bulletproofs.batch_verify(&mut thread_rng(), verifier, (), commitments)
+            }) {
                 Ok(())
             } else {
                 Err(RingCTError::BulletproofsRangeInvalid)
@@ -121,7 +122,7 @@ fn check_output_range_proofs(
 pub(crate) fn ring_ct_semantic_checks(
     tx: &Transaction,
     tx_hash: &[u8; 32],
-    verifier: &mut BatchVerifier<(), dalek_ff_group::EdwardsPoint>,
+    verifier: impl BatchVerifier,
     hf: &HardFork,
 ) -> Result<(), RingCTError> {
     let rct_type = tx.rct_signatures.rct_type();

diff --git a/consensus/src/batch_verifier.rs b/consensus/src/batch_verifier.rs
@@ -4,8 +4,6 @@ use multiexp::BatchVerifier as InternalBatchVerifier;
 use rayon::prelude::*;
 use thread_local::ThreadLocal;
 
-use crate::ConsensusError;
-
 /// A multithreaded batch verifier.
 pub struct MultiThreadedBatchVerifier {
     internal: ThreadLocal<RefCell<InternalBatchVerifier<(), dalek_ff_group::EdwardsPoint>>>,
@@ -19,19 +17,6 @@ impl MultiThreadedBatchVerifier {
         }
     }
 
-    pub fn queue_statement<R>(
-        &self,
-        stmt: impl FnOnce(
-            &mut InternalBatchVerifier<(), dalek_ff_group::EdwardsPoint>,
-        ) -> Result<R, ConsensusError>,
-    ) -> Result<R, ConsensusError> {
-        let verifier_cell = self
-            .internal
-            .get_or(|| RefCell::new(InternalBatchVerifier::new(8)));
-        // TODO: this is not ok as a rayon par_iter could be called in stmt.
-        stmt(verifier_cell.borrow_mut().deref_mut())
-    }
-
     pub fn verify(self) -> bool {
         self.internal
             .into_iter()
@@ -41,3 +26,17 @@ impl MultiThreadedBatchVerifier {
             .is_none()
     }
 }
+
+impl cuprate_consensus_rules::batch_verifier::BatchVerifier for &'_ MultiThreadedBatchVerifier {
+    fn queue_statement<R>(
+        &mut self,
+        stmt: impl FnOnce(&mut InternalBatchVerifier<(), dalek_ff_group::EdwardsPoint>) -> R,
+    ) -> R {
+        let mut verifier = self
+            .internal
+            .get_or(|| RefCell::new(InternalBatchVerifier::new(32)))
+            .borrow_mut();
+
+        stmt(verifier.deref_mut())
+    }
+}
diff --git a/consensus/src/transactions.rs b/consensus/src/transactions.rs
@@ -484,26 +484,23 @@ where
         batch_get_ring_member_info(txs.iter().map(|(tx, _)| tx), &hf, database).await?;
 
     rayon_spawn_async(move || {
-        let batch_veriifier = MultiThreadedBatchVerifier::new(rayon::current_num_threads());
+        let batch_verifier = MultiThreadedBatchVerifier::new(rayon::current_num_threads());
 
         txs.par_iter()
             .zip(txs_ring_member_info.par_iter())
             .try_for_each(|((tx, verification_needed), ring)| {
                 // do semantic validation if needed.
                 if *verification_needed == VerificationNeeded::SemanticAndContextual {
-                    batch_veriifier.queue_statement(|verifier| {
-                        let fee = check_transaction_semantic(
-                            &tx.tx,
-                            tx.tx_blob.len(),
-                            tx.tx_weight,
-                            &tx.tx_hash,
-                            &hf,
-                            verifier,
-                        )?;
-                        // make sure monero-serai calculated the same fee.
-                        assert_eq!(fee, tx.fee);
-                        Ok(())
-                    })?;
+                    let fee = check_transaction_semantic(
+                        &tx.tx,
+                        tx.tx_blob.len(),
+                        tx.tx_weight,
+                        &tx.tx_hash,
+                        &hf,
+                        &batch_verifier,
+                    )?;
+                    // make sure monero-serai calculated the same fee.
+                    assert_eq!(fee, tx.fee);
                 }
 
                 // Both variants of `VerificationNeeded` require contextual validation.
@@ -518,7 +515,7 @@ where
                 Ok::<_, ConsensusError>(())
             })?;
 
-        if !batch_veriifier.verify() {
+        if !batch_verifier.verify() {
             return Err(ExtendedConsensusError::OneOrMoreBatchVerificationStatementsInvalid);
         }