vsock_service_lib/
port.rs

1// Copyright 2018 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::collections::HashSet;
6use std::ops::Range;
7
8// https://www.iana.org/assignments/service-names-port-numbers/service-names-port-numbers.xhtml
9const 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    // TODO: use a bit-vec instead
17    used: HashSet<u32>,
18    // Track the next port we should attempt to begin allocating from as a
19    // heuristic.
20    next_allocation: u32,
21}
22
23impl Tracker {
24    pub fn allocate(&mut self) -> Option<u32> {
25        // Search for a port starting from `next_allocation` and wrapping around if none is
26        // found to try all the ports up until one before next_allocation.
27        (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}