Skip to main content

forest/libp2p/hello/
behaviour.rs

1// Copyright 2019-2026 ChainSafe Systems
2// SPDX-License-Identifier: Apache-2.0, MIT
3
4use 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            // Remove a peer from `pending_inbound_hello_peers` when its hello request is received.
177            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        // Disconnect peers whose hello request are not received after a TIMEOUT
195        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}