feat: verify expected PeerId as part of security handshake

core/src/transport/upgrade.rs

@@ -31,11 +31,11 @@ use crate::{
     },
     upgrade::{
         self, apply_inbound, apply_outbound, InboundConnectionUpgrade, InboundUpgradeApply,
-        OutboundConnectionUpgrade, OutboundUpgradeApply, UpgradeError,
+        OutboundConnectionUpgrade, OutboundUpgradeApply, SecurityUpgrade, UpgradeError,
     },
     Negotiated,
 };
-use futures::{prelude::*, ready};
+use futures::{future::LocalBoxFuture, prelude::*, ready};
 use libp2p_identity::PeerId;
 use multiaddr::Multiaddr;
 use std::{
@@ -99,16 +99,15 @@ where
     ) -> Authenticated<AndThen<T, impl FnOnce(C, ConnectedPoint) -> Authenticate<C, U> + Clone>>
     where
         T: Transport<Output = C>,
-        C: AsyncRead + AsyncWrite + Unpin,
+        C: AsyncRead + AsyncWrite + Unpin + 'static,
         D: AsyncRead + AsyncWrite + Unpin,
-        U: InboundConnectionUpgrade<Negotiated<C>, Output = (PeerId, D), Error = E>,
-        U: OutboundConnectionUpgrade<Negotiated<C>, Output = (PeerId, D), Error = E> + Clone,
+        U: SecurityUpgrade<Negotiated<C>, Output = (PeerId, D), Error = E> + Clone + 'static,
         E: Error + 'static,
     {
         let version = self.version;
         Authenticated(Builder::new(
             self.inner.and_then(move |conn, endpoint| Authenticate {
-                inner: upgrade::apply(conn, upgrade, endpoint, version),
+                inner: upgrade::secure(conn, upgrade, endpoint, version).boxed_local(),
             }),
             version,
         ))
@@ -123,23 +122,18 @@ where
 pub struct Authenticate<C, U>
 where
     C: AsyncRead + AsyncWrite + Unpin,
-    U: InboundConnectionUpgrade<Negotiated<C>> + OutboundConnectionUpgrade<Negotiated<C>>,
+    U: SecurityUpgrade<Negotiated<C>>,
 {
     #[pin]
-    inner: EitherUpgrade<C, U>,
+    inner: LocalBoxFuture<'static, Result<U::Output, UpgradeError<U::Error>>>,
 }
 
 impl<C, U> Future for Authenticate<C, U>
 where
     C: AsyncRead + AsyncWrite + Unpin,
-    U: InboundConnectionUpgrade<Negotiated<C>>
-        + OutboundConnectionUpgrade<
-            Negotiated<C>,
-            Output = <U as InboundConnectionUpgrade<Negotiated<C>>>::Output,
-            Error = <U as InboundConnectionUpgrade<Negotiated<C>>>::Error,
-        >,
+    U: SecurityUpgrade<Negotiated<C>>,
 {
-    type Output = <EitherUpgrade<C, U> as Future>::Output;
+    type Output = Result<U::Output, UpgradeError<U::Error>>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = self.project();

core/src/upgrade.rs

@@ -63,19 +63,21 @@ mod either;
 mod error;
 mod pending;
 mod ready;
+mod secure;
 mod select;
 
+pub use self::{
+    denied::DeniedUpgrade, pending::PendingUpgrade, ready::ReadyUpgrade, select::SelectUpgrade,
+};
+pub use crate::Negotiated;
 pub(crate) use apply::{
     apply, apply_inbound, apply_outbound, InboundUpgradeApply, OutboundUpgradeApply,
 };
 pub(crate) use error::UpgradeError;
 use futures::future::Future;
-
-pub use self::{
-    denied::DeniedUpgrade, pending::PendingUpgrade, ready::ReadyUpgrade, select::SelectUpgrade,
-};
-pub use crate::Negotiated;
+use libp2p_identity::PeerId;
 pub use multistream_select::{NegotiatedComplete, NegotiationError, ProtocolError, Version};
+pub(crate) use secure::secure;
 
 /// Common trait for upgrades that can be applied on inbound substreams, outbound substreams,
 /// or both.
@@ -152,3 +154,22 @@ pub trait OutboundConnectionUpgrade<T>: UpgradeInfo {
     /// The `info` is the identifier of the protocol, as produced by `protocol_info`.
     fn upgrade_outbound(self, socket: T, info: Self::Info) -> Self::Future;
 }
+
+/// Possible security upgrade on a connection
+pub trait SecurityUpgrade<T>: UpgradeInfo {
+    /// Output after the upgrade has been successfully negotiated and the handshake performed.
+    type Output;
+    /// Possible error during the handshake.
+    type Error;
+    /// Future that performs the handshake with the remote.
+    type Future: Future<Output = Result<Self::Output, Self::Error>>;
+
+    /// After we have determined that the remote supports one of the protocols we support, this
+    /// method is called to start the handshake.
+    ///
+    /// The `info` is the identifier of the protocol, as produced by `protocol_info`. Security
+    /// transports use the optional `peer_id` parameter on outgoing upgrades to validate validate
+    /// the expected `PeerId`.
+    fn upgrade_security(self, socket: T, info: Self::Info, peer_id: Option<PeerId>)
+        -> Self::Future;
+}

core/src/upgrade/secure.rs

@@ -0,0 +1,88 @@
+//! After a successful protocol negotiation as part of the upgrade process, the `SecurityUpgrade::upgrade_security`
+//! method is called and a [`Future`] that performs a handshake is returned.
+
+use crate::upgrade::{SecurityUpgrade, UpgradeError};
+use crate::{connection::ConnectedPoint, Negotiated};
+use futures::prelude::*;
+
+use libp2p_identity::PeerId;
+use multiaddr::Protocol;
+use multistream_select::Version;
+
+/// Applies a security upgrade to the inbound and outbound direction of a connection or substream.
+pub(crate) async fn secure<C, U>(
+    conn: C,
+    up: U,
+    cp: ConnectedPoint,
+    v: Version,
+) -> Result<U::Output, UpgradeError<U::Error>>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    U: SecurityUpgrade<Negotiated<C>>,
+{
+    match cp {
+        ConnectedPoint::Dialer { role_override, .. } if role_override.is_dialer() => {
+            let peer_id = cp
+                .get_remote_address()
+                .iter()
+                .find_map(|protocol| match protocol {
+                    Protocol::P2p(peer_id) => Some(peer_id),
+                    _ => None,
+                })
+                .expect("It should have /p2p as part of the address");
+            secure_outbound(conn, up, Some(peer_id), v).await
+        }
+        _ => secure_inbound(conn, up, None).await,
+    }
+}
+
+/// Tries to perform a security upgrade on an inbound connection or substream.
+async fn secure_inbound<C, U>(
+    conn: C,
+    up: U,
+    peer_id: Option<PeerId>,
+) -> Result<U::Output, UpgradeError<U::Error>>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    U: SecurityUpgrade<Negotiated<C>>,
+{
+    let (info, stream) =
+        multistream_select::listener_select_proto(conn, up.protocol_info()).await?;
+    let name = info.as_ref().to_owned();
+    match up.upgrade_security(stream, info, peer_id).await {
+        Ok(x) => {
+            tracing::trace!(up=%name, "Secured inbound stream");
+            Ok(x)
+        }
+        Err(e) => {
+            tracing::trace!(up=%name, "Failed to secure inbound stream");
+            Err(UpgradeError::Apply(e))
+        }
+    }
+}
+
+/// Tries to perform a security upgrade on an outbound connection or substream.
+async fn secure_outbound<C, U>(
+    conn: C,
+    up: U,
+    peer_id: Option<PeerId>,
+    v: Version,
+) -> Result<U::Output, UpgradeError<U::Error>>
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+    U: SecurityUpgrade<Negotiated<C>>,
+{
+    let (info, stream) =
+        multistream_select::dialer_select_proto(conn, up.protocol_info(), v).await?;
+    let name = info.as_ref().to_owned();
+    match up.upgrade_security(stream, info, peer_id).await {
+        Ok(x) => {
+            tracing::trace!(up=%name, "Secured outbound stream");
+            Ok(x)
+        }
+        Err(e) => {
+            tracing::trace!(up=%name, "Failed to secure outbound stream");
+            Err(UpgradeError::Apply(e))
+        }
+    }
+}