aramperes/onetun
1
0
Fork 0
mirror of https://github.com/aramperes/onetun.git synced 2025-09-09 06:58:31 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

udp: use tokio select instead of 1ms loop

Browse source
This commit is contained in:
Aram 🍐 2021-10-26 00:38:22 -04:00
parent faf157cfeb
commit 0da6fa51de
1 changed files with 92 additions and 74 deletions
Download patch file Download diff file Expand all files Collapse all files

26
src/virtual_iface/udp.rs
View file

@ -93,9 +93,19 @@ impl VirtualInterfacePoll for UdpVirtualInterface {
// A map of virtual port to client socket.
let mut client_sockets: HashMap<VirtualPort, SocketHandle> = HashMap::new();
// The next instant required to poll the virtual interface
// None means "immediate poll required".
let mut next_poll: Option<tokio::time::Instant> = None;
loop {
let loop_start = smoltcp::time::Instant::now();
let wg = self.wg.clone();
tokio::select! {
// Wait the recommended amount of time by smoltcp, and poll again.
_ = match next_poll {
None => tokio::time::sleep(Duration::ZERO),
Some(until) => tokio::time::sleep_until(until)
} => {
let loop_start = smoltcp::time::Instant::now();
match virtual_interface.poll(&mut socket_set, loop_start) {
Ok(processed) if processed => {
@ -132,7 +142,15 @@ impl VirtualInterfacePoll for UdpVirtualInterface {
}
}
if let Ok((client_port, data)) = data_to_virtual_server_rx.try_recv() {
next_poll = match virtual_interface.poll_delay(&socket_set, loop_start) {
Some(smoltcp::time::Duration::ZERO) => None,
Some(delay) => Some(tokio::time::Instant::now() + Duration::from_millis(delay.millis())),
None => None,
}
}
// Wait for data to be received from the real client
data_recv_result = data_to_virtual_server_rx.recv() => {
if let Some((client_port, data)) = data_recv_result {
// Register the socket in WireGuard Tunnel (overrides any previous registration as well)
wg.register_virtual_interface(client_port, base_ip_dispatch_tx.clone())
.unwrap_or_else(|e| {
@ -176,8 +194,8 @@ impl VirtualInterfacePoll for UdpVirtualInterface {
);
});
}
tokio::time::sleep(Duration::from_millis(1)).await;
}
}
}
}
}