feat(relay): allow enabling and disabling of relay HOP advertisement

protocols/relay/CHANGELOG.md

@@ -1,5 +1,9 @@
 ## 0.21.1
 
+- Automatically configure HOP protocol advertisement based on external addresses, with the ability to override this
+  functionality using `Behaviour::set_status` to explicitly set `Status::{Enable,Disable}` to enable or disable
+  protocol advertisement.
+  See [PR 6154](https://github.com/libp2p/rust-libp2p/pull/6154).
 - reduce allocations by replacing `get_or_insert` with `get_or_insert_with`
   See [PR 6136](https://github.com/libp2p/rust-libp2p/pull/6136)
 

protocols/relay/src/behaviour.rs

@@ -26,7 +26,7 @@ use std::{
     collections::{hash_map, HashMap, HashSet, VecDeque},
     num::NonZeroU32,
     ops::Add,
-    task::{Context, Poll},
+    task::{Context, Poll, Waker},
     time::Duration,
 };
 
@@ -35,6 +35,7 @@ use libp2p_core::{multiaddr::Protocol, transport::PortUse, ConnectedPoint, Endpo
 use libp2p_identity::PeerId;
 use libp2p_swarm::{
     behaviour::{ConnectionClosed, FromSwarm},
+    derive_prelude::ConnectionEstablished,
     dummy, ConnectionDenied, ConnectionId, ExternalAddresses, NetworkBehaviour, NotifyHandler,
     THandler, THandlerInEvent, THandlerOutEvent, ToSwarm,
 };
@@ -253,27 +254,132 @@ pub struct Behaviour {
 
     local_peer_id: PeerId,
 
+    connections: HashMap<PeerId, HashSet<ConnectionId>>,
     reservations: HashMap<PeerId, HashSet<ConnectionId>>,
     circuits: CircuitsTracker,
 
     /// Queue of actions to return when polled.
     queued_actions: VecDeque<ToSwarm<Event, THandlerInEvent<Self>>>,
 
     external_addresses: ExternalAddresses,
+
+    status: Status,
+
+    auto_status_change: bool,
+
+    waker: Option<Waker>,
+}
+
+#[derive(PartialEq, Copy, Clone, Debug)]
+pub enum Status {
+    /// Enables advertisement of the HOP protocol
+    Enable,
+
+    /// Disables advertisement of the HOP protocol
+    Disable,
 }
 
 impl Behaviour {
     pub fn new(local_peer_id: PeerId, config: Config) -> Self {
         Self {
             config,
             local_peer_id,
+            connections: Default::default(),
             reservations: Default::default(),
             circuits: Default::default(),
             queued_actions: Default::default(),
             external_addresses: Default::default(),
+            status: Status::Disable,
+            auto_status_change: true,
+            waker: None,
+        }
+    }
+
+    pub fn set_status(&mut self, status: Option<Status>) {
+        match status {
+            Some(status) => {
+                self.auto_status_change = false;
+                if self.status != status {
+                    self.status = status;
+                    self.reconfigure_relay_status();
+                }
+            }
+            None => {
+                self.auto_status_change = true;
+                self.determine_relay_status_from_external_address();
+            }
+        }
+        if let Some(waker) = self.waker.take() {
+            waker.wake();
         }
     }
 
+    fn reconfigure_relay_status(&mut self) {
+        if self.connections.is_empty() {
+            return;
+        }
+
+        for (peer_id, connections) in self.connections.iter() {
+            self.queued_actions
+                .extend(connections.iter().map(|id| ToSwarm::NotifyHandler {
+                    peer_id: *peer_id,
+                    handler: NotifyHandler::One(*id),
+                    event: Either::Left(handler::In::SetStatus {
+                        status: self.status,
+                    }),
+                }));
+        }
+    }
+
+    fn determine_relay_status_from_external_address(&mut self) {
+        let old = self.status;
+
+        self.status = match (self.external_addresses.as_slice(), self.status) {
+            ([], Status::Enable) => {
+                tracing::debug!("disabling protocol advertisment because we no longer have any confirmed external addresses");
+                Status::Disable
+            }
+            ([], Status::Disable) => {
+                // Previously disabled because of no external addresses.
+                Status::Disable
+            }
+            (confirmed_external_addresses, Status::Disable) => {
+                debug_assert!(
+                    !confirmed_external_addresses.is_empty(),
+                    "Previous match arm handled empty list"
+                );
+                tracing::debug!("advertising protcol because we are now externally reachable");
+                Status::Enable
+            }
+            (confirmed_external_addresses, Status::Enable) => {
+                debug_assert!(
+                    !confirmed_external_addresses.is_empty(),
+                    "Previous match arm handled empty list"
+                );
+
+                Status::Enable
+            }
+        };
+
+        if self.status != old {
+            self.reconfigure_relay_status();
+        }
+    }
+
+    fn on_connection_established(
+        &mut self,
+        ConnectionEstablished {
+            peer_id,
+            connection_id,
+            ..
+        }: ConnectionEstablished,
+    ) {
+        self.connections
+            .entry(peer_id)
+            .or_default()
+            .insert(connection_id);
+    }
+
     fn on_connection_closed(
         &mut self,
         ConnectionClosed {
@@ -294,6 +400,13 @@ impl Behaviour {
             }
         }
 
+        if let hash_map::Entry::Occupied(mut peer) = self.connections.entry(peer_id) {
+            peer.get_mut().remove(&connection_id);
+            if peer.get().is_empty() {
+                peer.remove();
+            }
+        }
+
         for circuit in self
             .circuits
             .remove_by_connection(peer_id, connection_id)
@@ -337,6 +450,7 @@ impl NetworkBehaviour for Behaviour {
                 local_addr: local_addr.clone(),
                 send_back_addr: remote_addr.clone(),
             },
+            self.status,
         )))
     }
 
@@ -364,14 +478,25 @@ impl NetworkBehaviour for Behaviour {
                 role_override,
                 port_use,
             },
+            self.status,
         )))
     }
 
     fn on_swarm_event(&mut self, event: FromSwarm) {
-        self.external_addresses.on_swarm_event(&event);
+        let changed = self.external_addresses.on_swarm_event(&event);
 
-        if let FromSwarm::ConnectionClosed(connection_closed) = event {
-            self.on_connection_closed(connection_closed)
+        if self.auto_status_change && changed {
+            self.determine_relay_status_from_external_address();
+        }
+
+        match event {
+            FromSwarm::ConnectionEstablished(connection_established) => {
+                self.on_connection_established(connection_established)
+            }
+            FromSwarm::ConnectionClosed(connection_closed) => {
+                self.on_connection_closed(connection_closed)
+            }
+            _ => {}
         }
     }
 
@@ -718,11 +843,16 @@ impl NetworkBehaviour for Behaviour {
     }
 
     #[tracing::instrument(level = "trace", name = "NetworkBehaviour::poll", skip(self))]
-    fn poll(&mut self, _: &mut Context<'_>) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
+    fn poll(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
         if let Some(to_swarm) = self.queued_actions.pop_front() {
             return Poll::Ready(to_swarm);
         }
 
+        self.waker = Some(cx.waker().clone());
+
         Poll::Pending
     }
 }

protocols/relay/src/behaviour/handler.rs

@@ -33,7 +33,10 @@ use futures::{
     stream::{FuturesUnordered, StreamExt},
 };
 use futures_timer::Delay;
-use libp2p_core::{upgrade::ReadyUpgrade, ConnectedPoint, Multiaddr};
+use libp2p_core::{
+    upgrade::{DeniedUpgrade, ReadyUpgrade},
+    ConnectedPoint, Multiaddr,
+};
 use libp2p_identity::PeerId;
 use libp2p_swarm::{
     handler::{ConnectionEvent, DialUpgradeError, FullyNegotiatedInbound, FullyNegotiatedOutbound},
@@ -43,7 +46,7 @@ use libp2p_swarm::{
 use web_time::Instant;
 
 use crate::{
-    behaviour::CircuitId,
+    behaviour::{self, CircuitId},
     copy_future::CopyFuture,
     proto,
     protocol::{inbound_hop, outbound_stop},
@@ -87,6 +90,9 @@ pub enum In {
         dst_stream: Stream,
         dst_pending_data: Bytes,
     },
+    SetStatus {
+        status: behaviour::Status,
+    },
 }
 
 impl fmt::Debug for In {
@@ -137,6 +143,10 @@ impl fmt::Debug for In {
                 .field("circuit_id", circuit_id)
                 .field("dst_peer_id", dst_peer_id)
                 .finish(),
+            In::SetStatus { status } => f
+                .debug_struct("In::SetStatus")
+                .field("status", status)
+                .finish(),
         }
     }
 }
@@ -385,10 +395,12 @@ pub struct Handler {
         CircuitId,
         Result<outbound_stop::Circuit, outbound_stop::Error>,
     >,
+
+    status: behaviour::Status,
 }
 
 impl Handler {
-    pub fn new(config: Config, endpoint: ConnectedPoint) -> Handler {
+    pub fn new(config: Config, endpoint: ConnectedPoint, status: behaviour::Status) -> Handler {
         Handler {
             inbound_workers: futures_bounded::FuturesSet::new(
                 STREAM_TIMEOUT,
@@ -409,6 +421,7 @@ impl Handler {
             active_reservation: Default::default(),
             pending_connect_requests: Default::default(),
             active_connect_requests: Default::default(),
+            status,
         }
     }
 
@@ -496,13 +509,18 @@ type Futures<T> = FuturesUnordered<BoxFuture<'static, T>>;
 impl ConnectionHandler for Handler {
     type FromBehaviour = In;
     type ToBehaviour = Event;
-    type InboundProtocol = ReadyUpgrade<StreamProtocol>;
+    type InboundProtocol = Either<ReadyUpgrade<StreamProtocol>, DeniedUpgrade>;
     type InboundOpenInfo = ();
     type OutboundProtocol = ReadyUpgrade<StreamProtocol>;
     type OutboundOpenInfo = ();
 
     fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol> {
-        SubstreamProtocol::new(ReadyUpgrade::new(HOP_PROTOCOL_NAME), ())
+        match self.status {
+            behaviour::Status::Enable => {
+                SubstreamProtocol::new(Either::Left(ReadyUpgrade::new(HOP_PROTOCOL_NAME)), ())
+            }
+            behaviour::Status::Disable => SubstreamProtocol::new(Either::Right(DeniedUpgrade), ()),
+        }
     }
 
     fn on_behaviour_event(&mut self, event: Self::FromBehaviour) {
@@ -594,6 +612,7 @@ impl ConnectionHandler for Handler {
                         .boxed(),
                 );
             }
+            In::SetStatus { status } => self.status = status,
         }
     }
 
@@ -890,7 +909,7 @@ impl ConnectionHandler for Handler {
     ) {
         match event {
             ConnectionEvent::FullyNegotiatedInbound(FullyNegotiatedInbound {
-                protocol: stream,
+                protocol: futures::future::Either::Left(stream),
                 ..
             }) => {
                 self.on_fully_negotiated_inbound(stream);

protocols/relay/src/lib.rs

@@ -39,7 +39,9 @@ mod proto {
     };
 }
 
-pub use behaviour::{rate_limiter::RateLimiter, Behaviour, CircuitId, Config, Event, StatusCode};
+pub use behaviour::{
+    rate_limiter::RateLimiter, Behaviour, CircuitId, Config, Event, Status, StatusCode,
+};
 pub use protocol::{HOP_PROTOCOL_NAME, STOP_PROTOCOL_NAME};
 
 /// Types related to the relay protocol inbound.