netsim_embed_nat/
port_allocator.rs

1use std::collections::hash_map::{Entry, HashMap};
2use std::net::SocketAddrV4;
3
4pub trait PortAllocator: std::fmt::Debug + Send {
5    fn next_port(&mut self, local_endpoint: SocketAddrV4) -> u16;
6}
7
8#[derive(Clone, Debug)]
9pub struct SequentialPortAllocator {
10    next_original_port: u16,
11    next_for_local_endpoint: HashMap<SocketAddrV4, u16>,
12}
13
14impl Default for SequentialPortAllocator {
15    fn default() -> Self {
16        Self {
17            next_original_port: 49152,
18            next_for_local_endpoint: HashMap::new(),
19        }
20    }
21}
22
23impl PortAllocator for SequentialPortAllocator {
24    fn next_port(&mut self, local_endpoint: SocketAddrV4) -> u16 {
25        match self.next_for_local_endpoint.entry(local_endpoint) {
26            Entry::Occupied(mut entry) => {
27                let port = *entry.get();
28                *entry.get_mut() = entry.get().checked_add(1).unwrap_or(49152);
29                port
30            }
31            Entry::Vacant(entry) => {
32                let port = self.next_original_port;
33                self.next_original_port = self.next_original_port.wrapping_add(16);
34                if self.next_original_port < 49152 {
35                    self.next_original_port += 49153
36                };
37                entry.insert(port);
38                port
39            }
40        }
41    }
42}
43
44#[derive(Clone, Debug, Default)]
45pub struct RandomPortAllocator;
46
47impl PortAllocator for RandomPortAllocator {
48    fn next_port(&mut self, _local_endpoint: SocketAddrV4) -> u16 {
49        loop {
50            let port = rand::random();
51            if port >= 1000 {
52                return port;
53            }
54        }
55    }
56}
netsim_embed_nat/port_allocator.rs

netsim_embed_nat/
port_allocator.rs
