net-io 0.4.0

Network IO substrate for Rust: TCP, UDP and local-socket listeners, connection lifecycle, limits, graceful drain, PROXY protocol, and a TLS seam. Protocol-free and dependency-free.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120

//! The per-source concurrency cap (decision D-025).
//!
//! Keyed by `PrefixKey` with the standard library's keyed hasher (SipHash with a
//! per-process random key), so an attacker cannot choose addresses that
//! collide. Sharded to keep accept threads and workers off one lock; every map
//! is sized at build for the global cap, so it never grows (no allocation after
//! build) and it can never fill: each live entry belongs to an open connection.

use std::collections::HashMap;
use std::collections::hash_map::RandomState;
use std::hash::BuildHasher;
use std::sync::{Mutex, PoisonError};

use crate::addr::PrefixKey;

const SHARDS: usize = 16;

pub(crate) struct SourceTable {
    shards: Vec<Mutex<HashMap<PrefixKey, u32, RandomState>>>,
    hasher: RandomState,
    cap: u32,
}

impl SourceTable {
    pub(crate) fn new(max_connections: u32, cap_per_source: u32) -> SourceTable {
        let per_shard = (max_connections as usize).div_ceil(SHARDS).max(4) * 2;
        let hasher = RandomState::new();
        let shards = (0..SHARDS)
            .map(|_| Mutex::new(HashMap::with_capacity_and_hasher(per_shard, hasher.clone())))
            .collect();
        SourceTable {
            shards,
            hasher,
            cap: cap_per_source,
        }
    }

    fn shard(&self, key: &PrefixKey) -> &Mutex<HashMap<PrefixKey, u32, RandomState>> {
        let h = self.hasher.hash_one(key);
        &self.shards[(h as usize) % SHARDS]
    }

    /// Counts one more connection from `key`; `false` when the source is at its cap.
    pub(crate) fn acquire(&self, key: PrefixKey) -> bool {
        let mut map = self
            .shard(&key)
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        let count = map.entry(key).or_insert(0);
        if *count >= self.cap {
            if *count == 0 {
                map.remove(&key);
            }
            return false;
        }
        *count += 1;
        true
    }

    /// One connection from `key` closed.
    pub(crate) fn release(&self, key: PrefixKey) {
        let mut map = self
            .shard(&key)
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        if let Some(count) = map.get_mut(&key) {
            *count = count.saturating_sub(1);
            if *count == 0 {
                map.remove(&key);
            }
        }
    }

    #[cfg(test)]
    pub(crate) fn live(&self) -> usize {
        self.shards
            .iter()
            .map(|s| s.lock().unwrap_or_else(PoisonError::into_inner).len())
            .sum()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::addr::{PrefixKey, PrefixPolicy};

    fn key(s: &str) -> PrefixKey {
        PrefixKey::of(s.parse().unwrap(), PrefixPolicy::DEFAULT)
    }

    #[test]
    fn test_source_cap_enforced_and_released() {
        let t = SourceTable::new(100, 2);
        assert!(t.acquire(key("10.0.0.1")));
        assert!(t.acquire(key("10.0.0.1")));
        assert!(!t.acquire(key("10.0.0.1")));
        assert!(t.acquire(key("10.0.0.2")));
        t.release(key("10.0.0.1"));
        assert!(t.acquire(key("10.0.0.1")));
        t.release(key("10.0.0.1"));
        t.release(key("10.0.0.1"));
        t.release(key("10.0.0.2"));
        assert_eq!(t.live(), 0);
    }

    #[test]
    fn test_source_ipv6_grouped_by_64() {
        let t = SourceTable::new(100, 1);
        assert!(t.acquire(key("2001:db8::1")));
        assert!(!t.acquire(key("2001:db8::2")));
    }

    #[test]
    fn test_source_release_unknown_is_noop() {
        let t = SourceTable::new(10, 1);
        t.release(key("1.1.1.1"));
        assert_eq!(t.live(), 0);
    }
}