1use reth_eth_wire_types::{
4 message::RequestPair, BlockBodies, BlockHeaders, Capabilities, DisconnectReason, EthMessage,
5 EthNetworkPrimitives, EthVersion, GetBlockBodies, GetBlockHeaders, GetNodeData,
6 GetPooledTransactions, GetReceipts, NetworkPrimitives, NodeData, PooledTransactions, Receipts,
7 Status,
8};
9use reth_ethereum_forks::ForkId;
10use reth_network_p2p::error::{RequestError, RequestResult};
11use reth_network_peers::PeerId;
12use reth_network_types::{PeerAddr, PeerKind};
13use reth_tokio_util::EventStream;
14use std::{
15 fmt,
16 net::SocketAddr,
17 pin::Pin,
18 sync::Arc,
19 task::{Context, Poll},
20};
21use tokio::sync::{mpsc, oneshot};
22use tokio_stream::{wrappers::UnboundedReceiverStream, Stream, StreamExt};
23
24pub struct PeerEventStream(Pin<Box<dyn Stream<Item = PeerEvent> + Send + Sync>>);
26
27impl fmt::Debug for PeerEventStream {
28 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
29 f.debug_struct("PeerEventStream").finish_non_exhaustive()
30 }
31}
32
33impl PeerEventStream {
34 pub fn new<S, T>(stream: S) -> Self
37 where
38 S: Stream<Item = T> + Send + Sync + 'static,
39 T: Into<PeerEvent> + 'static,
40 {
41 let mapped_stream = stream.map(Into::into);
42 Self(Box::pin(mapped_stream))
43 }
44}
45
46impl Stream for PeerEventStream {
47 type Item = PeerEvent;
48
49 fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
50 self.0.as_mut().poll_next(cx)
51 }
52}
53
54#[derive(Debug, Clone)]
56pub struct SessionInfo {
57 pub peer_id: PeerId,
59 pub remote_addr: SocketAddr,
61 pub client_version: Arc<str>,
63 pub capabilities: Arc<Capabilities>,
65 pub status: Arc<Status>,
67 pub version: EthVersion,
69 pub peer_kind: PeerKind,
71}
72
73#[derive(Debug, Clone)]
79pub enum PeerEvent {
80 SessionClosed {
82 peer_id: PeerId,
84 reason: Option<DisconnectReason>,
86 },
87 SessionEstablished(SessionInfo),
89 PeerAdded(PeerId),
91 PeerRemoved(PeerId),
93}
94
95#[derive(Debug)]
97pub enum NetworkEvent<R = PeerRequest> {
98 Peer(PeerEvent),
100 ActivePeerSession {
102 info: SessionInfo,
104 messages: PeerRequestSender<R>,
106 },
107}
108
109impl<R> Clone for NetworkEvent<R> {
110 fn clone(&self) -> Self {
111 match self {
112 Self::Peer(event) => Self::Peer(event.clone()),
113 Self::ActivePeerSession { info, messages } => {
114 Self::ActivePeerSession { info: info.clone(), messages: messages.clone() }
115 }
116 }
117 }
118}
119
120impl<R> From<NetworkEvent<R>> for PeerEvent {
121 fn from(event: NetworkEvent<R>) -> Self {
122 match event {
123 NetworkEvent::Peer(peer_event) => peer_event,
124 NetworkEvent::ActivePeerSession { info, .. } => Self::SessionEstablished(info),
125 }
126 }
127}
128
129#[auto_impl::auto_impl(&, Arc)]
131pub trait NetworkPeersEvents: Send + Sync {
132 fn peer_events(&self) -> PeerEventStream;
134}
135
136#[auto_impl::auto_impl(&, Arc)]
138pub trait NetworkEventListenerProvider: NetworkPeersEvents {
139 type Primitives: NetworkPrimitives;
141
142 fn event_listener(&self) -> EventStream<NetworkEvent<PeerRequest<Self::Primitives>>>;
144 fn discovery_listener(&self) -> UnboundedReceiverStream<DiscoveryEvent>;
148}
149
150#[derive(Debug, Clone, PartialEq, Eq)]
152pub enum DiscoveryEvent {
153 NewNode(DiscoveredEvent),
155 EnrForkId(PeerId, ForkId),
157}
158
159#[derive(Debug, Clone, PartialEq, Eq)]
161pub enum DiscoveredEvent {
162 EventQueued {
174 peer_id: PeerId,
176 addr: PeerAddr,
178 fork_id: Option<ForkId>,
181 },
182}
183
184#[derive(Debug)]
186pub enum PeerRequest<N: NetworkPrimitives = EthNetworkPrimitives> {
187 GetBlockHeaders {
191 request: GetBlockHeaders,
193 response: oneshot::Sender<RequestResult<BlockHeaders<N::BlockHeader>>>,
195 },
196 GetBlockBodies {
200 request: GetBlockBodies,
202 response: oneshot::Sender<RequestResult<BlockBodies<N::BlockBody>>>,
204 },
205 GetPooledTransactions {
209 request: GetPooledTransactions,
211 response: oneshot::Sender<RequestResult<PooledTransactions<N::PooledTransaction>>>,
213 },
214 GetNodeData {
218 request: GetNodeData,
220 response: oneshot::Sender<RequestResult<NodeData>>,
222 },
223 GetReceipts {
227 request: GetReceipts,
229 response: oneshot::Sender<RequestResult<Receipts<N::Receipt>>>,
231 },
232}
233
234impl<N: NetworkPrimitives> PeerRequest<N> {
237 pub fn send_bad_response(self) {
239 self.send_err_response(RequestError::BadResponse)
240 }
241
242 pub fn send_err_response(self, err: RequestError) {
244 let _ = match self {
245 Self::GetBlockHeaders { response, .. } => response.send(Err(err)).ok(),
246 Self::GetBlockBodies { response, .. } => response.send(Err(err)).ok(),
247 Self::GetPooledTransactions { response, .. } => response.send(Err(err)).ok(),
248 Self::GetNodeData { response, .. } => response.send(Err(err)).ok(),
249 Self::GetReceipts { response, .. } => response.send(Err(err)).ok(),
250 };
251 }
252
253 pub fn create_request_message(&self, request_id: u64) -> EthMessage<N> {
255 match self {
256 Self::GetBlockHeaders { request, .. } => {
257 EthMessage::GetBlockHeaders(RequestPair { request_id, message: *request })
258 }
259 Self::GetBlockBodies { request, .. } => {
260 EthMessage::GetBlockBodies(RequestPair { request_id, message: request.clone() })
261 }
262 Self::GetPooledTransactions { request, .. } => {
263 EthMessage::GetPooledTransactions(RequestPair {
264 request_id,
265 message: request.clone(),
266 })
267 }
268 Self::GetNodeData { request, .. } => {
269 EthMessage::GetNodeData(RequestPair { request_id, message: request.clone() })
270 }
271 Self::GetReceipts { request, .. } => {
272 EthMessage::GetReceipts(RequestPair { request_id, message: request.clone() })
273 }
274 }
275 }
276
277 pub fn into_get_pooled_transactions(self) -> Option<GetPooledTransactions> {
279 match self {
280 Self::GetPooledTransactions { request, .. } => Some(request),
281 _ => None,
282 }
283 }
284}
285
286pub struct PeerRequestSender<R = PeerRequest> {
288 pub peer_id: PeerId,
290 pub to_session_tx: mpsc::Sender<R>,
292}
293
294impl<R> Clone for PeerRequestSender<R> {
295 fn clone(&self) -> Self {
296 Self { peer_id: self.peer_id, to_session_tx: self.to_session_tx.clone() }
297 }
298}
299
300impl<R> PeerRequestSender<R> {
303 pub const fn new(peer_id: PeerId, to_session_tx: mpsc::Sender<R>) -> Self {
305 Self { peer_id, to_session_tx }
306 }
307
308 pub fn try_send(&self, req: R) -> Result<(), mpsc::error::TrySendError<R>> {
310 self.to_session_tx.try_send(req)
311 }
312
313 pub const fn peer_id(&self) -> &PeerId {
315 &self.peer_id
316 }
317}
318
319impl<R> fmt::Debug for PeerRequestSender<R> {
320 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
321 f.debug_struct("PeerRequestSender").field("peer_id", &self.peer_id).finish_non_exhaustive()
322 }
323}