[QuorumStore] Introduce BatchExtInfo to support Encrypted Transactions

config/src/config/quorum_store_config.rs

@@ -99,6 +99,8 @@ pub struct QuorumStoreConfig {
     pub enable_opt_quorum_store: bool,
     pub opt_qs_minimum_batch_age_usecs: u64,
     pub enable_payload_v2: bool,
+    /// Boolean flag that controls the usage of `BatchInfoExt::V1`
+    pub enable_proof_v2: bool,
 }
 
 impl Default for QuorumStoreConfig {
@@ -140,6 +142,7 @@ impl Default for QuorumStoreConfig {
             enable_opt_quorum_store: true,
             opt_qs_minimum_batch_age_usecs: Duration::from_millis(50).as_micros() as u64,
             enable_payload_v2: false,
+            enable_proof_v2: false,
         }
     }
 }

consensus/consensus-types/src/block.rs

@@ -6,6 +6,7 @@ use crate::{
     block_data::{BlockData, BlockType},
     common::{Author, Payload, Round},
     opt_block_data::OptBlockData,
+    payload::{OptQuorumStorePayload, TDataInfo},
     quorum_cert::QuorumCert,
 };
 use anyhow::{bail, ensure, format_err, Result};
@@ -143,11 +144,19 @@ impl Block {
                     proof_with_data.num_txns(),
                     proof_with_data.num_bytes(),
                 ),
-                Payload::OptQuorumStore(opt_quorum_store_payload) => (
-                    opt_quorum_store_payload.proof_with_data().num_proofs(),
-                    opt_quorum_store_payload.proof_with_data().num_txns(),
-                    opt_quorum_store_payload.proof_with_data().num_bytes(),
-                ),
+                Payload::OptQuorumStore(opt_quorum_store_payload) => match opt_quorum_store_payload
+                {
+                    OptQuorumStorePayload::V1(p) => (
+                        p.proof_with_data().num_proofs(),
+                        p.proof_with_data().num_txns(),
+                        p.proof_with_data().num_bytes(),
+                    ),
+                    OptQuorumStorePayload::V2(p) => (
+                        p.proof_with_data().num_proofs(),
+                        p.proof_with_data().num_txns(),
+                        p.proof_with_data().num_bytes(),
+                    ),
+                },
             },
         }
     }
@@ -169,11 +178,19 @@ impl Block {
                         .map(|(b, _)| b.num_bytes() as usize)
                         .sum(),
                 ),
-                Payload::OptQuorumStore(opt_quorum_store_payload) => (
-                    opt_quorum_store_payload.inline_batches().num_batches(),
-                    opt_quorum_store_payload.inline_batches().num_txns(),
-                    opt_quorum_store_payload.inline_batches().num_bytes(),
-                ),
+                Payload::OptQuorumStore(opt_quorum_store_payload) => match opt_quorum_store_payload
+                {
+                    OptQuorumStorePayload::V1(p) => (
+                        p.inline_batches().num_batches(),
+                        p.inline_batches().num_txns(),
+                        p.inline_batches().num_bytes(),
+                    ),
+                    OptQuorumStorePayload::V2(p) => (
+                        p.inline_batches().num_batches(),
+                        p.inline_batches().num_txns(),
+                        p.inline_batches().num_bytes(),
+                    ),
+                },
                 _ => (0, 0, 0),
             },
         }
@@ -184,19 +201,19 @@ impl Block {
         match self.block_data.payload() {
             None => (0, 0, 0),
             Some(payload) => match payload {
-                Payload::OptQuorumStore(opt_quorum_store_payload) => (
-                    opt_quorum_store_payload.opt_batches().len(),
-                    opt_quorum_store_payload
-                        .opt_batches()
-                        .iter()
-                        .map(|b| b.num_txns() as usize)
-                        .sum(),
-                    opt_quorum_store_payload
-                        .opt_batches()
-                        .iter()
-                        .map(|b| b.num_bytes() as usize)
-                        .sum(),
-                ),
+                Payload::OptQuorumStore(opt_quorum_store_payload) => match opt_quorum_store_payload
+                {
+                    OptQuorumStorePayload::V1(p) => (
+                        p.opt_batches().len(),
+                        p.opt_batches().iter().map(|b| b.num_txns() as usize).sum(),
+                        p.opt_batches().iter().map(|b| b.num_bytes() as usize).sum(),
+                    ),
+                    OptQuorumStorePayload::V2(p) => (
+                        p.opt_batches().len(),
+                        p.opt_batches().iter().map(|b| b.num_txns() as usize).sum(),
+                        p.opt_batches().iter().map(|b| b.num_bytes() as usize).sum(),
+                    ),
+                },
                 _ => (0, 0, 0),
             },
         }

consensus/consensus-types/src/common.rs

@@ -4,7 +4,7 @@
 
 use crate::{
     payload::{OptBatches, OptQuorumStorePayload, PayloadExecutionLimit, TxnAndGasLimits},
-    proof_of_store::{BatchInfo, ProofCache, ProofOfStore},
+    proof_of_store::{BatchInfo, BatchInfoExt, ProofCache, ProofOfStore, TBatchInfo},
 };
 use anyhow::ensure;
 use aptos_crypto::{
@@ -127,11 +127,11 @@ pub struct RejectedTransactionSummary {
 
 #[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
 pub struct ProofWithData {
-    pub proofs: Vec<ProofOfStore>,
+    pub proofs: Vec<ProofOfStore<BatchInfo>>,
 }
 
 impl ProofWithData {
-    pub fn new(proofs: Vec<ProofOfStore>) -> Self {
+    pub fn new(proofs: Vec<ProofOfStore<BatchInfo>>) -> Self {
         Self { proofs }
     }
 
@@ -515,16 +515,20 @@ impl Payload {
         }
     }
 
-    fn verify_with_cache(
-        proofs: &[ProofOfStore],
+    fn verify_with_cache<T>(
+        proofs: &[ProofOfStore<T>],
         validator: &ValidatorVerifier,
         proof_cache: &ProofCache,
-    ) -> anyhow::Result<()> {
+    ) -> anyhow::Result<()>
+    where
+        T: TBatchInfo + Send + Sync + 'static,
+        BatchInfoExt: From<T>,
+    {
         let unverified: Vec<_> = proofs
             .iter()
             .filter(|proof| {
                 proof_cache
-                    .get(proof.info())
+                    .get(&BatchInfoExt::from(proof.info().clone()))
                     .is_none_or(|cached_proof| cached_proof != *proof.multi_signature())
             })
             .collect();
@@ -535,15 +539,15 @@ impl Payload {
         Ok(())
     }
 
-    pub fn verify_inline_batches<'a>(
-        inline_batches: impl Iterator<Item = (&'a BatchInfo, &'a Vec<SignedTransaction>)>,
+    pub fn verify_inline_batches<'a, T: TBatchInfo + 'a>(
+        inline_batches: impl Iterator<Item = (&'a T, &'a Vec<SignedTransaction>)>,
     ) -> anyhow::Result<()> {
         for (batch, payload) in inline_batches {
             // TODO: Can cloning be avoided here?
             let computed_digest = BatchPayload::new(batch.author(), payload.clone()).hash();
             ensure!(
                 computed_digest == *batch.digest(),
-                "Hash of the received inline batch doesn't match the digest value for batch {}: {} != {}",
+                "Hash of the received inline batch doesn't match the digest value for batch {:?}: {} != {}",
                 batch,
                 computed_digest,
                 batch.digest()
@@ -552,9 +556,9 @@ impl Payload {
         Ok(())
     }
 
-    pub fn verify_opt_batches(
+    pub fn verify_opt_batches<T: TBatchInfo>(
         verifier: &ValidatorVerifier,
-        opt_batches: &OptBatches,
+        opt_batches: &OptBatches<T>,
     ) -> anyhow::Result<()> {
         let authors = verifier.address_to_validator_index();
         for batch in &opt_batches.batch_summary {
@@ -592,16 +596,26 @@ impl Payload {
                 )?;
                 Ok(())
             },
-            (true, Payload::OptQuorumStore(opt_quorum_store)) => {
-                let proof_with_data = opt_quorum_store.proof_with_data();
+            (true, Payload::OptQuorumStore(OptQuorumStorePayload::V1(p))) => {
+                let proof_with_data = p.proof_with_data();
+                Self::verify_with_cache(&proof_with_data.batch_summary, verifier, proof_cache)?;
+                Self::verify_inline_batches(
+                    p.inline_batches()
+                        .iter()
+                        .map(|batch| (batch.info(), batch.transactions())),
+                )?;
+                Self::verify_opt_batches(verifier, p.opt_batches())?;
+                Ok(())
+            },
+            (true, Payload::OptQuorumStore(OptQuorumStorePayload::V2(p))) => {
+                let proof_with_data = p.proof_with_data();
                 Self::verify_with_cache(&proof_with_data.batch_summary, verifier, proof_cache)?;
                 Self::verify_inline_batches(
-                    opt_quorum_store
-                        .inline_batches()
+                    p.inline_batches()
                         .iter()
                         .map(|batch| (batch.info(), batch.transactions())),
                 )?;
-                Self::verify_opt_batches(verifier, opt_quorum_store.opt_batches())?;
+                Self::verify_opt_batches(verifier, p.opt_batches())?;
                 Ok(())
             },
             (_, _) => Err(anyhow::anyhow!(
@@ -741,7 +755,7 @@ impl BatchPayload {
 #[derive(Deserialize, Serialize, Clone, Debug, PartialEq, Eq)]
 pub enum PayloadFilter {
     DirectMempool(Vec<TransactionSummary>),
-    InQuorumStore(HashSet<BatchInfo>),
+    InQuorumStore(HashSet<BatchInfoExt>),
     Empty,
 }
 
@@ -772,34 +786,45 @@ impl From<&Vec<&Payload>> for PayloadFilter {
                 match payload {
                     Payload::InQuorumStore(proof_with_status) => {
                         for proof in &proof_with_status.proofs {
-                            exclude_batches.insert(proof.info().clone());
+                            exclude_batches.insert(proof.info().clone().into());
                         }
                     },
                     Payload::InQuorumStoreWithLimit(proof_with_status) => {
                         for proof in &proof_with_status.proof_with_data.proofs {
-                            exclude_batches.insert(proof.info().clone());
+                            exclude_batches.insert(proof.info().clone().into());
                         }
                     },
                     Payload::QuorumStoreInlineHybrid(inline_batches, proof_with_data, _)
                     | Payload::QuorumStoreInlineHybridV2(inline_batches, proof_with_data, _) => {
                         for proof in &proof_with_data.proofs {
-                            exclude_batches.insert(proof.info().clone());
+                            exclude_batches.insert(proof.info().clone().into());
                         }
                         for (batch_info, _) in inline_batches {
-                            exclude_batches.insert(batch_info.clone());
+                            exclude_batches.insert(batch_info.clone().into());
                         }
                     },
                     Payload::DirectMempool(_) => {
                         error!("DirectMempool payload in InQuorumStore filter");
                     },
-                    Payload::OptQuorumStore(opt_qs_payload) => {
-                        for batch in opt_qs_payload.inline_batches().iter() {
+                    Payload::OptQuorumStore(OptQuorumStorePayload::V1(p)) => {
+                        for batch in p.inline_batches().iter() {
+                            exclude_batches.insert(batch.info().clone().into());
+                        }
+                        for batch_info in &p.opt_batches().batch_summary {
+                            exclude_batches.insert(batch_info.clone().into());
+                        }
+                        for proof in &p.proof_with_data().batch_summary {
+                            exclude_batches.insert(proof.info().clone().into());
+                        }
+                    },
+                    Payload::OptQuorumStore(OptQuorumStorePayload::V2(p)) => {
+                        for batch in p.inline_batches().iter() {
                             exclude_batches.insert(batch.info().clone());
                         }
-                        for batch_info in &opt_qs_payload.opt_batches().batch_summary {
+                        for batch_info in &p.opt_batches().batch_summary {
                             exclude_batches.insert(batch_info.clone());
                         }
-                        for proof in &opt_qs_payload.proof_with_data().batch_summary {
+                        for proof in &p.proof_with_data().batch_summary {
                             exclude_batches.insert(proof.info().clone());
                         }
                     },

consensus/consensus-types/src/opt_proposal_msg.rs

@@ -3,7 +3,10 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use crate::{
-    common::Author, opt_block_data::OptBlockData, proof_of_store::ProofCache, sync_info::SyncInfo,
+    common::Author,
+    opt_block_data::OptBlockData,
+    proof_of_store::{BatchInfo, ProofCache},
+    sync_info::SyncInfo,
 };
 use anyhow::{ensure, Context, Result};
 use aptos_types::validator_verifier::ValidatorVerifier;