refactor(node/router): improve disconnect error propagation in the router

Author: kaimast

Cargo.lock

@@ -29,6 +29,8 @@ pub enum DisconnectReason {
     ProtocolViolation,
     /// The peer's client is outdated, judging by its version.
     OutdatedClientVersion,
+    /// The two validators are the same node.
+    SelfConnect,
 }
 
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
@@ -64,6 +66,7 @@ impl FromBytes for Disconnect {
             Ok(1) => DisconnectReason::NoReasonGiven,
             Ok(2) => DisconnectReason::ProtocolViolation,
             Ok(3) => DisconnectReason::OutdatedClientVersion,
+            Ok(4) => DisconnectReason::SelfConnect,
             _ => return Err(io::Error::other("Invalid 'Disconnect' event")),
         };
 
@@ -86,6 +89,7 @@ mod tests {
             DisconnectReason::NoReasonGiven,
             DisconnectReason::InvalidChallengeResponse,
             DisconnectReason::OutdatedClientVersion,
+            DisconnectReason::SelfConnect,
         ];
 
         for reason in all_reasons.iter() {

node/bft/events/src/disconnect.rs

@@ -104,6 +104,9 @@ features = [ "codec" ]
 [dependencies.tokio-stream]
 version = "=0.1"
 
+[dependencies.thiserror]
+workspace = true
+
 [dependencies.tracing]
 workspace = true
 
@@ -136,7 +139,3 @@ features = [ "test" ]
 [dev-dependencies.snarkvm]
 workspace = true
 features = [ "test-helpers" ]
-
-[dev-dependencies.tracing-subscriber]
-workspace = true
-features = [ "env-filter", "fmt" ]

node/bft/src/gateway.rs

@@ -18,6 +18,7 @@ use snarkvm::prelude::{FromBytes, ToBytes, error};
 use std::io;
 
 /// The reason behind the node disconnecting from a peer.
+// TODO(kaimast): implement Display
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum DisconnectReason {
     /// The fork length limit was exceeded.
@@ -48,8 +49,16 @@ pub enum DisconnectReason {
     TooManyPeers,
     /// The peer is a sync node that's behind our node, and it needs to sync itself first.
     YouNeedToSyncFirst,
-    /// The peer's listening port is closed.
+    /// The peer's listening port is closed
     YourPortIsClosed(u16),
+    /// The two peers are the same node.
+    SelfConnect,
+    /// No untrusted external peers are allowed.
+    NoExternalPeersAllowed,
+    /// Already connecting to the same node (through another TCP channel).
+    AlreadyConnecting,
+    /// Already connected to the same node (through another TCP channel).
+    AlreadyConnected,
 }
 
 impl ToBytes for DisconnectReason {
@@ -73,6 +82,10 @@ impl ToBytes for DisconnectReason {
                 14u8.write_le(&mut writer)?;
                 port.write_le(writer)
             }
+            Self::SelfConnect => 15u8.write_le(writer),
+            Self::NoExternalPeersAllowed => 16u8.write_le(writer),
+            Self::AlreadyConnecting => 17u8.write_le(writer),
+            Self::AlreadyConnected => 18u8.write_le(writer),
         }
     }
 }
@@ -98,6 +111,10 @@ impl FromBytes for DisconnectReason {
                 let port = u16::read_le(reader)?;
                 Ok(Self::YourPortIsClosed(port))
             }
+            15 => Ok(Self::SelfConnect),
+            16 => Ok(Self::NoExternalPeersAllowed),
+            17 => Ok(Self::AlreadyConnecting),
+            18 => Ok(Self::AlreadyConnected),
             _ => Err(error("Invalid disconnect reason")),
         }
     }

node/router/Cargo.toml

@@ -14,18 +14,19 @@
 // limitations under the License.
 
 use crate::{
+    AddPeerError,
     NodeType,
     PeerPoolHandling,
     Router,
     messages::{ChallengeRequest, ChallengeResponse, DisconnectReason, Message, MessageCodec, MessageTrait},
 };
-use snarkos_node_tcp::{ConnectionSide, P2P, Tcp};
+use snarkos_node_tcp::{ConnectError, ConnectionSide, P2P, Tcp};
 use snarkvm::{
     ledger::narwhal::Data,
-    prelude::{Address, Field, Network, block::Header, error},
+    prelude::{Address, Field, Network, block::Header},
 };
 
-use anyhow::{Result, bail};
+use anyhow::{Result, anyhow};
 use futures::SinkExt;
 use rand::{Rng, rngs::OsRng};
 use std::{io, net::SocketAddr};
@@ -52,24 +53,19 @@ macro_rules! expect_message {
             }
             // Received a disconnect message, abort.
             Some(Message::Disconnect(reason)) => {
-                return Err(error(format!("'{}' disconnected: {reason:?}", $peer_addr)))
+                return Err(ConnectError::other(format!("'{}' disconnected with reason {reason:?}", $peer_addr)));
             }
             // Received an unexpected message, abort.
             Some(ty) => {
-                return Err(error(format!(
+                return Err(ConnectError::other(format!(
                     "'{}' did not follow the handshake protocol: received {:?} instead of {}",
                     $peer_addr,
                     ty.name(),
                     stringify!($msg_ty),
-                )))
+                )));
             }
             // Received nothing.
-            None => {
-                return Err(error(format!(
-                    "the peer disconnected before sending {:?}, likely due to peer saturation or shutdown",
-                    stringify!($msg_ty),
-                )))
-            }
+            None => return Err(ConnectError::IoError(io::ErrorKind::BrokenPipe.into())),
         }
     };
 }
@@ -93,7 +89,7 @@ impl<N: Network> Router<N> {
         peer_side: ConnectionSide,
         genesis_header: Header<N>,
         restrictions_id: Field<N>,
-    ) -> io::Result<Option<ChallengeRequest<N>>> {
+    ) -> Result<ChallengeRequest<N>, ConnectError> {
         // If this is an inbound connection, we log it, but don't know the listening address yet.
         // Otherwise, we can immediately register the listening address.
         let mut listener_addr = if peer_side == ConnectionSide::Initiator {
@@ -110,7 +106,7 @@ impl<N: Network> Router<N> {
             // If the IP is already banned reject the connection.
             if self.is_ip_banned(peer_addr.ip()) {
                 trace!("Rejected a connection request from banned IP '{}'", peer_addr.ip());
-                return Err(error(format!("'{}' is a banned IP address", peer_addr.ip())));
+                return Err(ConnectError::other(anyhow!("'{}' is a banned IP address", peer_addr.ip())));
             }
 
             let num_attempts =
@@ -120,30 +116,32 @@ impl<N: Network> Router<N> {
             if num_attempts > Router::<N>::MAX_CONNECTION_ATTEMPTS {
                 self.update_ip_ban(peer_addr.ip());
                 trace!("Rejected a consecutive connection request from IP '{}'", peer_addr.ip());
-                return Err(error(format!("'{}' appears to be spamming connections", peer_addr.ip())));
+                return Err(ConnectError::other(anyhow!("'{}' appears to be spamming connections", peer_addr.ip())));
             }
         }
 
         // Perform the handshake; we pass on a mutable reference to listener_addr in case the process is broken at any point in time.
         let handshake_result = if peer_side == ConnectionSide::Responder {
-            self.handshake_inner_initiator(peer_addr, stream, genesis_header, restrictions_id).await
+            self.handshake_inner_initiator(peer_addr, stream, genesis_header, restrictions_id)
+                .await
+                .map_err(|err| ConnectError::Other(err.into()))
         } else {
             self.handshake_inner_responder(peer_addr, &mut listener_addr, stream, genesis_header, restrictions_id).await
         };
 
         if let Some(addr) = listener_addr {
             match handshake_result {
-                Ok(Some(ref cr)) => {
+                Ok(ref cr) => {
                     if let Some(peer) = self.peer_pool.write().get_mut(&addr) {
                         self.resolver.write().insert_peer(peer.listener_addr(), peer_addr, None);
                         peer.upgrade_to_connected(peer_addr, cr.listener_port, cr.address, cr.node_type, cr.version);
+                        debug!("Completed the handshake with '{peer_addr}'");
+                    } else {
+                        debug!("Completed the handshake with '{peer_addr}', but it is not in the peer pool");
                     }
+
                     #[cfg(feature = "metrics")]
                     self.update_metrics();
-                    debug!("Completed the handshake with '{peer_addr}'");
-                }
-                Ok(None) => {
-                    return Err(error("Duplicate handshake attempt with '{addr}'"));
                 }
                 Err(_) => {
                     if let Some(peer) = self.peer_pool.write().get_mut(&addr) {
@@ -163,12 +161,13 @@ impl<N: Network> Router<N> {
         stream: &'a mut TcpStream,
         genesis_header: Header<N>,
         restrictions_id: Field<N>,
-    ) -> io::Result<Option<ChallengeRequest<N>>> {
+    ) -> Result<ChallengeRequest<N>, ConnectError> {
         // Introduce the peer into the peer pool.
-        if !self.add_connecting_peer(peer_addr) {
-            // Return early if already being connected to.
-            return Ok(None);
-        }
+        // If we are connecting, the peer and listener address are identical.
+        self.add_connecting_peer(peer_addr).map_err(|err| match err {
+            AddPeerError::AlreadyConnected => ConnectError::AlreadyConnected { address: peer_addr },
+            AddPeerError::AlreadyConnecting => ConnectError::AlreadyConnecting { address: peer_addr },
+        })?;
 
         // Construct the stream.
         let mut framed = Framed::new(stream, MessageCodec::<N>::handshake());
@@ -205,12 +204,13 @@ impl<N: Network> Router<N> {
             .await
         {
             send(&mut framed, peer_addr, reason.into()).await?;
-            return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
+            return Err(ConnectError::other(anyhow!("{reason:?}")));
         }
+
         // Verify the challenge request. If a disconnect reason was returned, send the disconnect message and abort.
         if let Some(reason) = self.verify_challenge_request(peer_addr, &peer_request) {
             send(&mut framed, peer_addr, reason.into()).await?;
-            return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
+            return Err(ConnectError::other(anyhow!("{reason:?}")));
         }
 
         /* Step 3: Send the challenge response. */
@@ -219,7 +219,7 @@ impl<N: Network> Router<N> {
         let data = [peer_request.nonce.to_le_bytes(), response_nonce.to_le_bytes()].concat();
         // Sign the counterparty nonce.
         let Ok(our_signature) = self.account.sign_bytes(&data, rng) else {
-            return Err(error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")));
+            return Err(ConnectError::other(anyhow!("Failed to sign the challenge request nonce")));
         };
         // Send the challenge response.
         let our_response = ChallengeResponse {
@@ -230,7 +230,7 @@ impl<N: Network> Router<N> {
         };
         send(&mut framed, peer_addr, Message::ChallengeResponse(our_response)).await?;
 
-        Ok(Some(peer_request))
+        Ok(peer_request)
     }
 
     /// The connection responder side of the handshake.
@@ -241,34 +241,42 @@ impl<N: Network> Router<N> {
         stream: &'a mut TcpStream,
         genesis_header: Header<N>,
         restrictions_id: Field<N>,
-    ) -> io::Result<Option<ChallengeRequest<N>>> {
+    ) -> Result<ChallengeRequest<N>, ConnectError> {
         // Construct the stream.
         let mut framed = Framed::new(stream, MessageCodec::<N>::handshake());
 
         /* Step 1: Receive the challenge request. */
 
-        // Listen for the challenge request message.
+        // Wait for the challenge request message.
         let peer_request = expect_message!(Message::ChallengeRequest, framed, peer_addr);
 
         // Obtain the peer's listening address.
         *listener_addr = Some(SocketAddr::new(peer_addr.ip(), peer_request.listener_port));
         let listener_addr = listener_addr.unwrap();
 
+        // Knowing the peers listening address, ensure we are not already connecting/connected to it.
+        if self.is_connecting_or_connected(listener_addr) {
+            return Err(ConnectError::AlreadyConnected { address: listener_addr });
+        }
+
         // Knowing the peer's listening address, ensure it is allowed to connect.
-        if let Err(forbidden_message) = self.ensure_peer_is_allowed(listener_addr) {
-            return Err(error(format!("{forbidden_message}")));
+        if let Err(reason) = self.ensure_peer_is_allowed(listener_addr) {
+            send(&mut framed, peer_addr, reason.into()).await?;
+            return Err(ConnectError::Other(
+                anyhow!("Dropping connection request from {listener_addr} ({reason:?})").into(),
+            ));
         }
 
         // Introduce the peer into the peer pool.
-        if !self.add_connecting_peer(listener_addr) {
-            // Return early if already being connected to.
-            return Ok(None);
-        }
+        self.add_connecting_peer(listener_addr).map_err(|err| match err {
+            AddPeerError::AlreadyConnected => ConnectError::AlreadyConnected { address: peer_addr },
+            AddPeerError::AlreadyConnecting => ConnectError::AlreadyConnecting { address: peer_addr },
+        })?;
 
         // Verify the challenge request. If a disconnect reason was returned, send the disconnect message and abort.
         if let Some(reason) = self.verify_challenge_request(peer_addr, &peer_request) {
             send(&mut framed, peer_addr, reason.into()).await?;
-            return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
+            return Err(ConnectError::Other(anyhow!("Dropped '{peer_addr}' for reason: {reason:?}").into()));
         }
 
         /* Step 2: Send the challenge response followed by own challenge request. */
@@ -280,7 +288,9 @@ impl<N: Network> Router<N> {
         let response_nonce: u64 = rng.r#gen();
         let data = [peer_request.nonce.to_le_bytes(), response_nonce.to_le_bytes()].concat();
         let Ok(our_signature) = self.account.sign_bytes(&data, rng) else {
-            return Err(error(format!("Failed to sign the challenge request nonce from '{peer_addr}'")));
+            return Err(ConnectError::Other(
+                anyhow!("Failed to sign the challenge request nonce from '{peer_addr}'").into(),
+            ));
         };
         // Send the challenge response.
         let our_response = ChallengeResponse {
@@ -299,8 +309,9 @@ impl<N: Network> Router<N> {
 
         /* Step 3: Receive the challenge response. */
 
-        // Listen for the challenge response message.
+        // Wait for the challenge response message.
         let peer_response = expect_message!(Message::ChallengeResponse, framed, peer_addr);
+
         // Verify the challenge response. If a disconnect reason was returned, send the disconnect message and abort.
         if let Some(reason) = self
             .verify_challenge_response(
@@ -315,21 +326,21 @@ impl<N: Network> Router<N> {
             .await
         {
             send(&mut framed, peer_addr, reason.into()).await?;
-            return Err(error(format!("Dropped '{peer_addr}' for reason: {reason:?}")));
+            Err(ConnectError::Other(anyhow!("Dropped '{peer_addr}' for reason: {reason:?}").into()))
+        } else {
+            Ok(peer_request)
         }
-
-        Ok(Some(peer_request))
     }
 
     /// Ensure the peer is allowed to connect.
-    fn ensure_peer_is_allowed(&self, listener_addr: SocketAddr) -> Result<()> {
+    fn ensure_peer_is_allowed(&self, listener_addr: SocketAddr) -> Result<(), DisconnectReason> {
         // Ensure that it's not a self-connect attempt.
         if self.is_local_ip(listener_addr) {
-            bail!("Dropping connection request from '{listener_addr}' (attempted to self-connect)");
+            return Err(DisconnectReason::SelfConnect);
         }
         // Unknown peers are untrusted, so check if `allow_external_peers` is true.
         if !self.allow_external_peers() && !self.is_trusted(listener_addr) {
-            bail!("Dropping connection request from '{listener_addr}' (untrusted)");
+            return Err(DisconnectReason::NoExternalPeersAllowed);
         }
         Ok(())
     }