--wip-- [skip ci]

src/tftp/server.rs

@@ -1,9 +1,9 @@
 use std::io;
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 use std::time::Duration;
 
 use thiserror::Error;
-use tokio::net::UdpSocket;
+use udp_sas_async::tokio::UdpSocketSas;
 
 use crate::tftp::packet::{ErrorCode, Mode};
 use crate::tftp::{self, Packet};
@@ -29,23 +29,25 @@ pub enum Error {
 }
 
 pub async fn handle_connection<F, Fut>(
-    src: SocketAddr,
     packet: Packet,
+    src: SocketAddr,
+    dest: IpAddr,
     callback: F,
 ) -> Result<(), Error>
 where
     F: Fn(&str) -> Fut,
     Fut: Future<Output = Option<Vec<u8>>>,
 {
-    let socket = UdpSocket::bind("0.0.0.0:0").await?;
-    println!(
-        "{src}: Connecting from {}",
-        socket.local_addr().expect("Socket is bound")
-    );
+    let socket = UdpSocketSas::bind(
+        format!("{dest}:0")
+            .parse()
+            .expect("Address should be valid"),
+    )?;
+    println!("{src}: Connecting from {}", socket.local_addr());
 
     let Packet::Read(request) = packet else {
         Packet::error(ErrorCode::NotDefined, "Operation not supported")
-            .send_to(&socket, &src)
+            .send_to(&socket, src)
             .await?;
 
         return Err(Error::UnsupportedOperation(packet));
@@ -56,7 +58,7 @@ where
 
     let Some(data) = callback(filename).await else {
         Packet::error(ErrorCode::FileNotFound, "Unknown file")
-            .send_to(&socket, &src)
+            .send_to(&socket, src)
             .await?;
 
         return Err(Error::FileNotFound(filename.into()));
@@ -72,7 +74,7 @@ where
                 ErrorCode::NotDefined,
                 format!("Mode {mode:?} is not supported"),
             )
-            .send_to(&socket, &src)
+            .send_to(&socket, src)
             .await?;
 
             return Err(Error::UnsupportedMode(mode));
@@ -88,10 +90,9 @@ where
         // Send the first chunk
         if let Some((index, chunk)) = chunk {
             Packet::data((index + 1) as u16, chunk)
-                .send_to(&socket, &src)
+                .send_to(&socket, src)
                 .await?;
         }
-
         tokio::select! {
             recv = Packet::recv_from(&socket) => {
                 timeout = 0;
@@ -101,7 +102,7 @@ where
                 // really know how to test it properly so we don't support it.
                 if new_src != src {
                     Packet::error(ErrorCode::Illegal, "Source changed")
-                        .send_to(&socket, &src)
+                        .send_to(&socket, src)
                         .await?;
 
                     return Err(Error::SourceChanged(new_src));
@@ -122,7 +123,7 @@ where
                     }
                     _ => {
                         Packet::error(ErrorCode::Illegal, "Unexpected packet")
-                            .send_to(&socket, &src)
+                            .send_to(&socket, src)
                             .await?;
 
 
@@ -150,16 +151,21 @@ where
     F: Fn(&str) -> Fut + Clone + Send + 'static,
     Fut: Future<Output = Option<Vec<u8>>> + Send + 'static,
 {
-    let socket = UdpSocket::bind("0.0.0.0:69").await?;
+    let socket = UdpSocketSas::bind("0.0.0.0:69".parse().expect("Address should be valid"))?;
 
     println!("Server started");
     loop {
+        let mut buf = Vec::with_capacity(0);
+        let (bytes, src, dest) = socket.recv_sas(&mut buf).await?;
+        println!("{src} -> {dest}");
+        println!("{bytes}, {}", buf.len());
+
         let (packet, src) = Packet::recv_from(&socket).await?;
 
         let callback = callback.clone();
         tokio::spawn({
             async move {
-                if let Err(err) = handle_connection(src, packet, callback).await {
+                if let Err(err) = handle_connection(packet, src, dest, callback).await {
                     println!("{src}: {err}")
                 }
             }