apollo_propeller: add nonce per peer to prevent replays (#13264)

Author: guy-starkware

Cargo.lock

@@ -16,6 +16,7 @@ apollo_protobuf.workspace = true
 asynchronous-codec.workspace = true
 futures.workspace = true
 libp2p.workspace = true
+lru.workspace = true
 prost.workspace = true
 rand.workspace = true
 reed-solomon-simd.workspace = true

crates/apollo_propeller/Cargo.toml

@@ -5,9 +5,12 @@
 //! `NetworkBehaviour` adapter in `behaviour.rs` via command/output channels.
 
 use std::collections::{HashMap, HashSet};
+use std::num::NonZeroUsize;
 use std::sync::Arc;
 
+use apollo_infra_utils::warn_every_n_ms;
 use libp2p::identity::{Keypair, PeerId, PublicKey};
+use lru::LruCache;
 use starknet_api::staking::StakingWeight;
 use tokio::sync::{mpsc, oneshot};
 use tracing::{debug, error, trace, warn};
@@ -23,6 +26,14 @@ use crate::tree::PropellerScheduleManager;
 use crate::types::{CommitteeId, CommitteeSetupError, Event, MessageRoot, ShardPublishError};
 use crate::unit::PropellerUnit;
 
+#[cfg(test)]
+#[path = "engine_test.rs"]
+mod engine_test;
+
+// TODO(guyn): move this to the propeller Config.
+// Must be much bigger than the number of peers we expect to work with (2*committee_size).
+const PEER_NONCE_CACHE_SIZE: usize = 1000;
+
 type BroadcastResponseTx = oneshot::Sender<Result<(), ShardPublishError>>;
 type BroadcastResult = (Result<Vec<PropellerUnit>, ShardPublishError>, BroadcastResponseTx);
 
@@ -93,6 +104,9 @@ pub struct Engine {
     message_to_unit_tx: HashMap<MessageKey, mpsc::UnboundedSender<UnitToValidate>>,
     /// Messages that have already been encountered and processed (for deduplication).
     messages_to_ignore_shards_from: TimeCache<MessageKey>,
+    // TODO(guyn): track nonces separately for each committee.
+    /// Nonce per peer. LRU cache is used as a passive garbage collection mechanism.
+    peer_nonce: LruCache<PeerId, u64>,
     /// Receiver for messages from state manager tasks.
     state_manager_rx: mpsc::UnboundedReceiver<EventStateManagerToEngine>,
     state_manager_tx: mpsc::UnboundedSender<EventStateManagerToEngine>,
@@ -126,6 +140,9 @@ impl Engine {
             local_peer_id,
             message_to_unit_tx: HashMap::new(),
             messages_to_ignore_shards_from,
+            peer_nonce: LruCache::new(
+                NonZeroUsize::new(PEER_NONCE_CACHE_SIZE).expect("Cache size must be non-zero"),
+            ),
             state_manager_rx,
             state_manager_tx,
             prepared_units_rx: broadcaster_results_rx,
@@ -243,15 +260,22 @@ impl Engine {
             root: claimed_root,
         };
 
-        // TODO(guyn): Add timestamps to message key to avoid replay attacks or issues with very
-        // late shards.
         if self.messages_to_ignore_shards_from.contains(&message_key) {
             trace!(?message_key, "Message already finalized, dropping unit");
             return;
         }
 
         // Spawn tasks if this is a new message.
         if !self.message_to_unit_tx.contains_key(&message_key) {
+            let nonce = self.peer_nonce.get(&claimed_publisher).copied().unwrap_or(0);
+            if nonce >= claimed_nonce {
+                warn_every_n_ms!(
+                    2000,
+                    "Message nonce is too old, dropping unit to prevent replay attacks"
+                );
+                return;
+            }
+
             debug!(?message_key, "[ENGINE] Spawning new message processor");
 
             let schedule_manager = committee_data.schedule_manager.clone();
@@ -376,6 +400,16 @@ impl Engine {
 
                 if !expired_keys.is_empty() {
                     trace!(?expired_keys, "[ENGINE] Removed expired messages from TTL cache");
+                    for key in expired_keys {
+                        // Update the nonce to the latest timestamp if it is bigger.
+                        let new_nonce = self
+                            .peer_nonce
+                            .peek(&key.publisher)
+                            .copied()
+                            .unwrap_or(0)
+                            .max(key.nonce);
+                        self.peer_nonce.put(key.publisher, new_nonce);
+                    }
                 }
 
                 // Clean up task handles

crates/apollo_propeller/src/engine.rs

@@ -0,0 +1,168 @@
+// TODO(andrew): test non-nonce scenarios.
+use std::time::Duration;
+
+use libp2p::identity::Keypair;
+use libp2p::PeerId;
+use starknet_api::staking::StakingWeight;
+use tokio::sync::mpsc;
+
+use crate::config::Config;
+use crate::engine::{Engine, EngineCommand, EngineOutput, MessageKey};
+use crate::message_processor::EventStateManagerToEngine;
+use crate::types::{CommitteeId, MessageRoot};
+use crate::{MerkleProof, PropellerUnit, Shard, ShardIndex, ShardsOfPeer};
+
+const TEST_COMMITTEE_ID: CommitteeId = CommitteeId([1; 32]);
+const BASE_NONCE: u64 = 1_000_000;
+
+fn test_config() -> Config {
+    Config { stale_message_timeout: Duration::from_millis(200), ..Config::default() }
+}
+
+struct TestSetup {
+    engine: Engine,
+    publisher: PeerId,
+    _output_rx: mpsc::UnboundedReceiver<EngineOutput>,
+    _cmd_tx: mpsc::UnboundedSender<EngineCommand>,
+}
+
+fn setup() -> TestSetup {
+    let local_keypair = Keypair::generate_ed25519();
+    let publisher_keypair = Keypair::generate_ed25519();
+    let publisher = PeerId::from(publisher_keypair.public());
+
+    let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
+    let (output_tx, output_rx) = mpsc::unbounded_channel();
+
+    let mut engine = Engine::new(local_keypair, test_config(), cmd_rx, output_tx, None);
+    engine
+        .register_committee(
+            TEST_COMMITTEE_ID,
+            vec![
+                (engine.local_peer_id, StakingWeight(10), None),
+                (publisher, StakingWeight(10), Some(publisher_keypair.public())),
+            ],
+        )
+        .unwrap();
+
+    TestSetup { engine, publisher, _output_rx: output_rx, _cmd_tx: cmd_tx }
+}
+
+fn make_unit(publisher: PeerId, nonce: u64, root: MessageRoot) -> PropellerUnit {
+    PropellerUnit::new(
+        TEST_COMMITTEE_ID,
+        publisher,
+        root,
+        vec![0; 64],
+        ShardIndex(0),
+        ShardsOfPeer(vec![Shard(vec![1, 2, 3])]),
+        MerkleProof { siblings: vec![] },
+        nonce,
+    )
+}
+
+fn message_key(publisher: PeerId, nonce: u64, root: MessageRoot) -> MessageKey {
+    MessageKey { committee_id: TEST_COMMITTEE_ID, publisher, nonce, root }
+}
+
+fn finalize_message(engine: &mut Engine, publisher: PeerId, nonce: u64, root: MessageRoot) {
+    engine.handle_state_manager_message(EventStateManagerToEngine::Finalized {
+        committee_id: TEST_COMMITTEE_ID,
+        publisher,
+        nonce,
+        message_root: root,
+    });
+}
+
+#[tokio::test]
+async fn reject_unit_of_new_message_with_old_nonce() {
+    let mut s = setup();
+    s.engine.peer_nonce.put(s.publisher, BASE_NONCE);
+
+    // Nonce earlier than BASE_NONCE are rejected.
+    let root_old = MessageRoot([1u8; 32]);
+    s.engine.handle_unit(PeerId::random(), make_unit(s.publisher, BASE_NONCE - 1, root_old));
+    assert!(
+        !s.engine.message_to_unit_tx.contains_key(&message_key(
+            s.publisher,
+            BASE_NONCE - 1,
+            root_old
+        )),
+        "unit with nonce < BASE_NONCE must be rejected"
+    );
+
+    // Nonce equal to BASE_NONCE: rejected.
+    let root_eq = MessageRoot([0u8; 32]);
+    s.engine.handle_unit(PeerId::random(), make_unit(s.publisher, BASE_NONCE, root_eq));
+    assert!(
+        !s.engine.message_to_unit_tx.contains_key(&message_key(s.publisher, BASE_NONCE, root_eq)),
+        "unit with nonce == BASE_NONCE must be rejected"
+    );
+}
+
+#[tokio::test]
+async fn allow_unit_of_new_message_with_fresh_nonce() {
+    let mut s = setup();
+    s.engine.peer_nonce.put(s.publisher, BASE_NONCE);
+
+    let root = MessageRoot([0u8; 32]);
+    s.engine.handle_unit(PeerId::random(), make_unit(s.publisher, BASE_NONCE + 1, root));
+
+    assert!(
+        s.engine.message_to_unit_tx.contains_key(&message_key(s.publisher, BASE_NONCE + 1, root)),
+        "unit with nonce > BASE_NONCE must be accepted"
+    );
+}
+
+#[tokio::test]
+async fn nonce_updated_on_cache_expiry() {
+    tokio::time::pause();
+    let mut s = setup();
+
+    // Finalize two messages from the same publisher with different nonces.
+    finalize_message(&mut s.engine, s.publisher, BASE_NONCE + 100, MessageRoot([2u8; 32]));
+    tokio::time::advance(Duration::from_millis(100)).await;
+    finalize_message(&mut s.engine, s.publisher, BASE_NONCE, MessageRoot([1u8; 32]));
+
+    // Before expiry: no nonce is cached (effective nonce = 0), so a unit with any nonce should be
+    // accepted.
+    let root_before = MessageRoot([10u8; 32]);
+    s.engine.handle_unit(PeerId::random(), make_unit(s.publisher, BASE_NONCE - 50, root_before));
+    assert!(
+        s.engine.message_to_unit_tx.contains_key(&message_key(
+            s.publisher,
+            BASE_NONCE - 50,
+            root_before
+        )),
+        "unit must be accepted when nonce cache has no entry for this publisher"
+    );
+
+    // Advance past TTL so both messages expire, then trigger cleanup via a third finalization.
+    tokio::time::advance(test_config().stale_message_timeout + Duration::from_millis(50)).await;
+    finalize_message(&mut s.engine, s.publisher, BASE_NONCE + 200, MessageRoot([3u8; 32]));
+
+    // After expiry: nonce advances to max(BASE_NONCE, BASE_NONCE+100) = BASE_NONCE+100.
+    // A unit with nonce <= BASE_NONCE+100 must now be rejected.
+    let root_stale = MessageRoot([11u8; 32]);
+    s.engine.handle_unit(PeerId::random(), make_unit(s.publisher, BASE_NONCE + 50, root_stale));
+    assert!(
+        !s.engine.message_to_unit_tx.contains_key(&message_key(
+            s.publisher,
+            BASE_NONCE + 50,
+            root_stale
+        )),
+        "unit with nonce <= max expired nonce must be rejected after cache expiry"
+    );
+
+    // A unit with nonce strictly above BASE_NONCE+100 must still be accepted.
+    let root_fresh = MessageRoot([12u8; 32]);
+    s.engine.handle_unit(PeerId::random(), make_unit(s.publisher, BASE_NONCE + 101, root_fresh));
+    assert!(
+        s.engine.message_to_unit_tx.contains_key(&message_key(
+            s.publisher,
+            BASE_NONCE + 101,
+            root_fresh
+        )),
+        "unit with nonce > max expired nonce must be accepted after cache expiry"
+    );
+}

crates/apollo_propeller/src/engine_test.rs

@@ -35,6 +35,8 @@ pub enum Event {
     MessageTimeout { committee_id: CommitteeId, publisher: PeerId, message_root: MessageRoot },
 }
 
+// TODO(andrew): add the epoch number to committee ID, so it doesn't repeat if the same members are
+// in different epochs.
 #[derive(Debug, Default, PartialEq, Clone, Copy, Ord, PartialOrd, Eq, Hash)]
 pub struct CommitteeId(pub [u8; 32]);