kad: Add global_addresses_only option

src/protocol/libp2p/kademlia/config.rs

@@ -23,7 +23,7 @@ use crate::{
     protocol::libp2p::kademlia::{
         handle::{
             IncomingRecordValidationMode, KademliaCommand, KademliaEvent, KademliaHandle,
-            RoutingTableUpdateMode,
+            RemoteAddressMode, RoutingTableUpdateMode,
         },
         store::MemoryStoreConfig,
     },
@@ -102,6 +102,9 @@ pub struct Config {
     /// Default record TTL.
     pub(super) record_ttl: Duration,
 
+    /// Whether to accept non-global addresses from remote peers.
+    pub(super) remote_address_mode: RemoteAddressMode,
+
     /// Provider record TTL.
     pub(super) memory_store_config: MemoryStoreConfig,
 
@@ -134,6 +137,7 @@ impl Config {
         mode: KademliaMode,
         validation_mode: IncomingRecordValidationMode,
         record_ttl: Duration,
+        remote_address_mode: RemoteAddressMode,
         memory_store_config: MemoryStoreConfig,
         max_message_size: usize,
     ) -> (Self, KademliaHandle) {
@@ -153,6 +157,7 @@ impl Config {
                 mode,
                 validation_mode,
                 record_ttl,
+                remote_address_mode,
                 memory_store_config,
                 codec: ProtocolCodec::UnsignedVarint(Some(max_message_size)),
                 replication_factor,
@@ -175,6 +180,7 @@ impl Config {
             KademliaMode::Server,
             IncomingRecordValidationMode::Automatic,
             DEFAULT_TTL,
+            RemoteAddressMode::All,
             Default::default(),
             DEFAULT_MAX_MESSAGE_SIZE,
         )
@@ -205,6 +211,9 @@ pub struct ConfigBuilder {
     /// Default TTL for the records.
     pub(super) record_ttl: Duration,
 
+    /// Whether to accept non-global addresses from remote peers.
+    pub(super) remote_address_mode: RemoteAddressMode,
+
     /// Memory store configuration.
     pub(super) memory_store_config: MemoryStoreConfig,
 
@@ -229,6 +238,7 @@ impl ConfigBuilder {
             mode: KademliaMode::Server,
             validation_mode: IncomingRecordValidationMode::Automatic,
             record_ttl: DEFAULT_TTL,
+            remote_address_mode: RemoteAddressMode::All,
             memory_store_config: Default::default(),
             max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
         }
@@ -294,6 +304,14 @@ impl ConfigBuilder {
         self
     }
 
+    /// Set whether non-global addresses received from remote peers are accepted.
+    ///
+    /// If unspecified, all addresses are accepted ([`RemoteAddressMode::All`]).
+    pub fn with_remote_address_mode(mut self, mode: RemoteAddressMode) -> Self {
+        self.remote_address_mode = mode;
+        self
+    }
+
     /// Set maximum number of records in the memory store.
     ///
     /// If unspecified, the default maximum number of records is 1024.
@@ -372,6 +390,7 @@ impl ConfigBuilder {
             self.mode,
             self.validation_mode,
             self.record_ttl,
+            self.remote_address_mode,
             self.memory_store_config,
             self.max_message_size,
         )

src/protocol/libp2p/kademlia/handle.rs

@@ -76,6 +76,20 @@ pub enum IncomingRecordValidationMode {
     Automatic,
 }
 
+/// Which addresses received from remote peers are accepted.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub enum RemoteAddressMode {
+    /// Accept all addresses.
+    All,
+
+    /// Drop non-global IP addresses of peers and providers received from remote peers. Closer
+    /// peers left without addresses are not used as query candidates. DNS addresses are kept.
+    ///
+    /// Locally supplied addresses (known peers, [`KademliaHandle::add_known_peer()`], addresses
+    /// of connections we dialed) are not affected.
+    GlobalOnly,
+}
+
 /// Kademlia commands.
 #[derive(Debug)]
 #[cfg_attr(feature = "fuzz", derive(serde::Serialize, serde::Deserialize))]

src/protocol/libp2p/kademlia/message.rs

@@ -357,6 +357,33 @@ impl KademliaMessage {
             }
         }
     }
+
+    /// Remove non-global addresses of the peers carried in the message.
+    ///
+    /// Closer peers left without addresses are dropped, as they are unreachable. Providers are
+    /// kept: the provider identity is the payload and can be resolved with `FIND_NODE`.
+    pub fn retain_global_addresses(&mut self) {
+        fn retain_closer_peers(peers: &mut Vec<KademliaPeer>) {
+            peers.retain_mut(|peer| {
+                peer.retain_global_addresses();
+                !peer.address_store.is_empty()
+            });
+        }
+
+        match self {
+            Self::FindNode { peers, .. } | Self::GetRecord { peers, .. } =>
+                retain_closer_peers(peers),
+            Self::GetProviders {
+                peers, providers, ..
+            } => {
+                retain_closer_peers(peers);
+                providers.iter_mut().for_each(KademliaPeer::retain_global_addresses);
+            }
+            Self::AddProvider { providers, .. } =>
+                providers.iter_mut().for_each(KademliaPeer::retain_global_addresses),
+            Self::PutValue { .. } => {}
+        }
+    }
 }
 
 fn record_to_schema(record: Record) -> schema::kademlia::Record {
@@ -399,6 +426,7 @@ fn record_from_schema(record: schema::kademlia::Record) -> Option<Record> {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use multiaddr::Multiaddr;
 
     #[test]
     fn non_empty_publisher_and_ttl_are_preserved() {
@@ -436,4 +464,54 @@ mod tests {
 
         assert_eq!(got_record, record);
     }
+
+    #[test]
+    fn retain_global_addresses() {
+        let peer = |addresses: &[&str]| {
+            KademliaPeer::new(
+                PeerId::random(),
+                addresses.iter().map(|address| address.parse().unwrap()).collect(),
+                ConnectionType::NotConnected,
+            )
+        };
+        let private_only = peer(&["/ip4/10.0.0.1/tcp/30333", "/ip6/fd00::1/tcp/30333"]);
+        let mixed = peer(&["/ip4/192.168.1.1/tcp/30333", "/ip4/8.8.8.8/tcp/30333"]);
+        let dns = peer(&["/dns/example.com/tcp/30333"]);
+        let no_addresses = peer(&[]);
+
+        let mut message = KademliaMessage::GetProviders {
+            key: None,
+            peers: vec![
+                private_only.clone(),
+                mixed.clone(),
+                dns.clone(),
+                no_addresses.clone(),
+            ],
+            providers: vec![private_only.clone(), mixed.clone()],
+        };
+        message.retain_global_addresses();
+
+        let KademliaMessage::GetProviders {
+            peers, providers, ..
+        } = message
+        else {
+            panic!("message type changed");
+        };
+
+        // Closer peers left without addresses are dropped.
+        assert_eq!(
+            peers.iter().map(|peer| peer.peer).collect::<Vec<_>>(),
+            vec![mixed.peer, dns.peer],
+        );
+        assert_eq!(
+            peers[0].addresses(),
+            vec!["/ip4/8.8.8.8/tcp/30333".parse::<Multiaddr>().unwrap()],
+        );
+        assert_eq!(peers[1].addresses(), dns.addresses());
+
+        // Providers are kept even if left without addresses.
+        assert_eq!(providers[0].peer, private_only.peer);
+        assert!(providers[0].addresses().is_empty());
+        assert_eq!(providers[1].addresses(), peers[0].addresses());
+    }
 }

src/protocol/libp2p/kademlia/mod.rs

@@ -57,7 +57,7 @@ use std::{
 pub use config::{Config, ConfigBuilder, KademliaMode};
 pub use handle::{
     IncomingRecordValidationMode, KademliaCommand, KademliaEvent, KademliaHandle, Quorum,
-    RoutingTableUpdateMode,
+    RemoteAddressMode, RoutingTableUpdateMode,
 };
 pub use query::QueryId;
 pub use record::{ContentProvider, Key as RecordKey, PeerRecord, Record};
@@ -178,6 +178,9 @@ pub(crate) struct Kademlia {
     /// Default record TTL.
     record_ttl: Duration,
 
+    /// Remote address mode.
+    remote_address_mode: RemoteAddressMode,
+
     /// Query engine.
     engine: QueryEngine,
 
@@ -217,6 +220,7 @@ impl Kademlia {
             mode: config.mode,
             validation_mode: config.validation_mode,
             record_ttl: config.record_ttl,
+            remote_address_mode: config.remote_address_mode,
             replication_factor: config.replication_factor,
             engine: QueryEngine::new(local_peer_id, config.replication_factor, PARALLELISM_FACTOR),
         }
@@ -485,9 +489,13 @@ impl Kademlia {
     ) -> crate::Result<()> {
         tracing::trace!(target: LOG_TARGET, ?peer, query = ?query_id, "handle message from peer");
 
-        match KademliaMessage::from_bytes(message, self.replication_factor)
-            .ok_or(Error::InvalidData)?
-        {
+        let mut message = KademliaMessage::from_bytes(message, self.replication_factor)
+            .ok_or(Error::InvalidData)?;
+        if self.remote_address_mode == RemoteAddressMode::GlobalOnly {
+            message.retain_global_addresses();
+        }
+
+        match message {
             KademliaMessage::FindNode { target, peers } => {
                 match query_id {
                     Some(query_id) => {
@@ -1466,10 +1474,13 @@ mod tests {
     }
 
     fn make_kademlia() -> (Kademlia, Context, TransportManager) {
-        make_kademlia_with_mode(KademliaMode::Server)
+        make_kademlia_with(KademliaMode::Server, RemoteAddressMode::All)
     }
 
-    fn make_kademlia_with_mode(mode: KademliaMode) -> (Kademlia, Context, TransportManager) {
+    fn make_kademlia_with(
+        mode: KademliaMode,
+        remote_address_mode: RemoteAddressMode,
+    ) -> (Kademlia, Context, TransportManager) {
         let manager = TransportManagerBuilder::new().build();
 
         let peer = PeerId::random();
@@ -1495,6 +1506,7 @@ mod tests {
             mode,
             validation_mode: IncomingRecordValidationMode::Automatic,
             record_ttl: Duration::from_secs(36 * 60 * 60),
+            remote_address_mode,
             memory_store_config: Default::default(),
             event_tx,
             cmd_rx,
@@ -1684,7 +1696,8 @@ mod tests {
 
     #[tokio::test]
     async fn client_mode_ignores_inbound_substreams() {
-        let (mut kademlia, _context, _manager) = make_kademlia_with_mode(KademliaMode::Client);
+        let (mut kademlia, _context, _manager) =
+            make_kademlia_with(KademliaMode::Client, RemoteAddressMode::All);
 
         let peer = PeerId::random();
         let mut substream = MockSubstream::new();
@@ -1867,4 +1880,91 @@ mod tests {
             )
             .is_none());
     }
+
+    #[tokio::test]
+    async fn non_global_addresses_dropped_from_responses() {
+        let (mut kademlia, mut context, _manager) =
+            make_kademlia_with(KademliaMode::Server, RemoteAddressMode::GlobalOnly);
+
+        let private_peer = PeerId::random();
+        let private = KademliaPeer::new(
+            private_peer,
+            vec![Multiaddr::from_str("/ip4/10.0.0.1/tcp/30333").unwrap()],
+            ConnectionType::NotConnected,
+        );
+        let public_peer = PeerId::random();
+        let public = KademliaPeer::new(
+            public_peer,
+            vec![
+                Multiaddr::from_str("/ip4/10.0.0.2/tcp/30333").unwrap(),
+                Multiaddr::from_str("/ip4/8.8.8.8/tcp/30333").unwrap(),
+            ],
+            ConnectionType::NotConnected,
+        );
+
+        let responder = PeerId::random();
+        let mut substream = MockSubstream::new();
+        substream.expect_poll_close().times(1).return_once(|_| Poll::Ready(Ok(())));
+        let message =
+            KademliaMessage::find_node_response(PeerId::random().to_bytes(), vec![private, public]);
+        kademlia
+            .on_message_received(
+                responder,
+                Some(QueryId(0)),
+                BytesMut::from(&message[..]),
+                Substream::new_mock(responder, SubstreamId::from(0usize), Box::new(substream)),
+            )
+            .await
+            .unwrap();
+
+        // The peer left without addresses is not discovered.
+        match context.event_rx.try_recv() {
+            Ok(KademliaEvent::PeersDiscovered { peers }) => assert_eq!(peers, vec![public_peer]),
+            event => panic!("unexpected event: {event:?}"),
+        }
+    }
+
+    #[tokio::test]
+    async fn non_global_addresses_dropped_from_incoming_provider() {
+        let (mut kademlia, mut context, _manager) =
+            make_kademlia_with(KademliaMode::Server, RemoteAddressMode::GlobalOnly);
+
+        let key = RecordKey::from(vec![1, 2, 3]);
+        let provider = PeerId::random();
+        let message = KademliaMessage::add_provider(
+            key.clone(),
+            ContentProvider {
+                peer: provider,
+                addresses: vec![
+                    Multiaddr::from_str("/ip4/10.0.0.1/tcp/30333").unwrap(),
+                    Multiaddr::from_str("/ip4/8.8.8.8/tcp/30333").unwrap(),
+                ],
+            },
+        );
+        kademlia
+            .on_message_received(
+                provider,
+                None,
+                BytesMut::from(&message[..]),
+                Substream::new_mock(
+                    provider,
+                    SubstreamId::from(0usize),
+                    Box::new(MockSubstream::new()),
+                ),
+            )
+            .await
+            .unwrap();
+
+        let global = vec![Multiaddr::from_str("/ip4/8.8.8.8/tcp/30333").unwrap()];
+        let providers = kademlia.store.get_providers(&key);
+        assert_eq!(providers.len(), 1);
+        assert_eq!(providers[0].addresses, global);
+        match context.event_rx.try_recv() {
+            Ok(KademliaEvent::IncomingProvider { provider: p, .. }) => {
+                assert_eq!(p.peer, provider);
+                assert_eq!(p.addresses, global);
+            }
+            event => panic!("unexpected event: {event:?}"),
+        }
+    }
 }

src/protocol/libp2p/kademlia/types.rs

@@ -26,7 +26,7 @@
 
 use crate::{
     protocol::libp2p::kademlia::schema,
-    transport::manager::address::{AddressRecord, AddressStore},
+    transport::manager::address::{is_global_multiaddr, AddressRecord, AddressStore},
     PeerId,
 };
 
@@ -288,6 +288,11 @@ impl KademliaPeer {
     pub fn addresses(&self) -> Vec<Multiaddr> {
         self.address_store.addresses(MAX_ADDRESSES)
     }
+
+    /// Remove non-global addresses of the peer.
+    pub(super) fn retain_global_addresses(&mut self) {
+        self.address_store.retain(is_global_multiaddr);
+    }
 }
 
 impl TryFrom<&schema::kademlia::Peer> for KademliaPeer {