Allow customizing the delay before closing a Kademlia connection (#1477)

* Reduce the delay before closing a Kademlia connection

* Rework pull request

* Update protocols/kad/src/behaviour.rs

Co-Authored-By: Roman Borschel <romanb@users.noreply.github.com>

* Update protocols/kad/src/behaviour.rs

Co-Authored-By: Roman Borschel <romanb@users.noreply.github.com>

* Rework the pull request

Co-authored-by: Roman Borschel <romanb@users.noreply.github.com>

protocols/kad/src/handler.rs

@@ -37,7 +37,7 @@ use libp2p_core::{
     upgrade::{self, InboundUpgrade, OutboundUpgrade}
 };
 use log::trace;
-use std::{borrow::Cow, error, fmt, io, pin::Pin, task::Context, task::Poll, time::Duration};
+use std::{error, fmt, io, pin::Pin, task::Context, task::Poll, time::Duration};
 use wasm_timer::Instant;
 
 /// Protocol handler that handles Kademlia communications with the remote.
@@ -48,10 +48,7 @@ use wasm_timer::Instant;
 /// It also handles requests made by the remote.
 pub struct KademliaHandler<TUserData> {
     /// Configuration for the Kademlia protocol.
-    config: KademliaProtocolConfig,
-
-    /// If false, we always refuse incoming Kademlia substreams.
-    allow_listening: bool,
+    config: KademliaHandlerConfig,
 
     /// Next unique ID of a connection.
     next_connec_unique_id: UniqueConnecId,
@@ -63,6 +60,19 @@ pub struct KademliaHandler<TUserData> {
     keep_alive: KeepAlive,
 }
 
+/// Configuration of a [`KademliaHandler`].
+#[derive(Debug, Clone)]
+pub struct KademliaHandlerConfig {
+    /// Configuration of the wire protocol.
+    pub protocol_config: KademliaProtocolConfig,
+
+    /// If false, we deny incoming requests.
+    pub allow_listening: bool,
+
+    /// Time after which we close an idle connection.
+    pub idle_timeout: Duration,
+}
+
 /// State of an active substream, opened either by us or by the remote.
 enum SubstreamState<TUserData> {
     /// We haven't started opening the outgoing substream yet.
@@ -369,42 +379,22 @@ pub struct KademliaRequestId {
 struct UniqueConnecId(u64);
 
 impl<TUserData> KademliaHandler<TUserData> {
-    /// Create a `KademliaHandler` that only allows sending messages to the remote but denying
-    /// incoming connections.
-    pub fn dial_only() -> Self {
-        KademliaHandler::with_allow_listening(false)
-    }
+    /// Create a [`KademliaHandler`] using the given configuration.
+    pub fn new(config: KademliaHandlerConfig) -> Self {
+        let keep_alive = KeepAlive::Until(Instant::now() + config.idle_timeout);
 
-    /// Create a `KademliaHandler` that only allows sending messages but also receive incoming
-    /// requests.
-    ///
-    /// The `Default` trait implementation wraps around this function.
-    pub fn dial_and_listen() -> Self {
-        KademliaHandler::with_allow_listening(true)
-    }
-
-    fn with_allow_listening(allow_listening: bool) -> Self {
         KademliaHandler {
-            config: Default::default(),
-            allow_listening,
+            config,
             next_connec_unique_id: UniqueConnecId(0),
             substreams: Vec::new(),
-            keep_alive: KeepAlive::Until(Instant::now() + Duration::from_secs(10)),
+            keep_alive,
         }
     }
-
-    /// Modifies the protocol name used on the wire. Can be used to create incompatibilities
-    /// between networks on purpose.
-    pub fn with_protocol_name(mut self, name: impl Into<Cow<'static, [u8]>>) -> Self {
-        self.config = self.config.with_protocol_name(name);
-        self
-    }
 }
 
 impl<TUserData> Default for KademliaHandler<TUserData> {
-    #[inline]
     fn default() -> Self {
-        KademliaHandler::dial_and_listen()
+        KademliaHandler::new(Default::default())
     }
 }
 
@@ -422,8 +412,8 @@ where
 
     #[inline]
     fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol> {
-        if self.allow_listening {
-            SubstreamProtocol::new(self.config.clone()).map_upgrade(upgrade::EitherUpgrade::A)
+        if self.config.allow_listening {
+            SubstreamProtocol::new(self.config.protocol_config.clone()).map_upgrade(upgrade::EitherUpgrade::A)
         } else {
             SubstreamProtocol::new(upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade))
         }
@@ -449,7 +439,7 @@ where
             EitherOutput::Second(p) => void::unreachable(p),
         };
 
-        debug_assert!(self.allow_listening);
+        debug_assert!(self.config.allow_listening);
         let connec_unique_id = self.next_connec_unique_id;
         self.next_connec_unique_id.0 += 1;
         self.substreams
@@ -635,7 +625,7 @@ where
             let mut substream = self.substreams.swap_remove(n);
 
             loop {
-                match advance_substream(substream, self.config.clone(), cx) {
+                match advance_substream(substream, self.config.protocol_config.clone(), cx) {
                     (Some(new_state), Some(event), _) => {
                         self.substreams.push(new_state);
                         return Poll::Ready(event);
@@ -672,6 +662,16 @@ where
     }
 }
 
+impl Default for KademliaHandlerConfig {
+    fn default() -> Self {
+        KademliaHandlerConfig {
+            protocol_config: Default::default(),
+            allow_listening: true,
+            idle_timeout: Duration::from_secs(10),
+        }
+    }
+}
+
 /// Advances one substream.
 ///
 /// Returns the new state for that substream, an event to generate, and whether the substream