test: add mainnet connection test (#306)

crates/net/eth-wire/src/error.rs

@@ -112,6 +112,8 @@ pub enum P2PHandshakeError {
     NoResponse,
     #[error("handshake timed out")]
     Timeout,
+    #[error("Disconnected by peer: {0}")]
+    Disconnected(DisconnectReason),
 }
 
 /// An error that can occur when interacting with a [`Pinger`].

crates/net/eth-wire/src/p2pstream.rs

@@ -100,18 +100,35 @@ where
 
         // get the message id
         let id = *hello_bytes.first().ok_or(P2PStreamError::EmptyProtocolMessage)?;
+        let id = P2PMessageID::try_from(id)?;
 
-        // the first message sent MUST be the hello message
-        if id != P2PMessageID::Hello as u8 {
+        // the first message sent MUST be the hello OR disconnect message
+        match id {
+            P2PMessageID::Hello => {}
+            P2PMessageID::Disconnect => {
+                return match P2PMessage::decode(&mut &hello_bytes[..])? {
+                    P2PMessage::Disconnect(reason) => {
+                        tracing::error!("Disconnected by peer during handshake: {}", reason);
+                        Err(P2PStreamError::HandshakeError(P2PHandshakeError::Disconnected(reason)))
+                    }
+                    msg => Err(P2PStreamError::HandshakeError(
+                        P2PHandshakeError::NonHelloMessageInHandshake,
+                    )),
+                }
+            }
+            id => {
+                tracing::error!("expected hello message but received: {:?}", id);
                 return Err(P2PStreamError::HandshakeError(
                     P2PHandshakeError::NonHelloMessageInHandshake,
                 ))
             }
+        }
 
         let their_hello = match P2PMessage::decode(&mut &hello_bytes[..])? {
             P2PMessage::Hello(hello) => Ok(hello),
-            _ => {
-                // TODO: this should never occur due to the id check
+            msg => {
+                // Note: this should never occur due to the id check
+                tracing::error!("expected hello message but received: {:?}", msg);
                 Err(P2PStreamError::HandshakeError(P2PHandshakeError::NonHelloMessageInHandshake))
             }
         }?;
@@ -575,6 +592,7 @@ impl Decodable for P2PMessage {
 }
 
 /// Message IDs for `p2p` subprotocol messages.
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
 pub enum P2PMessageID {
     /// Message ID for the [`P2PMessage::Hello`] message.
     Hello = 0x00,
@@ -642,13 +660,13 @@ pub enum ProtocolVersion {
 }
 
 impl Encodable for ProtocolVersion {
+    fn encode(&self, out: &mut dyn bytes::BufMut) {
+        (*self as u8).encode(out)
+    }
     fn length(&self) -> usize {
         // the version should be a single byte
         (*self as u8).length()
     }
-    fn encode(&self, out: &mut dyn bytes::BufMut) {
-        (*self as u8).encode(out)
-    }
 }
 
 impl Decodable for ProtocolVersion {
@@ -748,12 +766,6 @@ impl TryFrom<u8> for DisconnectReason {
 }
 
 impl Encodable for DisconnectReason {
-    fn length(&self) -> usize {
-        // disconnect reasons are snappy encoded as follows:
-        // [0x01, 0x00, reason as u8]
-        // this is 3 bytes
-        3
-    }
     fn encode(&self, out: &mut dyn bytes::BufMut) {
         // disconnect reasons are snappy encoded as follows:
         // [0x01, 0x00, reason as u8]
@@ -762,6 +774,12 @@ impl Encodable for DisconnectReason {
         out.put_u8(0x00);
         out.put_u8(*self as u8);
     }
+    fn length(&self) -> usize {
+        // disconnect reasons are snappy encoded as follows:
+        // [0x01, 0x00, reason as u8]
+        // this is 3 bytes
+        3
+    }
 }
 
 impl Decodable for DisconnectReason {

crates/net/network/src/config.rs

@@ -147,6 +147,12 @@ impl<C> NetworkConfigBuilder<C> {
         self
     }
 
+    /// Sets the discv4 config to use.
+    pub fn discovery(mut self, builder: Discv4ConfigBuilder) -> Self {
+        self.discovery_v4_builder = builder;
+        self
+    }
+
     /// Consumes the type and creates the actual [`NetworkConfig`]
     pub fn build(self) -> NetworkConfig<C> {
         let Self {

crates/net/network/tests/it/connect.rs

@@ -2,8 +2,11 @@
 
 use super::testnet::Testnet;
 use futures::StreamExt;
-use reth_network::NetworkEvent;
-use std::collections::HashSet;
+use reth_discv4::{bootnodes::mainnet_nodes, Discv4Config};
+use reth_interfaces::test_utils::TestApi;
+use reth_network::{NetworkConfig, NetworkEvent, NetworkManager};
+use secp256k1::SecretKey;
+use std::{collections::HashSet, sync::Arc};
 
 #[tokio::test(flavor = "multi_thread")]
 async fn test_establish_connections() {
@@ -56,3 +59,24 @@ async fn test_establish_connections() {
         assert_eq!(net.peers()[2].num_peers(), 1);
     }
 }
+
+#[tokio::test(flavor = "multi_thread")]
+#[ignore]
+async fn test_connect_with_boot_nodes() {
+    reth_tracing::init_tracing();
+    let secret_key = SecretKey::new(&mut rand::thread_rng());
+    let mut discv4 = Discv4Config::builder();
+    discv4.add_boot_nodes(mainnet_nodes());
+
+    let config =
+        NetworkConfig::builder(Arc::new(TestApi::default()), secret_key).discovery(discv4).build();
+    let network = NetworkManager::new(config).await.unwrap();
+
+    let handle = network.handle().clone();
+    let mut events = handle.event_listener();
+    tokio::task::spawn(network);
+
+    while let Some(ev) = events.next().await {
+        dbg!(ev);
+    }
+}