Use tokio select for polling client socket

src/main.rs

@@ -10,7 +10,7 @@ use anyhow::Context;
 use smoltcp::iface::InterfaceBuilder;
 use smoltcp::socket::{SocketSet, TcpSocket, TcpSocketBuffer};
 use smoltcp::wire::{IpAddress, IpCidr};
-use tokio::io::Interest;
+use tokio::io::{AsyncReadExt, Interest};
 use tokio::net::{TcpListener, TcpStream};
 use tokio::sync::mpsc::error::TryRecvError;
 
@@ -158,19 +158,16 @@ async fn handle_tcp_proxy_connection(
     }
 
     loop {
-        let ready = socket
-            .ready(Interest::READABLE | Interest::WRITABLE)
-            .await
-            .with_context(|| "Failed to wait for TCP proxy socket readiness")?;
-
         if abort.load(Ordering::Relaxed) {
             break;
         }
 
-        if ready.is_readable() {
-            let mut buffer = [0u8; MAX_PACKET];
-
-            match socket.try_read(&mut buffer) {
+        tokio::select! {
+            readable_result = socket.readable() => {
+                match readable_result {
+                    Ok(_) => {
+                        let mut buffer = vec![];
+                        match socket.try_read_buf(&mut buffer) {
                             Ok(size) if size > 0 => {
                                 let data = &buffer[..size];
                                 debug!(
@@ -197,14 +194,20 @@ async fn handle_tcp_proxy_connection(
                                 );
                                 break;
                             }
-                _ => {}
+                            _ => {
+                                break;
                             }
                         }
-
-        if ready.is_writable() {
-            // Flush the data_to_real_client_rx channel
-            match data_to_real_client_rx.try_recv() {
-                Ok(data) => match socket.try_write(&data) {
+                    }
+                    Err(e) => {
+                        error!("[{}] Failed to check if readable: {:?}", virtual_port, e);
+                        break;
+                    }
+                }
+            }
+            data_recv_result = data_to_real_client_rx.recv() => {
+                match data_recv_result {
+                    Some(data) => match socket.try_write(&data) {
                         Ok(size) => {
                             debug!(
                                 "[{}] Wrote {} bytes of TCP data to real client",
@@ -221,22 +224,10 @@ async fn handle_tcp_proxy_connection(
                             );
                         }
                     },
-                Err(e) => match e {
-                    TryRecvError::Empty => {
-                        // Nothing else to consume in the data channel.
-                    }
-                    TryRecvError::Disconnected => {
-                        // Channel is broken, probably terminated.
-                    }
-                },
+                    None => continue,
                 }
             }
-
-        if ready.is_read_closed() || ready.is_write_closed() {
-            break;
         }
-
-        tokio::time::sleep(Duration::from_millis(1)).await;
     }
 
     trace!("[{}] TCP socket handler task terminated", virtual_port);

src/port_pool.rs

@@ -14,9 +14,7 @@ impl PortPool {
     pub fn new() -> Self {
         let inner = lockfree::queue::Queue::default();
         PORT_RANGE.for_each(|p| inner.push(p) as ());
-        Self {
-            inner,
-        }
+        Self { inner }
     }
 
     pub fn next(&self) -> anyhow::Result<u16> {