crete/src/tftp/server.rs

use std::io;
use std::net::{IpAddr, SocketAddr};
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, PACKET_SIZE};
use crate::tftp::{self, Packet};

#[derive(Debug, Error)]
pub enum Error {
    #[error("{0}")]
    Tftp(#[from] tftp::Error),
    #[error("Io error: {0}")]
    Io(#[from] io::Error),
    #[error("Unsupported operation: {0:#?}")]
    UnsupportedOperation(Packet),
    #[error("File not found: {0}")]
    FileNotFound(String),
    #[error("Unsupported mode: {0:?}")]
    UnsupportedMode(Mode),
    #[error("Source changed: {0}")]
    SourceChanged(SocketAddr),
    #[error("UnexpectedPacket: {0:#?}")]
    UnexpectedPacket(Packet),
    #[error("Timeout")]
    Timeout,
}

pub async fn handle_connection<F, Fut>(
    info: PktInfo,
    packet: Packet,
    callback: F,
) -> Result<(), Error>
where
    F: Fn(&str, IpAddr) -> Fut,
    Fut: Future<Output = Option<Vec<u8>>>,
{
    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, &info.addr_src)
            .await?;

        return Err(Error::UnsupportedOperation(packet));
    };

    let filename = request.filename();
    println!("{src}: {filename} [{:?}]", request.mode());

    let Some(data) = callback(filename, info.addr_dst).await else {
        Packet::error(ErrorCode::FileNotFound, "Unknown file")
            .send_to(&socket, &src)
            .await?;

        return Err(Error::FileNotFound(filename.into()));
    };

    let mode = request.mode();
    let data = match mode {
        // In case of netascii mode we need to re-encode to escape certain sequences
        Mode::NetAscii => netascii::Netascii::from_bytes(data).collect(),
        Mode::Octet => data,
        _ => {
            Packet::error(
                ErrorCode::NotDefined,
                format!("Mode {mode:?} is not supported"),
            )
            .send_to(&socket, &src)
            .await?;

            return Err(Error::UnsupportedMode(mode));
        }
    };

    let mut chunks = data.chunks(512).enumerate();
    let mut chunk = chunks.next();

    let mut timeout = 0;

    loop {
        // Send the first chunk
        if let Some((index, chunk)) = chunk {
            Packet::data((index + 1) as u16, chunk)
                .send_to(&socket, &src)
                .await?;
        }

        tokio::select! {
            recv = Packet::recv_from(&socket) => {
                timeout = 0;
                let (packet, new_src) = recv?;

                // I think this is might technically be possible and fine, but I don't
                // 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)
                        .await?;

                    return Err(Error::SourceChanged(new_src));
                }

                match packet {
                    Packet::Ack(ack) => {
                        // If we receive an ack for the current chunk go to the next one
                        if let Some((index, _)) = chunk
                            && (index + 1) as u16 == ack.block()
                        {
                            chunk = chunks.next();
                        }

                        if chunk.is_none() {
                            break;
                        }
                    }
                    _ => {
                        Packet::error(ErrorCode::Illegal, "Unexpected packet")
                            .send_to(&socket, &src)
                            .await?;


                        return Err(Error::UnexpectedPacket(packet));
                    }
                }
            },
            _ = tokio::time::sleep(Duration::from_secs(1)) => {
                timeout += 1;

                if timeout == 5 {
                    return Err(Error::Timeout);
                }
            }
        }
    }

    println!("{src}: Transmission complete");

    Ok(())
}

pub async fn serve<F, Fut>(ip: IpAddr, callback: F) -> Result<!, Error>
where
    F: Fn(&str, IpAddr) -> Fut + Clone + Send + 'static,
    Fut: Future<Output = Option<Vec<u8>>> + Send + 'static,
{
    let socket = PktInfoUdpSocket::new(Domain::IPV4)?;
    socket.set_nonblocking(true)?;

    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 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(info, packet, callback).await {
                    println!("{src}: {err}")
                }
            }
        });
    }
}