Add Mode enum for kademlia client/server mode

libp2p · Aug 10, 2021 · 02a0a74 · 02a0a74
1 parent 1a042f8
commit 02a0a74
Show file tree

Hide file tree

Showing 4 changed files with 39 additions and 20 deletions.
diff --git a/protocols/kad/src/behaviour.rs b/protocols/kad/src/behaviour.rs
@@ -33,7 +33,7 @@ use crate::handler::{
 };
 use crate::jobs::*;
 use crate::kbucket::{self, Distance, KBucketsTable, NodeStatus};
-use crate::protocol::{KademliaProtocolConfig, KadConnectionType, KadPeer};
+use crate::protocol::{Mode, KademliaProtocolConfig, KadConnectionType, KadPeer};
 use crate::query::{Query, QueryId, QueryPool, QueryConfig, QueryPoolState};
 use crate::record::{self, store::{self, RecordStore}, Record, ProviderRecord};
 use fnv::{FnvHashMap, FnvHashSet};
@@ -148,7 +148,7 @@ pub struct KademliaConfig {
     kbucket_inserts: KademliaBucketInserts,
     caching: KademliaCaching,
     // Set the peer mode.
-    client: bool,
+    client: Mode,
 }
 
 /// The configuration for Kademlia "write-back" caching after successful
@@ -183,7 +183,7 @@ impl Default for KademliaConfig {
             kbucket_inserts: KademliaBucketInserts::OnConnected,
             caching: KademliaCaching::Enabled { max_peers: 1 },
             // By default, a peer will not be in client mode.
-            client: false,
+            client: Mode::default(),
         }
     }
 }
@@ -361,12 +361,12 @@ impl KademliaConfig {
     ///
     /// Client mode is disabled by default.
     ///
-    /// In `client` mode, the peer would not be added to any other peer's 
+    /// In `client` mode, the peer would not be added to any other peers 
     /// routing table. The client peer can send queries over the Kademlia
-    /// network, but will not receive any queries from the other peer's in
+    /// network, but will not receive any queries from the other peers in
     /// the network.
-    pub fn set_client_mode(&mut self, t: bool) -> &mut Self {
-        self.client = t;
+    pub fn set_mode(&mut self, mode: Mode) -> &mut Self {
+        self.client = mode;
         self
     }
 }

diff --git a/protocols/kad/src/dht.proto b/protocols/kad/src/dht.proto
@@ -26,9 +26,6 @@ message Record {
     // The remaining TTL of the record, in seconds.
     // Currently specific to rust-libp2p.
     uint32 ttl = 777;
-
-	// Client mode or server mode.
-	bool client = false;
 };
 
 message Message {

diff --git a/protocols/kad/src/handler.rs b/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::*;
@@ -63,10 +63,15 @@ impl<T: Clone + Send + 'static> IntoProtocolsHandler for KademliaHandlerProto<T>
     }
 
     fn inbound_protocol(&self) -> <Self::Handler as ProtocolsHandler>::InboundProtocol {
-        if self.config.allow_listening || !self.config.client {
-            upgrade::EitherUpgrade::A(self.config.protocol_config.clone())
-        } else {
-            upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade)
+        match self.config.client {
+            Mode::Client => upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade),
+            Mode::Server => {
+                if self.config.allow_listening {
+                    upgrade::EitherUpgrade::A(self.config.protocol_config.clone())
+                } else {
+                    upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade)
+                }
+            }
         }
     }
 }
@@ -126,7 +131,7 @@ pub struct KademliaHandlerConfig {
     pub idle_timeout: Duration,
 
     // If true, node will act in `client` mode in the Kademlia network.
-    pub client: bool,
+    pub client: Mode,
 }
 
 /// State of an active substream, opened either by us or by the remote.
@@ -472,10 +477,15 @@ where
     type InboundOpenInfo = ();
 
     fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol, Self::InboundOpenInfo> {
-        if self.config.allow_listening || !self.config.client {
-            SubstreamProtocol::new(self.config.protocol_config.clone(), ()).map_upgrade(upgrade::EitherUpgrade::A)
-        } else {
-            SubstreamProtocol::new(upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade), ())
+        match self.config.client {
+            Mode::Server => {
+                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), ())
+                }
+            },
+            Mode::Client => SubstreamProtocol::new(upgrade::EitherUpgrade::B(upgrade::DeniedUpgrade), ()),
         }
     }
 

diff --git a/protocols/kad/src/protocol.rs b/protocols/kad/src/protocol.rs
@@ -46,6 +46,18 @@ 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)]
+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 {