Support client mode in kademlia

protocols/kad/Cargo.toml

@@ -33,6 +33,7 @@ futures-timer = "3.0"
 libp2p-noise = { path = "../../transports/noise" }
 libp2p-yamux = { path = "../../muxers/yamux" }
 quickcheck = "0.9.0"
+futures-test = "0.3.17"
 
 [build-dependencies]
 prost-build = "0.8"

protocols/kad/src/behaviour.rs

@@ -29,7 +29,7 @@ use crate::handler::{
 };
 use crate::jobs::*;
 use crate::kbucket::{self, Distance, KBucketsTable, NodeStatus};
-use crate::protocol::{KadConnectionType, KadPeer, KademliaProtocolConfig};
+use crate::protocol::{KadConnectionType, KadPeer, KademliaProtocolConfig, Mode};
 use crate::query::{Query, QueryConfig, QueryId, QueryPool, QueryPoolState};
 use crate::record::{
     self,
@@ -175,6 +175,8 @@ pub struct KademliaConfig {
     connection_idle_timeout: Duration,
     kbucket_inserts: KademliaBucketInserts,
     caching: KademliaCaching,
+    // Set the peer mode.
+    mode: Mode,
 }
 
 /// The configuration for Kademlia "write-back" caching after successful
@@ -209,6 +211,8 @@ impl Default for KademliaConfig {
             connection_idle_timeout: Duration::from_secs(10),
             kbucket_inserts: KademliaBucketInserts::OnConnected,
             caching: KademliaCaching::Enabled { max_peers: 1 },
+            // By default, a peer will not be in client mode.
+            mode: Mode::default(),
         }
     }
 }
@@ -390,6 +394,21 @@ impl KademliaConfig {
         self.caching = c;
         self
     }
+
+    /// Sets the [`Mode`] the node is operating in.
+    ///
+    /// The default is [`Mode::Server`].
+    ///
+    /// In [`Mode::Client`] the node does not advertise support for the Kademlia
+    /// protocol, nor does it accept incoming Kademlia streams. Peers will thus not
+    /// include the local node in their routing table. The node can still make
+    /// outbound requests.
+    ///
+    /// Use [`Mode::Client`] when not publicly reachable.
+    pub fn set_mode(&mut self, mode: Mode) -> &mut Self {
+        self.mode = mode;
+        self
+    }
 }
 
 impl<TStore> Kademlia<TStore>
@@ -1715,8 +1734,8 @@ where
     fn new_handler(&mut self) -> Self::ProtocolsHandler {
         KademliaHandlerProto::new(KademliaHandlerConfig {
             protocol_config: self.protocol_config.clone(),
-            allow_listening: true,
             idle_timeout: self.connection_idle_timeout,
+            mode: Mode::default(),
         })
     }
 

protocols/kad/src/behaviour/test.rs

@@ -1318,6 +1318,119 @@ fn network_behaviour_inject_address_change() {
     );
 }
 
+#[futures_test::test]
+async fn client_mode_test() {
+    let mut addrs = Vec::new();
+    let mut peers = Vec::new();
+
+    // Create server peers.
+    let server1 = build_node();
+    addrs.push(server1.0);
+    let mut server1 = server1.1;
+
+    let server2 = build_node();
+    addrs.push(server2.0);
+    let mut server2 = server2.1;
+
+    // Create a client peer.
+    let mut cfg = KademliaConfig::default();
+    cfg.set_mode(Mode::Server);
+    let client = build_node_with_config(cfg);
+    addrs.push(client.0);
+    let mut client = client.1;
+
+    // Fitler out peer Ids.
+    peers.push(*server1.local_peer_id());
+    peers.push(*server2.local_peer_id());
+    peers.push(*client.local_peer_id());
+
+    // Connect server1 and server2.
+    server1
+        .behaviour_mut()
+        .add_address(&peers[1], addrs[1].clone());
+
+    server2
+        .behaviour_mut()
+        .add_address(&peers[0], addrs[0].clone());
+
+    // Put a record from a server peer.
+    let key = Key::new(&"123".to_string());
+    let record = Record::new(key.clone(), vec![4, 5, 6]);
+
+    let put_qid = server1
+        .behaviour_mut()
+        .put_record(record.clone(), Quorum::One)
+        .unwrap();
+
+    loop {
+        futures::select! {
+            event = server1.next() => if let Some(SwarmEvent::Behaviour(KademliaEvent::OutboundQueryCompleted {
+                result: QueryResult::PutRecord(result),
+                id: qid,
+                ..
+            })) = event {
+                assert_eq!(qid, put_qid);
+                match result {
+                    Ok(PutRecordOk { key: actual_key }) => {
+                        if actual_key == key {
+                            // Dial a server peer from the client peer.
+                            // client.dial_addr(addrs[1].clone());
+                            client
+                                .behaviour_mut()
+                                .add_address(&peers[1], addrs[1].clone());
+
+                            client.behaviour_mut().get_record(&key, Quorum::One);
+                        }
+                    },
+                    Err(e) => panic!("{:?}", e),
+                }
+            },
+            _event = server2.next() => {},
+            event = client.next() => if let Some(SwarmEvent::Behaviour(KademliaEvent::OutboundQueryCompleted{
+                result: QueryResult::GetRecord(result),
+                ..
+            })) = event {
+                match result {
+                    Ok(GetRecordOk { records: actual_record, .. }) => {
+                        assert_eq!(&record.key, &actual_record[0].record.key);
+
+                        let mut records = Vec::new();
+
+                        // Read the routing table.
+                        for bucket in client.behaviour_mut().kbuckets() {
+                            for record in bucket.iter() {
+                                records.push(*record.node.key.preimage());
+                            }
+                        }
+
+                        assert!(records.contains(&peers[1]));
+
+                        records.clear();
+                        for bucket in server1.behaviour_mut().kbuckets() {
+                            for record in bucket.iter() {
+                                records.push(*record.node.key.preimage());
+                            }
+                        }
+
+                        assert!(!records.contains(&peers[2]));
+
+                        records.clear();
+                        for bucket in server2.behaviour_mut().kbuckets() {
+                            for record in bucket.iter() {
+                                records.push(*record.node.key.preimage());
+                            }
+                        }
+
+                        assert!(!records.contains(&peers[2]));
+                        break;
+                    },
+                    Err(e) => panic!("{:?}", e),
+                }
+            }
+        }
+    }
+}
+
 #[test]
 fn get_providers() {
     fn prop(key: record::Key) {

protocols/kad/src/handler.rs

@@ -20,7 +20,7 @@
 
 use crate::protocol::{
     KadInStreamSink, KadOutStreamSink, KadPeer, KadRequestMsg, KadResponseMsg,
-    KademliaProtocolConfig,
+    KademliaProtocolConfig, Mode,
 };
 use crate::record::{self, Record};
 use futures::prelude::*;
@@ -62,10 +62,9 @@ impl<T: Clone + fmt::Debug + Send + 'static> IntoProtocolsHandler for KademliaHa
     }
 
     fn inbound_protocol(&self) -> <Self::Handler as ProtocolsHandler>::InboundProtocol {
-        if self.config.allow_listening {
-            upgrade::EitherUpgrade::A(self.config.protocol_config.clone())
-        } else {
-            upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade)
+        match self.config.mode {
+            Mode::Client => upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade),
+            Mode::Server => upgrade::EitherUpgrade::A(self.config.protocol_config.clone()),
         }
     }
 }
@@ -118,11 +117,11 @@ 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,
+
+    /// [`Mode`] the handler is operating in.
+    pub mode: Mode,
 }
 
 /// State of an active substream, opened either by us or by the remote.
@@ -480,11 +479,14 @@ where
     type InboundOpenInfo = ();
 
     fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol, Self::InboundOpenInfo> {
-        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), ())
+        match self.config.mode {
+            Mode::Server => SubstreamProtocol::new(self.config.protocol_config.clone(), ())
+                .map_upgrade(upgrade::EitherUpgrade::A),
+            // If we are in client mode, I don't want to advertise Kademlia so that other peers will
+            // not send me Kademlia requests.
+            Mode::Client => {
+                SubstreamProtocol::new(upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade), ())
+            }
         }
     }
 
@@ -508,24 +510,17 @@ where
         protocol: <Self::InboundProtocol as InboundUpgrade<NegotiatedSubstream>>::Output,
         (): Self::InboundOpenInfo,
     ) {
-        // If `self.allow_listening` is false, then we produced a `DeniedUpgrade` and `protocol`
-        // is a `Void`.
         let protocol = match protocol {
             EitherOutput::First(p) => p,
             EitherOutput::Second(p) => void::unreachable(p),
         };
 
-        debug_assert!(self.config.allow_listening);
         let connec_unique_id = self.next_connec_unique_id;
         self.next_connec_unique_id.0 += 1;
         self.substreams
             .push(SubstreamState::InWaitingMessage(connec_unique_id, protocol));
-        if let ProtocolStatus::Unconfirmed = self.protocol_status {
-            // Upon the first successfully negotiated substream, we know that the
-            // remote is configured with the same protocol name and we want
-            // the behaviour to add this peer to the routing table, if possible.
-            self.protocol_status = ProtocolStatus::Confirmed;
-        }
+        // Just because another peer is sending us Kademlia requests, doesn't necessarily
+        // mean that it will answer the Kademlia requests that we send to it.
     }
 
     fn inject_event(&mut self, message: KademliaHandlerIn<TUserData>) {
@@ -761,8 +756,8 @@ impl Default for KademliaHandlerConfig {
     fn default() -> Self {
         KademliaHandlerConfig {
             protocol_config: Default::default(),
-            allow_listening: true,
             idle_timeout: Duration::from_secs(10),
+            mode: Mode::default(),
         }
     }
 }

protocols/kad/src/protocol.rs

@@ -46,6 +46,19 @@ pub const DEFAULT_PROTO_NAME: &[u8] = b"/ipfs/kad/1.0.0";
 /// The default maximum size for a varint length-delimited packet.
 pub const DEFAULT_MAX_PACKET_SIZE: usize = 16 * 1024;
 
+#[derive(Debug, Clone)]
+/// See [`crate::KademliaConfig::set_mode`].
+pub enum Mode {
+    Client,
+    Server,
+}
+
+impl Default for Mode {
+    fn default() -> Self {
+        Mode::Server
+    }
+}
+
 /// Status of our connection to a node reported by the Kademlia protocol.
 #[derive(Copy, Clone, PartialEq, Eq, Debug, Hash)]
 pub enum KadConnectionType {