From 5ba111002f3ce1373cde090383f98926d7f0f8af Mon Sep 17 00:00:00 2001
From: Aram Peres <aram.peres@wavy.fm>
Date: Wed, 13 Oct 2021 22:22:17 -0400
Subject: [PATCH] Reduce sleepers

---
 src/main.rs | 10 +++-------
 src/wg.rs   | 15 +++++++++------
 2 files changed, 12 insertions(+), 13 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 9bbac59..02ff45e 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -29,7 +29,7 @@ pub const MAX_PACKET: usize = 65536;
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    pretty_env_logger::init_custom_env("ONETUN_LOG");
+    pretty_env_logger::try_init_timed_custom_env("ONETUN_LOG").unwrap();
     let config = Config::from_args().with_context(|| "Failed to read config")?;
     let port_pool = Arc::new(PortPool::new());
 
@@ -236,7 +236,7 @@ async fn handle_tcp_proxy_connection(
             break;
         }
 
-        tokio::time::sleep(Duration::from_millis(5)).await;
+        tokio::time::sleep(Duration::from_millis(1)).await;
     }
 
     trace!("[{}] TCP socket handler task terminated", virtual_port);
@@ -370,11 +370,7 @@ async fn virtual_tcp_interface(
             }
         }
 
-        match virtual_interface.poll_delay(&socket_set, loop_start) {
-            None => tokio::time::sleep(Duration::from_millis(1)).await,
-            Some(smoltcp::time::Duration::ZERO) => {}
-            Some(delay) => tokio::time::sleep(Duration::from_millis(delay.millis())).await,
-        };
+        tokio::time::sleep(Duration::from_millis(1)).await;
     }
     trace!("[{}] Virtual interface task terminated", virtual_port);
     Ok(())
diff --git a/src/wg.rs b/src/wg.rs
index ed92a77..56852ab 100644
--- a/src/wg.rs
+++ b/src/wg.rs
@@ -46,6 +46,7 @@ impl WireGuardTunnel {
 
     /// Encapsulates and sends an IP packet through to the WireGuard endpoint.
     pub async fn send_ip_packet(&self, packet: &[u8]) -> anyhow::Result<()> {
+        trace_ip_packet("Sending IP packet", packet);
         let mut send_buf = [0u8; MAX_PACKET];
         match self.peer.encapsulate(packet, &mut send_buf) {
             TunnResult::WriteToNetwork(packet) => {
@@ -109,7 +110,7 @@ impl WireGuardTunnel {
                 }
                 TunnResult::Done => {
                     // Sleep for a bit
-                    tokio::time::sleep(Duration::from_millis(100)).await;
+                    tokio::time::sleep(Duration::from_millis(1)).await;
                 }
                 other => {
                     warn!("Unexpected WireGuard routine task state: {:?}", other);
@@ -132,7 +133,7 @@ impl WireGuardTunnel {
                 Err(e) => {
                     error!("Failed to read from WireGuard endpoint: {:?}", e);
                     // Sleep a little bit and try again
-                    tokio::time::sleep(Duration::from_millis(100)).await;
+                    tokio::time::sleep(Duration::from_millis(1)).await;
                     continue;
                 }
             };
@@ -172,7 +173,7 @@ impl WireGuardTunnel {
                     );
 
                     // For debugging purposes: parse packet
-                    trace_ip_packet(packet);
+                    trace_ip_packet("Received IP packet", packet);
 
                     // Broadcast IP packet
                     match self.ip_broadcast_tx.send(packet.to_vec()) {
@@ -206,17 +207,19 @@ impl WireGuardTunnel {
     }
 }
 
-fn trace_ip_packet(packet: &[u8]) {
+fn trace_ip_packet(message: &str, packet: &[u8]) {
     if log_enabled!(Level::Trace) {
         use smoltcp::wire::*;
 
         match IpVersion::of_packet(&packet) {
             Ok(IpVersion::Ipv4) => trace!(
-                "IPv4 packet received: {}",
+                "{}: {}",
+                message,
                 PrettyPrinter::<Ipv4Packet<&mut [u8]>>::new("", &packet)
             ),
             Ok(IpVersion::Ipv6) => trace!(
-                "IPv6 packet received: {}",
+                "{}: {}",
+                message,
                 PrettyPrinter::<Ipv6Packet<&mut [u8]>>::new("", &packet)
             ),
             _ => {}