refactor rust udp

Author: sword-jin

src/client/client.rs

@@ -1,6 +1,7 @@
-use anyhow::{Context, Result};
+use anyhow::{Context as _, Result};
 use async_shutdown::{ShutdownManager, ShutdownSignal};
-use std::net::SocketAddr;
+use std::task::{Context, Poll};
+use std::{net::SocketAddr, pin::Pin};
 use tokio_stream::{wrappers::ReceiverStream, StreamExt};
 use tonic::{transport::Channel, Response, Status, Streaming};
 use tracing::{debug, error, info, instrument, span};
@@ -12,6 +13,7 @@ use tokio::{
     sync::{mpsc, oneshot},
 };
 
+use crate::io::AsyncUdpSocket;
 use crate::{
     constant,
     io::{StreamingReader, StreamingWriter, TrafficToServerWrapper},
@@ -331,7 +333,7 @@ async fn handle_work_traffic(
 
     // write the data streaming response to transfer_tx,
     // then forward_traffic_to_local can read the data from transfer_rx
-    let (transfer_tx, mut transfer_rx) = mpsc::channel::<TrafficToClient>(64);
+    let (transfer_tx, transfer_rx) = mpsc::channel::<TrafficToClient>(64);
 
     let (local_conn_established_tx, local_conn_established_rx) = mpsc::channel::<()>(1);
     let mut local_conn_established_rx = Some(local_conn_established_rx);
@@ -375,7 +377,7 @@ async fn handle_work_traffic(
     });
 
     let wrapper = TrafficToServerWrapper::new(connection_id.clone());
-    let mut writer = StreamingWriter::new(streaming_tx.clone(), wrapper);
+    let writer = StreamingWriter::new(streaming_tx.clone(), wrapper);
 
     if is_udp {
         tokio::spawn(async move {
@@ -410,30 +412,16 @@ async fn handle_work_traffic(
 
             local_conn_established_tx.send(()).await.unwrap();
 
-            let read_transfer_send_to_local = async {
-                while let Some(buf) = transfer_rx.recv().await {
-                    socket.send(&buf.data).await.unwrap();
-                }
-            };
-
-            let read_local_send_to_server = async {
-                loop {
-                    let mut buf = vec![0u8; 65507];
-                    let result = socket.recv(&mut buf).await;
-                    match result {
-                        Ok(n) => {
-                            writer.write_all(&buf[..n]).await.unwrap();
-                        }
-                        Err(err) => {
-                            error!(err = ?err, "failed to read from local endpoint");
-                            break;
-                        }
-                    }
-                }
-                writer.shutdown().await.unwrap();
-            };
-
-            tokio::join!(read_transfer_send_to_local, read_local_send_to_server);
+            if let Err(err) = forward_traffic_to_local(
+                AsyncUdpSocket::new(&socket),
+                AsyncUdpSocket::new(&socket),
+                StreamingReader::new(transfer_rx),
+                writer,
+            )
+            .await
+            {
+                debug!("failed to forward traffic to local: {:?}", err);
+            }
         });
     } else {
         tokio::spawn(async move {
@@ -485,8 +473,8 @@ async fn handle_work_traffic(
 async fn forward_traffic_to_local(
     local_r: impl AsyncRead + Unpin,
     mut local_w: impl AsyncWrite + Unpin,
-    remote_r: StreamingReader<TrafficToClient>,
-    mut remote_w: StreamingWriter<TrafficToServer>,
+    remote_r: impl AsyncRead + Unpin,
+    mut remote_w: impl AsyncWrite + Unpin,
 ) -> Result<()> {
     let remote_to_me_to_local = async {
         // read from remote, write to local

src/io.rs

@@ -8,6 +8,7 @@ use std::fmt::Debug;
 use std::task::{Context, Poll};
 use tokio::io::AsyncRead;
 use tokio::io::AsyncWrite;
+use tokio::net::UdpSocket;
 use tokio::sync::mpsc;
 use tokio::{io, sync::mpsc::Sender};
 use tokio_util::sync::CancellationToken;
@@ -263,3 +264,51 @@ macro_rules! generate_async_write_impl {
 
 generate_async_write_impl!(TrafficToServer);
 generate_async_write_impl!(Vec<u8>);
+
+pub(crate) struct AsyncUdpSocket<'a> {
+    socket: &'a UdpSocket,
+}
+
+impl<'a> AsyncUdpSocket<'a> {
+    pub(crate) fn new(socket: &'a UdpSocket) -> Self {
+        Self { socket }
+    }
+}
+
+impl<'a> AsyncRead for AsyncUdpSocket<'a> {
+    fn poll_read(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut io::ReadBuf<'_>,
+    ) -> Poll<io::Result<()>> {
+        match self.get_mut().socket.poll_recv_from(cx, buf) {
+            Poll::Ready(Ok(_addr)) => Poll::Ready(Ok(())),
+            Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+impl<'a> AsyncWrite for AsyncUdpSocket<'a> {
+    fn poll_write(
+        self: std::pin::Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::result::Result<usize, std::io::Error>> {
+        self.get_mut().socket.poll_send(cx, buf)
+    }
+
+    fn poll_flush(
+        self: std::pin::Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<std::result::Result<(), std::io::Error>> {
+        Poll::Ready(Ok(())) // No-op for UDP
+    }
+
+    fn poll_shutdown(
+        self: std::pin::Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<std::result::Result<(), std::io::Error>> {
+        Poll::Ready(Ok(())) // No-op for UDP
+    }
+}