feat: move Identify I/O from NetworkBehaviour to ConnectionHandler (#3208)

Addresses #2885

protocols/identify/src/handler.rs

@@ -19,14 +19,15 @@
 // DEALINGS IN THE SOFTWARE.
 
 use crate::protocol::{
-    InboundPush, Info, OutboundPush, Protocol, PushProtocol, ReplySubstream, UpgradeError,
+    self, Identify, InboundPush, Info, OutboundPush, Protocol, Push, UpgradeError,
 };
 use futures::future::BoxFuture;
 use futures::prelude::*;
+use futures::stream::FuturesUnordered;
 use futures_timer::Delay;
 use libp2p_core::either::{EitherError, EitherOutput};
 use libp2p_core::upgrade::{EitherUpgrade, SelectUpgrade};
-use libp2p_core::{ConnectedPoint, PeerId};
+use libp2p_core::{ConnectedPoint, Multiaddr, PeerId, PublicKey};
 use libp2p_swarm::handler::{
     ConnectionEvent, DialUpgradeError, FullyNegotiatedInbound, FullyNegotiatedOutbound,
 };
@@ -36,18 +37,31 @@ use libp2p_swarm::{
 };
 use log::warn;
 use smallvec::SmallVec;
+use std::collections::VecDeque;
 use std::{io, pin::Pin, task::Context, task::Poll, time::Duration};
 
 pub struct Proto {
     initial_delay: Duration,
     interval: Duration,
+    public_key: PublicKey,
+    protocol_version: String,
+    agent_version: String,
 }
 
 impl Proto {
-    pub fn new(initial_delay: Duration, interval: Duration) -> Self {
+    pub fn new(
+        initial_delay: Duration,
+        interval: Duration,
+        public_key: PublicKey,
+        protocol_version: String,
+        agent_version: String,
+    ) -> Self {
         Proto {
             initial_delay,
             interval,
+            public_key,
+            protocol_version,
+            agent_version,
         }
     }
 }
@@ -55,12 +69,25 @@ impl Proto {
 impl IntoConnectionHandler for Proto {
     type Handler = Handler;
 
-    fn into_handler(self, remote_peer_id: &PeerId, _endpoint: &ConnectedPoint) -> Self::Handler {
-        Handler::new(self.initial_delay, self.interval, *remote_peer_id)
+    fn into_handler(self, remote_peer_id: &PeerId, endpoint: &ConnectedPoint) -> Self::Handler {
+        let observed_addr = match endpoint {
+            ConnectedPoint::Dialer { address, .. } => address,
+            ConnectedPoint::Listener { send_back_addr, .. } => send_back_addr,
+        };
+
+        Handler::new(
+            self.initial_delay,
+            self.interval,
+            *remote_peer_id,
+            self.public_key,
+            self.protocol_version,
+            self.agent_version,
+            observed_addr.clone(),
+        )
     }
 
     fn inbound_protocol(&self) -> <Self::Handler as ConnectionHandler>::InboundProtocol {
-        SelectUpgrade::new(Protocol, PushProtocol::inbound())
+        SelectUpgrade::new(Identify, Push::inbound())
     }
 }
 
@@ -74,14 +101,16 @@ pub struct Handler {
     inbound_identify_push: Option<BoxFuture<'static, Result<Info, UpgradeError>>>,
     /// Pending events to yield.
     events: SmallVec<
-        [ConnectionHandlerEvent<
-            EitherUpgrade<Protocol, PushProtocol<OutboundPush>>,
-            (),
-            Event,
-            io::Error,
-        >; 4],
+        [ConnectionHandlerEvent<EitherUpgrade<Identify, Push<OutboundPush>>, (), Event, io::Error>;
+            4],
     >,
 
+    /// Streams awaiting `BehaviourInfo` to then send identify requests.
+    reply_streams: VecDeque<NegotiatedSubstream>,
+
+    /// Pending identification replies, awaiting being sent.
+    pending_replies: FuturesUnordered<BoxFuture<'static, Result<PeerId, UpgradeError>>>,
+
     /// Future that fires when we need to identify the node again.
     trigger_next_identify: Delay,
 
@@ -90,36 +119,75 @@ pub struct Handler {
 
     /// The interval of `trigger_next_identify`, i.e. the recurrent delay.
     interval: Duration,
+
+    /// The public key of the local peer.
+    public_key: PublicKey,
+
+    /// Application-specific version of the protocol family used by the peer,
+    /// e.g. `ipfs/1.0.0` or `polkadot/1.0.0`.
+    protocol_version: String,
+
+    /// Name and version of the peer, similar to the `User-Agent` header in
+    /// the HTTP protocol.
+    agent_version: String,
+
+    /// Address observed by or for the remote.
+    observed_addr: Multiaddr,
 }
 
-/// Event produced by the `IdentifyHandler`.
+/// An event from `Behaviour` with the information requested by the `Handler`.
+#[derive(Debug)]
+pub struct InEvent {
+    /// The addresses that the peer is listening on.
+    pub listen_addrs: Vec<Multiaddr>,
+
+    /// The list of protocols supported by the peer, e.g. `/ipfs/ping/1.0.0`.
+    pub supported_protocols: Vec<String>,
+
+    /// The protocol w.r.t. the information requested.
+    pub protocol: Protocol,
+}
+
+/// Event produced by the `Handler`.
 #[derive(Debug)]
 #[allow(clippy::large_enum_variant)]
 pub enum Event {
     /// We obtained identification information from the remote.
     Identified(Info),
+    /// We replied to an identification request from the remote.
+    Identification(PeerId),
     /// We actively pushed our identification information to the remote.
     IdentificationPushed,
     /// We received a request for identification.
-    Identify(ReplySubstream<NegotiatedSubstream>),
-    /// Failed to identify the remote.
+    Identify,
+    /// Failed to identify the remote, or to reply to an identification request.
     IdentificationError(ConnectionHandlerUpgrErr<UpgradeError>),
 }
 
-/// Identifying information of the local node that is pushed to a remote.
-#[derive(Debug)]
-pub struct Push(pub Info);
-
 impl Handler {
-    /// Creates a new `IdentifyHandler`.
-    pub fn new(initial_delay: Duration, interval: Duration, remote_peer_id: PeerId) -> Self {
+    /// Creates a new `Handler`.
+    pub fn new(
+        initial_delay: Duration,
+        interval: Duration,
+        remote_peer_id: PeerId,
+        public_key: PublicKey,
+        protocol_version: String,
+        agent_version: String,
+        observed_addr: Multiaddr,
+    ) -> Self {
         Self {
             remote_peer_id,
             inbound_identify_push: Default::default(),
             events: SmallVec::new(),
+            reply_streams: VecDeque::new(),
+            pending_replies: FuturesUnordered::new(),
             trigger_next_identify: Delay::new(initial_delay),
             keep_alive: KeepAlive::Yes,
             interval,
+            public_key,
+            protocol_version,
+            agent_version,
+            observed_addr,
         }
     }
 
@@ -133,9 +201,18 @@ impl Handler {
         >,
     ) {
         match output {
-            EitherOutput::First(substream) => self
-                .events
-                .push(ConnectionHandlerEvent::Custom(Event::Identify(substream))),
+            EitherOutput::First(substream) => {
+                self.events
+                    .push(ConnectionHandlerEvent::Custom(Event::Identify));
+                if !self.reply_streams.is_empty() {
+                    warn!(
+                        "New inbound identify request from {} while a previous one \
+                         is still pending. Queueing the new one.",
+                        self.remote_peer_id,
+                    );
+                }
+                self.reply_streams.push_back(substream);
+            }
             EitherOutput::Second(fut) => {
                 if self.inbound_identify_push.replace(fut).is_some() {
                     warn!(
@@ -195,26 +272,58 @@ impl Handler {
 }
 
 impl ConnectionHandler for Handler {
-    type InEvent = Push;
+    type InEvent = InEvent;
     type OutEvent = Event;
     type Error = io::Error;
-    type InboundProtocol = SelectUpgrade<Protocol, PushProtocol<InboundPush>>;
-    type OutboundProtocol = EitherUpgrade<Protocol, PushProtocol<OutboundPush>>;
+    type InboundProtocol = SelectUpgrade<Identify, Push<InboundPush>>;
+    type OutboundProtocol = EitherUpgrade<Identify, Push<OutboundPush>>;
     type OutboundOpenInfo = ();
     type InboundOpenInfo = ();
 
     fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol, Self::InboundOpenInfo> {
-        SubstreamProtocol::new(SelectUpgrade::new(Protocol, PushProtocol::inbound()), ())
+        SubstreamProtocol::new(SelectUpgrade::new(Identify, Push::inbound()), ())
     }
 
-    fn on_behaviour_event(&mut self, Push(push): Self::InEvent) {
-        self.events
-            .push(ConnectionHandlerEvent::OutboundSubstreamRequest {
-                protocol: SubstreamProtocol::new(
-                    EitherUpgrade::B(PushProtocol::outbound(push)),
-                    (),
-                ),
-            });
+    fn on_behaviour_event(
+        &mut self,
+        InEvent {
+            listen_addrs,
+            supported_protocols,
+            protocol,
+        }: Self::InEvent,
+    ) {
+        let info = Info {
+            public_key: self.public_key.clone(),
+            protocol_version: self.protocol_version.clone(),
+            agent_version: self.agent_version.clone(),
+            listen_addrs,
+            protocols: supported_protocols,
+            observed_addr: self.observed_addr.clone(),
+        };
+
+        match protocol {
+            Protocol::Push => {
+                self.events
+                    .push(ConnectionHandlerEvent::OutboundSubstreamRequest {
+                        protocol: SubstreamProtocol::new(
+                            EitherUpgrade::B(Push::outbound(info)),
+                            (),
+                        ),
+                    });
+            }
+            Protocol::Identify(_) => {
+                let substream = self
+                    .reply_streams
+                    .pop_front()
+                    .expect("A BehaviourInfo reply should have a matching substream.");
+                let peer = self.remote_peer_id;
+                let fut = Box::pin(async move {
+                    protocol::send(substream, info).await?;
+                    Ok(peer)
+                });
+                self.pending_replies.push(fut);
+            }
+        }
     }
 
     fn connection_keep_alive(&self) -> KeepAlive {
@@ -237,7 +346,7 @@ impl ConnectionHandler for Handler {
             Poll::Ready(()) => {
                 self.trigger_next_identify.reset(self.interval);
                 let ev = ConnectionHandlerEvent::OutboundSubstreamRequest {
-                    protocol: SubstreamProtocol::new(EitherUpgrade::A(Protocol), ()),
+                    protocol: SubstreamProtocol::new(EitherUpgrade::A(Identify), ()),
                 };
                 return Poll::Ready(ev);
             }
@@ -255,7 +364,18 @@ impl ConnectionHandler for Handler {
             }
         }
 
-        Poll::Pending
+        // Check for pending replies to send.
+        match self.pending_replies.poll_next_unpin(cx) {
+            Poll::Ready(Some(Ok(peer_id))) => Poll::Ready(ConnectionHandlerEvent::Custom(
+                Event::Identification(peer_id),
+            )),
+            Poll::Ready(Some(Err(err))) => Poll::Ready(ConnectionHandlerEvent::Custom(
+                Event::IdentificationError(ConnectionHandlerUpgrErr::Upgrade(
+                    libp2p_core::upgrade::UpgradeError::Apply(err),
+                )),
+            )),
+            Poll::Ready(None) | Poll::Pending => Poll::Pending,
+        }
     }
 
     fn on_connection_event(