1use std::{
5 sync::Arc,
6 task::Poll,
7 time::{Duration, Instant},
8};
9
10use ahash::HashMap;
11use libp2p::{
12 PeerId,
13 request_response::{self, OutboundRequestId, ProtocolSupport, ResponseChannel},
14 swarm::{CloseConnection, NetworkBehaviour, THandlerOutEvent, derive_prelude::*},
15};
16use tracing::warn;
17
18use super::*;
19use crate::libp2p::{PeerManager, service::metrics};
20
21type InnerBehaviour = request_response::Behaviour<HelloCodec>;
22
23pub struct HelloBehaviour {
24 inner: InnerBehaviour,
25 response_channels: HashMap<OutboundRequestId, flume::Sender<HelloResponse>>,
26 pending_inbound_hello_peers: HashMap<PeerId, Instant>,
27 peer_manager: Arc<PeerManager>,
28}
29
30impl HelloBehaviour {
31 pub fn new(cfg: request_response::Config, peer_manager: Arc<PeerManager>) -> Self {
32 Self {
33 inner: InnerBehaviour::new([(HELLO_PROTOCOL_NAME, ProtocolSupport::Full)], cfg),
34 response_channels: Default::default(),
35 pending_inbound_hello_peers: Default::default(),
36 peer_manager,
37 }
38 }
39
40 pub fn send_request(
41 &mut self,
42 peer: &PeerId,
43 request: HelloRequest,
44 response_channel: flume::Sender<HelloResponse>,
45 ) -> OutboundRequestId {
46 let request_id = self.inner.send_request(peer, request);
47 self.response_channels.insert(request_id, response_channel);
48 self.track_metrics();
49 request_id
50 }
51
52 pub fn send_response(
53 &mut self,
54 channel: ResponseChannel<HelloResponse>,
55 response: HelloResponse,
56 ) -> Result<(), HelloResponse> {
57 self.inner.send_response(channel, response)
58 }
59
60 pub async fn handle_response(
61 &mut self,
62 request_id: &OutboundRequestId,
63 response: HelloResponse,
64 ) {
65 if let Some(channel) = self.response_channels.remove(request_id) {
66 self.response_channels.shrink_to_fit();
67 self.track_metrics();
68 if let Err(err) = channel.send_async(response).await {
69 warn!("{err}");
70 }
71 }
72 }
73
74 pub fn on_outbound_failure(&mut self, request_id: &OutboundRequestId) {
75 if self.response_channels.remove(request_id).is_some() {
76 self.response_channels.shrink_to_fit();
77 self.track_metrics();
78 }
79 }
80
81 fn track_metrics(&self) {
82 metrics::NETWORK_CONTAINER_CAPACITIES
83 .get_or_create(&metrics::values::HELLO_REQUEST_TABLE)
84 .set(self.response_channels.capacity() as _);
85 }
86}
87
88impl NetworkBehaviour for HelloBehaviour {
89 type ConnectionHandler = <InnerBehaviour as NetworkBehaviour>::ConnectionHandler;
90
91 type ToSwarm = <InnerBehaviour as NetworkBehaviour>::ToSwarm;
92
93 fn handle_established_inbound_connection(
94 &mut self,
95 connection_id: ConnectionId,
96 peer: PeerId,
97 local_addr: &libp2p::Multiaddr,
98 remote_addr: &libp2p::Multiaddr,
99 ) -> Result<THandler<Self>, ConnectionDenied> {
100 self.inner.handle_established_inbound_connection(
101 connection_id,
102 peer,
103 local_addr,
104 remote_addr,
105 )
106 }
107
108 fn handle_established_outbound_connection(
109 &mut self,
110 connection_id: ConnectionId,
111 peer: PeerId,
112 addr: &Multiaddr,
113 role_override: libp2p::core::Endpoint,
114 port_use: PortUse,
115 ) -> Result<THandler<Self>, ConnectionDenied> {
116 self.inner.handle_established_outbound_connection(
117 connection_id,
118 peer,
119 addr,
120 role_override,
121 port_use,
122 )
123 }
124
125 fn handle_pending_inbound_connection(
126 &mut self,
127 connection_id: ConnectionId,
128 local_addr: &libp2p::Multiaddr,
129 remote_addr: &libp2p::Multiaddr,
130 ) -> Result<(), ConnectionDenied> {
131 self.inner
132 .handle_pending_inbound_connection(connection_id, local_addr, remote_addr)
133 }
134
135 fn handle_pending_outbound_connection(
136 &mut self,
137 connection_id: ConnectionId,
138 maybe_peer: Option<PeerId>,
139 addresses: &[libp2p::Multiaddr],
140 effective_role: libp2p::core::Endpoint,
141 ) -> Result<Vec<libp2p::Multiaddr>, ConnectionDenied> {
142 self.inner.handle_pending_outbound_connection(
143 connection_id,
144 maybe_peer,
145 addresses,
146 effective_role,
147 )
148 }
149
150 fn on_connection_handler_event(
151 &mut self,
152 peer_id: PeerId,
153 connection_id: ConnectionId,
154 event: THandlerOutEvent<Self>,
155 ) {
156 self.inner
157 .on_connection_handler_event(peer_id, connection_id, event)
158 }
159
160 fn on_swarm_event(&mut self, event: FromSwarm) {
161 if let FromSwarm::ConnectionEstablished(e) = &event
162 && e.other_established == 0
163 {
164 self.pending_inbound_hello_peers
165 .insert(e.peer_id, Instant::now());
166 }
167
168 self.inner.on_swarm_event(event)
169 }
170
171 fn poll(
172 &mut self,
173 cx: &mut std::task::Context<'_>,
174 ) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
175 if let Poll::Ready(ev) = self.inner.poll(cx) {
176 if let ToSwarm::GenerateEvent(request_response::Event::Message {
178 peer,
179 message:
180 request_response::Message::Request {
181 request: HelloRequest { .. },
182 ..
183 },
184 ..
185 }) = &ev
186 && self.pending_inbound_hello_peers.remove(peer).is_some()
187 {
188 self.pending_inbound_hello_peers.shrink_to_fit();
189 }
190
191 return Poll::Ready(ev);
192 }
193
194 const INBOUND_HELLO_WAIT_TIMEOUT: Duration = Duration::from_secs(30);
196 let now = Instant::now();
197 if let Some((&peer_to_disconnect, _)) =
198 self.pending_inbound_hello_peers
199 .iter()
200 .find(|&(_, &connected_instant)| {
201 now.duration_since(connected_instant) > INBOUND_HELLO_WAIT_TIMEOUT
202 })
203 {
204 if self
205 .pending_inbound_hello_peers
206 .remove(&peer_to_disconnect)
207 .is_some()
208 {
209 self.pending_inbound_hello_peers.shrink_to_fit();
210 }
211 if !self.peer_manager.is_peer_protected(&peer_to_disconnect) {
212 tracing::debug!(peer=%peer_to_disconnect, "Disconnecting peer for not receiving hello in 30s");
213 return Poll::Ready(ToSwarm::CloseConnection {
214 peer_id: peer_to_disconnect,
215 connection: CloseConnection::All,
216 });
217 }
218 }
219
220 Poll::Pending
221 }
222}