quinn-proto: harden congestion arithmetic and resolve token cache panic surface

quinn-proto/src/bloom_token_log.rs

@@ -3,7 +3,7 @@ use std::{
     f64::consts::LN_2,
     hash::{BuildHasher, Hasher},
     mem::{size_of, take},
-    sync::Mutex,
+    sync::{Mutex, PoisonError},
 };
 
 use fastbloom::BloomFilter;
@@ -67,11 +67,22 @@ impl TokenLog for BloomTokenLog {
             return Err(TokenReuseError);
         }
 
-        let mut guard = self.0.lock().unwrap();
+        // Calculate the token's expiry *before* taking the lock, and via checked
+        // arithmetic: an extreme `issued + lifetime` would otherwise panic inside the
+        // `SystemTime` addition while the guard is held, poisoning the mutex and making
+        // every subsequent validation panic on `lock()` (quinn-rs/quinn#2702).
+        let Some(expires_at) = issued.checked_add(lifetime) else {
+            warn!("BloomTokenLog presented with token whose lifetime overflows its issue time");
+            return Err(TokenReuseError);
+        };
+
+        // Recover a poisoned lock rather than propagating the panic: a poisoned guard
+        // only means a past holder panicked, and the bloom/hash-set state remains a
+        // valid (if possibly stale) token log.
+        let mut guard = self.0.lock().unwrap_or_else(PoisonError::into_inner);
         let state = &mut *guard;
 
         // calculate how many periods past period 1 the token expires
-        let expires_at = issued + lifetime;
         let Ok(periods_forward) = expires_at
             .duration_since(state.period_1_start)
             .map(|duration| duration.as_nanos() / lifetime.as_nanos())
@@ -365,4 +376,22 @@ mod test {
     fn k_num_zero() {
         test_doesnt_panic(BloomTokenLog::new(100, 0));
     }
+
+    #[test]
+    fn extreme_lifetime_rejected_without_poisoning() {
+        // An `issued + lifetime` that overflows `SystemTime` must be reported as a
+        // rejected token, not panic while the mutex is held (which would poison it and
+        // make every later validation panic on `lock()`). See quinn-rs/quinn#2702.
+        let log = BloomTokenLog::new_expected_items(1024, 16);
+
+        let extreme = log.check_and_insert(1, SystemTime::now(), Duration::from_secs(u64::MAX));
+        assert!(extreme.is_err(), "overflowing lifetime should be rejected");
+
+        // The lock must not be poisoned: a subsequent normal validation still works.
+        let normal = log.check_and_insert(2, SystemTime::now(), Duration::from_secs(1));
+        assert!(
+            normal.is_ok(),
+            "a normal token must still validate after an extreme one"
+        );
+    }
 }

quinn-proto/src/congestion/bbr/bw_estimation.rs

@@ -20,7 +20,7 @@ pub(crate) struct BandwidthEstimation {
 impl BandwidthEstimation {
     pub(crate) fn on_sent(&mut self, now: Instant, bytes: u64) {
         self.prev_total_sent = self.total_sent;
-        self.total_sent += bytes;
+        self.total_sent = self.total_sent.saturating_add(bytes);
         self.prev_sent_time = self.sent_time;
         self.sent_time = Some(now);
     }
@@ -34,7 +34,7 @@ impl BandwidthEstimation {
         app_limited: bool,
     ) {
         self.prev_total_acked = self.total_acked;
-        self.total_acked += bytes;
+        self.total_acked = self.total_acked.saturating_add(bytes);
         self.prev_acked_time = self.acked_time;
         self.acked_time = Some(now);
 
@@ -83,7 +83,7 @@ impl BandwidthEstimation {
         if window_duration_ns == 0 {
             return None;
         }
-        let b_ns = bytes * 1_000_000_000;
+        let b_ns = bytes.saturating_mul(1_000_000_000);
         let bytes_per_second = b_ns / (window_duration_ns as u64);
         Some(bytes_per_second)
     }

quinn-proto/src/congestion/bbr/mod.rs

@@ -313,21 +313,21 @@ impl Bbr {
         let mut target_window = self.get_target_cwnd(self.cwnd_gain);
         if self.is_at_full_bandwidth {
             // Add the max recently measured ack aggregation to CWND.
-            target_window += self.ack_aggregation.max_ack_height.get();
+            target_window = target_window.saturating_add(self.ack_aggregation.max_ack_height.get());
         } else {
             // Add the most recent excess acked.  Because CWND never decreases in
             // STARTUP, this will automatically create a very localized max filter.
-            target_window += excess_acked;
+            target_window = target_window.saturating_add(excess_acked);
         }
         // Instead of immediately setting the target CWND as the new one, BBR grows
         // the CWND towards |target_window| by only increasing it |bytes_acked| at a
         // time.
         if self.is_at_full_bandwidth {
-            self.cwnd = target_window.min(self.cwnd + bytes_acked);
+            self.cwnd = target_window.min(self.cwnd.saturating_add(bytes_acked));
         } else if (self.cwnd_gain < target_window as f32) || (self.acked_bytes < self.init_cwnd) {
             // If the connection is not yet out of startup phase, do not decrease
             // the window.
-            self.cwnd += bytes_acked;
+            self.cwnd = self.cwnd.saturating_add(bytes_acked);
         }
 
         // Enforce the limits on the congestion window.
@@ -342,7 +342,7 @@ impl Bbr {
         }
         // Set up the initial recovery window.
         if self.recovery_window == 0 {
-            self.recovery_window = self.min_cwnd.max(in_flight + bytes_acked);
+            self.recovery_window = self.min_cwnd.max(in_flight.saturating_add(bytes_acked));
             return;
         }
 
@@ -357,14 +357,14 @@ impl Bbr {
         // In CONSERVATION mode, just subtracting losses is sufficient.  In GROWTH,
         // release additional |bytes_acked| to achieve a slow-start-like behavior.
         if self.recovery_state == RecoveryState::Growth {
-            self.recovery_window += bytes_acked;
+            self.recovery_window = self.recovery_window.saturating_add(bytes_acked);
         }
 
         // Sanity checks.  Ensure that we always allow to send at least an MSS or
         // |bytes_acked| in response, whichever is larger.
         self.recovery_window = self
             .recovery_window
-            .max(in_flight + bytes_acked)
+            .max(in_flight.saturating_add(bytes_acked))
             .max(self.min_cwnd);
     }
 
@@ -407,7 +407,7 @@ impl Controller for Bbr {
     ) {
         self.max_bandwidth
             .on_ack(now, sent, bytes, self.round_count, app_limited);
-        self.acked_bytes += bytes;
+        self.acked_bytes = self.acked_bytes.saturating_add(bytes);
         if self.is_min_rtt_expired(now, app_limited) || self.min_rtt > rtt.min() {
             self.min_rtt = rtt.min();
         }
@@ -560,11 +560,10 @@ impl AckAggregationState {
     ) -> u64 {
         // Compute how many bytes are expected to be delivered, assuming max
         // bandwidth is correct.
-        let expected_bytes_acked = max_bandwidth
-            * now
-                .saturating_duration_since(self.aggregation_epoch_start_time.unwrap_or(now))
-                .as_micros() as u64
-            / 1_000_000;
+        let expected_bytes_acked = max_bandwidth.saturating_mul(
+            now.saturating_duration_since(self.aggregation_epoch_start_time.unwrap_or(now))
+                .as_micros() as u64,
+        ) / 1_000_000;
 
         // Reset the current aggregation epoch as soon as the ack arrival rate is
         // less than or equal to the max bandwidth.
@@ -577,7 +576,9 @@ impl AckAggregationState {
 
         // Compute how many extra bytes were delivered vs max bandwidth.
         // Include the bytes most recently acknowledged to account for stretch acks.
-        self.aggregation_epoch_bytes += newly_acked_bytes;
+        self.aggregation_epoch_bytes = self
+            .aggregation_epoch_bytes
+            .saturating_add(newly_acked_bytes);
         let diff = self.aggregation_epoch_bytes - expected_bytes_acked;
         self.max_ack_height.update_max(round, diff);
         diff
@@ -650,3 +651,29 @@ const K_MAX_INITIAL_CONGESTION_WINDOW: u64 = 200;
 
 const PROBE_RTT_BASED_ON_BDP: bool = true;
 const DRAIN_TO_TARGET: bool = true;
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// Extreme/forged ACK byte counts must be contained by saturating arithmetic
+    /// rather than overflowing BBR's internal counters (which panics in debug
+    /// builds). See quinn-rs/quinn#2702. Before the fix this panicked on the first
+    /// `on_end_acks` via `aggregation_epoch_bytes += newly_acked_bytes`.
+    #[test]
+    fn extreme_ack_accounting_does_not_panic() {
+        let now = Instant::now();
+        let sent = now + Duration::from_micros(1);
+        let rtt = RttEstimator::new(Duration::from_millis(100));
+        let mut controller = Arc::new(BbrConfig::default()).build(now, 1200);
+
+        for _ in 0..4 {
+            controller.on_ack(now, sent, u64::MAX, false, &rtt);
+            controller.on_end_acks(now, u64::MAX, false, Some(u64::MAX));
+        }
+
+        // Reaching here without an arithmetic-overflow panic is the assertion; the
+        // window must also remain a sane, non-zero value.
+        assert!(controller.window() > 0);
+    }
+}

quinn-proto/src/congestion/new_reno.rs

@@ -55,7 +55,9 @@ impl Controller for NewReno {
 
         if self.window < self.ssthresh {
             // Slow start
-            self.window += bytes;
+            // Saturate: an extreme (e.g. forged/aggregated) ACK byte count must not
+            // overflow the window and panic in debug builds (quinn-rs/quinn#2702).
+            self.window = self.window.saturating_add(bytes);
 
             if self.window >= self.ssthresh {
                 // Exiting slow start
@@ -73,11 +75,11 @@ impl Controller for NewReno {
             // for every round trip.
             // This mechanism is called Appropriate Byte Counting in
             // https://tools.ietf.org/html/rfc3465
-            self.bytes_acked += bytes;
+            self.bytes_acked = self.bytes_acked.saturating_add(bytes);
 
             if self.bytes_acked >= self.window {
                 self.bytes_acked -= self.window;
-                self.window += self.current_mtu;
+                self.window = self.window.saturating_add(self.current_mtu);
             }
         }
     }
@@ -171,3 +173,28 @@ impl ControllerFactory for NewRenoConfig {
         Box::new(NewReno::new(self, now, current_mtu))
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::Duration;
+
+    /// A forged/aggregated ACK with an extreme byte count must saturate the window
+    /// rather than overflow-panicking in debug builds. See quinn-rs/quinn#2702.
+    #[test]
+    fn extreme_ack_saturates_window_without_panic() {
+        let now = Instant::now();
+        // `sent` must be after `recovery_start_time` (== now) or on_ack early-returns.
+        let sent = now + Duration::from_micros(1);
+        let rtt = RttEstimator::new(Duration::from_millis(100));
+        let mut controller = NewReno::new(Arc::new(NewRenoConfig::default()), now, 1200);
+
+        // First ACK: slow-start path (`window += bytes`).
+        controller.on_ack(now, sent, u64::MAX, false, &rtt);
+        assert_eq!(controller.window(), u64::MAX);
+
+        // Second ACK: congestion-avoidance path (`bytes_acked += bytes`, `window += mtu`).
+        controller.on_ack(now, sent, u64::MAX, false, &rtt);
+        assert_eq!(controller.window(), u64::MAX);
+    }
+}

quinn-proto/src/token_memory_cache.rs

@@ -2,7 +2,7 @@
 
 use std::{
     collections::{HashMap, VecDeque, hash_map},
-    sync::{Arc, Mutex},
+    sync::{Arc, Mutex, PoisonError},
 };
 
 use bytes::Bytes;
@@ -29,11 +29,20 @@ impl TokenMemoryCache {
 impl TokenStore for TokenMemoryCache {
     fn insert(&self, server_name: &str, token: Bytes) {
         trace!(%server_name, "storing token");
-        self.0.lock().unwrap().store(server_name, token)
+        // A poisoned lock only means a previous holder panicked; the cache state is
+        // still a valid token store, so recover the guard rather than panicking.
+        self.0
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .store(server_name, token)
     }
 
     fn take(&self, server_name: &str) -> Option<Bytes> {
-        let token = self.0.lock().unwrap().take(server_name);
+        let token = self
+            .0
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+            .take(server_name);
         trace!(%server_name, found=%token.is_some(), "taking token");
         token
     }