network/ip/
node_limit.rs

1use crate::{
2    ip::{bucket::NodeBucket, sample::SampleHashMap, util::SubnetType},
3    node_database::NodeDatabase,
4    node_table::NodeId,
5};
6use std::{
7    collections::{HashMap, HashSet},
8    net::IpAddr,
9    time::Duration,
10};
11
12/// Default maximum duration of last node contact time that used to evict a node
13/// when `NodeBucket` is full.
14const DEFAULT_EVICT_TIMEOUT: Duration = Duration::from_secs(7 * 24 * 3600); // 7 days
15
16/// Validation result before adding a node.
17#[derive(Debug, PartialEq)]
18pub enum ValidateInsertResult {
19    /// Node already exists and IP address not changed
20    AlreadyExists,
21    /// Node will be new added, and occupy the IP address used by existing
22    /// node. In this case, the existing node will be evicted before adding
23    /// the new node.
24    OccupyIp(NodeId),
25    /// Node is allowed to add or update with new IP address, because the
26    /// corresponding `NodeBucket` has enough quota.
27    QuotaEnough,
28    /// Node is allowed to add or update, but need to evict the specified
29    /// existing node first.
30    Evict(NodeId),
31    /// Node is not allowed to add or update, because the corresponding
32    /// `NodeBucket` is full, and no node could be evicted. E.g. all nodes in
33    /// the bucket are in connecting status.
34    QuotaNotEnough,
35}
36
37/// NodeLimit is used to limit the number of nodes stored in database via IP
38/// address and subnet. Basically, one IP address only allow one node stored
39/// in database, and a subnet allow configured N nodes stored in database.
40///
41/// When adding a new node, and the number of nodes for the subnet reaches the
42/// quota limitation, any node may be evicted from database according to
43/// a pre-defined rule.
44#[derive(Debug)]
45pub struct NodeIpLimit {
46    subnet_type: SubnetType,
47    subnet_quota: usize,     // quota for a subnet
48    evict_timeout: Duration, // used to evict out-of-date node
49
50    // all trusted nodes grouped by subnet
51    trusted_buckets: SampleHashMap<u32, NodeBucket>,
52    // all untrusted nodes grouped by subnet
53    untrusted_buckets: SampleHashMap<u32, NodeBucket>,
54
55    // helpful indices
56    ip_index: HashMap<IpAddr, NodeId>,
57    node_index: HashMap<NodeId, IpAddr>,
58}
59
60impl NodeIpLimit {
61    pub fn new(subnet_quota: usize) -> Self {
62        NodeIpLimit {
63            subnet_type: SubnetType::C,
64            subnet_quota,
65            evict_timeout: DEFAULT_EVICT_TIMEOUT,
66            trusted_buckets: Default::default(),
67            untrusted_buckets: Default::default(),
68            ip_index: HashMap::new(),
69            node_index: HashMap::new(),
70        }
71    }
72
73    #[inline]
74    pub fn is_enabled(&self) -> bool { self.subnet_quota > 0 }
75
76    /// Get the subnet of specified node `id`.
77    pub fn subnet(&self, id: &NodeId) -> Option<u32> {
78        let ip = self.node_index.get(id)?;
79        Some(self.subnet_type.subnet(ip))
80    }
81
82    /// Remove the specified node `id` and return `true` if removed
83    /// successfully. If not found, return `false`.
84    pub fn remove(&mut self, id: &NodeId) -> bool {
85        if !self.is_enabled() {
86            return true;
87        }
88
89        let ip = match self.node_index.remove(id) {
90            Some(ip) => ip,
91            None => return false,
92        };
93
94        self.ip_index.remove(&ip);
95
96        let subnet = self.subnet_type.subnet(&ip);
97        if !Self::remove_with_buckets(&mut self.trusted_buckets, subnet, id) {
98            Self::remove_with_buckets(&mut self.untrusted_buckets, subnet, id);
99        }
100
101        true
102    }
103
104    /// Remove node from specified buckets.
105    fn remove_with_buckets(
106        buckets: &mut SampleHashMap<u32, NodeBucket>, subnet: u32, id: &NodeId,
107    ) -> bool {
108        let bucket = match buckets.get_mut(&subnet) {
109            Some(bucket) => bucket,
110            None => return false,
111        };
112
113        if !bucket.remove(id) {
114            return false;
115        }
116
117        // remove bucket on empty
118        if bucket.count() == 0 {
119            buckets.remove(&subnet);
120        }
121
122        true
123    }
124
125    /// Randomly select `n` trusted nodes. Note, it may return less than `n`
126    /// nodes. Note, the time complexity is O(n), where n is the number of
127    /// sampled nodes.
128    pub fn sample_trusted(&self, n: u32) -> HashSet<NodeId> {
129        if !self.is_enabled() {
130            return HashSet::new();
131        }
132
133        let mut sampled = HashSet::new();
134        if self.trusted_buckets.is_empty() {
135            return sampled;
136        }
137
138        let mut rng = rand::rng();
139
140        for _ in 0..n {
141            if let Some(bucket) = self.trusted_buckets.sample(&mut rng) {
142                if let Some(id) = bucket.sample(&mut rng) {
143                    sampled.insert(id);
144                }
145            }
146        }
147
148        sampled
149    }
150
151    /// Validate before inserting a node with specified `id` and `ip`.
152    /// The returned result indicates whether insertion is allowed,
153    /// and possible evictee before insertion.
154    ///
155    /// When subnet quota is not enough before insertion:
156    ///   1. If node IP changed and still in the same subnet, just evict the
157    ///      node itself.
158    /// 2. Otherwise, someone in the subnet of `ip` may be evicted.
159    ///
160    /// There are 2 cases that insertion is not allowed:
161    /// 1. Node already exists and ip not changed;
162    /// 2. Subnet quota is not enough, and no evictee found.
163    pub fn validate_insertion(
164        &self, id: &NodeId, ip: &IpAddr, db: &NodeDatabase,
165    ) -> ValidateInsertResult {
166        if !self.is_enabled() {
167            return ValidateInsertResult::QuotaEnough;
168        }
169
170        // node exists and ip not changed.
171        let maybe_cur_ip = self.node_index.get(id);
172        if let Some(cur_ip) = maybe_cur_ip {
173            if cur_ip == ip {
174                return ValidateInsertResult::AlreadyExists;
175            }
176        }
177
178        // ip already in use by other node
179        if let Some(old_id) = self.ip_index.get(ip) {
180            return ValidateInsertResult::OccupyIp(*old_id);
181        }
182
183        // quota enough
184        if self.is_quota_allowed(ip) {
185            return ValidateInsertResult::QuotaEnough;
186        }
187
188        // Node ip changed, but still in the same subnet.
189        // So, just evict the node itself.
190        if let Some(cur_ip) = maybe_cur_ip {
191            let cur_subnet = self.subnet_type.subnet(cur_ip);
192            let new_subnet = self.subnet_type.subnet(ip);
193            if cur_subnet == new_subnet {
194                return ValidateInsertResult::Evict(*id);
195            }
196        }
197
198        // quota not enough, try to evict one.
199        if let Some(evictee) = self.select_evictee(ip, db) {
200            return ValidateInsertResult::Evict(evictee);
201        }
202
203        ValidateInsertResult::QuotaNotEnough
204    }
205
206    /// Insert a node with specified `id` and `ip` as trusted or untrusted.
207    /// If evictee specified, then remove that one before insert new one.
208    /// Returns `true` if insert successfully, otherwise `false`.
209    pub fn insert(
210        &mut self, id: NodeId, ip: IpAddr, trusted: bool,
211        evictee: Option<NodeId>,
212    ) -> bool {
213        if !self.is_enabled() {
214            return true;
215        }
216
217        // node exists and ip not changed.
218        if let Some(cur_ip) = self.node_index.get(&id) {
219            if *cur_ip == ip {
220                return false;
221            }
222        }
223
224        // remove evictee before insertion
225        if let Some(id) = evictee {
226            self.remove(&id);
227        }
228
229        // ip already in use by other node
230        if self.ip_index.contains_key(&ip) {
231            return false;
232        }
233
234        if self.is_quota_allowed(&ip) {
235            self.add_or_update(ip, id, trusted);
236            return true;
237        }
238
239        false
240    }
241
242    /// Demote `id` from trusted to untrusted
243    pub fn demote(&mut self, id: &NodeId) {
244        if let Some(ip) = self.node_index.get(id) {
245            let subnet = self.subnet_type.subnet(ip);
246            if let Some(b) = self.trusted_buckets.get_mut(&subnet) {
247                if b.remove(id) {
248                    if let Some(b) = self.untrusted_buckets.get_mut(&subnet) {
249                        b.add(*id);
250                    }
251                }
252            }
253        }
254    }
255
256    /// Add or update a node with new IP address. It assumes that the bucket
257    /// of corresponding subnet have enough quota.
258    fn add_or_update(&mut self, ip: IpAddr, id: NodeId, trusted: bool) {
259        // clear the old ip information at first
260        self.remove(&id);
261
262        self.node_index.insert(id, ip);
263        self.ip_index.insert(ip, id);
264
265        let subnet = self.subnet_type.subnet(&ip);
266        if trusted {
267            self.trusted_buckets
268                .get_mut_or_insert_with(subnet, NodeBucket::default)
269                .add(id);
270        } else {
271            self.untrusted_buckets
272                .get_mut_or_insert_with(subnet, NodeBucket::default)
273                .add(id);
274        }
275    }
276
277    /// Check whether the subnet quota is enough for the specified IP address .
278    fn is_quota_allowed(&self, ip: &IpAddr) -> bool {
279        let subnet = self.subnet_type.subnet(ip);
280
281        let num_trusted = self
282            .trusted_buckets
283            .get(&subnet)
284            .map_or(0, |bucket| bucket.count());
285
286        let num_untrusted = self
287            .untrusted_buckets
288            .get(&subnet)
289            .map_or(0, |bucket| bucket.count());
290
291        num_trusted + num_untrusted < self.subnet_quota
292    }
293
294    /// Select a node to evict.
295    fn select_evictee(&self, ip: &IpAddr, db: &NodeDatabase) -> Option<NodeId> {
296        let subnet = self.subnet_type.subnet(ip);
297
298        // evict untrusted node prior to trusted node
299        self.untrusted_buckets
300            .get(&subnet)
301            .and_then(|bucket| bucket.select_evictee(db, self.evict_timeout))
302            .or_else(|| {
303                self.trusted_buckets.get(&subnet).and_then(|bucket| {
304                    bucket.select_evictee(db, self.evict_timeout)
305                })
306            })
307    }
308}
309
310#[cfg(test)]
311mod tests {
312    use super::{NodeDatabase, NodeId, NodeIpLimit, ValidateInsertResult};
313    use std::{net::IpAddr, str::FromStr};
314
315    fn new_ip(ip: &'static str) -> IpAddr { IpAddr::from_str(ip).unwrap() }
316
317    #[test]
318    fn test_remove() {
319        let mut limit = NodeIpLimit::new(2);
320
321        // remove non-exist node
322        assert_eq!(limit.remove(&NodeId::random()), false);
323
324        // add 2 new nodes
325        let n1 = NodeId::random();
326        let ip1 = new_ip("127.0.0.1");
327        assert_eq!(limit.insert(n1.clone(), ip1, true, None), true);
328
329        let n2 = NodeId::random();
330        let ip2 = new_ip("127.0.0.2");
331        assert_eq!(limit.insert(n2.clone(), ip2, true, None), true);
332
333        // remove those 2 nodes
334        validate_node(&limit, &n1, &ip1, true /* exists */);
335        assert_eq!(limit.remove(&n1), true);
336        validate_node(&limit, &n1, &ip1, false /* exists */);
337
338        validate_node(&limit, &n2, &ip2, true /* exists */);
339        assert_eq!(limit.remove(&n2), true);
340        validate_node(&limit, &n2, &ip2, false /* exists */);
341    }
342
343    #[test]
344    fn test_sample() {
345        let mut limit = NodeIpLimit::new(2);
346
347        // empty case
348        assert_eq!(limit.sample_trusted(3).is_empty(), true);
349
350        // only untrusted nodes case
351        let n1 = NodeId::random();
352        let ip1 = new_ip("127.0.0.1");
353        assert_eq!(limit.insert(n1, ip1, false, None), true);
354        assert_eq!(limit.sample_trusted(3).is_empty(), true);
355
356        // trusted nodes case
357        let n2 = NodeId::random();
358        let ip2 = new_ip("127.0.0.2");
359        assert_eq!(limit.insert(n2, ip2, true, None), true);
360        assert_eq!(limit.sample_trusted(0).len(), 0);
361        assert_eq!(limit.sample_trusted(1).len(), 1);
362        assert_eq!(limit.sample_trusted(3).len(), 1);
363    }
364
365    fn validate_node(
366        limit: &NodeIpLimit, id: &NodeId, ip: &IpAddr, exists: bool,
367    ) {
368        if !exists {
369            assert_eq!(limit.node_index.contains_key(id), false);
370        } else {
371            assert_eq!(limit.node_index.contains_key(id), true);
372            assert_eq!(limit.node_index[id], *ip);
373        }
374    }
375
376    #[test]
377    fn test_insert_duplicate_id_ip() {
378        let mut limit = NodeIpLimit::new(2);
379        let db = NodeDatabase::new(None, 2);
380
381        // quota is enough
382        let n = NodeId::random();
383        let ip = new_ip("127.0.0.1");
384        assert_eq!(
385            limit.validate_insertion(&n, &ip, &db),
386            ValidateInsertResult::QuotaEnough
387        );
388        assert_eq!(limit.insert(n.clone(), ip, true, None), true);
389        validate_node(&limit, &n, &ip, true /* exists */);
390
391        // cannot insert with same id and ip as trusted or untrusted
392        assert_eq!(
393            limit.validate_insertion(&n, &ip, &db),
394            ValidateInsertResult::AlreadyExists
395        );
396        assert_eq!(limit.insert(n.clone(), ip, true, None), false);
397        assert_eq!(limit.insert(n.clone(), ip, false, None), false);
398        validate_node(&limit, &n, &ip, true /* exists */);
399    }
400
401    #[test]
402    fn test_insert_occupy_ip_new_node() {
403        let mut limit = NodeIpLimit::new(2);
404        let db = NodeDatabase::new(None, 2);
405
406        // insert n1
407        let n1 = NodeId::random();
408        let ip = new_ip("127.0.0.1");
409        assert_eq!(
410            limit.validate_insertion(&n1, &ip, &db),
411            ValidateInsertResult::QuotaEnough
412        );
413        assert_eq!(limit.insert(n1.clone(), ip, true, None), true);
414        validate_node(&limit, &n1, &ip, true /* exists */);
415
416        // add n2 with existing ip which need to evict the n1
417        let n2 = NodeId::random();
418        assert_eq!(
419            limit.validate_insertion(&n2, &ip, &db),
420            ValidateInsertResult::OccupyIp(n1.clone())
421        );
422
423        // add n2 without evicting n1
424        assert_eq!(limit.insert(n2.clone(), ip, true, None), false);
425        validate_node(&limit, &n1, &ip, true /* exists */); // n1 not evicted
426        validate_node(&limit, &n2, &ip, false /* exists */); // n2 not inserted
427
428        // add n2 with evicting n1
429        assert_eq!(limit.insert(n2.clone(), ip, true, Some(n1.clone())), true);
430        validate_node(&limit, &n1, &ip, false /* exists */); // n1 evicted
431        validate_node(&limit, &n2, &ip, true /* exists */); // n2 inserted
432    }
433
434    #[test]
435    fn test_insert_occupy_ip_update_node() {
436        let mut limit = NodeIpLimit::new(2);
437        let db = NodeDatabase::new(None, 2);
438
439        // insert n1 and n2
440        let n1 = NodeId::random();
441        let ip1 = new_ip("127.0.0.1");
442        assert_eq!(
443            limit.validate_insertion(&n1, &ip1, &db),
444            ValidateInsertResult::QuotaEnough
445        );
446        assert_eq!(limit.insert(n1.clone(), ip1, true, None), true);
447        validate_node(&limit, &n1, &ip1, true /* exists */);
448
449        let n2 = NodeId::random();
450        let ip2 = new_ip("127.0.0.2");
451        assert_eq!(
452            limit.validate_insertion(&n2, &ip2, &db),
453            ValidateInsertResult::QuotaEnough
454        );
455        assert_eq!(limit.insert(n2.clone(), ip2, true, None), true);
456        validate_node(&limit, &n2, &ip2, true /* exists */);
457
458        // change n2's ip from ip2 to ip1
459        assert_eq!(
460            limit.validate_insertion(&n2, &ip1, &db),
461            ValidateInsertResult::OccupyIp(n1.clone())
462        );
463
464        // update n2 without evicting n1
465        assert_eq!(limit.insert(n2.clone(), ip1, true, None), false);
466        validate_node(&limit, &n1, &ip1, true /* exists */); // n1 not evicted
467        validate_node(&limit, &n2, &ip2, true /* exists */); // n2 not updated
468
469        // update n2 with evicting n1
470        assert_eq!(limit.insert(n2.clone(), ip1, true, Some(n1.clone())), true);
471        validate_node(&limit, &n1, &ip1, false /* exists */); // n1 evicted
472        validate_node(&limit, &n2, &ip1, true /* exists */); // n2 updated
473    }
474
475    #[test]
476    fn test_is_quota_allowed() {
477        let mut limit = NodeIpLimit::new(2);
478
479        // add n1
480        let n1 = NodeId::random();
481        let ip1 = new_ip("127.0.0.1");
482        assert_eq!(limit.insert(n1, ip1, true, None), true);
483
484        // add n2
485        let n2 = NodeId::random();
486        let ip2 = new_ip("127.0.0.2");
487        assert_eq!(limit.insert(n2, ip2, true, None), true);
488
489        // same subnet
490        assert_eq!(limit.is_quota_allowed(&new_ip("127.0.0.3")), false);
491
492        // different subnet
493        assert_eq!(limit.is_quota_allowed(&new_ip("127.0.1.1")), true);
494    }
495
496    #[test]
497    fn test_select_evictee() {
498        let limit = NodeIpLimit::new(2);
499        let db = NodeDatabase::new(None, 2);
500
501        // select from empty bucket
502        assert_eq!(limit.select_evictee(&new_ip("127.0.0.1"), &db), None);
503    }
504}
network/ip/node_limit.rs

network/ip/
node_limit.rs
