feat: Bind to address that was connected to

src/tftp/server.rs

@@ -1,11 +1,14 @@
 use std::io;
-use std::net::SocketAddr;
+use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
 use std::time::Duration;
 
+use socket_pktinfo::{PktInfo, PktInfoUdpSocket};
+use socket2::Domain;
 use thiserror::Error;
+use tokio::io::unix::AsyncFd;
 use tokio::net::UdpSocket;
 
-use crate::tftp::packet::{ErrorCode, Mode};
+use crate::tftp::packet::{ErrorCode, Mode, PACKET_SIZE};
 use crate::tftp::{self, Packet};
 
 #[derive(Debug, Error)]
@@ -29,7 +32,7 @@ pub enum Error {
 }
 
 pub async fn handle_connection<F, Fut>(
-    src: SocketAddr,
+    info: PktInfo,
     packet: Packet,
     callback: F,
 ) -> Result<(), Error>
@@ -37,15 +40,17 @@ where
     F: Fn(&str) -> Fut,
     Fut: Future<Output = Option<Vec<u8>>>,
 {
-    let socket = UdpSocket::bind("0.0.0.0:0").await?;
+    let src = info.addr_src;
+    let socket = UdpSocket::bind((info.addr_dst, 0)).await?;
     println!(
-        "{src}: Connecting from {}",
-        socket.local_addr().expect("Socket is bound")
+        "{}: Connecting back from {}",
+        src,
+        socket.local_addr().expect("Local addr should be set")
     );
 
     let Packet::Read(request) = packet else {
         Packet::error(ErrorCode::NotDefined, "Operation not supported")
-            .send_to(&socket, &src)
+            .send_to(&socket, &info.addr_src)
             .await?;
 
         return Err(Error::UnsupportedOperation(packet));
@@ -150,16 +155,30 @@ 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 = PktInfoUdpSocket::new(Domain::IPV4)?;
+    socket.set_nonblocking(true)?;
+
+    socket.bind(&SocketAddrV4::new(Ipv4Addr::UNSPECIFIED, 69).into())?;
+
+    let async_fd = AsyncFd::new(socket)?;
+    let mut buf = [0u8; PACKET_SIZE];
 
     println!("Server started");
     loop {
-        let (packet, src) = Packet::recv_from(&socket).await?;
+        let mut guard = async_fd.readable().await?;
+
+        let (packet, info) = match guard.try_io(|inner| inner.get_ref().recv(&mut buf)) {
+            Ok(Ok((_, info))) => (buf.try_into()?, info),
+            Ok(Err(e)) => return Err(e)?,
+            // Socket was not ready
+            Err(_would_block) => continue,
+        };
 
         let callback = callback.clone();
+        let src = info.addr_src;
         tokio::spawn({
             async move {
-                if let Err(err) = handle_connection(src, packet, callback).await {
+                if let Err(err) = handle_connection(info, packet, callback).await {
                     println!("{src}: {err}")
                 }
             }