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};
@@ -375,7 +376,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 +411,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 +472,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
@@ -528,3 +515,51 @@ async fn forward_traffic_to_local(
 
     Ok(())
 }
+
+struct AsyncUdpSocket<'a> {
+    socket: &'a UdpSocket,
+}
+
+impl<'a> AsyncUdpSocket<'a> {
+    fn new(socket: &'a UdpSocket) -> Self {
+        Self { socket }
+    }
+}
+
+impl<'a> AsyncRead for AsyncUdpSocket<'a> {
+    fn poll_read(
+        self: 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: 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: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<std::result::Result<(), std::io::Error>> {
+        Poll::Ready(Ok(())) // No-op for UDP
+    }
+
+    fn poll_shutdown(
+        self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+    ) -> Poll<std::result::Result<(), std::io::Error>> {
+        Poll::Ready(Ok(())) // No-op for UDP
+    }
+}