Make some networking parts deterministic

beacon_node/lighthouse_network/Cargo.toml

@@ -56,6 +56,7 @@ unsigned-varint = { version = "0.8", features = ["codec"] }
 async-channel = { workspace = true }
 logging = { workspace = true }
 proptest = { workspace = true }
+rand_chacha = "0.9.0"
 tempfile = { workspace = true }
 
 [[test]]

beacon_node/lighthouse_network/src/peer_manager/peerdb.rs

@@ -472,7 +472,11 @@ impl<E: EthSpec> PeerDB<E> {
     /// NOTE: Peer scores cannot be penalized during the update, they can only increase. Therefore
     /// it not possible to ban peers when updating scores.
     #[must_use = "The unbanned peers must be sent to libp2p"]
-    pub(super) fn update_scores(&mut self) -> Vec<(PeerId, ScoreUpdateResult)> {
+    /// Applies score decay and ban-expiry to all peers using the supplied `now`.
+    ///
+    /// The production path (the peer manager heartbeat) passes `Instant::now()`; tests can advance
+    /// logical time deterministically by supplying their own `now`.
+    pub(super) fn update_scores_at(&mut self, now: Instant) -> Vec<(PeerId, ScoreUpdateResult)> {
         // Peer can be unbanned in this process.
         // We return the result, such that the peer manager can inform the swarm to lift the libp2p
         // ban on these peers.
@@ -482,7 +486,7 @@ impl<E: EthSpec> PeerDB<E> {
         for (peer_id, info) in self.peers.iter_mut() {
             let previous_state = info.score_state();
             // Update scores
-            info.score_update();
+            info.score_update_at(now);
 
             match Self::handle_score_transition(previous_state, peer_id, info) {
                 // A peer should not be able to be banned from a score update.

beacon_node/lighthouse_network/src/peer_manager/peerdb/peer_info.rs

@@ -458,10 +458,13 @@ impl<E: EthSpec> PeerInfo<E> {
         self.partial_message_subnets.clear()
     }
 
-    /// Applies decay rates to a non-trusted peer's score.
-    pub(super) fn score_update(&mut self) {
+    /// Applies decay rates to a non-trusted peer's score using the supplied `now`.
+    ///
+    /// The production path passes `Instant::now()` (via the peer manager heartbeat); tests can
+    /// advance logical time deterministically instead of relying on the wall clock.
+    pub(super) fn score_update_at(&mut self, now: std::time::Instant) {
         if !self.is_trusted {
-            self.score.update()
+            self.score.update_at(now)
         }
     }
 

beacon_node/lighthouse_network/src/peer_manager/peerdb/score.rs

@@ -219,15 +219,12 @@ impl RealScore {
         self.update_state();
     }
 
-    /// Applies time-based logic such as decay rates to the score.
-    /// This function should be called periodically.
-    pub fn update(&mut self) {
-        self.update_at(Instant::now())
-    }
-
     /// Applies time-based logic such as decay rates to the score with the given now value.
-    /// This private sub function is mainly used for testing.
-    fn update_at(&mut self, now: Instant) {
+    /// This function should be called periodically.
+    ///
+    /// The production path passes `Instant::now()` (via the peer manager heartbeat). Tests can
+    /// drive decay/ban-expiry deterministically by supplying their own `now`.
+    pub(crate) fn update_at(&mut self, now: Instant) {
         // Decay the current score
         // Using exponential decay based on a constant half life.
 
@@ -288,7 +285,7 @@ macro_rules! apply {
 }
 
 apply!(apply_peer_action, peer_action: PeerAction);
-apply!(update);
+apply!(update_at, now: Instant);
 apply!(update_gossipsub_score, new_score: f64, ignore: bool);
 #[cfg(test)]
 apply!(test_add, score: f64);
@@ -348,6 +345,14 @@ impl std::fmt::Display for Score {
     }
 }
 
+/// Test-only re-exports of internal scoring constants so deterministic tests in other modules can
+/// reason about ban-expiry timing without duplicating magic numbers.
+#[cfg(test)]
+pub mod testing {
+    /// The number of seconds we ban a peer for before their score begins to decay.
+    pub const BANNED_BEFORE_DECAY: super::Duration = super::BANNED_BEFORE_DECAY;
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

beacon_node/lighthouse_network/src/rpc/rate_limiter.rs

@@ -396,12 +396,33 @@ impl RPCRateLimiter {
         RPCRateLimiterBuilder::default()
     }
 
+    /// The instant the rate limiter was created. Logical time offsets used by the `*_at` methods
+    /// are measured relative to this. Exposed for deterministic tests.
+    #[cfg(test)]
+    pub(crate) fn init_time(&self) -> Instant {
+        self.init_time
+    }
+
     pub fn allows<Item: RateLimiterItem>(
         &mut self,
         peer_id: &PeerId,
         request: &Item,
     ) -> Result<(), RateLimitedErr> {
-        let time_since_start = self.init_time.elapsed();
+        self.allows_at(Instant::now(), peer_id, request)
+    }
+
+    /// As [`RPCRateLimiter::allows`], but uses the supplied `now` to compute the elapsed time
+    /// since the limiter was created.
+    ///
+    /// The production path (via [`RPCRateLimiter::allows`]) passes `Instant::now()`; tests can
+    /// supply their own `now` to advance logical time deterministically without sleeping.
+    pub(crate) fn allows_at<Item: RateLimiterItem>(
+        &mut self,
+        now: Instant,
+        peer_id: &PeerId,
+        request: &Item,
+    ) -> Result<(), RateLimitedErr> {
+        let time_since_start = now.saturating_duration_since(self.init_time);
         let tokens = request
             .max_responses(
                 self.fork_context.current_fork_epoch(),
@@ -434,7 +455,14 @@ impl RPCRateLimiter {
     }
 
     pub fn prune(&mut self) {
-        let time_since_start = self.init_time.elapsed();
+        self.prune_at(Instant::now());
+    }
+
+    /// As [`RPCRateLimiter::prune`], but uses the supplied `now` to compute the elapsed time since
+    /// the limiter was created. The production path passes `Instant::now()`; tests can advance
+    /// logical time deterministically.
+    pub(crate) fn prune_at(&mut self, now: Instant) {
+        let time_since_start = now.saturating_duration_since(self.init_time);
 
         let Self {
             prune_interval: _,
@@ -679,4 +707,161 @@ mod tests {
         let result = limiter.allows(Duration::from_secs_f32(0.0), &10, tokens);
         assert!(matches!(result, Err(RateLimitedErr::TooLarge)));
     }
+
+    /// A partial refill after a partial interval should make exactly the refilled tokens
+    /// available, and no more.
+    ///
+    /// Quota: 4 tokens per 2s => 1 token replenished every 0.5s. After consuming all 4 tokens at
+    /// t=0, at t=1.0s exactly 2 tokens have been replenished: two 1-token requests succeed and a
+    /// third fails. This is a pure-math (logical time) test of the token-bucket refill edge case.
+    #[test]
+    fn partial_refill_after_partial_interval() {
+        let mut limiter = Limiter::from_quota(Quota {
+            replenish_all_every: Duration::from_secs(2),
+            max_tokens: NonZeroU64::new(4).unwrap(),
+        })
+        .unwrap();
+        let key = 1;
+
+        // Drain the full bucket at t=0.
+        assert!(limiter.allows(Duration::from_secs(0), &key, 4).is_ok());
+        // Immediately after, the bucket is empty.
+        assert!(limiter.allows(Duration::from_secs(0), &key, 1).is_err());
+
+        // After 1.0s, exactly 2 tokens are replenished (1 token / 0.5s).
+        assert!(
+            limiter
+                .allows(Duration::from_secs_f32(1.0), &key, 1)
+                .is_ok()
+        );
+        assert!(
+            limiter
+                .allows(Duration::from_secs_f32(1.0), &key, 1)
+                .is_ok()
+        );
+        // The third token at t=1.0s is not yet available.
+        assert!(
+            limiter
+                .allows(Duration::from_secs_f32(1.0), &key, 1)
+                .is_err()
+        );
+    }
+
+    /// Deterministic ("Tier A") tests at the `RPCRateLimiter` level, exercising the injectable
+    /// time seam (`allows_at` / `init_time`). No real time elapses: all time is logical, expressed
+    /// as offsets from the limiter's `init_time`.
+    mod rpc_rate_limiter_deterministic {
+        use crate::rpc::config::RateLimiterConfig;
+        use crate::rpc::rate_limiter::{RPCRateLimiter, RateLimitedErr};
+        use crate::rpc::{Ping, RequestType};
+        use libp2p::PeerId;
+        use std::num::NonZeroU64;
+        use std::sync::Arc;
+        use std::time::Duration;
+        use types::{EthSpec, ForkContext, Hash256, MainnetEthSpec, Slot};
+
+        fn fork_context() -> Arc<ForkContext> {
+            Arc::new(ForkContext::new::<MainnetEthSpec>(
+                Slot::new(0),
+                Hash256::ZERO,
+                &MainnetEthSpec::default_spec(),
+            ))
+        }
+
+        /// Build a limiter whose Ping protocol allows `max_tokens` tokens every `secs` seconds.
+        fn limiter_with_ping_quota(max_tokens: u64, secs: u64) -> RPCRateLimiter {
+            let config = RateLimiterConfig {
+                ping_quota: crate::rpc::rate_limiter::Quota::n_every(
+                    NonZeroU64::new(max_tokens).unwrap(),
+                    secs,
+                ),
+                ..Default::default()
+            };
+            RPCRateLimiter::new_with_config(config, fork_context()).unwrap()
+        }
+
+        fn ping(data: u64) -> RequestType<MainnetEthSpec> {
+            RequestType::Ping(Ping { data })
+        }
+
+        /// The limiter allows up to capacity, then rejects once the bucket is empty.
+        #[tokio::test]
+        async fn allows_up_to_capacity_then_rejects() {
+            // 3 tokens (each Ping = 1 token) per 30s.
+            let mut limiter = limiter_with_ping_quota(3, 30);
+            let peer = PeerId::random();
+            let t0 = limiter.init_time();
+
+            // All requests happen at the same logical instant t0.
+            assert!(limiter.allows_at(t0, &peer, &ping(0)).is_ok());
+            assert!(limiter.allows_at(t0, &peer, &ping(1)).is_ok());
+            assert!(limiter.allows_at(t0, &peer, &ping(2)).is_ok());
+
+            // The 4th request at the same instant is rejected (bucket empty).
+            assert!(matches!(
+                limiter.allows_at(t0, &peer, &ping(3)),
+                Err(RateLimitedErr::TooSoon(_))
+            ));
+        }
+
+        /// After advancing logical time enough to refill, requests are allowed again.
+        #[tokio::test]
+        async fn refills_after_logical_time_advance() {
+            // 1 token per 10s, so a full refill takes 10s.
+            let mut limiter = limiter_with_ping_quota(1, 10);
+            let peer = PeerId::random();
+            let t0 = limiter.init_time();
+
+            // Consume the single token.
+            assert!(limiter.allows_at(t0, &peer, &ping(0)).is_ok());
+            // Immediately rejected.
+            assert!(matches!(
+                limiter.allows_at(t0, &peer, &ping(1)),
+                Err(RateLimitedErr::TooSoon(_))
+            ));
+
+            // Still too soon after 9 logical seconds.
+            assert!(matches!(
+                limiter.allows_at(t0 + Duration::from_secs(9), &peer, &ping(2)),
+                Err(RateLimitedErr::TooSoon(_))
+            ));
+
+            // After 10 logical seconds the token is replenished and the request is allowed.
+            assert!(
+                limiter
+                    .allows_at(t0 + Duration::from_secs(10), &peer, &ping(3))
+                    .is_ok()
+            );
+        }
+
+        /// `prune_at` removes keys whose bucket is full by the given logical time, but keeps keys
+        /// that are still rate limited.
+        #[tokio::test]
+        async fn prune_at_respects_logical_time() {
+            // 1 token per 10s.
+            let mut limiter = limiter_with_ping_quota(1, 10);
+            let peer = PeerId::random();
+            let t0 = limiter.init_time();
+
+            // Consume the token, putting the peer's TAT into the future (bucket not yet full).
+            assert!(limiter.allows_at(t0, &peer, &ping(0)).is_ok());
+
+            // Pruning at t0 must NOT remove the peer: their bucket is not full yet, so a new
+            // request at t0 is still rate limited.
+            limiter.prune_at(t0);
+            assert!(matches!(
+                limiter.allows_at(t0, &peer, &ping(1)),
+                Err(RateLimitedErr::TooSoon(_))
+            ));
+
+            // Pruning well after the bucket has refilled removes the key; a subsequent request is
+            // treated as a fresh (full-bucket) peer and is allowed.
+            limiter.prune_at(t0 + Duration::from_secs(100));
+            assert!(
+                limiter
+                    .allows_at(t0 + Duration::from_secs(100), &peer, &ping(2))
+                    .is_ok()
+            );
+        }
+    }
 }