From 04105ec82bbd9eeeab288a1de644a1dc2e4c9103 Mon Sep 17 00:00:00 2001 From: Matthias Seitz Date: Thu, 1 Dec 2022 17:59:52 +0100 Subject: [PATCH] test: add mainnet connection test (#306) --- crates/net/eth-wire/src/error.rs | 2 ++ crates/net/eth-wire/src/p2pstream.rs | 50 +++++++++++++++++--------- crates/net/network/src/config.rs | 6 ++++ crates/net/network/tests/it/connect.rs | 28 +++++++++++++-- 4 files changed, 68 insertions(+), 18 deletions(-) diff --git a/crates/net/eth-wire/src/error.rs b/crates/net/eth-wire/src/error.rs index fbfbd22fbd..333ed9f0e6 100644 --- a/crates/net/eth-wire/src/error.rs +++ b/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`]. diff --git a/crates/net/eth-wire/src/p2pstream.rs b/crates/net/eth-wire/src/p2pstream.rs index 0f28d50091..fced14eee6 100644 --- a/crates/net/eth-wire/src/p2pstream.rs +++ b/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 { - return Err(P2PStreamError::HandshakeError( - P2PHandshakeError::NonHelloMessageInHandshake, - )) + // 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 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 { diff --git a/crates/net/network/src/config.rs b/crates/net/network/src/config.rs index 16c53d64de..35b8ffe680 100644 --- a/crates/net/network/src/config.rs +++ b/crates/net/network/src/config.rs @@ -147,6 +147,12 @@ impl NetworkConfigBuilder { 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 { let Self { diff --git a/crates/net/network/tests/it/connect.rs b/crates/net/network/tests/it/connect.rs index af7a807098..775d39c087 100644 --- a/crates/net/network/tests/it/connect.rs +++ b/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); + } +}