subspace_networking/protocols/
subspace_connection_limits.rs

1use libp2p::connection_limits::{Behaviour as ConnectionLimitsBehaviour, ConnectionLimits};
2use libp2p::core::transport::PortUse;
3use libp2p::core::Endpoint;
4use libp2p::multiaddr::Protocol;
5use libp2p::swarm::{
6    ConnectionDenied, ConnectionId, FromSwarm, NetworkBehaviour, THandler, THandlerInEvent,
7    THandlerOutEvent, ToSwarm,
8};
9use libp2p::{Multiaddr, PeerId};
10use std::collections::hash_map::Entry;
11use std::collections::{HashMap, HashSet};
12use std::net::IpAddr;
13use std::task::{Context, Poll};
14
15// TODO: Upstream these capabilities
16pub(crate) struct Behaviour {
17    inner: ConnectionLimitsBehaviour,
18    /// For every peer ID store both their expected IP addresses as well as number of incoming connection attempts
19    /// allowed before this allow list entry no longer has an effect
20    incoming_allow_list: HashMap<PeerId, (HashSet<IpAddr>, usize)>,
21    /// For every peer ID store number of outgoing connection attempts allowed before this allow list entry no longer
22    /// has an effect
23    outgoing_allow_list: HashMap<PeerId, usize>,
24}
25
26impl Behaviour {
27    pub(crate) fn new(limits: ConnectionLimits) -> Self {
28        Self {
29            inner: ConnectionLimitsBehaviour::new(limits),
30            incoming_allow_list: HashMap::default(),
31            outgoing_allow_list: HashMap::default(),
32        }
33    }
34
35    /// Add to allow list some attempts of incoming connections from specified peer ID that will bypass global limits
36    pub(crate) fn add_to_incoming_allow_list<IpAddresses>(
37        &mut self,
38        peer: PeerId,
39        ip_addresses: IpAddresses,
40        add_attempts: usize,
41    ) where
42        IpAddresses: Iterator<Item = IpAddr>,
43    {
44        match self.incoming_allow_list.entry(peer) {
45            Entry::Occupied(mut entry) => {
46                let (existing_ip_addresses, attempts) = entry.get_mut();
47                existing_ip_addresses.extend(ip_addresses);
48                *attempts = attempts.saturating_add(add_attempts);
49            }
50            Entry::Vacant(entry) => {
51                entry.insert((ip_addresses.collect(), add_attempts));
52            }
53        }
54    }
55
56    /// Remove some (or all) attempts of incoming connections from specified peer ID
57    pub(crate) fn remove_from_incoming_allow_list(
58        &mut self,
59        peer: &PeerId,
60        remove_attempts: Option<usize>,
61    ) {
62        if let Some(remove_attempts) = remove_attempts {
63            if let Some((_ip_addresses, attempts)) = self.incoming_allow_list.get_mut(peer) {
64                *attempts = attempts.saturating_sub(remove_attempts);
65
66                if *attempts == 0 {
67                    self.incoming_allow_list.remove(peer);
68                }
69            }
70        } else {
71            self.incoming_allow_list.remove(peer);
72        }
73    }
74}
75
76impl NetworkBehaviour for Behaviour {
77    type ConnectionHandler = <ConnectionLimitsBehaviour as NetworkBehaviour>::ConnectionHandler;
78    type ToSwarm = <ConnectionLimitsBehaviour as NetworkBehaviour>::ToSwarm;
79
80    fn handle_pending_inbound_connection(
81        &mut self,
82        connection_id: ConnectionId,
83        local_addr: &Multiaddr,
84        remote_addr: &Multiaddr,
85    ) -> Result<(), ConnectionDenied> {
86        // PeerId is not yet known at this point, so we check against IP address instead
87        if let Some(ip_address) = remote_addr.iter().find_map(|protocol| match protocol {
88            Protocol::Ip4(ip) => Some(IpAddr::V4(ip)),
89            Protocol::Ip6(ip) => Some(IpAddr::V6(ip)),
90            _ => None,
91        }) {
92            if self
93                .incoming_allow_list
94                .values()
95                .any(|(ip_addresses, _attempts)| ip_addresses.contains(&ip_address))
96            {
97                return Ok(());
98            }
99        }
100
101        self.inner
102            .handle_pending_inbound_connection(connection_id, local_addr, remote_addr)
103    }
104
105    fn handle_established_inbound_connection(
106        &mut self,
107        connection_id: ConnectionId,
108        peer: PeerId,
109        local_addr: &Multiaddr,
110        remote_addr: &Multiaddr,
111    ) -> Result<THandler<Self>, ConnectionDenied> {
112        if let Some((_ip_addresses, attempts)) = self.incoming_allow_list.get_mut(&peer) {
113            *attempts -= 1;
114
115            if *attempts == 0 {
116                self.incoming_allow_list.remove(&peer);
117            }
118
119            return Ok(Self::ConnectionHandler {});
120        }
121
122        self.inner.handle_established_inbound_connection(
123            connection_id,
124            peer,
125            local_addr,
126            remote_addr,
127        )
128    }
129
130    fn handle_pending_outbound_connection(
131        &mut self,
132        connection_id: ConnectionId,
133        maybe_peer: Option<PeerId>,
134        addresses: &[Multiaddr],
135        effective_role: Endpoint,
136    ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
137        if let Some(peer) = &maybe_peer {
138            if self.incoming_allow_list.contains_key(peer) {
139                return Ok(Vec::new());
140            }
141        }
142
143        self.inner.handle_pending_outbound_connection(
144            connection_id,
145            maybe_peer,
146            addresses,
147            effective_role,
148        )
149    }
150
151    fn handle_established_outbound_connection(
152        &mut self,
153        connection_id: ConnectionId,
154        peer: PeerId,
155        addr: &Multiaddr,
156        role_override: Endpoint,
157        port_use: PortUse,
158    ) -> Result<THandler<Self>, ConnectionDenied> {
159        if let Some(attempts) = self.outgoing_allow_list.get_mut(&peer) {
160            *attempts -= 1;
161
162            if *attempts == 0 {
163                self.outgoing_allow_list.remove(&peer);
164            }
165
166            return Ok(Self::ConnectionHandler {});
167        }
168
169        self.inner.handle_established_outbound_connection(
170            connection_id,
171            peer,
172            addr,
173            role_override,
174            port_use,
175        )
176    }
177
178    fn on_swarm_event(&mut self, event: FromSwarm) {
179        self.inner.on_swarm_event(event)
180    }
181
182    fn on_connection_handler_event(
183        &mut self,
184        id: PeerId,
185        connection_id: ConnectionId,
186        event: THandlerOutEvent<Self>,
187    ) {
188        self.inner
189            .on_connection_handler_event(id, connection_id, event)
190    }
191
192    fn poll(
193        &mut self,
194        cx: &mut Context<'_>,
195    ) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
196        self.inner.poll(cx)
197    }
198}
subspace_networking/protocols/subspace_connection_limits.rs

subspace_networking/protocols/
subspace_connection_limits.rs
