Verify commit message author matches with the sender

aptos-labs · Nov 24, 2024 · 11b5b0a · 11b5b0a
1 parent 0747561
commit 11b5b0a
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 0 deletions.
diff --git a/consensus/src/epoch_manager.rs b/consensus/src/epoch_manager.rs
@@ -1446,6 +1446,7 @@ impl<P: OnChainConfigProvider> EpochManager<P> {
                 BlockStage::EPOCH_MANAGER_RECEIVED,
             );
         }
+        self.check_author(peer_id, &consensus_msg)?;
         // we can't verify signatures from a different epoch
         let maybe_unverified_event = self.check_epoch(peer_id, consensus_msg).await?;
 
@@ -1512,6 +1513,46 @@ impl<P: OnChainConfigProvider> EpochManager<P> {
         Ok(())
     }
 
+    fn check_author(&mut self, peer_id: AccountAddress, msg: &ConsensusMsg) -> anyhow::Result<()> {
+        let author = match msg {
+            ConsensusMsg::CommitMessage(commit) => commit.author(),
+            ConsensusMsg::ProposalMsg(proposal) => proposal.proposal().author(),
+            ConsensusMsg::VoteMsg(vote) => Some(vote.vote().author()),
+            ConsensusMsg::OrderVoteMsg(order_vote) => Some(order_vote.order_vote().author()),
+            ConsensusMsg::CommitVoteMsg(commit_vote) => Some(commit_vote.author()),
+            ConsensusMsg::BatchMsg(batch) => Some(batch.author()),
+            ConsensusMsg::RoundTimeoutMsg(round_timeout) => Some(round_timeout.author()),
+            ConsensusMsg::BatchResponse(batch_response) => Some(batch_response.author()),
+            ConsensusMsg::BatchRequestMsg(batch_request) => Some(batch_request.source()),
+
+            ConsensusMsg::CommitDecisionMsg(_)
+            | ConsensusMsg::DAGMessage(_)
+            | ConsensusMsg::EpochChangeProof(_)
+            | ConsensusMsg::EpochRetrievalRequest(_)
+            | ConsensusMsg::ProofOfStoreMsg(_)
+            | ConsensusMsg::SyncInfo(_)
+            | ConsensusMsg::RandGenMessage(_)
+            | ConsensusMsg::BatchResponseV2(_)
+            | ConsensusMsg::BlockRetrievalRequest(_)
+            |  ConsensusMsg::BlockRetrievalResponse(_)
+            // For SignedBatchInfo, the verify function will check the author
+            |  ConsensusMsg::SignedBatchInfo(_) => None,
+        };
+
+        if let Some(author) = author {
+            if author != peer_id {
+                bail!(
+                    "Received {:?} message from peer {} with different author {}",
+                    discriminant(msg),
+                    peer_id,
+                    author
+                );
+            }
+        }
+
+        Ok(())
+    }
+
     async fn check_epoch(
         &mut self,
         peer_id: AccountAddress,

diff --git a/consensus/src/pipeline/commit_reliable_broadcast.rs b/consensus/src/pipeline/commit_reliable_broadcast.rs
@@ -60,6 +60,13 @@ impl CommitMessage {
             _ => None,
         }
     }
+
+    pub fn author(&self) -> Option<PeerId> {
+        match self {
+            CommitMessage::Vote(vote) => Some(vote.author()),
+            _ => None,
+        }
+    }
 }
 
 impl RBMessage for CommitMessage {}