feat(swarm)!: Allow NetworkBehaviours to manage incoming connections

protocols/autonat/src/behaviour.rs

@@ -36,9 +36,8 @@ use libp2p_request_response::{
     handler::RequestResponseHandlerEvent, ProtocolSupport, RequestId, RequestResponse,
     RequestResponseConfig, RequestResponseEvent, RequestResponseMessage, ResponseChannel,
 };
-use libp2p_swarm::{
-    DialError, IntoConnectionHandler, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
-};
+use libp2p_swarm::behaviour::THandlerInEvent;
+use libp2p_swarm::{DialError, NetworkBehaviour, NetworkBehaviourAction, PollParameters};
 use std::{
     collections::{HashMap, VecDeque},
     iter,
@@ -356,7 +355,7 @@ impl NetworkBehaviour for Behaviour {
         peer: &PeerId,
         conn: &ConnectionId,
         endpoint: &ConnectedPoint,
-        handler: <Self::ConnectionHandler as IntoConnectionHandler>::Handler,
+        handler: Self::ConnectionHandler,
         remaining_established: usize,
     ) {
         self.inner
@@ -369,14 +368,9 @@ impl NetworkBehaviour for Behaviour {
         }
     }
 
-    fn inject_dial_failure(
-        &mut self,
-        peer: Option<PeerId>,
-        handler: Self::ConnectionHandler,
-        error: &DialError,
-    ) {
-        self.inner.inject_dial_failure(peer, handler, error);
-        if let Some(event) = self.as_server().on_outbound_dial_error(peer, error) {
+    fn inject_dial_failure(&mut self, peer_id: Option<PeerId>, error: &DialError) {
+        self.inner.inject_dial_failure(peer_id, error);
+        if let Some(event) = self.as_server().on_outbound_dial_error(peer_id, error) {
             self.pending_out_events
                 .push_back(Event::InboundProbe(event));
         }
@@ -470,8 +464,12 @@ impl NetworkBehaviour for Behaviour {
         }
     }
 
-    fn new_handler(&mut self) -> Self::ConnectionHandler {
-        self.inner.new_handler()
+    fn new_handler(
+        &mut self,
+        peer: &PeerId,
+        connected_point: &ConnectedPoint,
+    ) -> Self::ConnectionHandler {
+        self.inner.new_handler(peer, connected_point)
     }
 
     fn addresses_of_peer(&mut self, peer: &PeerId) -> Vec<Multiaddr> {
@@ -487,14 +485,8 @@ impl NetworkBehaviour for Behaviour {
         self.inner.inject_event(peer_id, conn, event)
     }
 
-    fn inject_listen_failure(
-        &mut self,
-        local_addr: &Multiaddr,
-        send_back_addr: &Multiaddr,
-        handler: Self::ConnectionHandler,
-    ) {
-        self.inner
-            .inject_listen_failure(local_addr, send_back_addr, handler)
+    fn inject_listen_failure(&mut self, local_addr: &Multiaddr, send_back_addr: &Multiaddr) {
+        self.inner.inject_listen_failure(local_addr, send_back_addr)
     }
 
     fn inject_new_listener(&mut self, id: ListenerId) {
@@ -512,7 +504,7 @@ impl NetworkBehaviour for Behaviour {
 
 type Action = NetworkBehaviourAction<
     <Behaviour as NetworkBehaviour>::OutEvent,
-    <Behaviour as NetworkBehaviour>::ConnectionHandler,
+    THandlerInEvent<<Behaviour as NetworkBehaviour>::ConnectionHandler>,
 >;
 
 // Trait implemented for `AsClient` as `AsServer` to handle events from the inner [`RequestResponse`] Protocol.

protocols/autonat/src/behaviour/as_server.rs

@@ -30,7 +30,7 @@ use libp2p_request_response::{
 };
 use libp2p_swarm::{
     dial_opts::{DialOpts, PeerCondition},
-    DialError, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
+    DialError, NetworkBehaviourAction, PollParameters,
 };
 use std::{
     collections::{HashMap, HashSet, VecDeque},
@@ -137,7 +137,6 @@ impl<'a> HandleInnerEvent for AsServer<'a> {
                                 .override_dial_concurrency_factor(NonZeroU8::new(1).expect("1 > 0"))
                                 .addresses(addrs)
                                 .build(),
-                            handler: self.inner.new_handler(),
                         });
                     }
                     Err((status_text, error)) => {

protocols/dcutr/src/behaviour.rs

@@ -26,11 +26,13 @@ use either::Either;
 use libp2p_core::connection::{ConnectedPoint, ConnectionId};
 use libp2p_core::multiaddr::Protocol;
 use libp2p_core::{Multiaddr, PeerId};
+use libp2p_swarm::behaviour::THandlerInEvent;
 use libp2p_swarm::dial_opts::{self, DialOpts};
 use libp2p_swarm::{
-    ConnectionHandler, ConnectionHandlerUpgrErr, DialError, IntoConnectionHandler,
-    NetworkBehaviour, NetworkBehaviourAction, NotifyHandler, PollParameters,
+    dummy, ConnectionHandler, ConnectionHandlerUpgrErr, DialError, NetworkBehaviour,
+    NetworkBehaviourAction, NotifyHandler, PollParameters,
 };
+use std::collections::hash_map::Entry;
 use std::collections::{HashMap, HashSet, VecDeque};
 use std::task::{Context, Poll};
 use thiserror::Error;
@@ -72,23 +74,82 @@ pub struct Behaviour {
 
     /// All direct (non-relayed) connections.
     direct_connections: HashMap<PeerId, HashSet<ConnectionId>>,
+
+    awaiting_direct_inbound_connections: HashMap<PeerId, ConnectionId>,
+
+    awaiting_direct_outbound_connections: HashMap<PeerId, (ConnectionId, u8)>,
 }
 
+type Handler =
+    Either<handler::relayed::Handler, Either<handler::direct::Handler, dummy::ConnectionHandler>>;
+
 impl Behaviour {
     pub fn new() -> Self {
         Behaviour {
             queued_actions: Default::default(),
             direct_connections: Default::default(),
+            awaiting_direct_inbound_connections: Default::default(),
+            awaiting_direct_outbound_connections: Default::default(),
         }
     }
 }
 
 impl NetworkBehaviour for Behaviour {
-    type ConnectionHandler = handler::Prototype;
+    type ConnectionHandler = Handler;
     type OutEvent = Event;
 
-    fn new_handler(&mut self) -> Self::ConnectionHandler {
-        handler::Prototype::UnknownConnection
+    fn new_handler(
+        &mut self,
+        peer: &PeerId,
+        connected_point: &ConnectedPoint,
+    ) -> Self::ConnectionHandler {
+        match (
+            self.awaiting_direct_inbound_connections.entry(*peer),
+            self.awaiting_direct_outbound_connections.entry(*peer),
+        ) {
+            (Entry::Vacant(_), Entry::Occupied(occupied)) => {
+                if connected_point.is_relayed() {
+                    log::debug!(
+                        "Unexpected relay connection whilst awaiting direct outbound connection"
+                    );
+
+                    return Either::Left(handler::relayed::Handler::new(connected_point.clone()));
+                }
+
+                let (relayed_connection_id, _) = occupied.remove();
+
+                Either::Right(Either::Left(handler::direct::Handler::new(
+                    relayed_connection_id,
+                )))
+            }
+            (Entry::Occupied(occupied), Entry::Vacant(_)) => {
+                if connected_point.is_relayed() {
+                    log::debug!(
+                        "Unexpected relay connection whilst awaiting direct inbound connection"
+                    );
+
+                    return Either::Left(handler::relayed::Handler::new(connected_point.clone()));
+                }
+
+                let relayed_connection_id = occupied.remove();
+
+                Either::Right(Either::Left(handler::direct::Handler::new(
+                    relayed_connection_id,
+                )))
+            }
+            (Entry::Vacant(_), Entry::Vacant(_)) => {
+                if connected_point.is_relayed() {
+                    Either::Left(handler::relayed::Handler::new(connected_point.clone()))
+                } else {
+                    Either::Right(Either::Right(dummy::ConnectionHandler))
+                }
+            }
+            (Entry::Occupied(_), Entry::Occupied(_)) => {
+                debug_assert!(false, "Should not have a pending outbound and pending inbound connection to same peer simultaneously");
+
+                Either::Right(Either::Right(dummy::ConnectionHandler))
+            }
+        }
     }
 
     fn addresses_of_peer(&mut self, _peer_id: &PeerId) -> Vec<Multiaddr> {
@@ -140,41 +201,38 @@ impl NetworkBehaviour for Behaviour {
         }
     }
 
-    fn inject_dial_failure(
-        &mut self,
-        peer_id: Option<PeerId>,
-        handler: Self::ConnectionHandler,
-        _error: &DialError,
-    ) {
-        if let handler::Prototype::DirectConnection {
-            relayed_connection_id,
-            role: handler::Role::Initiator { attempt },
-        } = handler
-        {
-            let peer_id = peer_id.expect("Peer of `Prototype::DirectConnection` is always known.");
-            if attempt < MAX_NUMBER_OF_UPGRADE_ATTEMPTS {
-                self.queued_actions.push_back(ActionBuilder::Connect {
+    fn inject_dial_failure(&mut self, maybe_peer_id: Option<PeerId>, _error: &DialError) {
+        let peer_id = match maybe_peer_id {
+            Some(peer_id) => peer_id,
+            None => return,
+        };
+
+        let (relayed_connection_id, attempt) =
+            match self.awaiting_direct_outbound_connections.remove(&peer_id) {
+                Some((relayed_connection_id, attempt)) => (relayed_connection_id, attempt),
+                None => return,
+            };
+
+        if attempt < MAX_NUMBER_OF_UPGRADE_ATTEMPTS {
+            self.queued_actions.push_back(ActionBuilder::Connect {
+                peer_id,
+                handler: NotifyHandler::One(relayed_connection_id),
+                attempt: attempt + 1,
+            });
+        } else {
+            self.queued_actions.extend([
+                NetworkBehaviourAction::NotifyHandler {
                     peer_id,
                     handler: NotifyHandler::One(relayed_connection_id),
-                    attempt: attempt + 1,
-                });
-            } else {
-                self.queued_actions.extend([
-                    NetworkBehaviourAction::NotifyHandler {
-                        peer_id,
-                        handler: NotifyHandler::One(relayed_connection_id),
-                        event: Either::Left(
-                            handler::relayed::Command::UpgradeFinishedDontKeepAlive,
-                        ),
-                    }
-                    .into(),
-                    NetworkBehaviourAction::GenerateEvent(Event::DirectConnectionUpgradeFailed {
-                        remote_peer_id: peer_id,
-                        error: UpgradeError::Dial,
-                    })
-                    .into(),
-                ]);
-            }
+                    event: Either::Left(handler::relayed::Command::UpgradeFinishedDontKeepAlive),
+                }
+                .into(),
+                NetworkBehaviourAction::GenerateEvent(Event::DirectConnectionUpgradeFailed {
+                    remote_peer_id: peer_id,
+                    error: UpgradeError::Dial,
+                })
+                .into(),
+            ]);
         }
     }
 
@@ -183,7 +241,7 @@ impl NetworkBehaviour for Behaviour {
         peer_id: &PeerId,
         connection_id: &ConnectionId,
         connected_point: &ConnectedPoint,
-        _handler: <<Self as NetworkBehaviour>::ConnectionHandler as IntoConnectionHandler>::Handler,
+        _handler: Self::ConnectionHandler,
         _remaining_established: usize,
     ) {
         if !connected_point.is_relayed() {
@@ -205,7 +263,7 @@ impl NetworkBehaviour for Behaviour {
         &mut self,
         event_source: PeerId,
         connection: ConnectionId,
-        handler_event: <<Self::ConnectionHandler as IntoConnectionHandler>::Handler as ConnectionHandler>::OutEvent,
+        handler_event: <Self::ConnectionHandler as ConnectionHandler>::OutEvent,
     ) {
         match handler_event {
             Either::Left(handler::relayed::Event::InboundConnectRequest {
@@ -243,13 +301,11 @@ impl NetworkBehaviour for Behaviour {
                             .addresses(remote_addrs)
                             .condition(dial_opts::PeerCondition::Always)
                             .build(),
-                        handler: handler::Prototype::DirectConnection {
-                            relayed_connection_id: connection,
-                            role: handler::Role::Listener,
-                        },
                     }
                     .into(),
                 );
+                self.awaiting_direct_inbound_connections
+                    .insert(event_source, connection);
             }
             Either::Left(handler::relayed::Event::OutboundNegotiationFailed { error }) => {
                 self.queued_actions.push_back(
@@ -271,13 +327,11 @@ impl NetworkBehaviour for Behaviour {
                             .addresses(remote_addrs)
                             .override_role()
                             .build(),
-                        handler: handler::Prototype::DirectConnection {
-                            relayed_connection_id: connection,
-                            role: handler::Role::Initiator { attempt },
-                        },
                     }
                     .into(),
                 );
+                self.awaiting_direct_outbound_connections
+                    .insert(event_source, (connection, attempt));
             }
             Either::Right(Either::Left(
                 handler::direct::Event::DirectConnectionUpgradeSucceeded {
@@ -309,7 +363,8 @@ impl NetworkBehaviour for Behaviour {
         &mut self,
         _cx: &mut Context<'_>,
         poll_parameters: &mut impl PollParameters,
-    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ConnectionHandler>> {
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, THandlerInEvent<Self::ConnectionHandler>>>
+    {
         if let Some(action) = self.queued_actions.pop_front() {
             return Poll::Ready(action.build(poll_parameters));
         }
@@ -321,7 +376,7 @@ impl NetworkBehaviour for Behaviour {
 /// A [`NetworkBehaviourAction`], either complete, or still requiring data from [`PollParameters`]
 /// before being returned in [`Behaviour::poll`].
 enum ActionBuilder {
-    Done(NetworkBehaviourAction<Event, handler::Prototype>),
+    Done(NetworkBehaviourAction<Event, THandlerInEvent<Handler>>),
     Connect {
         attempt: u8,
         handler: NotifyHandler,
@@ -334,8 +389,8 @@ enum ActionBuilder {
     },
 }
 
-impl From<NetworkBehaviourAction<Event, handler::Prototype>> for ActionBuilder {
-    fn from(action: NetworkBehaviourAction<Event, handler::Prototype>) -> Self {
+impl From<NetworkBehaviourAction<Event, THandlerInEvent<Handler>>> for ActionBuilder {
+    fn from(action: NetworkBehaviourAction<Event, THandlerInEvent<Handler>>) -> Self {
         Self::Done(action)
     }
 }
@@ -344,7 +399,7 @@ impl ActionBuilder {
     fn build(
         self,
         poll_parameters: &mut impl PollParameters,
-    ) -> NetworkBehaviourAction<Event, handler::Prototype> {
+    ) -> NetworkBehaviourAction<Event, THandlerInEvent<Handler>> {
         let obs_addrs = || {
             poll_parameters
                 .external_addresses()