litep2p/transport/manager/
limits.rs

1// Copyright 2024 litep2p developers
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21//! Limits for the transport manager.
22
23use crate::types::ConnectionId;
24
25use std::collections::HashSet;
26
27/// Configuration for the connection limits.
28#[derive(Debug, Clone, Default)]
29pub struct ConnectionLimitsConfig {
30    /// Maximum number of incoming connections that can be established.
31    max_incoming_connections: Option<usize>,
32    /// Maximum number of outgoing connections that can be established.
33    max_outgoing_connections: Option<usize>,
34}
35
36impl ConnectionLimitsConfig {
37    /// Configures the maximum number of incoming connections that can be established.
38    pub fn max_incoming_connections(mut self, limit: Option<usize>) -> Self {
39        self.max_incoming_connections = limit;
40        self
41    }
42
43    /// Configures the maximum number of outgoing connections that can be established.
44    pub fn max_outgoing_connections(mut self, limit: Option<usize>) -> Self {
45        self.max_outgoing_connections = limit;
46        self
47    }
48}
49
50/// Error type for connection limits.
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
52pub enum ConnectionLimitsError {
53    /// Maximum number of incoming connections exceeded.
54    MaxIncomingConnectionsExceeded,
55    /// Maximum number of outgoing connections exceeded.
56    MaxOutgoingConnectionsExceeded,
57}
58
59/// Connection limits.
60#[derive(Debug, Clone)]
61pub struct ConnectionLimits {
62    /// Configuration for the connection limits.
63    config: ConnectionLimitsConfig,
64
65    /// Established incoming connections.
66    incoming_connections: HashSet<ConnectionId>,
67    /// Established outgoing connections.
68    outgoing_connections: HashSet<ConnectionId>,
69}
70
71impl ConnectionLimits {
72    /// Creates a new connection limits instance.
73    pub fn new(config: ConnectionLimitsConfig) -> Self {
74        let max_incoming_connections = config.max_incoming_connections.unwrap_or(0);
75        let max_outgoing_connections = config.max_outgoing_connections.unwrap_or(0);
76
77        Self {
78            config,
79            incoming_connections: HashSet::with_capacity(max_incoming_connections),
80            outgoing_connections: HashSet::with_capacity(max_outgoing_connections),
81        }
82    }
83
84    /// Called when dialing an address.
85    ///
86    /// Returns the number of outgoing connections permitted to be established.
87    /// It is guaranteed that at least one connection can be established if the method returns `Ok`.
88    /// The number of available outgoing connections can influence the maximum parallel dials to a
89    /// single address.
90    ///
91    /// If the maximum number of outgoing connections is not set, `Ok(usize::MAX)` is returned.
92    pub fn on_dial_address(&mut self) -> Result<usize, ConnectionLimitsError> {
93        if let Some(max_outgoing_connections) = self.config.max_outgoing_connections {
94            if self.outgoing_connections.len() >= max_outgoing_connections {
95                return Err(ConnectionLimitsError::MaxOutgoingConnectionsExceeded);
96            }
97
98            return Ok(max_outgoing_connections - self.outgoing_connections.len());
99        }
100
101        Ok(usize::MAX)
102    }
103
104    /// Called before accepting a new incoming connection.
105    pub fn on_incoming(&mut self) -> Result<(), ConnectionLimitsError> {
106        if let Some(max_incoming_connections) = self.config.max_incoming_connections {
107            if self.incoming_connections.len() >= max_incoming_connections {
108                return Err(ConnectionLimitsError::MaxIncomingConnectionsExceeded);
109            }
110        }
111
112        Ok(())
113    }
114
115    /// Called when a new connection is established.
116    ///
117    /// Returns an error if the connection cannot be accepted due to connection limits.
118    pub fn can_accept_connection(
119        &mut self,
120        is_listener: bool,
121    ) -> Result<(), ConnectionLimitsError> {
122        // Check connection limits.
123        if is_listener {
124            if let Some(max_incoming_connections) = self.config.max_incoming_connections {
125                if self.incoming_connections.len() >= max_incoming_connections {
126                    return Err(ConnectionLimitsError::MaxIncomingConnectionsExceeded);
127                }
128            }
129        } else if let Some(max_outgoing_connections) = self.config.max_outgoing_connections {
130            if self.outgoing_connections.len() >= max_outgoing_connections {
131                return Err(ConnectionLimitsError::MaxOutgoingConnectionsExceeded);
132            }
133        }
134
135        Ok(())
136    }
137
138    /// Accept an established connection.
139    ///
140    /// # Note
141    ///
142    /// This method should be called after the `Self::can_accept_connection` method
143    /// to ensure that the connection can be accepted.
144    pub fn accept_established_connection(
145        &mut self,
146        connection_id: ConnectionId,
147        is_listener: bool,
148    ) {
149        if is_listener {
150            if self.config.max_incoming_connections.is_some() {
151                self.incoming_connections.insert(connection_id);
152            }
153        } else if self.config.max_outgoing_connections.is_some() {
154            self.outgoing_connections.insert(connection_id);
155        }
156    }
157
158    /// Called when a connection is closed.
159    pub fn on_connection_closed(&mut self, connection_id: ConnectionId) {
160        self.incoming_connections.remove(&connection_id);
161        self.outgoing_connections.remove(&connection_id);
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::*;
168    use crate::types::ConnectionId;
169
170    #[test]
171    fn connection_limits() {
172        let config = ConnectionLimitsConfig::default()
173            .max_incoming_connections(Some(3))
174            .max_outgoing_connections(Some(2));
175        let mut limits = ConnectionLimits::new(config);
176
177        let connection_id_in_1 = ConnectionId::random();
178        let connection_id_in_2 = ConnectionId::random();
179        let connection_id_out_1 = ConnectionId::random();
180        let connection_id_out_2 = ConnectionId::random();
181        let connection_id_in_3 = ConnectionId::random();
182
183        // Establish incoming connection.
184        assert!(limits.can_accept_connection(true).is_ok());
185        limits.accept_established_connection(connection_id_in_1, true);
186        assert_eq!(limits.incoming_connections.len(), 1);
187
188        assert!(limits.can_accept_connection(true).is_ok());
189        limits.accept_established_connection(connection_id_in_2, true);
190        assert_eq!(limits.incoming_connections.len(), 2);
191
192        assert!(limits.can_accept_connection(true).is_ok());
193        limits.accept_established_connection(connection_id_in_3, true);
194        assert_eq!(limits.incoming_connections.len(), 3);
195
196        assert_eq!(
197            limits.can_accept_connection(true).unwrap_err(),
198            ConnectionLimitsError::MaxIncomingConnectionsExceeded
199        );
200        assert_eq!(limits.incoming_connections.len(), 3);
201
202        // Establish outgoing connection.
203        assert!(limits.can_accept_connection(false).is_ok());
204        limits.accept_established_connection(connection_id_out_1, false);
205        assert_eq!(limits.incoming_connections.len(), 3);
206        assert_eq!(limits.outgoing_connections.len(), 1);
207
208        assert!(limits.can_accept_connection(false).is_ok());
209        limits.accept_established_connection(connection_id_out_2, false);
210        assert_eq!(limits.incoming_connections.len(), 3);
211        assert_eq!(limits.outgoing_connections.len(), 2);
212
213        assert_eq!(
214            limits.can_accept_connection(false).unwrap_err(),
215            ConnectionLimitsError::MaxOutgoingConnectionsExceeded
216        );
217
218        // Close connections with peer a.
219        limits.on_connection_closed(connection_id_in_1);
220        assert_eq!(limits.incoming_connections.len(), 2);
221        assert_eq!(limits.outgoing_connections.len(), 2);
222
223        limits.on_connection_closed(connection_id_out_1);
224        assert_eq!(limits.incoming_connections.len(), 2);
225        assert_eq!(limits.outgoing_connections.len(), 1);
226    }
227}
litep2p/transport/manager/limits.rs

litep2p/transport/manager/
limits.rs
