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

protocols/autonat/src/behaviour.rs

@@ -36,6 +36,7 @@ use libp2p_request_response::{
     handler::RequestResponseHandlerEvent, ProtocolSupport, RequestId, RequestResponse,
     RequestResponseConfig, RequestResponseEvent, RequestResponseMessage, ResponseChannel,
 };
+use libp2p_swarm::behaviour::THandlerInEvent;
 use libp2p_swarm::{
     DialError, IntoConnectionHandler, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
 };
@@ -303,6 +304,7 @@ impl Behaviour {
 impl NetworkBehaviour for Behaviour {
     type ConnectionHandler = <RequestResponse<AutoNatCodec> as NetworkBehaviour>::ConnectionHandler;
     type OutEvent = Event;
+    type DialPayload = ();
 
     fn inject_connection_established(
         &mut self,
@@ -372,10 +374,10 @@ impl NetworkBehaviour for Behaviour {
     fn inject_dial_failure(
         &mut self,
         peer: Option<PeerId>,
-        handler: Self::ConnectionHandler,
+        initial_event: Option<Self::DialPayload>,
         error: &DialError,
     ) {
-        self.inner.inject_dial_failure(peer, handler, error);
+        self.inner.inject_dial_failure(peer, initial_event, error);
         if let Some(event) = self.as_server().on_outbound_dial_error(peer, error) {
             self.pending_out_events
                 .push_back(Event::InboundProbe(event));
@@ -487,14 +489,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 +508,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,7 @@ impl<'a> HandleInnerEvent for AsServer<'a> {
                                 .override_dial_concurrency_factor(NonZeroU8::new(1).expect("1 > 0"))
                                 .addresses(addrs)
                                 .build(),
-                            handler: self.inner.new_handler(),
+                            dial_payload: (),
                         });
                     }
                     Err((status_text, error)) => {

protocols/dcutr/src/behaviour.rs

@@ -26,6 +26,7 @@ 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,
@@ -86,11 +87,26 @@ impl Behaviour {
 impl NetworkBehaviour for Behaviour {
     type ConnectionHandler = handler::Prototype;
     type OutEvent = Event;
+    type DialPayload = handler::Prototype;
 
     fn new_handler(&mut self) -> Self::ConnectionHandler {
         handler::Prototype::UnknownConnection
     }
 
+    fn new_inbound_handler(&mut self) -> Self::ConnectionHandler {
+        handler::Prototype::UnknownConnection
+    }
+
+    fn new_outbound_handler(
+        &mut self,
+        dial_payload: Option<Self::DialPayload>,
+    ) -> Self::ConnectionHandler {
+        match dial_payload {
+            Some(prototype) => prototype,
+            None => handler::Prototype::UnknownConnection,
+        }
+    }
+
     fn addresses_of_peer(&mut self, _peer_id: &PeerId) -> Vec<Multiaddr> {
         vec![]
     }
@@ -143,13 +159,13 @@ impl NetworkBehaviour for Behaviour {
     fn inject_dial_failure(
         &mut self,
         peer_id: Option<PeerId>,
-        handler: Self::ConnectionHandler,
+        dial_payload: Option<Self::DialPayload>,
         _error: &DialError,
     ) {
-        if let handler::Prototype::DirectConnection {
+        if let Some(handler::Prototype::DirectConnection {
             relayed_connection_id,
             role: handler::Role::Initiator { attempt },
-        } = handler
+        }) = dial_payload
         {
             let peer_id = peer_id.expect("Peer of `Prototype::DirectConnection` is always known.");
             if attempt < MAX_NUMBER_OF_UPGRADE_ATTEMPTS {
@@ -243,7 +259,7 @@ impl NetworkBehaviour for Behaviour {
                             .addresses(remote_addrs)
                             .condition(dial_opts::PeerCondition::Always)
                             .build(),
-                        handler: handler::Prototype::DirectConnection {
+                        dial_payload: handler::Prototype::DirectConnection {
                             relayed_connection_id: connection,
                             role: handler::Role::Listener,
                         },
@@ -271,7 +287,7 @@ impl NetworkBehaviour for Behaviour {
                             .addresses(remote_addrs)
                             .override_role()
                             .build(),
-                        handler: handler::Prototype::DirectConnection {
+                        dial_payload: handler::Prototype::DirectConnection {
                             relayed_connection_id: connection,
                             role: handler::Role::Initiator { attempt },
                         },
@@ -309,7 +325,13 @@ 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>,
+            handler::Prototype,
+        >,
+    > {
         if let Some(action) = self.queued_actions.pop_front() {
             return Poll::Ready(action.build(poll_parameters));
         }
@@ -321,7 +343,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::Prototype>, handler::Prototype>),
     Connect {
         attempt: u8,
         handler: NotifyHandler,
@@ -334,8 +356,16 @@ enum ActionBuilder {
     },
 }
 
-impl From<NetworkBehaviourAction<Event, handler::Prototype>> for ActionBuilder {
-    fn from(action: NetworkBehaviourAction<Event, handler::Prototype>) -> Self {
+impl From<NetworkBehaviourAction<Event, THandlerInEvent<handler::Prototype>, handler::Prototype>>
+    for ActionBuilder
+{
+    fn from(
+        action: NetworkBehaviourAction<
+            Event,
+            THandlerInEvent<handler::Prototype>,
+            handler::Prototype,
+        >,
+    ) -> Self {
         Self::Done(action)
     }
 }
@@ -344,7 +374,8 @@ impl ActionBuilder {
     fn build(
         self,
         poll_parameters: &mut impl PollParameters,
-    ) -> NetworkBehaviourAction<Event, handler::Prototype> {
+    ) -> NetworkBehaviourAction<Event, THandlerInEvent<handler::Prototype>, handler::Prototype>
+    {
         let obs_addrs = || {
             poll_parameters
                 .external_addresses()

protocols/floodsub/src/layer.rs

@@ -28,6 +28,7 @@ use cuckoofilter::{CuckooError, CuckooFilter};
 use fnv::FnvHashSet;
 use libp2p_core::{connection::ConnectionId, PeerId};
 use libp2p_core::{ConnectedPoint, Multiaddr};
+use libp2p_swarm::behaviour::THandlerInEvent;
 use libp2p_swarm::{
     dial_opts::DialOpts, NetworkBehaviour, NetworkBehaviourAction, NotifyHandler, OneShotHandler,
     PollParameters,
@@ -41,12 +42,7 @@ use std::{collections::VecDeque, iter};
 /// Network behaviour that handles the floodsub protocol.
 pub struct Floodsub {
     /// Events that need to be yielded to the outside when polling.
-    events: VecDeque<
-        NetworkBehaviourAction<
-            FloodsubEvent,
-            OneShotHandler<FloodsubProtocol, FloodsubRpc, InnerMessage>,
-        >,
-    >,
+    events: VecDeque<NetworkBehaviourAction<FloodsubEvent, FloodsubRpc>>,
 
     config: FloodsubConfig,
 
@@ -107,10 +103,9 @@ impl Floodsub {
         }
 
         if self.target_peers.insert(peer_id) {
-            let handler = self.new_handler();
             self.events.push_back(NetworkBehaviourAction::Dial {
                 opts: DialOpts::peer_id(peer_id).build(),
-                handler,
+                dial_payload: (),
             });
         }
     }
@@ -281,6 +276,7 @@ impl Floodsub {
 impl NetworkBehaviour for Floodsub {
     type ConnectionHandler = OneShotHandler<FloodsubProtocol, FloodsubRpc, InnerMessage>;
     type OutEvent = FloodsubEvent;
+    type DialPayload = ();
 
     fn new_handler(&mut self) -> Self::ConnectionHandler {
         Default::default()
@@ -339,10 +335,9 @@ impl NetworkBehaviour for Floodsub {
         // We can be disconnected by the remote in case of inactivity for example, so we always
         // try to reconnect.
         if self.target_peers.contains(id) {
-            let handler = self.new_handler();
             self.events.push_back(NetworkBehaviourAction::Dial {
                 opts: DialOpts::peer_id(*id).build(),
-                handler,
+                dial_payload: (),
             });
         }
     }
@@ -470,7 +465,8 @@ impl NetworkBehaviour for Floodsub {
         &mut self,
         _: &mut Context<'_>,
         _: &mut impl PollParameters,
-    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ConnectionHandler>> {
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, THandlerInEvent<Self::ConnectionHandler>>>
+    {
         if let Some(event) = self.events.pop_front() {
             return Poll::Ready(event);
         }

protocols/gossipsub/src/behaviour.rs

@@ -200,7 +200,7 @@ impl From<MessageAuthenticity> for PublishConfig {
 }
 
 type GossipsubNetworkBehaviourAction =
-    NetworkBehaviourAction<GossipsubEvent, GossipsubHandler, Arc<GossipsubHandlerIn>>;
+    NetworkBehaviourAction<GossipsubEvent, Arc<GossipsubHandlerIn>>;
 
 /// Network behaviour that handles the gossipsub protocol.
 ///
@@ -1139,10 +1139,9 @@ where
         if !self.peer_topics.contains_key(peer_id) {
             // Connect to peer
             debug!("Connecting to explicit peer {:?}", peer_id);
-            let handler = self.new_handler();
             self.events.push_back(NetworkBehaviourAction::Dial {
                 opts: DialOpts::peer_id(*peer_id).build(),
-                handler,
+                dial_payload: (),
             });
         }
     }
@@ -1637,11 +1636,9 @@ where
                 // mark as px peer
                 self.px_peers.insert(peer_id);
 
-                // dial peer
-                let handler = self.new_handler();
                 self.events.push_back(NetworkBehaviourAction::Dial {
                     opts: DialOpts::peer_id(peer_id).build(),
-                    handler,
+                    dial_payload: (),
                 });
             }
         }
@@ -3045,6 +3042,7 @@ where
 {
     type ConnectionHandler = GossipsubHandler;
     type OutEvent = GossipsubEvent;
+    type DialPayload = ();
 
     fn new_handler(&mut self) -> Self::ConnectionHandler {
         let protocol_config = ProtocolConfig::new(
@@ -3437,7 +3435,7 @@ where
         &mut self,
         cx: &mut Context<'_>,
         _: &mut impl PollParameters,
-    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ConnectionHandler>> {
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, GossipsubHandlerIn>> {
         if let Some(event) = self.events.pop_front() {
             return Poll::Ready(event.map_in(|e: Arc<GossipsubHandlerIn>| {
                 // clone send event reference if others references are present

protocols/gossipsub/src/behaviour/tests.rs

@@ -1356,7 +1356,10 @@ fn test_explicit_peer_gets_connected() {
         .events
         .iter()
         .filter(|e| match e {
-            NetworkBehaviourAction::Dial { opts, handler: _ } => opts.get_peer_id() == Some(peer),
+            NetworkBehaviourAction::Dial {
+                opts,
+                initial_in_event: _,
+            } => opts.get_peer_id() == Some(peer),
             _ => false,
         })
         .count();
@@ -1397,8 +1400,10 @@ fn test_explicit_peer_reconnects() {
         gs.events
             .iter()
             .filter(|e| match e {
-                NetworkBehaviourAction::Dial { opts, handler: _ } =>
-                    opts.get_peer_id() == Some(*peer),
+                NetworkBehaviourAction::Dial {
+                    opts,
+                    initial_in_event: _,
+                } => opts.get_peer_id() == Some(*peer),
                 _ => false,
             })
             .count(),
@@ -1413,8 +1418,10 @@ fn test_explicit_peer_reconnects() {
         gs.events
             .iter()
             .filter(|e| match e {
-                NetworkBehaviourAction::Dial { opts, handler: _ } =>
-                    opts.get_peer_id() == Some(*peer),
+                NetworkBehaviourAction::Dial {
+                    opts,
+                    initial_in_event: _,
+                } => opts.get_peer_id() == Some(*peer),
                 _ => false,
             })
             .count()
@@ -1794,7 +1801,10 @@ fn test_connect_to_px_peers_on_handle_prune() {
         .events
         .iter()
         .filter_map(|e| match e {
-            NetworkBehaviourAction::Dial { opts, handler: _ } => opts.get_peer_id(),
+            NetworkBehaviourAction::Dial {
+                opts,
+                initial_in_event: _,
+            } => opts.get_peer_id(),
             _ => None,
         })
         .collect();