fix(swarm): prevent overflow in keep-alive computation

swarm/CHANGELOG.md

@@ -6,6 +6,8 @@
   See [PR 4120].
 - Make the `Debug` implementation of `StreamProtocol` more concise.
   See [PR 4631](https://github.com/libp2p/rust-libp2p/pull/4631).
+- Fix overflow in `KeepAlive` computation that could occur panic at `Delay::new` if `SwarmBuilder::idle_connection_timeout` is configured too large.
+  See [PR 4644](https://github.com/libp2p/rust-libp2p/pull/4644).
 - Deprecate `KeepAlive::Until`.
   Individual protocols should not keep connections alive for longer than necessary.
   Users should use `swarm::Config::idle_connection_timeout` instead.

swarm/src/connection.rs

@@ -344,49 +344,12 @@
                 }
             }
 
-            // Ask the handler whether it wants the connection (and the handler itself)
-            // to be kept alive, which determines the planned shutdown, if any.
-            let keep_alive = handler.connection_keep_alive();
-            #[allow(deprecated)]
-            match (&mut *shutdown, keep_alive) {
-                (Shutdown::Later(timer, deadline), KeepAlive::Until(t)) => {
-                    if *deadline != t {
-                        *deadline = t;
-                        if let Some(new_duration) = deadline.checked_duration_since(Instant::now())
-                        {
-                            let effective_keep_alive = max(new_duration, *idle_timeout);
-
-                            timer.reset(effective_keep_alive)
-                        }
-                    }
-                }
-                (_, KeepAlive::Until(earliest_shutdown)) => {
-                    let now = Instant::now();
-
-                    if let Some(requested) = earliest_shutdown.checked_duration_since(now) {
-                        let effective_keep_alive = max(requested, *idle_timeout);
-
-                        let safe_keep_alive = checked_add_fraction(now, effective_keep_alive);
-
-                        // Important: We store the _original_ `Instant` given by the `ConnectionHandler` in the `Later` instance to ensure we can compare it in the above branch.
-                        // This is quite subtle but will hopefully become simpler soon once `KeepAlive::Until` is fully deprecated. See <https://github.com/libp2p/rust-libp2p/issues/3844>/
-                        *shutdown = Shutdown::Later(Delay::new(safe_keep_alive), earliest_shutdown)
-                    }
-                }
-                (_, KeepAlive::No) if idle_timeout == &Duration::ZERO => {
-                    *shutdown = Shutdown::Asap;
-                }
-                (Shutdown::Later(_, _), KeepAlive::No) => {
-                    // Do nothing, i.e. let the shutdown timer continue to tick.
-                }
-                (_, KeepAlive::No) => {
-                    let now = Instant::now();
-                    let safe_keep_alive = checked_add_fraction(now, *idle_timeout);
-
-                    *shutdown = Shutdown::Later(Delay::new(safe_keep_alive), now + safe_keep_alive);
-                }
-                (_, KeepAlive::Yes) => *shutdown = Shutdown::None,
-            };
+            // Compute new shutdown
+            if let Some(new_shutdown) =
+                compute_new_shutdown(handler.connection_keep_alive(), shutdown, *idle_timeout)
+            {
+                *shutdown = new_shutdown;
+            }
 
             // Check if the connection (and handler) should be shut down.
             // As long as we're still negotiating substreams, shutdown is always postponed.
@@ -482,6 +445,62 @@
         .collect()
 }
 
+fn compute_new_shutdown(
+    handler_keep_alive: KeepAlive,
+    current_shutdown: &Shutdown,
+    idle_timeout: Duration,
+) -> Option<Shutdown> {
+    #[allow(deprecated)]
+    match (current_shutdown, handler_keep_alive) {
+        (Shutdown::Later(_, deadline), KeepAlive::Until(t)) => {
+            let now = Instant::now();
+
+            if *deadline != t {
+                let deadline = t;
+                if let Some(new_duration) = deadline.checked_duration_since(Instant::now()) {
+                    let effective_keep_alive = max(new_duration, idle_timeout);
+
+                    let safe_keep_alive = checked_add_fraction(now, effective_keep_alive);
+                    return Some(Shutdown::Later(Delay::new(safe_keep_alive), deadline));
+                }
+            }
+            None
+        }
+        (_, KeepAlive::Until(earliest_shutdown)) => {
+            let now = Instant::now();
+
+            if let Some(requested) = earliest_shutdown.checked_duration_since(now) {
+                let effective_keep_alive = max(requested, idle_timeout);
+
+                let safe_keep_alive = checked_add_fraction(now, effective_keep_alive);
+
+                // Important: We store the _original_ `Instant` given by the `ConnectionHandler` in the `Later` instance to ensure we can compare it in the above branch.
+                // This is quite subtle but will hopefully become simpler soon once `KeepAlive::Until` is fully deprecated. See <https://github.com/libp2p/rust-libp2p/issues/3844>/
+                return Some(Shutdown::Later(
+                    Delay::new(safe_keep_alive),
+                    earliest_shutdown,
+                ));
+            }
+            None
+        }
+        (_, KeepAlive::No) if idle_timeout == Duration::ZERO => Some(Shutdown::Asap),
+        (Shutdown::Later(_, _), KeepAlive::No) => {
+            // Do nothing, i.e. let the shutdown timer continue to tick.
+            None
+        }
+        (_, KeepAlive::No) => {
+            let now = Instant::now();
+            let safe_keep_alive = checked_add_fraction(now, idle_timeout);
+
+            Some(Shutdown::Later(
+                Delay::new(safe_keep_alive),
+                now + safe_keep_alive,
+            ))
+        }
+        (_, KeepAlive::Yes) => Some(Shutdown::None),
+    }
+}
+
 /// Repeatedly halves and adds the [`Duration`] to the [`Instant`] until [`Instant::checked_add`] succeeds.
 ///
 /// [`Instant`] depends on the underlying platform and has a limit of which points in time it can represent.
@@ -986,6 +1005,59 @@
         assert!(start.checked_add(duration).is_some())
     }
 
+    #[test]
+    fn compute_new_shutdown_does_not_panic() {
+        let _ = env_logger::try_init();
+
+        #[derive(Debug)]
+        struct ArbitraryShutdown(Shutdown);
+
+        impl Clone for ArbitraryShutdown {
+            fn clone(&self) -> Self {
+                let shutdown = match self.0 {
+                    Shutdown::None => Shutdown::None,
+                    Shutdown::Asap => Shutdown::Asap,
+                    Shutdown::Later(_, instant) => Shutdown::Later(
+                        // compute_new_shutdown does not touch the delay. Delay does not
+                        // implement Clone. Thus use a placeholder delay.
+                        Delay::new(Duration::from_secs(1)),
+                        instant.clone(),
+                    ),
+                };
+
+                ArbitraryShutdown(shutdown)
+            }
+        }
+
+        impl Arbitrary for ArbitraryShutdown {
+            fn arbitrary(g: &mut Gen) -> Self {
+                let shutdown = match g.gen_range(1u8..4) {
+                    1 => Shutdown::None,
+                    2 => Shutdown::Asap,
+                    3 => Shutdown::Later(
+                        Delay::new(Duration::from_secs(u32::arbitrary(g) as u64)),
+                        Instant::now()
+                            .checked_add(Duration::arbitrary(g))
+                            .unwrap_or(Instant::now()),
+                    ),
+                    _ => unreachable!(),
+                };
+
+                Self(shutdown)
+            }
+        }
+
+        fn prop(
+            handler_keep_alive: KeepAlive,
+            current_shutdown: ArbitraryShutdown,
+            idle_timeout: Duration,
+        ) {
+            compute_new_shutdown(handler_keep_alive, &current_shutdown.0, idle_timeout);
+        }
+
+        QuickCheck::new().quickcheck(prop as fn(_, _, _));
+    }
+
     struct KeepAliveUntilConnectionHandler {
         until: Instant,
     }

swarm/src/handler.rs

@@ -765,6 +765,26 @@ impl Ord for KeepAlive {
     }
 }
 
+#[cfg(test)]
+impl quickcheck::Arbitrary for KeepAlive {
+    fn arbitrary(g: &mut quickcheck::Gen) -> Self {
+        match quickcheck::GenRange::gen_range(g, 1u8..4) {
+            1 =>
+            {
+                #[allow(deprecated)]
+                KeepAlive::Until(
+                    Instant::now()
+                        .checked_add(Duration::arbitrary(g))
+                        .unwrap_or(Instant::now()),
+                )
+            }
+            2 => KeepAlive::Yes,
+            3 => KeepAlive::No,
+            _ => unreachable!(),
+        }
+    }
+}
+
 /// A statically declared, empty [`HashSet`] allows us to work around borrow-checker rules for
 /// [`ProtocolsAdded::from_set`]. The lifetimes don't work unless we have a [`HashSet`] with a `'static' lifetime.
 static EMPTY_HASHSET: Lazy<HashSet<StreamProtocol>> = Lazy::new(HashSet::new);