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

use std::collections::hash_map::{Entry, HashMap};
use std::net::SocketAddrV4;

pub trait PortAllocator: std::fmt::Debug + Send {
    fn next_port(&mut self, local_endpoint: SocketAddrV4) -> u16;
}

#[derive(Clone, Debug)]
pub struct SequentialPortAllocator {
    next_original_port: u16,
    next_for_local_endpoint: HashMap<SocketAddrV4, u16>,
}

impl Default for SequentialPortAllocator {
    fn default() -> Self {
        Self {
            next_original_port: 49152,
            next_for_local_endpoint: HashMap::new(),
        }
    }
}

impl PortAllocator for SequentialPortAllocator {
    fn next_port(&mut self, local_endpoint: SocketAddrV4) -> u16 {
        match self.next_for_local_endpoint.entry(local_endpoint) {
            Entry::Occupied(mut entry) => {
                let port = *entry.get();
                *entry.get_mut() = entry.get().checked_add(1).unwrap_or(49152);
                port
            }
            Entry::Vacant(entry) => {
                let port = self.next_original_port;
                self.next_original_port = self.next_original_port.wrapping_add(16);
                if self.next_original_port < 49152 {
                    self.next_original_port += 49153
                };
                entry.insert(port);
                port
            }
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct RandomPortAllocator;

impl PortAllocator for RandomPortAllocator {
    fn next_port(&mut self, _local_endpoint: SocketAddrV4) -> u16 {
        loop {
            let port = rand::random();
            if port >= 1000 {
                return port;
            }
        }
    }
}