From e989225564a2618190a5ae02521fa7ffdc139bc3 Mon Sep 17 00:00:00 2001
From: Balaji Arun <balaji@aptoslabs.com>
Date: Mon, 16 Dec 2024 09:40:01 -0800
Subject: [PATCH] optqs inline batch verification

---
 consensus/consensus-types/src/common.rs  | 35 ++++++++++++++++++------
 consensus/consensus-types/src/payload.rs |  4 +++
 consensus/src/round_manager_test.rs      |  1 +
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/consensus/consensus-types/src/common.rs b/consensus/consensus-types/src/common.rs
index 6f10dc0b29d7f..f57c5c50cbf5d 100644
--- a/consensus/consensus-types/src/common.rs
+++ b/consensus/consensus-types/src/common.rs
@@ -488,6 +488,23 @@ impl Payload {
         Ok(())
     }
 
+    pub fn verify_inline_batches<'a>(
+        inline_batches: impl Iterator<Item = (&'a BatchInfo, &'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 {}: {} != {}",
+                batch,
+                computed_digest,
+                batch.digest()
+            );
+        }
+        Ok(())
+    }
+
     pub fn verify(
         &self,
         validator: &ValidatorVerifier,
@@ -506,20 +523,20 @@ impl Payload {
             ),
             (true, Payload::QuorumStoreInlineHybrid(inline_batches, proof_with_data, _)) => {
                 Self::verify_with_cache(&proof_with_data.proofs, validator, proof_cache)?;
-                for (batch, payload) in inline_batches.iter() {
-                    // TODO: Can cloning be avoided here?
-                    if BatchPayload::new(batch.author(), payload.clone()).hash() != *batch.digest()
-                    {
-                        return Err(anyhow::anyhow!(
-                            "Hash of the received inline batch doesn't match the digest value",
-                        ));
-                    }
-                }
+                Self::verify_inline_batches(
+                    inline_batches.iter().map(|(info, txns)| (info, txns)),
+                )?;
                 Ok(())
             },
             (true, Payload::OptQuorumStore(opt_quorum_store)) => {
                 let proof_with_data = opt_quorum_store.proof_with_data();
                 Self::verify_with_cache(&proof_with_data.batch_summary, validator, proof_cache)?;
+                Self::verify_inline_batches(
+                    opt_quorum_store
+                        .inline_batches()
+                        .iter()
+                        .map(|batch| (batch.info(), batch.transactions())),
+                )?;
                 Ok(())
             },
             (_, _) => Err(anyhow::anyhow!(
diff --git a/consensus/consensus-types/src/payload.rs b/consensus/consensus-types/src/payload.rs
index 08ad201788e2d..d4dd7b26db86a 100644
--- a/consensus/consensus-types/src/payload.rs
+++ b/consensus/consensus-types/src/payload.rs
@@ -188,6 +188,10 @@ impl InlineBatch {
     pub fn info(&self) -> &BatchInfo {
         &self.batch_info
     }
+
+    pub fn transactions(&self) -> &Vec<SignedTransaction> {
+        &self.transactions
+    }
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
diff --git a/consensus/src/round_manager_test.rs b/consensus/src/round_manager_test.rs
index 1123f33352919..292ddfe269ea8 100644
--- a/consensus/src/round_manager_test.rs
+++ b/consensus/src/round_manager_test.rs
@@ -482,6 +482,7 @@ impl NodeSetup {
         self.vote_queue.pop_front().unwrap()
     }
 
+    #[allow(unused)]
     pub async fn next_order_vote(&mut self) -> OrderVoteMsg {
         while self.order_vote_queue.is_empty() {
             self.next_network_message().await;