use crate::{
builder::ETH_REQUEST_CHANNEL_CAPACITY,
error::NetworkError,
eth_requests::EthRequestHandler,
protocol::IntoRlpxSubProtocol,
transactions::{TransactionsHandle, TransactionsManager, TransactionsManagerConfig},
NetworkConfig, NetworkConfigBuilder, NetworkHandle, NetworkManager,
};
use futures::{FutureExt, StreamExt};
use pin_project::pin_project;
use reth_chainspec::{ChainSpecProvider, Hardforks, MAINNET};
use reth_eth_wire::{
protocol::Protocol, DisconnectReason, EthNetworkPrimitives, HelloMessageWithProtocols,
};
use reth_network_api::{
events::{PeerEvent, SessionInfo},
test_utils::{PeersHandle, PeersHandleProvider},
NetworkEvent, NetworkEventListenerProvider, NetworkInfo, Peers,
};
use reth_network_peers::PeerId;
use reth_primitives::{PooledTransaction, TransactionSigned};
use reth_storage_api::{
noop::NoopProvider, BlockReader, BlockReaderIdExt, HeaderProvider, StateProviderFactory,
};
use reth_tasks::TokioTaskExecutor;
use reth_tokio_util::EventStream;
use reth_transaction_pool::{
blobstore::InMemoryBlobStore,
test_utils::{TestPool, TestPoolBuilder},
EthTransactionPool, PoolTransaction, TransactionPool, TransactionValidationTaskExecutor,
};
use secp256k1::SecretKey;
use std::{
fmt,
future::Future,
net::{Ipv4Addr, SocketAddr, SocketAddrV4},
pin::Pin,
task::{Context, Poll},
};
use tokio::{
sync::{
mpsc::{channel, unbounded_channel},
oneshot,
},
task::JoinHandle,
};
pub struct Testnet<C, Pool> {
peers: Vec<Peer<C, Pool>>,
}
impl<C> Testnet<C, TestPool>
where
C: BlockReader + HeaderProvider + Clone + 'static + ChainSpecProvider<ChainSpec: Hardforks>,
{
pub async fn create_with(num_peers: usize, provider: C) -> Self {
Self::try_create_with(num_peers, provider).await.unwrap()
}
pub async fn try_create_with(num_peers: usize, provider: C) -> Result<Self, NetworkError> {
let mut this = Self { peers: Vec::with_capacity(num_peers) };
for _ in 0..num_peers {
let config = PeerConfig::new(provider.clone());
this.add_peer_with_config(config).await?;
}
Ok(this)
}
pub async fn extend_peer_with_config(
&mut self,
configs: impl IntoIterator<Item = PeerConfig<C>>,
) -> Result<(), NetworkError> {
let peers = configs.into_iter().map(|c| c.launch()).collect::<Vec<_>>();
let peers = futures::future::join_all(peers).await;
for peer in peers {
self.peers.push(peer?);
}
Ok(())
}
}
impl<C, Pool> Testnet<C, Pool>
where
C: BlockReader + HeaderProvider + Clone + 'static,
Pool: TransactionPool,
{
pub fn peers_mut(&mut self) -> &mut [Peer<C, Pool>] {
&mut self.peers
}
pub fn peers(&self) -> &[Peer<C, Pool>] {
&self.peers
}
pub fn remove_peer(&mut self, index: usize) -> Peer<C, Pool> {
self.peers.remove(index)
}
pub fn peers_iter_mut(&mut self) -> impl Iterator<Item = &mut Peer<C, Pool>> + '_ {
self.peers.iter_mut()
}
pub fn peers_iter(&self) -> impl Iterator<Item = &Peer<C, Pool>> + '_ {
self.peers.iter()
}
pub async fn add_peer_with_config(
&mut self,
config: PeerConfig<C>,
) -> Result<(), NetworkError> {
let PeerConfig { config, client, secret_key } = config;
let network = NetworkManager::new(config).await?;
let peer = Peer {
network,
client,
secret_key,
request_handler: None,
transactions_manager: None,
pool: None,
};
self.peers.push(peer);
Ok(())
}
pub fn handles(&self) -> impl Iterator<Item = NetworkHandle<EthNetworkPrimitives>> + '_ {
self.peers.iter().map(|p| p.handle())
}
pub fn map_pool<F, P>(self, f: F) -> Testnet<C, P>
where
F: Fn(Peer<C, Pool>) -> Peer<C, P>,
P: TransactionPool,
{
Testnet { peers: self.peers.into_iter().map(f).collect() }
}
pub fn for_each<F>(&self, f: F)
where
F: Fn(&Peer<C, Pool>),
{
self.peers.iter().for_each(f)
}
pub fn for_each_mut<F>(&mut self, f: F)
where
F: FnMut(&mut Peer<C, Pool>),
{
self.peers.iter_mut().for_each(f)
}
}
impl<C, Pool> Testnet<C, Pool>
where
C: StateProviderFactory + BlockReaderIdExt + HeaderProvider + Clone + 'static,
Pool: TransactionPool,
{
pub fn with_eth_pool(self) -> Testnet<C, EthTransactionPool<C, InMemoryBlobStore>> {
self.map_pool(|peer| {
let blob_store = InMemoryBlobStore::default();
let pool = TransactionValidationTaskExecutor::eth(
peer.client.clone(),
MAINNET.clone(),
blob_store.clone(),
TokioTaskExecutor::default(),
);
peer.map_transactions_manager(EthTransactionPool::eth_pool(
pool,
blob_store,
Default::default(),
))
})
}
pub fn with_eth_pool_config(
self,
tx_manager_config: TransactionsManagerConfig,
) -> Testnet<C, EthTransactionPool<C, InMemoryBlobStore>> {
self.map_pool(|peer| {
let blob_store = InMemoryBlobStore::default();
let pool = TransactionValidationTaskExecutor::eth(
peer.client.clone(),
MAINNET.clone(),
blob_store.clone(),
TokioTaskExecutor::default(),
);
peer.map_transactions_manager_with_config(
EthTransactionPool::eth_pool(pool, blob_store, Default::default()),
tx_manager_config.clone(),
)
})
}
}
impl<C, Pool> Testnet<C, Pool>
where
C: BlockReader<
Block = reth_primitives::Block,
Receipt = reth_primitives::Receipt,
Header = reth_primitives::Header,
> + HeaderProvider
+ Clone
+ Unpin
+ 'static,
Pool: TransactionPool<
Transaction: PoolTransaction<Consensus = TransactionSigned, Pooled = PooledTransaction>,
> + Unpin
+ 'static,
{
pub fn spawn(self) -> TestnetHandle<C, Pool> {
let (tx, rx) = oneshot::channel::<oneshot::Sender<Self>>();
let peers = self.peers.iter().map(|peer| peer.peer_handle()).collect::<Vec<_>>();
let mut net = self;
let handle = tokio::task::spawn(async move {
let mut tx = None;
tokio::select! {
_ = &mut net => {}
inc = rx => {
tx = inc.ok();
}
}
if let Some(tx) = tx {
let _ = tx.send(net);
}
});
TestnetHandle { _handle: handle, peers, terminate: tx }
}
}
impl Testnet<NoopProvider, TestPool> {
pub async fn create(num_peers: usize) -> Self {
Self::try_create(num_peers).await.unwrap()
}
pub async fn try_create(num_peers: usize) -> Result<Self, NetworkError> {
let mut this = Self::default();
this.extend_peer_with_config((0..num_peers).map(|_| Default::default())).await?;
Ok(this)
}
pub async fn add_peer(&mut self) -> Result<(), NetworkError> {
self.add_peer_with_config(Default::default()).await
}
}
impl<C, Pool> Default for Testnet<C, Pool> {
fn default() -> Self {
Self { peers: Vec::new() }
}
}
impl<C, Pool> fmt::Debug for Testnet<C, Pool> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Testnet {{}}").finish_non_exhaustive()
}
}
impl<C, Pool> Future for Testnet<C, Pool>
where
C: BlockReader<
Block = reth_primitives::Block,
Receipt = reth_primitives::Receipt,
Header = reth_primitives::Header,
> + HeaderProvider
+ Unpin
+ 'static,
Pool: TransactionPool<
Transaction: PoolTransaction<Consensus = TransactionSigned, Pooled = PooledTransaction>,
> + Unpin
+ 'static,
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
for peer in &mut this.peers {
let _ = peer.poll_unpin(cx);
}
Poll::Pending
}
}
#[derive(Debug)]
pub struct TestnetHandle<C, Pool> {
_handle: JoinHandle<()>,
peers: Vec<PeerHandle<Pool>>,
terminate: oneshot::Sender<oneshot::Sender<Testnet<C, Pool>>>,
}
impl<C, Pool> TestnetHandle<C, Pool> {
pub async fn terminate(self) -> Testnet<C, Pool> {
let (tx, rx) = oneshot::channel();
self.terminate.send(tx).unwrap();
rx.await.unwrap()
}
pub fn peers(&self) -> &[PeerHandle<Pool>] {
&self.peers
}
pub async fn connect_peers(&self) {
if self.peers.len() < 2 {
return
}
let streams =
self.peers.iter().map(|handle| NetworkEventStream::new(handle.event_listener()));
for (idx, handle) in self.peers.iter().enumerate().take(self.peers.len() - 1) {
for idx in (idx + 1)..self.peers.len() {
let neighbour = &self.peers[idx];
handle.network.add_peer(*neighbour.peer_id(), neighbour.local_addr());
}
}
let num_sessions_per_peer = self.peers.len() - 1;
let fut = streams.into_iter().map(|mut stream| async move {
stream.take_session_established(num_sessions_per_peer).await
});
futures::future::join_all(fut).await;
}
}
#[pin_project]
#[derive(Debug)]
pub struct Peer<C, Pool = TestPool> {
#[pin]
network: NetworkManager<EthNetworkPrimitives>,
#[pin]
request_handler: Option<EthRequestHandler<C, EthNetworkPrimitives>>,
#[pin]
transactions_manager: Option<TransactionsManager<Pool, EthNetworkPrimitives>>,
pool: Option<Pool>,
client: C,
secret_key: SecretKey,
}
impl<C, Pool> Peer<C, Pool>
where
C: BlockReader + HeaderProvider + Clone + 'static,
Pool: TransactionPool,
{
pub fn num_peers(&self) -> usize {
self.network.num_connected_peers()
}
pub fn add_rlpx_sub_protocol(&mut self, protocol: impl IntoRlpxSubProtocol) {
self.network.add_rlpx_sub_protocol(protocol);
}
pub fn peer_handle(&self) -> PeerHandle<Pool> {
PeerHandle {
network: self.network.handle().clone(),
pool: self.pool.clone(),
transactions: self.transactions_manager.as_ref().map(|mgr| mgr.handle()),
}
}
pub const fn local_addr(&self) -> SocketAddr {
self.network.local_addr()
}
pub fn peer_id(&self) -> PeerId {
*self.network.peer_id()
}
pub fn network_mut(&mut self) -> &mut NetworkManager<EthNetworkPrimitives> {
&mut self.network
}
pub fn handle(&self) -> NetworkHandle<EthNetworkPrimitives> {
self.network.handle().clone()
}
pub const fn pool(&self) -> Option<&Pool> {
self.pool.as_ref()
}
pub fn install_request_handler(&mut self) {
let (tx, rx) = channel(ETH_REQUEST_CHANNEL_CAPACITY);
self.network.set_eth_request_handler(tx);
let peers = self.network.peers_handle();
let request_handler = EthRequestHandler::new(self.client.clone(), peers, rx);
self.request_handler = Some(request_handler);
}
pub fn install_transactions_manager(&mut self, pool: Pool) {
let (tx, rx) = unbounded_channel();
self.network.set_transactions(tx);
let transactions_manager = TransactionsManager::new(
self.handle(),
pool.clone(),
rx,
TransactionsManagerConfig::default(),
);
self.transactions_manager = Some(transactions_manager);
self.pool = Some(pool);
}
pub fn map_transactions_manager<P>(self, pool: P) -> Peer<C, P>
where
P: TransactionPool,
{
let Self { mut network, request_handler, client, secret_key, .. } = self;
let (tx, rx) = unbounded_channel();
network.set_transactions(tx);
let transactions_manager = TransactionsManager::new(
network.handle().clone(),
pool.clone(),
rx,
TransactionsManagerConfig::default(),
);
Peer {
network,
request_handler,
transactions_manager: Some(transactions_manager),
pool: Some(pool),
client,
secret_key,
}
}
pub fn map_transactions_manager_with_config<P>(
self,
pool: P,
config: TransactionsManagerConfig,
) -> Peer<C, P>
where
P: TransactionPool,
{
let Self { mut network, request_handler, client, secret_key, .. } = self;
let (tx, rx) = unbounded_channel();
network.set_transactions(tx);
let transactions_manager = TransactionsManager::new(
network.handle().clone(),
pool.clone(),
rx,
config, );
Peer {
network,
request_handler,
transactions_manager: Some(transactions_manager),
pool: Some(pool),
client,
secret_key,
}
}
}
impl<C> Peer<C>
where
C: BlockReader + HeaderProvider + Clone + 'static,
{
pub fn install_test_pool(&mut self) {
self.install_transactions_manager(TestPoolBuilder::default().into())
}
}
impl<C, Pool> Future for Peer<C, Pool>
where
C: BlockReader<
Block = reth_primitives::Block,
Receipt = reth_primitives::Receipt,
Header = reth_primitives::Header,
> + HeaderProvider
+ Unpin
+ 'static,
Pool: TransactionPool<
Transaction: PoolTransaction<Consensus = TransactionSigned, Pooled = PooledTransaction>,
> + Unpin
+ 'static,
{
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
if let Some(request) = this.request_handler.as_pin_mut() {
let _ = request.poll(cx);
}
if let Some(tx_manager) = this.transactions_manager.as_pin_mut() {
let _ = tx_manager.poll(cx);
}
this.network.poll(cx)
}
}
#[derive(Debug)]
pub struct PeerConfig<C = NoopProvider> {
config: NetworkConfig<C>,
client: C,
secret_key: SecretKey,
}
#[derive(Debug)]
pub struct PeerHandle<Pool> {
network: NetworkHandle<EthNetworkPrimitives>,
transactions: Option<TransactionsHandle<EthNetworkPrimitives>>,
pool: Option<Pool>,
}
impl<Pool> PeerHandle<Pool> {
pub fn peer_id(&self) -> &PeerId {
self.network.peer_id()
}
pub fn peer_handle(&self) -> &PeersHandle {
self.network.peers_handle()
}
pub fn local_addr(&self) -> SocketAddr {
self.network.local_addr()
}
pub fn event_listener(&self) -> EventStream<NetworkEvent> {
self.network.event_listener()
}
pub const fn transactions(&self) -> Option<&TransactionsHandle> {
self.transactions.as_ref()
}
pub const fn pool(&self) -> Option<&Pool> {
self.pool.as_ref()
}
pub const fn network(&self) -> &NetworkHandle<EthNetworkPrimitives> {
&self.network
}
}
impl<C> PeerConfig<C>
where
C: BlockReader + HeaderProvider + Clone + 'static,
{
pub async fn launch(self) -> Result<Peer<C>, NetworkError> {
let Self { config, client, secret_key } = self;
let network = NetworkManager::new(config).await?;
let peer = Peer {
network,
client,
secret_key,
request_handler: None,
transactions_manager: None,
pool: None,
};
Ok(peer)
}
pub fn new(client: C) -> Self
where
C: ChainSpecProvider<ChainSpec: Hardforks>,
{
let secret_key = SecretKey::new(&mut rand::thread_rng());
let config = Self::network_config_builder(secret_key).build(client.clone());
Self { config, client, secret_key }
}
pub fn with_secret_key(client: C, secret_key: SecretKey) -> Self
where
C: ChainSpecProvider<ChainSpec: Hardforks>,
{
let config = Self::network_config_builder(secret_key).build(client.clone());
Self { config, client, secret_key }
}
pub fn with_protocols(client: C, protocols: impl IntoIterator<Item = Protocol>) -> Self
where
C: ChainSpecProvider<ChainSpec: Hardforks>,
{
let secret_key = SecretKey::new(&mut rand::thread_rng());
let builder = Self::network_config_builder(secret_key);
let hello_message =
HelloMessageWithProtocols::builder(builder.get_peer_id()).protocols(protocols).build();
let config = builder.hello_message(hello_message).build(client.clone());
Self { config, client, secret_key }
}
fn network_config_builder(secret_key: SecretKey) -> NetworkConfigBuilder {
NetworkConfigBuilder::new(secret_key)
.listener_addr(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)))
.discovery_addr(SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 0)))
.disable_dns_discovery()
.disable_discv4_discovery()
}
}
impl Default for PeerConfig {
fn default() -> Self {
Self::new(NoopProvider::default())
}
}
#[derive(Debug)]
pub struct NetworkEventStream {
inner: EventStream<NetworkEvent>,
}
impl NetworkEventStream {
pub const fn new(inner: EventStream<NetworkEvent>) -> Self {
Self { inner }
}
pub async fn next_session_closed(&mut self) -> Option<(PeerId, Option<DisconnectReason>)> {
while let Some(ev) = self.inner.next().await {
match ev {
NetworkEvent::Peer(PeerEvent::SessionClosed { peer_id, reason }) => {
return Some((peer_id, reason))
}
_ => continue,
}
}
None
}
pub async fn next_session_established(&mut self) -> Option<PeerId> {
while let Some(ev) = self.inner.next().await {
match ev {
NetworkEvent::ActivePeerSession { info, .. } |
NetworkEvent::Peer(PeerEvent::SessionEstablished(info)) => {
return Some(info.peer_id)
}
_ => continue,
}
}
None
}
pub async fn take_session_established(&mut self, mut num: usize) -> Vec<PeerId> {
if num == 0 {
return Vec::new();
}
let mut peers = Vec::with_capacity(num);
while let Some(ev) = self.inner.next().await {
match ev {
NetworkEvent::ActivePeerSession { info: SessionInfo { peer_id, .. }, .. } => {
peers.push(peer_id);
num -= 1;
if num == 0 {
return peers;
}
}
_ => continue,
}
}
peers
}
pub async fn peer_added_and_established(&mut self) -> Option<PeerId> {
let peer_id = match self.inner.next().await {
Some(NetworkEvent::Peer(PeerEvent::PeerAdded(peer_id))) => peer_id,
_ => return None,
};
match self.inner.next().await {
Some(NetworkEvent::ActivePeerSession {
info: SessionInfo { peer_id: peer_id2, .. },
..
}) => {
debug_assert_eq!(
peer_id, peer_id2,
"PeerAdded peer_id {peer_id} does not match SessionEstablished peer_id {peer_id2}"
);
Some(peer_id)
}
_ => None,
}
}
}