1use crate::{
2 listener::{ConnectionListener, ListenerEvent},
3 message::PeerMessage,
4 peers::{InboundConnectionError, PeersManager},
5 protocol::IntoRlpxSubProtocol,
6 session::{Direction, PendingSessionHandshakeError, SessionEvent, SessionId, SessionManager},
7 state::{NetworkState, SessionActivation, StateAction},
8};
9use futures::Stream;
10use reth_eth_wire::{
11 errors::EthStreamError, Capabilities, DisconnectReason, EthNetworkPrimitives, EthVersion,
12 NetworkPrimitives, UnifiedStatus,
13};
14use reth_network_api::{PeerRequest, PeerRequestSender};
15use reth_network_peers::PeerId;
16use std::{
17 io,
18 net::SocketAddr,
19 pin::Pin,
20 sync::Arc,
21 task::{Context, Poll},
22};
23use tracing::trace;
24
25#[cfg_attr(doc, aquamarine::aquamarine)]
26#[derive(Debug)]
50#[must_use = "Swarm does nothing unless polled"]
51pub(crate) struct Swarm<N: NetworkPrimitives = EthNetworkPrimitives> {
52 incoming: ConnectionListener,
54 sessions: SessionManager<N>,
56 state: NetworkState<N>,
58}
59
60impl<N: NetworkPrimitives> Swarm<N> {
63 pub(crate) const fn new(
65 incoming: ConnectionListener,
66 sessions: SessionManager<N>,
67 state: NetworkState<N>,
68 ) -> Self {
69 Self { incoming, sessions, state }
70 }
71
72 pub(crate) fn add_rlpx_sub_protocol(&mut self, protocol: impl IntoRlpxSubProtocol) {
74 self.sessions_mut().add_rlpx_sub_protocol(protocol);
75 }
76
77 pub(crate) const fn state(&self) -> &NetworkState<N> {
79 &self.state
80 }
81
82 pub(crate) const fn state_mut(&mut self) -> &mut NetworkState<N> {
84 &mut self.state
85 }
86
87 pub(crate) const fn listener(&self) -> &ConnectionListener {
89 &self.incoming
90 }
91
92 pub(crate) const fn sessions(&self) -> &SessionManager<N> {
94 &self.sessions
95 }
96
97 pub(crate) const fn sessions_mut(&mut self) -> &mut SessionManager<N> {
99 &mut self.sessions
100 }
101
102 pub(crate) const fn peers(&self) -> &PeersManager {
104 self.state.peers()
105 }
106
107 pub(crate) const fn peers_mut(&mut self) -> &mut PeersManager {
109 self.state.peers_mut()
110 }
111}
112
113impl<N: NetworkPrimitives> Swarm<N> {
114 pub(crate) fn dial_outbound(&mut self, remote_addr: SocketAddr, remote_id: PeerId) -> bool {
116 if let Err(err) = self.sessions.try_dial_outbound(remote_addr, remote_id) {
117 trace!(target: "net", %err, "Outgoing connection rejected, capacity already reached.");
118 self.peers_mut().on_outgoing_pending_session_gracefully_closed(&remote_id);
120 return false
121 }
122 true
123 }
124
125 fn on_session_event(&mut self, event: SessionEvent<N>) -> Option<SwarmEvent<N>> {
130 match event {
131 SessionEvent::SessionEstablished {
132 peer_id,
133 remote_addr,
134 client_version,
135 capabilities,
136 version,
137 status,
138 messages,
139 direction,
140 timeout,
141 range_info,
142 supports_snap,
143 } => {
144 self.state.on_session_activated(SessionActivation {
145 peer: peer_id,
146 capabilities: capabilities.clone(),
147 status: status.clone(),
148 request_tx: messages.clone(),
149 timeout,
150 range_info,
151 supports_snap,
152 });
153 Some(SwarmEvent::SessionEstablished {
154 peer_id,
155 remote_addr,
156 client_version,
157 capabilities,
158 version,
159 messages,
160 status,
161 direction,
162 })
163 }
164 SessionEvent::AlreadyConnected { peer_id, remote_addr, direction } => {
165 trace!(target: "net", ?peer_id, ?remote_addr, ?direction, "already connected");
166 self.state.peers_mut().on_already_connected(direction);
167 None
168 }
169 SessionEvent::ValidMessage { peer_id, message } => {
170 Some(SwarmEvent::ValidMessage { peer_id, message })
171 }
172 SessionEvent::IncomingPendingSessionClosed { remote_addr, error } => {
173 Some(SwarmEvent::IncomingPendingSessionClosed { remote_addr, error })
174 }
175 SessionEvent::OutgoingPendingSessionClosed { remote_addr, peer_id, error } => {
176 Some(SwarmEvent::OutgoingPendingSessionClosed { remote_addr, peer_id, error })
177 }
178 SessionEvent::Disconnected { peer_id, remote_addr } => {
179 self.state.on_session_closed(peer_id);
180 Some(SwarmEvent::SessionClosed { peer_id, remote_addr, error: None })
181 }
182 SessionEvent::SessionClosedOnConnectionError { peer_id, remote_addr, error } => {
183 self.state.on_session_closed(peer_id);
184 Some(SwarmEvent::SessionClosed { peer_id, remote_addr, error: Some(error) })
185 }
186 SessionEvent::OutgoingConnectionError { remote_addr, peer_id, error } => {
187 Some(SwarmEvent::OutgoingConnectionError { peer_id, remote_addr, error })
188 }
189 SessionEvent::BadMessage { peer_id } => Some(SwarmEvent::BadMessage { peer_id }),
190 SessionEvent::ProtocolBreach { peer_id } => {
191 Some(SwarmEvent::ProtocolBreach { peer_id })
192 }
193 }
194 }
195
196 fn on_connection(&mut self, event: ListenerEvent) -> Option<SwarmEvent<N>> {
200 match event {
201 ListenerEvent::Error(err) => return Some(SwarmEvent::TcpListenerError(err)),
202 ListenerEvent::ListenerClosed { local_address: address } => {
203 return Some(SwarmEvent::TcpListenerClosed { remote_addr: address })
204 }
205 ListenerEvent::Incoming { stream, remote_addr } => {
206 if self.is_shutting_down() {
208 return None
209 }
210 if let Err(err) = self.peers_mut().on_incoming_pending_session(remote_addr.ip()) {
212 match err {
213 InboundConnectionError::IpBanned => {
214 trace!(target: "net", ?remote_addr, "The incoming ip address is in the ban list");
215 }
216 InboundConnectionError::ExceedsCapacity => {
217 trace!(target: "net", ?remote_addr, "No capacity for incoming connection");
218 self.sessions.try_disconnect_incoming_connection(
219 stream,
220 DisconnectReason::TooManyPeers,
221 );
222 }
223 }
224 return None
225 }
226
227 match self.sessions.on_incoming(stream, remote_addr) {
228 Ok(session_id) => {
229 trace!(target: "net", ?remote_addr, "Incoming connection");
230 return Some(SwarmEvent::IncomingTcpConnection { session_id, remote_addr })
231 }
232 Err(err) => {
233 trace!(target: "net", %err, "Incoming connection rejected, capacity already reached.");
234 self.state_mut()
235 .peers_mut()
236 .on_incoming_pending_session_rejected_internally();
237 }
238 }
239 }
240 }
241 None
242 }
243
244 fn on_state_action(&mut self, event: StateAction<N>) -> Option<SwarmEvent<N>> {
246 match event {
247 StateAction::Connect { remote_addr, peer_id } => {
248 if self.dial_outbound(remote_addr, peer_id) {
249 return Some(SwarmEvent::OutgoingTcpConnection { remote_addr, peer_id })
250 }
251 }
252 StateAction::Disconnect { peer_id, reason } => {
253 self.sessions.disconnect(peer_id, reason);
254 }
255 StateAction::NewBlock { peer_id, block: msg } => {
256 let msg = PeerMessage::NewBlock(msg);
257 self.sessions.send_message(&peer_id, msg);
258 }
259 StateAction::NewBlockHashes { peer_id, hashes } => {
260 let msg = PeerMessage::NewBlockHashes(hashes);
261 self.sessions.send_message(&peer_id, msg);
262 }
263 StateAction::PeerAdded(peer_id) => return Some(SwarmEvent::PeerAdded(peer_id)),
264 StateAction::PeerRemoved(peer_id) => return Some(SwarmEvent::PeerRemoved(peer_id)),
265 StateAction::DiscoveredNode { peer_id, addr, fork_id } => {
266 if self.is_shutting_down() {
267 return None
268 }
269
270 let enforce = self.peers().enforce_enr_fork_id();
277 let allow = match fork_id {
278 Some(f) => self.sessions.is_valid_fork_id(f),
279 None => !enforce,
280 };
281 if allow {
282 self.peers_mut().add_peer(peer_id, addr, fork_id);
283 }
284 }
285 StateAction::DiscoveredEnrForkId { peer_id, addr, fork_id } => {
286 if self.sessions.is_valid_fork_id(fork_id) {
287 self.peers_mut().add_peer(peer_id, addr, Some(fork_id));
288 } else {
289 trace!(target: "net", ?peer_id, remote_fork_id=?fork_id, our_fork_id=?self.sessions.fork_id(), "fork id mismatch, removing peer");
290 self.peers_mut().remove_peer(peer_id);
291 }
292 }
293 }
294 None
295 }
296
297 pub(crate) const fn on_shutdown_requested(&mut self) {
299 self.peers_mut().on_shutdown();
300 }
301
302 #[inline]
304 pub(crate) const fn is_shutting_down(&self) -> bool {
305 self.peers().connection_state().is_shutting_down()
306 }
307
308 pub(crate) const fn on_network_state_change(&mut self, network_state: NetworkConnectionState) {
310 self.peers_mut().on_network_state_change(network_state);
311 }
312}
313
314impl<N: NetworkPrimitives> Stream for Swarm<N> {
315 type Item = SwarmEvent<N>;
316
317 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
325 let this = self.get_mut();
326
327 loop {
331 while let Poll::Ready(action) = this.state.poll(cx) {
332 if let Some(event) = this.on_state_action(action) {
333 return Poll::Ready(Some(event))
334 }
335 }
336
337 match this.sessions.poll(cx, |peer_id| this.state.peers().has_inbound_capacity(peer_id))
339 {
340 Poll::Pending => {}
341 Poll::Ready(event) => {
342 if let Some(event) = this.on_session_event(event) {
343 return Poll::Ready(Some(event))
344 }
345 continue
346 }
347 }
348
349 match Pin::new(&mut this.incoming).poll(cx) {
351 Poll::Pending => {}
352 Poll::Ready(event) => {
353 if let Some(event) = this.on_connection(event) {
354 return Poll::Ready(Some(event))
355 }
356 continue
357 }
358 }
359
360 return Poll::Pending
361 }
362 }
363}
364
365pub(crate) enum SwarmEvent<N: NetworkPrimitives = EthNetworkPrimitives> {
368 ValidMessage {
370 peer_id: PeerId,
372 message: PeerMessage<N>,
374 },
375 BadMessage {
377 peer_id: PeerId,
379 },
380 ProtocolBreach {
382 peer_id: PeerId,
384 },
385 TcpListenerClosed {
387 remote_addr: SocketAddr,
389 },
390 TcpListenerError(io::Error),
392 IncomingTcpConnection {
397 session_id: SessionId,
399 remote_addr: SocketAddr,
401 },
402 OutgoingTcpConnection {
404 peer_id: PeerId,
406 remote_addr: SocketAddr,
407 },
408 SessionEstablished {
409 peer_id: PeerId,
410 remote_addr: SocketAddr,
411 client_version: Arc<str>,
412 capabilities: Arc<Capabilities>,
413 version: EthVersion,
415 messages: PeerRequestSender<PeerRequest<N>>,
416 status: Arc<UnifiedStatus>,
417 direction: Direction,
418 },
419 SessionClosed {
420 peer_id: PeerId,
421 remote_addr: SocketAddr,
422 error: Option<EthStreamError>,
424 },
425 PeerAdded(PeerId),
427 PeerRemoved(PeerId),
429 IncomingPendingSessionClosed {
431 remote_addr: SocketAddr,
432 error: Option<PendingSessionHandshakeError>,
433 },
434 OutgoingPendingSessionClosed {
436 remote_addr: SocketAddr,
437 peer_id: PeerId,
438 error: Option<PendingSessionHandshakeError>,
439 },
440 OutgoingConnectionError { remote_addr: SocketAddr, peer_id: PeerId, error: io::Error },
442}
443
444#[derive(Debug, Default)]
449pub enum NetworkConnectionState {
450 #[default]
452 Active,
453 ShuttingDown,
455 Hibernate,
457}
458
459impl NetworkConnectionState {
460 pub(crate) const fn is_active(&self) -> bool {
462 matches!(self, Self::Active)
463 }
464
465 pub(crate) const fn is_shutting_down(&self) -> bool {
467 matches!(self, Self::ShuttingDown)
468 }
469}