net_srv/
port.rs

1use std::{collections::HashSet, sync::Mutex};
2pub struct PortAssignerInner {
3    unused_start: u16,
4    unused: HashSet<u16>,
5    used: HashSet<u16>,
6}
7
8pub struct PortAssigner {
9    inner: Mutex<PortAssignerInner>,
10}
11
12const EPHEMERAL_START: u16 = 49152;
13const EPHEMERAL_END: u16 = 65535;
14
15impl PortAssignerInner {
16    pub fn new() -> Self {
17        Self {
18            unused_start: EPHEMERAL_START,
19            unused: HashSet::new(),
20            used: HashSet::new(),
21        }
22    }
23
24    pub fn return_port(&mut self, port: u16) {
25        self.used.remove(&port);
26        self.unused.insert(port);
27    }
28
29    pub fn allocate_port(&mut self, port: u16) -> bool {
30        if self.used.contains(&port) {
31            return false;
32        }
33
34        self.unused.remove(&port);
35        self.used.insert(port);
36        true
37    }
38
39    pub fn get_ephemeral_port(&mut self) -> Option<u16> {
40        if self.unused.is_empty() {
41            if self.unused_start == EPHEMERAL_END {
42                return None;
43            }
44            let port = self.unused_start;
45            self.unused_start += 1;
46            if self.used.contains(&port) {
47                return self.get_ephemeral_port();
48            }
49            self.used.insert(port);
50            self.unused.remove(&port);
51            Some(port)
52        } else {
53            let port = self.unused.iter().next().unwrap();
54            self.used.insert(*port);
55            Some(*port)
56        }
57    }
58}
59
60impl PortAssigner {
61    pub fn new() -> Self {
62        Self {
63            inner: Mutex::new(PortAssignerInner::new()),
64        }
65    }
66
67    pub fn return_port(&self, port: u16) {
68        self.inner.lock().unwrap().return_port(port);
69    }
70
71    pub fn get_ephemeral_port(&self) -> Option<u16> {
72        self.inner.lock().unwrap().get_ephemeral_port()
73    }
74
75    pub fn allocate_port(&self, port: u16) -> bool {
76        self.inner.lock().unwrap().allocate_port(port)
77    }
78}
net_srv/port.rs

net_srv/
port.rs
