--wip-- [skip ci]

Cargo.lock

@@ -375,9 +375,10 @@ dependencies = [
  "serde",
  "serde_json",
  "serde_yaml",
+ "socket-pktinfo",
+ "socket2",
  "thiserror",
  "tokio",
- "udp_sas_async",
 ]
 
 [[package]]
@@ -1460,16 +1461,6 @@ dependencies = [
  "syn 2.0.119",
 ]
 
-[[package]]
-name = "os_socketaddr"
-version = "0.2.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fa149f3ec09cadeb41968390fcd71e6ea038a9a373d1d4eb070f4d9484d4b244"
-dependencies = [
- "libc",
- "winapi",
-]
-
 [[package]]
 name = "parking"
 version = "2.2.1"
@@ -1955,6 +1946,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"
@@ -2208,18 +2210,6 @@ version = "0.2.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b"
 
-[[package]]
-name = "udp_sas_async"
-version = "0.2.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "073390c0d7b1682a5f044db2f769b0564568f741b9a7e23b8494b79fc868d491"
-dependencies = [
- "cc",
- "libc",
- "os_socketaddr",
- "tokio",
-]
-
 [[package]]
 name = "unicode-ident"
 version = "1.0.26"
@@ -2362,22 +2352,6 @@ dependencies = [
  "wasm-bindgen",
 ]
 
-[[package]]
-name = "winapi"
-version = "0.3.9"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
-dependencies = [
- "winapi-i686-pc-windows-gnu",
- "winapi-x86_64-pc-windows-gnu",
-]
-
-[[package]]
-name = "winapi-i686-pc-windows-gnu"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
-
 [[package]]
 name = "winapi-util"
 version = "0.1.11"
@@ -2387,12 +2361,6 @@ dependencies = [
  "windows-sys",
 ]
 
-[[package]]
-name = "winapi-x86_64-pc-windows-gnu"
-version = "0.4.0"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
-
 [[package]]
 name = "windows-link"
 version = "0.2.1"

Cargo.toml

@@ -37,7 +37,8 @@ thiserror = "2.0.21"
 tokio = { version = "1.52.1", features = ["full"] }
 ipxe = { path = "./ipxe" }
 axum = { version = "0.8.9", features = ["macros"] }
-udp_sas_async = { version = "0.2.0", features = ["async_tokio"] }
+socket2 = "0.6.5"
+socket-pktinfo = "0.4.1"
 
 [lib]
 test = false

src/boot.ipxe

@@ -1,4 +1,3 @@
-{% set httpUrl = "http://192.168.1.1:8000" -%}
 #!ipxe
 echo Starting ${serial}
 

src/cli.rs

@@ -1,3 +1,4 @@
+use std::net::IpAddr;
 use std::path::PathBuf;
 
 use clap::{Args, Parser, Subcommand};
@@ -50,5 +51,10 @@ pub enum Commands {
     /// Generate completions for your current shell
     ShellCompletions,
     /// Serve config files over tftp
-    Serve,
+    Serve {
+        #[arg(short, long, default_value = "3000")]
+        port: u16,
+        #[arg(long, default_value = "0.0.0.0")]
+        ip: IpAddr,
+    },
 }

src/main.rs

@@ -1,9 +1,11 @@
 mod cli;
 
+use std::net::SocketAddr;
 use std::path::PathBuf;
 use std::pin::pin;
 
 use axum::extract::Request;
+use axum::http::Uri;
 use axum::response::IntoResponse;
 use axum::routing::get;
 use axum::{Router, debug_handler};
@@ -102,22 +104,25 @@ async fn generate(opts: &GlobalOpts) -> Result<(), Error> {
 #[debug_handler]
 async fn handler(req: Request) -> impl IntoResponse {
     println!("{req:?}");
-    let host = req
+    let uri: Uri = req
         .headers()
         .iter()
         .find(|(name, _)| header::HOST.eq(name))
         .map(|(_, value)| value)
-        .unwrap()
+        .expect("Host header should exist")
         .to_str()
-        .unwrap();
+        .expect("Header value should be valid ASCII")
+        .parse()
+        .expect("Host should be valid Uri");
 
-    println!("{host}");
+    let host = uri.host().expect("Uri should have host");
+    let port = uri.port().map(|port| port.as_u16()).unwrap_or(80);
 
     let Some(path) = req.uri().path().strip_prefix("/") else {
         return StatusCode::BAD_REQUEST.into_response();
     };
 
-    let Some((file, content_type)) = get_file(path).await else {
+    let Some((file, content_type)) = get_file(path, format!("http://{host}:{port}")).await else {
         return StatusCode::NOT_FOUND.into_response();
     };
 
@@ -135,12 +140,15 @@ async fn handler(req: Request) -> impl IntoResponse {
 }
 
 // TODO: Make this a proper async function
-fn get_file(filename: impl Into<String>) -> impl Future<Output = Option<(Vec<u8>, &'static str)>> {
+fn get_file(
+    filename: impl ToString,
+    http_url: impl ToString,
+) -> impl Future<Output = Option<(Vec<u8>, &'static str)>> {
     // TODO: Make this a static that is initialized once
     let re =
         Regex::new(r"^talos/(?<cluster>.+)/(?<node>.+)\.ya?ml$").expect("regex should be valid");
 
-    let filename: String = filename.into();
+    let filename: String = filename.to_string();
     let re = re.clone();
 
     async move {
@@ -158,6 +166,7 @@ fn get_file(filename: impl Into<String>) -> impl Future<Output = Option<(Vec<u8>
             let content = template
                 .render(context! {
                     clusters => clusters.values().collect::<Vec<_>>(),
+                    httpUrl => http_url.to_string()
                 })
                 .unwrap();
 
@@ -199,16 +208,23 @@ async fn main() -> Result<(), Error> {
             "crete",
             &mut std::io::stdout(),
         ),
-        Commands::Serve => {
-            tokio::spawn(async {
+        Commands::Serve { port, ip } => {
+            tokio::spawn(async move {
                 let app = Router::new().fallback_service(get(handler));
-                let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
+                let addr: SocketAddr = (ip, port).into();
+                let listener = tokio::net::TcpListener::bind(addr).await.unwrap();
+
+                println!("Http server started on {addr}");
                 axum::serve(listener, app).await.unwrap();
             });
 
-            tftp::serve(move |filename| {
+            tftp::serve(ip, move |filename, ip| {
                 let filename = filename.to_owned();
-                async move { get_file(filename).await.map(|(file, _)| file) }
+                async move {
+                    get_file(filename, format!("http://{ip}:{port}"))
+                        .await
+                        .map(|(file, _)| file)
+                }
             })
             .await?;
         }

src/tftp/packet.rs

@@ -3,7 +3,7 @@ use std::net::SocketAddr;
 
 use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
 use num_enum::{IntoPrimitive, TryFromPrimitive};
-use udp_sas_async::tokio::UdpSocketSas;
+use tokio::net::UdpSocket;
 
 use crate::tftp::Error;
 
@@ -309,15 +309,16 @@ impl Packet {
         Ok(())
     }
 
-    pub async fn recv_from(socket: &UdpSocketSas) -> Result<(Self, SocketAddr), Error> {
-        let mut buf = Vec::with_capacity(PACKET_SIZE);
+    pub async fn recv_from(socket: &UdpSocket) -> Result<(Self, SocketAddr), Error> {
+        let mut buf = [0u8; PACKET_SIZE];
         let (_, src) = socket.recv_from(&mut buf).await?;
+
         let packet = buf.try_into()?;
 
         Ok((packet, src))
     }
 
-    pub async fn send_to(&self, socket: &UdpSocketSas, src: SocketAddr) -> Result<(), Error> {
+    pub async fn send_to(&self, socket: &UdpSocket, src: &SocketAddr) -> Result<(), Error> {
         let buf: Vec<_> = self.try_into()?;
         socket.send_to(&buf, src).await?;
 
@@ -341,10 +342,12 @@ impl TryFrom<&Packet> for Vec<u8> {
     }
 }
 
-impl TryFrom<Vec<u8>> for Packet {
+impl TryFrom<[u8; PACKET_SIZE]> for Packet {
     type Error = Error;
 
-    fn try_from(buf: Vec<u8>) -> Result<Self, <Packet as TryFrom<Vec<u8>>>::Error> {
+    fn try_from(
+        buf: [u8; PACKET_SIZE],
+    ) -> Result<Self, <Packet as TryFrom<[u8; PACKET_SIZE]>>::Error> {
         let mut reader = BufReader::new(Cursor::new(buf));
 
         Packet::decode(&mut reader)

src/tftp/server.rs

@@ -2,10 +2,13 @@ use std::io;
 use std::net::{IpAddr, SocketAddr};
 use std::time::Duration;
 
+use socket_pktinfo::{PktInfo, PktInfoUdpSocket};
+use socket2::Domain;
 use thiserror::Error;
-use udp_sas_async::tokio::UdpSocketSas;
+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,25 +32,25 @@ pub enum Error {
 }
 
 pub async fn handle_connection<F, Fut>(
+    info: PktInfo,
     packet: Packet,
-    src: SocketAddr,
-    dest: IpAddr,
     callback: F,
 ) -> Result<(), Error>
 where
-    F: Fn(&str) -> Fut,
+    F: Fn(&str, IpAddr) -> Fut,
     Fut: Future<Output = Option<Vec<u8>>>,
 {
-    let socket = UdpSocketSas::bind(
-        format!("{dest}:0")
-            .parse()
-            .expect("Address should be valid"),
-    )?;
-    println!("{src}: Connecting from {}", socket.local_addr());
+    let src = info.addr_src;
+    let socket = UdpSocket::bind((info.addr_dst, 0)).await?;
+    println!(
+        "{}: 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));
@@ -56,9 +59,9 @@ where
     let filename = request.filename();
     println!("{src}: {filename} [{:?}]", request.mode());
 
-    let Some(data) = callback(filename).await else {
+    let Some(data) = callback(filename, info.addr_dst).await else {
         Packet::error(ErrorCode::FileNotFound, "Unknown file")
-            .send_to(&socket, src)
+            .send_to(&socket, &src)
             .await?;
 
         return Err(Error::FileNotFound(filename.into()));
@@ -74,7 +77,7 @@ where
                 ErrorCode::NotDefined,
                 format!("Mode {mode:?} is not supported"),
             )
-            .send_to(&socket, src)
+            .send_to(&socket, &src)
             .await?;
 
             return Err(Error::UnsupportedMode(mode));
@@ -90,9 +93,10 @@ 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;
@@ -102,7 +106,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));
@@ -123,7 +127,7 @@ where
                     }
                     _ => {
                         Packet::error(ErrorCode::Illegal, "Unexpected packet")
-                            .send_to(&socket, src)
+                            .send_to(&socket, &src)
                             .await?;
 
 
@@ -146,26 +150,36 @@ where
     Ok(())
 }
 
-pub async fn serve<F, Fut>(callback: F) -> Result<!, Error>
+pub async fn serve<F, Fut>(ip: IpAddr, callback: F) -> Result<!, Error>
 where
-    F: Fn(&str) -> Fut + Clone + Send + 'static,
+    F: Fn(&str, IpAddr) -> Fut + Clone + Send + 'static,
     Fut: Future<Output = Option<Vec<u8>>> + Send + 'static,
 {
-    let socket = UdpSocketSas::bind("0.0.0.0:69".parse().expect("Address should be valid"))?;
+    let socket = PktInfoUdpSocket::new(Domain::IPV4)?;
+    socket.set_nonblocking(true)?;
 
-    println!("Server started");
+    let addr: SocketAddr = (ip, 69).into();
+    socket.bind(&addr.into())?;
+
+    let async_fd = AsyncFd::new(socket)?;
+    let mut buf = [0u8; PACKET_SIZE];
+
+    println!("Tftp server started on {addr}");
     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 mut guard = async_fd.readable().await?;
 
-        let (packet, src) = Packet::recv_from(&socket).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(packet, src, dest, callback).await {
+                if let Err(err) = handle_connection(info, packet, callback).await {
                     println!("{src}: {err}")
                 }
             }