crete/src/tftp/server.rs

use std::io;
use std::net::{IpAddr, SocketAddr};
use std::time::Duration;

use thiserror::Error;
use udp_sas_async::tokio::UdpSocketSas;

use crate::tftp::packet::{ErrorCode, Mode};
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>(
    packet: Packet,
    src: SocketAddr,
    dest: IpAddr,
    callback: F,
) -> Result<(), Error>
where
    F: Fn(&str) -> 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 Packet::Read(request) = packet else {
        Packet::error(ErrorCode::NotDefined, "Operation not supported")
            .send_to(&socket, src)
            .await?;

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

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

    let Some(data) = callback(filename).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>(callback: F) -> Result<!, Error>
where
    F: Fn(&str) -> 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"))?;

    println!("Server started");
    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 (packet, src) = Packet::recv_from(&socket).await?;

        let callback = callback.clone();
        tokio::spawn({
            async move {
                if let Err(err) = handle_connection(packet, src, dest, callback).await {
                    println!("{src}: {err}")
                }
            }
        });
    }
}