From f78ea9d0a66b89a326fe95df817773ef568245e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20RIBEAU?=
 <francois.ribeau.external@stormshield.eu>
Date: Wed, 28 Aug 2024 17:18:57 +0200
Subject: [PATCH] feat(SwarmBuilder): add with_connection_timeout method

---
 libp2p/CHANGELOG.md               |  2 ++
 libp2p/src/builder/phase/build.rs | 20 +++++++++++++-------
 libp2p/src/builder/phase/swarm.rs |  1 +
 3 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/libp2p/CHANGELOG.md b/libp2p/CHANGELOG.md
index 72a624786d4..f6f6a9253c5 100644
--- a/libp2p/CHANGELOG.md
+++ b/libp2p/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 - Update individual crates.
     - Update to [`libp2p-metrics` `0.15.0`](misc/metrics/CHANGELOG.md#0150).
+- Add `with_connection_timeout` on `SwarmBuilder` to allow configuration of the connection_timeout parameter.
+  See [PR XXXX](https://github.com/libp2p/rust-libp2p/pull/XXXX).
 
 ## 0.54.0
 
diff --git a/libp2p/src/builder/phase/build.rs b/libp2p/src/builder/phase/build.rs
index 80a83994eeb..6f9e8990e14 100644
--- a/libp2p/src/builder/phase/build.rs
+++ b/libp2p/src/builder/phase/build.rs
@@ -1,28 +1,34 @@
+use std::time::Duration;
+
 #[allow(unused_imports)]
 use super::*;
 
 use crate::SwarmBuilder;
-use libp2p_core::Transport;
+use libp2p_core::{transport::timeout::TransportTimeout, Transport};
 use libp2p_swarm::Swarm;
 
 pub struct BuildPhase<T, B> {
     pub(crate) behaviour: B,
     pub(crate) transport: T,
     pub(crate) swarm_config: libp2p_swarm::Config,
+    pub(crate) connection_timeout: Duration,
 }
 
-const CONNECTION_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
+#[allow(unused)] // used in `../swarm.rs` but due to feature flag combinations, clippy gives an unnecessary warning.
+pub(crate) const DEFAULT_CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
 
 impl<Provider, T: AuthenticatedMultiplexedTransport, B: libp2p_swarm::NetworkBehaviour>
     SwarmBuilder<Provider, BuildPhase<T, B>>
 {
+    /// Timeout of the [`TransportTimeout`] wrapping the transport.
+    pub fn with_connection_timeout(mut self, connection_timeout: Duration) -> Self {
+        self.phase.connection_timeout = connection_timeout;
+        self
+    }
+
     pub fn build(self) -> Swarm<B> {
         Swarm::new(
-            libp2p_core::transport::timeout::TransportTimeout::new(
-                self.phase.transport,
-                CONNECTION_TIMEOUT,
-            )
-            .boxed(),
+            TransportTimeout::new(self.phase.transport, self.phase.connection_timeout).boxed(),
             self.phase.behaviour,
             self.keypair.public().to_peer_id(),
             self.phase.swarm_config,
diff --git a/libp2p/src/builder/phase/swarm.rs b/libp2p/src/builder/phase/swarm.rs
index ee456ced927..b8e4bd2b107 100644
--- a/libp2p/src/builder/phase/swarm.rs
+++ b/libp2p/src/builder/phase/swarm.rs
@@ -20,6 +20,7 @@ macro_rules! impl_with_swarm_config {
                         behaviour: self.phase.behaviour,
                         transport: self.phase.transport,
                         swarm_config: constructor($config),
+                        connection_timeout: DEFAULT_CONNECTION_TIMEOUT,
                     },
                     keypair: self.keypair,
                     phantom: std::marker::PhantomData,