vsock_service_lib/
port.rs
1use std::collections::HashSet;
6use std::ops::Range;
7
8const EPHEMERAL_PORT_RANGE: Range<u32> = 49152..65535;
10
11pub fn is_ephemeral(port: u32) -> bool {
12 port >= EPHEMERAL_PORT_RANGE.start && port < EPHEMERAL_PORT_RANGE.end
13}
14
15pub struct Tracker {
16 used: HashSet<u32>,
18 next_allocation: u32,
21}
22
23impl Tracker {
24 pub fn allocate(&mut self) -> Option<u32> {
25 (self.next_allocation..EPHEMERAL_PORT_RANGE.end)
28 .chain(EPHEMERAL_PORT_RANGE.start..self.next_allocation)
29 .find(|&p| self.used.insert(p))
30 .map(|x| {
31 self.next_allocation = x + 1;
32 x
33 })
34 }
35 pub fn free(&mut self, port: u32) {
36 if !self.used.remove(&port) {
37 panic!("Tried to free unallocated port.");
38 }
39 self.next_allocation = std::cmp::min(self.next_allocation, port);
40 }
41 pub fn new() -> Self {
42 Tracker { used: HashSet::new(), next_allocation: EPHEMERAL_PORT_RANGE.start }
43 }
44}
45
46impl Default for Tracker {
47 fn default() -> Self {
48 Self::new()
49 }
50}
51
52#[cfg(test)]
53mod tests {
54 use super::*;
55
56 #[test]
57 fn ports_reused() {
58 let mut t = Tracker::new();
59 let p1 = t.allocate().unwrap();
60 let p2 = t.allocate().unwrap();
61 let p3 = t.allocate().unwrap();
62 t.free(p2);
63 let p4 = t.allocate().unwrap();
64 assert_eq!(p2, p4);
65 let p5 = t.allocate().unwrap();
66 t.free(p4);
67 t.free(p3);
68 let p6 = t.allocate().unwrap();
69 let p7 = t.allocate().unwrap();
70 assert_eq!(p4, p6);
71 assert_eq!(p3, p7);
72 t.free(p1);
73 t.free(p5);
74 t.free(p6);
75 t.free(p7);
76 }
77}