ya_relay_stack/
port.rs

1use rand::distributions::{Distribution, Uniform};
2use std::collections::{BTreeMap, BTreeSet};
3use std::ops::RangeInclusive;
4
5use crate::{Error, Protocol, Result};
6
7#[derive(Default)]
8pub struct Allocator {
9    taken: BTreeMap<Protocol, BTreeSet<u16>>,
10}
11
12impl Allocator {
13    const RANGE: RangeInclusive<u16> = 1000..=65535;
14
15    pub fn next(&mut self, protocol: Protocol) -> Result<u16> {
16        let mut rng = rand::thread_rng();
17        let mut port = Uniform::from(Self::RANGE).sample(&mut rng);
18        let taken = self.taken.entry(protocol).or_insert_with(Default::default);
19
20        let range_start = *Self::RANGE.start();
21        let mut num = Self::RANGE.len() as i32;
22
23        while num > 0 {
24            if !taken.contains(&port) {
25                taken.insert(port);
26                return Ok(port);
27            }
28            port = range_start.max(port.overflowing_add(1).0);
29            num -= 1;
30        }
31
32        Err(Error::Other("no ports available".into()))
33    }
34
35    #[allow(unused)]
36    pub fn reserve(&mut self, protocol: Protocol, port: u16) -> Result<()> {
37        let entry = self.taken.entry(protocol).or_insert_with(Default::default);
38        if entry.contains(&port) {
39            return Err(Error::Other(format!("port {} is unavailable", port)));
40        }
41        entry.insert(port);
42        Ok(())
43    }
44
45    pub fn free(&mut self, protocol: Protocol, port: u16) {
46        self.taken
47            .entry(protocol)
48            .or_insert_with(Default::default)
49            .remove(&port);
50    }
51}
ya_relay_stack/port.rs

ya_relay_stack/
port.rs
