OpenPXE/crates/tftp/src/server.rs

//! TFTP server implementation.
//!
//! Design: the main socket on :69 accepts RRQ packets. For each RRQ we spawn
//! a task that creates a new ephemeral UDP socket and handles the full
//! transfer there (per RFC 1350 — each transfer uses its own port pair so
//! multiple clients can download concurrently). This matches exactly how
//! `tftpd`/`in.tftpd` works and is why TFTP is awkward behind stateful NAT:
//! the ephemeral ports must be reachable from the client.
//!
//! We only serve files from `pxeforge_ipxe_assets::asset_bytes` — that is,
//! the bundled iPXE binaries and wimboot. No filesystem is ever opened, so
//! `../` path traversal attempts simply return ENOENT.

use pxeforge_core::{ClientEvent, ClientRegistry};
use pxeforge_ipxe_assets::asset_bytes;
use socket2::{Domain, Protocol, Socket, Type};
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Duration;
use tokio::net::UdpSocket;

// TFTP opcodes.
const OP_RRQ: u16 = 1;
const OP_DATA: u16 = 3;
const OP_ACK: u16 = 4;
const OP_ERROR: u16 = 5;
const OP_OACK: u16 = 6;

// Error codes (RFC 1350).
const ERR_NOT_DEFINED: u16 = 0;
const ERR_FILE_NOT_FOUND: u16 = 1;
const ERR_ILLEGAL_OP: u16 = 4;

pub struct TftpServer {
    bind: IpAddr,
    port: u16,
    clients: Arc<ClientRegistry>,
    metrics: pxeforge_core::Metrics,
}

impl TftpServer {
    pub fn new(
        bind: IpAddr,
        port: u16,
        clients: Arc<ClientRegistry>,
        metrics: pxeforge_core::Metrics,
    ) -> Self {
        Self { bind, port, clients, metrics }
    }

    pub async fn run(self) -> anyhow::Result<()> {
        let sock = bind_udp(self.bind, self.port)?;
        tracing::info!(target: "pxeforge::tftp", "TFTP listening on {}:{}", self.bind, self.port);
        let clients = self.clients.clone();
        let metrics = self.metrics.clone();
        let mut buf = vec![0u8; 2048];
        loop {
            let (n, from) = match sock.recv_from(&mut buf).await {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!(target: "pxeforge::tftp", "recv error: {e}");
                    continue;
                }
            };
            let data = buf[..n].to_vec();
            let clients = clients.clone();
            let metrics = metrics.clone();
            let bind_ip = self.bind;
            tokio::spawn(async move {
                if let Err(e) = handle_rrq(data, from, bind_ip, clients, metrics.clone()).await {
                    metrics.record_tftp_err();
                    tracing::warn!(target: "pxeforge::tftp", peer=%from, "handler error: {e}");
                }
            });
        }
    }
}

async fn handle_rrq(
    packet: Vec<u8>,
    peer: SocketAddr,
    bind_ip: IpAddr,
    clients: Arc<ClientRegistry>,
    metrics: pxeforge_core::Metrics,
) -> anyhow::Result<()> {
    let Some(req) = parse_rrq(&packet) else {
        return Ok(());
    };
    let Request { filename, options, .. } = req;

    // Per-transfer ephemeral socket.
    let sock = bind_udp(bind_ip, 0)?;

    let Some(file_bytes) = asset_bytes(&filename) else {
        let _ = send_error(&sock, peer, ERR_FILE_NOT_FOUND, "no such file").await;
        tracing::info!(target: "pxeforge::tftp", peer=%peer, file=%filename, "404");
        clients.record(
            &peer.ip().to_string(),
            Some(peer.ip()),
            None,
            ClientEvent::TftpRead { file: filename.clone() },
        );
        return Ok(());
    };

    tracing::info!(
        target: "pxeforge::tftp",
        peer=%peer, file=%filename, size=file_bytes.len(),
        "serving"
    );
    clients.record(
        &peer.ip().to_string(),
        Some(peer.ip()),
        None,
        ClientEvent::TftpRead { file: filename.clone() },
    );

    // Negotiate options.
    let mut blksize: usize = 512;
    let mut window: u16 = 1;
    let mut accepted_opts: Vec<(String, String)> = Vec::new();

    for (k, v) in &options {
        match k.as_str() {
            "blksize" => {
                if let Ok(n) = v.parse::<usize>() {
                    blksize = n.clamp(8, 65464);
                    accepted_opts.push(("blksize".into(), blksize.to_string()));
                }
            }
            "tsize" => {
                accepted_opts.push(("tsize".into(), file_bytes.len().to_string()));
            }
            "windowsize" => {
                if let Ok(n) = v.parse::<u16>() {
                    window = n.clamp(1, 64);
                    accepted_opts.push(("windowsize".into(), window.to_string()));
                }
            }
            _ => {}
        }
    }

    if !accepted_opts.is_empty() {
        let oack = encode_oack(&accepted_opts);
        // Loop until client ACKs block 0 (the OACK).
        if !wait_for_ack(&sock, peer, 0, &oack).await? {
            return Ok(());
        }
    }

    // DATA transfer. Block numbers are u16 and may wrap at 65535 — we handle
    // that via `wrapping_add`. For each window we remember the starting
    // (offset, block_no) explicitly; on retransmit we replay from there
    // instead of trying to compute it back from `last_block_in_window` which
    // is wrong when a window short-sends at EOF (previous bug: window=8 but
    // only 3 blocks sent, then rewind subtracted 7 landing in the wrong id).
    let total = file_bytes.len();
    let mut offset: usize = 0;
    let mut block_no: u16 = 1;
    // Per RFC 1350: if the final data block is exactly blksize, the
    // server must follow up with a zero-length DATA so the client knows
    // the transfer has ended. The flag is set inside the loop and
    // tested at end-of-transfer.
    let needs_zero_final;

    'transfer: loop {
        let window_start_offset = offset;
        let window_start_block = block_no;
        let mut last_block_in_window = block_no;
        let mut window_reached_eof = false;
        let mut last_chunk_len = 0usize;

        // Send one window worth of DATA.
        for _ in 0..window {
            if offset >= total { break; }
            let end = (offset + blksize).min(total);
            let chunk = &file_bytes[offset..end];
            let pkt = encode_data(block_no, chunk);
            sock.send_to(&pkt, peer).await?;
            last_block_in_window = block_no;
            last_chunk_len = chunk.len();
            offset = end;
            block_no = block_no.wrapping_add(1);
            if end == total {
                window_reached_eof = true;
                break;
            }
        }

        // Wait for an ACK of the last block we actually sent (not the
        // nominal last block of a theoretical full window).
        let mut tries = 0u8;
        loop {
            match tokio::time::timeout(Duration::from_secs(3), recv_ack(&sock, peer)).await {
                Ok(Ok(acked)) if acked == last_block_in_window => break,
                Ok(Ok(_)) => {} // stale ACK from an earlier block — ignore
                Ok(Err(e)) => return Err(e),
                Err(_) => {
                    tries += 1;
                    if tries > 5 {
                        tracing::warn!(
                            target: "pxeforge::tftp",
                            peer=%peer, last_block=last_block_in_window,
                            "timeout after {tries} retries, aborting transfer"
                        );
                        return Ok(());
                    }
                    // Rewind to the start of this window and resend exactly
                    // the same blocks (same count, same block numbers). This
                    // is cheap and correct even for short final windows.
                    offset = window_start_offset;
                    block_no = window_start_block;
                    let mut resent = 0;
                    while resent < window && offset < total {
                        let end = (offset + blksize).min(total);
                        let pkt = encode_data(block_no, &file_bytes[offset..end]);
                        sock.send_to(&pkt, peer).await?;
                        last_block_in_window = block_no;
                        last_chunk_len = end - offset;
                        offset = end;
                        block_no = block_no.wrapping_add(1);
                        resent += 1;
                    }
                }
            }
        }

        if window_reached_eof || offset >= total {
            // If the very last DATA was exactly blksize, RFC 1350 requires a
            // following zero-length DATA to signal end-of-transfer. If it
            // was shorter, the short block already signals EOF.
            needs_zero_final = last_chunk_len == blksize;
            break 'transfer;
        }
    }

    if needs_zero_final {
        let pkt = encode_data(block_no, &[]);
        sock.send_to(&pkt, peer).await?;
        let _ = tokio::time::timeout(Duration::from_secs(3), recv_ack(&sock, peer)).await;
    }

    tracing::debug!(target: "pxeforge::tftp", peer=%peer, bytes=total, "transfer complete");
    metrics.record_tftp_ok(total as u64);
    Ok(())
}

#[derive(Debug)]
struct Request {
    filename: String,
    #[allow(dead_code)]
    mode: String,
    options: Vec<(String, String)>,
}

fn parse_rrq(pkt: &[u8]) -> Option<Request> {
    if pkt.len() < 4 { return None; }
    let op = u16::from_be_bytes([pkt[0], pkt[1]]);
    if op != OP_RRQ { return None; }
    let mut rest = &pkt[2..];
    let filename = read_cstr(&mut rest)?;
    let mode = read_cstr(&mut rest)?;
    let mut options = Vec::new();
    while !rest.is_empty() {
        let Some(k) = read_cstr(&mut rest) else { break };
        if k.is_empty() { break; }
        let v = read_cstr(&mut rest).unwrap_or_default();
        options.push((k.to_ascii_lowercase(), v));
    }
    Some(Request { filename, mode, options })
}

fn read_cstr(buf: &mut &[u8]) -> Option<String> {
    let pos = buf.iter().position(|b| *b == 0)?;
    let s = std::str::from_utf8(&buf[..pos]).ok()?.to_string();
    *buf = &buf[pos + 1..];
    Some(s)
}

fn encode_data(block: u16, chunk: &[u8]) -> Vec<u8> {
    let mut v = Vec::with_capacity(4 + chunk.len());
    v.extend_from_slice(&OP_DATA.to_be_bytes());
    v.extend_from_slice(&block.to_be_bytes());
    v.extend_from_slice(chunk);
    v
}

fn encode_oack(opts: &[(String, String)]) -> Vec<u8> {
    let mut v = Vec::new();
    v.extend_from_slice(&OP_OACK.to_be_bytes());
    for (k, val) in opts {
        v.extend_from_slice(k.as_bytes());
        v.push(0);
        v.extend_from_slice(val.as_bytes());
        v.push(0);
    }
    v
}

async fn send_error(
    sock: &UdpSocket,
    peer: SocketAddr,
    code: u16,
    msg: &str,
) -> std::io::Result<()> {
    let mut v = Vec::with_capacity(5 + msg.len());
    v.extend_from_slice(&OP_ERROR.to_be_bytes());
    v.extend_from_slice(&code.to_be_bytes());
    v.extend_from_slice(msg.as_bytes());
    v.push(0);
    sock.send_to(&v, peer).await.map(|_| ())
}

async fn recv_ack(sock: &UdpSocket, peer: SocketAddr) -> anyhow::Result<u16> {
    let mut buf = [0u8; 32];
    loop {
        let (n, from) = sock.recv_from(&mut buf).await?;
        if from.ip() != peer.ip() { continue; }
        if n < 4 { continue; }
        let op = u16::from_be_bytes([buf[0], buf[1]]);
        match op {
            OP_ACK => return Ok(u16::from_be_bytes([buf[2], buf[3]])),
            OP_ERROR => {
                let code = u16::from_be_bytes([buf[2], buf[3]]);
                anyhow::bail!("client error {code}");
            }
            _ => {}
        }
    }
}

async fn wait_for_ack(
    sock: &UdpSocket,
    peer: SocketAddr,
    expect_block: u16,
    to_retx: &[u8],
) -> anyhow::Result<bool> {
    let mut tries = 0;
    loop {
        sock.send_to(to_retx, peer).await?;
        match tokio::time::timeout(Duration::from_secs(3), recv_ack(sock, peer)).await {
            Ok(Ok(b)) if b == expect_block => return Ok(true),
            Ok(Ok(_)) => {}
            Ok(Err(_)) | Err(_) => {
                tries += 1;
                if tries > 5 { return Ok(false); }
            }
        }
    }
}

fn bind_udp(bind: IpAddr, port: u16) -> anyhow::Result<UdpSocket> {
    let domain = match bind { IpAddr::V4(_) => Domain::IPV4, IpAddr::V6(_) => Domain::IPV6 };
    let sock = Socket::new(domain, Type::DGRAM, Some(Protocol::UDP))?;
    sock.set_reuse_address(true)?;
    sock.set_nonblocking(true)?;
    let addr: SocketAddr = SocketAddr::new(bind, port);
    sock.bind(&addr.into())?;
    let std_sock: std::net::UdpSocket = sock.into();
    Ok(UdpSocket::from_std(std_sock)?)
}

#[allow(dead_code)]
const _UNUSED: (u16, u16) = (ERR_NOT_DEFINED, ERR_ILLEGAL_OP);

/// Pure-logic helper used by the unit tests below and (in a refactor) by
/// `handle_rrq`. Given a position in the file and the window, return the
/// (block_no, chunk_len) list this window will emit. Useful as a sanity
/// check that our windowing math matches the wire behavior the spec
/// requires — tested against edge cases (exact-blksize tail, short tail,
/// single-block window).
#[must_use]
pub fn plan_window(
    total: usize,
    offset: usize,
    blksize: usize,
    window: u16,
    starting_block: u16,
) -> Vec<(u16, usize)> {
    let mut out = Vec::new();
    let mut o = offset;
    let mut b = starting_block;
    for _ in 0..window {
        if o >= total { break; }
        let end = (o + blksize).min(total);
        out.push((b, end - o));
        o = end;
        b = b.wrapping_add(1);
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_rrq_with_options() {
        // RRQ "snponly.efi" mode "octet" blksize=1468 tsize=0
        let mut pkt = vec![0, OP_RRQ as u8];
        // The single-digit `\0` escapes here are NUL terminators between
        // TFTP option name/value pairs — using `\x00` to dodge clippy's
        // "octal-looking escape" lint.
        pkt.extend_from_slice(b"snponly.efi\x00octet\x00blksize\x001468\x00tsize\x000\x00");
        let r = parse_rrq(&pkt).unwrap();
        assert_eq!(r.filename, "snponly.efi");
        assert_eq!(r.mode, "octet");
        assert_eq!(r.options.len(), 2);
        assert_eq!(r.options[0].0, "blksize");
        assert_eq!(r.options[0].1, "1468");
    }

    #[test]
    fn encode_decode_data() {
        let p = encode_data(7, b"hello");
        assert_eq!(&p[0..2], &OP_DATA.to_be_bytes());
        assert_eq!(&p[2..4], &7u16.to_be_bytes());
        assert_eq!(&p[4..], b"hello");
    }

    #[test]
    fn plan_window_full_blocks() {
        // 4 KB file, 1024 blksize, window 4 → one window of 4 full blocks.
        let p = plan_window(4096, 0, 1024, 4, 1);
        assert_eq!(p, vec![(1, 1024), (2, 1024), (3, 1024), (4, 1024)]);
    }

    #[test]
    fn plan_window_short_tail_at_eof() {
        // 3.5 KB file, 1024 blksize, window 8 starting at offset 0.
        // Expect 3 full + 1 half, then stop (below 8 blocks).
        let p = plan_window(3584, 0, 1024, 8, 1);
        assert_eq!(p, vec![(1, 1024), (2, 1024), (3, 1024), (4, 512)]);
    }

    #[test]
    fn plan_window_exact_boundary_needs_zero_final() {
        // 2 KB file, 1024 blksize, window 8 — last block is exactly blksize.
        // `handle_rrq` checks `last_chunk_len == blksize` to decide whether to
        // emit the terminating zero-length DATA. Assert that condition here.
        let p = plan_window(2048, 0, 1024, 8, 1);
        assert_eq!(p, vec![(1, 1024), (2, 1024)]);
        let last = p.last().unwrap();
        assert_eq!(last.1, 1024); // => needs zero final per RFC 1350
    }

    #[test]
    fn plan_window_wraparound() {
        // Block number wraps from u16::MAX to 0 on next window — 2 blocks,
        // starting at 65534.
        let p = plan_window(2048, 0, 1024, 2, 65534);
        assert_eq!(p, vec![(65534, 1024), (65535, 1024)]);
        let p2 = plan_window(2048, 2048, 1024, 2, 0);
        assert!(p2.is_empty()); // nothing past EOF
        // And a cross-boundary case:
        let p3 = plan_window(3072, 0, 1024, 3, 65535);
        assert_eq!(p3, vec![(65535, 1024), (0, 1024), (1, 1024)]);
    }
}