Cargo.lock

@@ -310,6 +310,8 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_yaml",
+ "socket-pktinfo",
+ "socket2",
  "thiserror",
  "tokio",
 ]
@@ -1837,6 +1839,17 @@ version = "1.16.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "f9395f0f0eee849a9b707b2f06bb92a6a422090e2123bb2ef8e87a0e61892a8e"
 
+[[package]]
+name = "socket-pktinfo"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "612942246d0cc239cfd83af1dfd39be47f649208a3524e5e9da651910128e0ac"
+dependencies = [
+ "libc",
+ "socket2",
+ "windows-sys",
+]
+
 [[package]]
 name = "socket2"
 version = "0.6.5"

Cargo.toml

@@ -36,6 +36,8 @@ serde_yaml = "0.9.34"
 thiserror = "2.0.21"
 tokio = { version = "1.52.1", features = ["full"] }
 ipxe = { path = "./ipxe" }
+socket2 = "0.6.5"
+socket-pktinfo = "0.4.1"
 
 [lib]
 test = false

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 {}",
+        "{}: Connecting back from {}",
-        socket.local_addr().expect("Socket is bound")
+        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}")
                 }
             }