apollo_consensus: limit messages from future heights or rounds (#8683) (#10239)

Author: asmaastarkware

crates/apollo_consensus/src/manager.rs

@@ -22,7 +22,7 @@ use futures::channel::mpsc;
 use futures::stream::FuturesUnordered;
 use futures::{FutureExt, StreamExt};
 use starknet_api::block::BlockNumber;
-use tracing::{debug, error, info, instrument, trace};
+use tracing::{debug, error, info, instrument, trace, warn};
 
 use crate::metrics::{
     register_metrics,
@@ -407,27 +407,36 @@ impl<ContextT: ConsensusContext> MultiHeightManager<ContextT> {
 
         match proposal_init.height.cmp(&height) {
             std::cmp::Ordering::Greater => {
-                debug!("Received a proposal for a future height. {:?}", proposal_init);
-                // Note: new proposals with the same height/round will be ignored.
-                //
-                // TODO(matan): This only work for trusted peers. In the case of possibly malicious
-                // peers this is a possible DoS attack (malicious users can insert
-                // invalid/bad/malicious proposals before "good" nodes can propose).
-                //
-                // When moving to version 1.0 make sure this is addressed.
-                self.cached_proposals
-                    .entry(proposal_init.height.0)
-                    .or_default()
-                    .entry(proposal_init.round)
-                    .or_insert((proposal_init, content_receiver));
+                if self.should_cache_proposal(&height, 0, &proposal_init) {
+                    debug!("Received a proposal for a future height. {:?}", proposal_init);
+                    // Note: new proposals with the same height/round will be ignored.
+                    //
+                    // TODO(matan): This only work for trusted peers. In the case of possibly
+                    // malicious peers this is a possible DoS attack (malicious
+                    // users can insert invalid/bad/malicious proposals before
+                    // "good" nodes can propose).
+                    //
+                    // When moving to version 1.0 make sure this is addressed.
+                    self.cached_proposals
+                        .entry(proposal_init.height.0)
+                        .or_default()
+                        .entry(proposal_init.round)
+                        .or_insert((proposal_init, content_receiver));
+                }
                 Ok(ShcReturn::Tasks(Vec::new()))
             }
             std::cmp::Ordering::Less => {
                 trace!("Drop proposal from past height. {:?}", proposal_init);
                 Ok(ShcReturn::Tasks(Vec::new()))
             }
             std::cmp::Ordering::Equal => match shc {
-                Some(shc) => shc.handle_proposal(context, proposal_init, content_receiver).await,
+                Some(shc) => {
+                    if self.should_cache_proposal(&height, shc.current_round(), &proposal_init) {
+                        shc.handle_proposal(context, proposal_init, content_receiver).await
+                    } else {
+                        Ok(ShcReturn::Tasks(Vec::new()))
+                    }
+                }
                 None => {
                     trace!("Drop proposal from just completed height. {:?}", proposal_init);
                     Ok(ShcReturn::Tasks(Vec::new()))
@@ -481,16 +490,24 @@ impl<ContextT: ConsensusContext> MultiHeightManager<ContextT> {
         // 2. Parallel proposals - we may send/receive a proposal for (H+1, 0).
         match message.height.cmp(&height.0) {
             std::cmp::Ordering::Greater => {
-                trace!("Cache message for a future height. {:?}", message);
-                self.future_votes.entry(message.height).or_default().push(message);
+                if self.should_cache_vote(&height, 0, &message) {
+                    trace!("Cache message for a future height. {:?}", message);
+                    self.future_votes.entry(message.height).or_default().push(message);
+                }
                 Ok(ShcReturn::Tasks(Vec::new()))
             }
             std::cmp::Ordering::Less => {
                 trace!("Drop message from past height. {:?}", message);
                 Ok(ShcReturn::Tasks(Vec::new()))
             }
             std::cmp::Ordering::Equal => match shc {
-                Some(shc) => shc.handle_vote(context, message).await,
+                Some(shc) => {
+                    if self.should_cache_vote(&height, shc.current_round(), &message) {
+                        shc.handle_vote(context, message).await
+                    } else {
+                        Ok(ShcReturn::Tasks(Vec::new()))
+                    }
+                }
                 None => {
                     trace!("Drop message from just completed height. {:?}", message);
                     Ok(ShcReturn::Tasks(Vec::new()))
@@ -548,4 +565,76 @@ impl<ContextT: ConsensusContext> MultiHeightManager<ContextT> {
         let max_cached_block_number = self.cached_proposals.keys().max().unwrap_or(&height.0);
         CONSENSUS_MAX_CACHED_BLOCK_NUMBER.set_lossy(*max_cached_block_number);
     }
+
+    fn should_cache_msg(
+        &self,
+        current_height: &BlockNumber,
+        current_round: u32,
+        msg_height: u64,
+        msg_round: u32,
+        msg_description: &str,
+    ) -> bool {
+        let height_diff = msg_height.saturating_sub(current_height.0);
+        let round_diff = msg_round.saturating_sub(current_round);
+        let mut should_cache = true;
+
+        // Check height limits first
+        if height_diff > self.consensus_config.static_config.future_height_limit.into() {
+            should_cache = false;
+        }
+
+        // Check round limits based on height
+        if height_diff == 0 {
+            // For current height, check against current round + future_round_limit
+            if round_diff > self.consensus_config.static_config.future_round_limit {
+                should_cache = false;
+            }
+        } else {
+            // For future heights, check absolute round limit
+            if msg_round > self.consensus_config.static_config.future_height_round_limit {
+                should_cache = false;
+            }
+        }
+
+        if !should_cache {
+            warn!(
+                "Dropping {} for height={} round={} when current_height={} current_round={} - \
+                 limits: future_height={}, future_height_round={}, future_round={}",
+                msg_description,
+                msg_height,
+                msg_round,
+                current_height.0,
+                current_round,
+                self.consensus_config.static_config.future_height_limit,
+                self.consensus_config.static_config.future_height_round_limit,
+                self.consensus_config.static_config.future_round_limit
+            );
+        }
+
+        should_cache
+    }
+
+    fn should_cache_proposal(
+        &self,
+        current_height: &BlockNumber,
+        current_round: u32,
+        proposal: &ProposalInit,
+    ) -> bool {
+        self.should_cache_msg(
+            current_height,
+            current_round,
+            proposal.height.0,
+            proposal.round,
+            "proposal",
+        )
+    }
+
+    fn should_cache_vote(
+        &self,
+        current_height: &BlockNumber,
+        current_round: u32,
+        vote: &Vote,
+    ) -> bool {
+        self.should_cache_msg(current_height, current_round, vote.height, vote.round, "vote")
+    }
 }

crates/apollo_consensus/src/manager_test.rs

@@ -335,6 +335,125 @@ async fn timely_message_handling(consensus_config: ConsensusConfig) {
     assert!(vote_sender.send((vote.clone(), metadata.clone())).now_or_never().is_some());
 }
 
+#[rstest]
+#[tokio::test]
+async fn future_height_limit_caching_and_dropping(mut consensus_config: ConsensusConfig) {
+    // Use very low limit - only cache 1 height ahead with round 0.
+    consensus_config.static_config.future_height_limit = 1;
+    consensus_config.static_config.future_round_limit = 0;
+    consensus_config.static_config.future_height_round_limit = 0;
+
+    let TestSubscriberChannels { mock_network, subscriber_channels } =
+        mock_register_broadcast_topic().unwrap();
+    let mut sender = mock_network.broadcasted_messages_sender;
+
+    let (mut proposal_receiver_sender, mut proposal_receiver_receiver) =
+        mpsc::channel(CHANNEL_SIZE);
+
+    // Send proposal and votes for height 2 (should be dropped when processing height 0).
+    send_proposal(
+        &mut proposal_receiver_sender,
+        vec![TestProposalPart::Init(proposal_init(2, 0, *PROPOSER_ID))],
+    )
+    .await;
+    send(&mut sender, prevote(Some(Felt::TWO), 2, 0, *PROPOSER_ID)).await;
+    send(&mut sender, precommit(Some(Felt::TWO), 2, 0, *PROPOSER_ID)).await;
+
+    // Send proposal and votes for height 1 (should be cached when processing height 0).
+    send_proposal(
+        &mut proposal_receiver_sender,
+        vec![TestProposalPart::Init(proposal_init(1, 0, *PROPOSER_ID))],
+    )
+    .await;
+    send(&mut sender, prevote(Some(Felt::ONE), 1, 0, *PROPOSER_ID)).await;
+    send(&mut sender, precommit(Some(Felt::ONE), 1, 0, *PROPOSER_ID)).await;
+
+    // Send proposal and votes for height 0 (current height - needed to reach consensus).
+    send_proposal(
+        &mut proposal_receiver_sender,
+        vec![TestProposalPart::Init(proposal_init(0, 0, *PROPOSER_ID))],
+    )
+    .await;
+    send(&mut sender, prevote(Some(Felt::ZERO), 0, 0, *PROPOSER_ID)).await;
+    send(&mut sender, precommit(Some(Felt::ZERO), 0, 0, *PROPOSER_ID)).await;
+
+    let mut context = MockTestContext::new();
+    context.expect_try_sync().returning(|_| false);
+    expect_validate_proposal(&mut context, Felt::ZERO, 1); // Height 0 validation
+    expect_validate_proposal(&mut context, Felt::ONE, 1); // Height 1 validation
+    context.expect_validators().returning(move |_| vec![*PROPOSER_ID, *VALIDATOR_ID]);
+    context.expect_proposer().returning(move |_, _| *PROPOSER_ID);
+    context.expect_set_height_and_round().returning(move |_, _| ());
+    // Set up coordination to detect when node votes Nil for height 2 (indicating proposal was
+    // dropped, so the node didn't received the proposal and votes Nil).
+    let (height2_nil_vote_trigger, height2_nil_vote_wait) = oneshot::channel();
+    context
+        .expect_broadcast()
+        .withf(move |vote: &Vote| vote.height == 2 && vote.proposal_commitment.is_none())
+        .times(1)
+        .return_once(move |_| {
+            height2_nil_vote_trigger.send(()).unwrap();
+            Ok(())
+        });
+    // Handle all other broadcasts normally.
+    context.expect_broadcast().returning(move |_| Ok(()));
+
+    let mut manager = MultiHeightManager::new(consensus_config, QuorumType::Byzantine);
+    let mut subscriber_channels = subscriber_channels.into();
+
+    // Run height 0 - should drop height 2 messages, cache height 1 messages, and reach consensus.
+    let decision = manager
+        .run_height(
+            &mut context,
+            BlockNumber(0),
+            false,
+            &mut subscriber_channels,
+            &mut proposal_receiver_receiver,
+        )
+        .await
+        .unwrap();
+    assert_decision(decision, Felt::ZERO);
+
+    // Run height 1 - should succeed using cached proposal.
+    let decision = manager
+        .run_height(
+            &mut context,
+            BlockNumber(1),
+            false,
+            &mut subscriber_channels,
+            &mut proposal_receiver_receiver,
+        )
+        .await
+        .unwrap();
+    assert_decision(decision, Felt::ONE);
+
+    // Run height 2 in background - shouldn't reach consensus because proposal was dropped.
+    let manager_handle = tokio::spawn(async move {
+        manager
+            .run_height(
+                &mut context,
+                BlockNumber(2),
+                false,
+                &mut subscriber_channels,
+                &mut proposal_receiver_receiver,
+            )
+            .await
+    });
+
+    // Race between consensus completing and height2_nil_vote_trigger being fired.
+    tokio::select! {
+        _ = height2_nil_vote_wait => {
+            // SUCCESS: height2_nil_vote_trigger was fired - this means the proposal was dropped as
+            // expected, and the node didn't receive the proposal and votes Nil.
+        }
+        consensus_result = manager_handle => {
+            panic!("FAIL: Node should not reach consensus. {consensus_result:?}");
+        }
+    }
+}
+
+// TODO(Asmaa): Add test for current height future round limit.
+
 #[rstest]
 #[tokio::test]
 async fn run_consensus_dynamic_client_updates_validator_between_heights(

crates/apollo_consensus/src/single_height_consensus.rs

@@ -172,6 +172,10 @@ impl SingleHeightConsensus {
         }
     }
 
+    pub(crate) fn current_round(&self) -> Round {
+        self.state_machine.round()
+    }
+
     #[instrument(skip_all)]
     pub(crate) async fn start<ContextT: ConsensusContext>(
         &mut self,