From bfe99b4c9ca584af86505957f2c21e0721e51936 Mon Sep 17 00:00:00 2001
From: Luca Fulchir <luca.fulchir@runesauth.com>
Date: Sat, 25 Feb 2023 16:33:48 +0100
Subject: [PATCH] Strong types for server/receiver and send packet

Signed-off-by: Luca Fulchir <luca.fulchir@runesauth.com>
---
 src/lib.rs | 80 ++++++++++++++++++++++++++++++++++++++++--------------
 1 file changed, 59 insertions(+), 21 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index d8a6447..d667a46 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -192,22 +192,22 @@ type TokenChecker =
 //  so we have to handle `drop()` manually, and garbage-collect the ones we
 //  are no longer using in the background. sigh.
 // Just go with a ArcSwapAny<Arc<Vec<Arc< ...
+type SocketTracker = (Arc<UdpSocket>, Arc<JoinHandle<::std::io::Result<()>>>);
 struct SocketList {
-    sockets:
-        ArcSwap<Vec<(Arc<UdpSocket>, Arc<JoinHandle<::std::io::Result<()>>>)>>,
+    list: ArcSwap<Vec<SocketTracker>>,
 }
 impl SocketList {
     fn new() -> Self {
         Self {
-            sockets: ArcSwap::new(Arc::new(Vec::new())),
+            list: ArcSwap::new(Arc::new(Vec::new())),
         }
     }
     // TODO: fn rm_socket()
     fn rm_all(&self) -> Self {
         let new_list = Arc::new(Vec::new());
-        let old_list = self.sockets.swap(new_list);
+        let old_list = self.list.swap(new_list);
         Self {
-            sockets: old_list.into(),
+            list: old_list.into(),
         }
     }
     async fn add_socket(
@@ -217,24 +217,17 @@ impl SocketList {
     ) {
         let mut new_list;
         {
-            let old_list = self.sockets.load();
+            let old_list = self.list.load();
             new_list = Arc::new(Vec::with_capacity(old_list.len() + 1));
             new_list = old_list.to_vec().into();
         }
         Arc::get_mut(&mut new_list)
             .unwrap()
             .push((socket, Arc::new(handle)));
-        self.sockets.swap(new_list);
-    }
-    async fn find(&self, sock: SocketAddr) -> Option<Arc<UdpSocket>> {
-        let list = self.sockets.load();
-        match list.iter().find(|&(s, _)| s.local_addr().unwrap() == sock) {
-            Some((sock, _)) => Some(sock.clone()),
-            None => None,
-        }
+        self.list.swap(new_list);
     }
     async fn stop_all(mut self) {
-        let mut arc_list = self.sockets.into_inner();
+        let mut arc_list = self.list.into_inner();
         let list = loop {
             match Arc::try_unwrap(arc_list) {
                 Ok(list) => break list,
@@ -251,12 +244,38 @@ impl SocketList {
             Arc::get_mut(&mut handle).unwrap().await;
         }
     }
+    fn lock(&self) -> SocketListRef {
+        SocketListRef {
+            list: self.list.load_full(),
+        }
+    }
 }
+// TODO: impl Drop for SocketList
+struct SocketListRef {
+    list: Arc<Vec<SocketTracker>>,
+}
+impl SocketListRef {
+    fn find(&self, sock: UdpServer) -> Option<Arc<UdpSocket>> {
+        match self
+            .list
+            .iter()
+            .find(|&(s, _)| s.local_addr().unwrap() == sock.0)
+        {
+            Some((sock_srv, _)) => Some(sock_srv.clone()),
+            None => None,
+        }
+    }
+}
+
+#[derive(Debug, Copy, Clone)]
+struct UdpClient(SocketAddr);
+#[derive(Debug, Copy, Clone)]
+struct UdpServer(SocketAddr);
 
 struct RawUdp {
     data: Vec<u8>,
-    src: SocketAddr,
-    dst: SocketAddr,
+    src: UdpClient,
+    dst: UdpServer,
 }
 
 enum Work {
@@ -430,7 +449,7 @@ impl Fenrir {
         socket: Arc<UdpSocket>,
     ) -> ::std::io::Result<()> {
         // jumbo frames are 9K max
-        let sock_receiver = socket.local_addr()?;
+        let sock_receiver = UdpServer(socket.local_addr()?);
         let mut buffer: [u8; 9000] = [0; 9000];
         loop {
             let (bytes, sock_sender) = ::tokio::select! {
@@ -444,8 +463,8 @@ impl Fenrir {
             let data: Vec<u8> = buffer[..bytes].to_vec();
             work_queue.send(Work::Recv(RawUdp {
                 data,
-                src: sock_sender,
-                dst: sock_receiver.clone(),
+                src: UdpClient(sock_sender),
+                dst: sock_receiver,
             }));
         }
         Ok(())
@@ -614,7 +633,8 @@ impl Fenrir {
                                 let mut raw_out =
                                     Vec::<u8>::with_capacity(packet.len());
                                 packet.serialize(&mut raw_out);
-                                todo!()
+                                self.send_packet(raw_out, udp.src, udp.dst)
+                                    .await;
                             }
                             _ => {
                                 todo!()
@@ -628,4 +648,22 @@ impl Fenrir {
         // copy packet, spawn
         todo!();
     }
+    async fn send_packet(
+        &self,
+        data: Vec<u8>,
+        client: UdpClient,
+        server: UdpServer,
+    ) {
+        let sockets = self.sockets.lock();
+        let src_sock = match sockets.find(server) {
+            Some(src_sock) => src_sock,
+            None => {
+                ::tracing::error!(
+                    "Can't send packet: Server changed listening ip!"
+                );
+                return;
+            }
+        };
+        src_sock.send_to(&data, client.0);
+    }
 }