subspace_networking/protocols/
subspace_connection_limits.rs

1use libp2p::connection_limits::{Behaviour as ConnectionLimitsBehaviour, ConnectionLimits};
2use libp2p::core::Endpoint;
3use libp2p::core::transport::PortUse;
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        }) && self
92            .incoming_allow_list
93            .values()
94            .any(|(ip_addresses, _attempts)| ip_addresses.contains(&ip_address))
95        {
96            return Ok(());
97        }
98
99        self.inner
100            .handle_pending_inbound_connection(connection_id, local_addr, remote_addr)
101    }
102
103    fn handle_established_inbound_connection(
104        &mut self,
105        connection_id: ConnectionId,
106        peer: PeerId,
107        local_addr: &Multiaddr,
108        remote_addr: &Multiaddr,
109    ) -> Result<THandler<Self>, ConnectionDenied> {
110        if let Some((_ip_addresses, attempts)) = self.incoming_allow_list.get_mut(&peer) {
111            *attempts -= 1;
112
113            if *attempts == 0 {
114                self.incoming_allow_list.remove(&peer);
115            }
116
117            return Ok(Self::ConnectionHandler {});
118        }
119
120        self.inner.handle_established_inbound_connection(
121            connection_id,
122            peer,
123            local_addr,
124            remote_addr,
125        )
126    }
127
128    fn handle_pending_outbound_connection(
129        &mut self,
130        connection_id: ConnectionId,
131        maybe_peer: Option<PeerId>,
132        addresses: &[Multiaddr],
133        effective_role: Endpoint,
134    ) -> Result<Vec<Multiaddr>, ConnectionDenied> {
135        if let Some(peer) = &maybe_peer
136            && self.incoming_allow_list.contains_key(peer)
137        {
138            return Ok(Vec::new());
139        }
140
141        self.inner.handle_pending_outbound_connection(
142            connection_id,
143            maybe_peer,
144            addresses,
145            effective_role,
146        )
147    }
148
149    fn handle_established_outbound_connection(
150        &mut self,
151        connection_id: ConnectionId,
152        peer: PeerId,
153        addr: &Multiaddr,
154        role_override: Endpoint,
155        port_use: PortUse,
156    ) -> Result<THandler<Self>, ConnectionDenied> {
157        if let Some(attempts) = self.outgoing_allow_list.get_mut(&peer) {
158            *attempts -= 1;
159
160            if *attempts == 0 {
161                self.outgoing_allow_list.remove(&peer);
162            }
163
164            return Ok(Self::ConnectionHandler {});
165        }
166
167        self.inner.handle_established_outbound_connection(
168            connection_id,
169            peer,
170            addr,
171            role_override,
172            port_use,
173        )
174    }
175
176    fn on_swarm_event(&mut self, event: FromSwarm) {
177        self.inner.on_swarm_event(event)
178    }
179
180    fn on_connection_handler_event(
181        &mut self,
182        id: PeerId,
183        connection_id: ConnectionId,
184        event: THandlerOutEvent<Self>,
185    ) {
186        self.inner
187            .on_connection_handler_event(id, connection_id, event)
188    }
189
190    fn poll(
191        &mut self,
192        cx: &mut Context<'_>,
193    ) -> Poll<ToSwarm<Self::ToSwarm, THandlerInEvent<Self>>> {
194        self.inner.poll(cx)
195    }
196}
subspace_networking/protocols/subspace_connection_limits.rs

subspace_networking/protocols/
subspace_connection_limits.rs
