refactor udp task

src/server/frontend_task.rs

@@ -10,7 +10,11 @@ use tokio::net::UnixStream;
 use tokio::net::TcpStream;
 
 use anyhow::{anyhow, Result};
-use tokio::{io::ReadHalf, sync::mpsc::Sender, task::JoinHandle};
+use tokio::{
+    io::ReadHalf,
+    sync::mpsc::{Receiver, Sender},
+    task::JoinHandle,
+};
 
 use crate::{
     client::{ClientEvent, ClientHandle, Position},
@@ -23,18 +27,14 @@ use super::{
 
 pub(crate) fn new(
     mut frontend: FrontendListener,
+    mut notify_rx: Receiver<FrontendNotify>,
     server: Server,
     producer_notify: Sender<ProducerEvent>,
     consumer_notify: Sender<ConsumerEvent>,
     resolve_ch: Sender<DnsRequest>,
     port_tx: Sender<u16>,
-) -> (
-    JoinHandle<Result<()>>,
-    Sender<FrontendEvent>,
-    Sender<FrontendNotify>,
-) {
+) -> (JoinHandle<Result<()>>, Sender<FrontendEvent>) {
     let (event_tx, mut event_rx) = tokio::sync::mpsc::channel(32);
-    let (notify_tx, mut notify_rx) = tokio::sync::mpsc::channel(32);
     let event_tx_clone = event_tx.clone();
     let frontend_task = tokio::task::spawn_local(async move {
         loop {
@@ -63,7 +63,7 @@ pub(crate) fn new(
         }
         anyhow::Ok(())
     });
-    (frontend_task, event_tx, notify_tx)
+    (frontend_task, event_tx)
 }
 
 async fn handle_frontend_stream(

src/server/network_task.rs

@@ -0,0 +1,90 @@
+use std::net::SocketAddr;
+
+use anyhow::Result;
+use tokio::{
+    net::UdpSocket,
+    sync::mpsc::{Receiver, Sender},
+    task::JoinHandle,
+};
+
+use crate::{event::Event, frontend::FrontendNotify};
+
+use super::Server;
+
+pub async fn new(
+    server: Server,
+    frontend_notify_tx: Sender<FrontendNotify>,
+) -> Result<(
+    JoinHandle<()>,
+    Sender<(Event, SocketAddr)>,
+    Receiver<Result<(Event, SocketAddr)>>,
+    Sender<u16>,
+)> {
+    // bind the udp socket
+    let listen_addr = SocketAddr::new("0.0.0.0".parse().unwrap(), server.port.get());
+    let mut socket = UdpSocket::bind(listen_addr).await?;
+    let (receiver_tx, receiver_rx) = tokio::sync::mpsc::channel(32);
+    let (sender_tx, mut sender_rx) = tokio::sync::mpsc::channel(32);
+    let (port_tx, mut port_rx) = tokio::sync::mpsc::channel(32);
+
+    let udp_task = tokio::task::spawn_local(async move {
+        loop {
+            tokio::select! {
+                event = receive_event(&socket) => {
+                    let _ = receiver_tx.send(event).await;
+                }
+                event = sender_rx.recv() => {
+                    let Some((event, addr)) = event else {
+                        break;
+                    };
+                    if let Err(e) = send_event(&socket, event, addr) {
+                        log::warn!("udp send failed: {e}");
+                    };
+                }
+                port = port_rx.recv() => {
+                    let Some(port) = port else {
+                        break;
+                    };
+
+                    if socket.local_addr().unwrap().port() == port {
+                        continue;
+                    }
+
+                    let listen_addr = SocketAddr::new("0.0.0.0".parse().unwrap(), port);
+                    match UdpSocket::bind(listen_addr).await {
+                        Ok(new_socket) => {
+                            socket = new_socket;
+                            server.port.replace(port);
+                            let _ = frontend_notify_tx.send(FrontendNotify::NotifyPortChange(port, None)).await;
+                        }
+                        Err(e) => {
+                            log::warn!("could not change port: {e}");
+                            let port = socket.local_addr().unwrap().port();
+                            let _ = frontend_notify_tx.send(FrontendNotify::NotifyPortChange(
+                                    port,
+                                    Some(format!("could not change port: {e}")),
+                                )).await;
+                        }
+                    }
+
+                }
+            }
+        }
+    });
+    Ok((udp_task, sender_tx, receiver_rx, port_tx))
+}
+
+async fn receive_event(socket: &UdpSocket) -> Result<(Event, SocketAddr)> {
+    let mut buf = vec![0u8; 22];
+    let (_amt, src) = socket.recv_from(&mut buf).await?;
+    Ok((Event::try_from(buf)?, src))
+}
+
+fn send_event(sock: &UdpSocket, e: Event, addr: SocketAddr) -> Result<usize> {
+    log::trace!("{:20} ------>->->-> {addr}", e.to_string());
+    let data: Vec<u8> = (&e).into();
+    // When udp blocks, we dont want to block the event loop.
+    // Dropping events is better than potentially crashing the event
+    // producer.
+    Ok(sock.try_send_to(&data, addr)?)
+}