diff --git a/Cargo.lock b/Cargo.lock index 94224f9..feb6127 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 6e7eb7c..0a39bd7 100644 --- a/Cargo.toml +++ b/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 diff --git a/src/tftp/server.rs b/src/tftp/server.rs index 372e18f..eb26943 100644 --- a/src/tftp/server.rs +++ b/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( - src: SocketAddr, + info: PktInfo, packet: Packet, callback: F, ) -> Result<(), Error> @@ -37,15 +40,17 @@ where F: Fn(&str) -> Fut, Fut: Future>>, { - 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>> + 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}") } }