path: root/src/tftp.rs

use std::{
    io::{Cursor, Read as _, Result, Write},
    net::UdpSocket,
    path::{Path, PathBuf},
    str::FromStr,
};

use crate::wire;

pub const PORT: u16 = 69;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct InvalidTftpOp(u16);

impl std::fmt::Display for InvalidTftpOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "invalid tftp opcode '{}'", self.0)
    }
}

impl std::error::Error for InvalidTftpOp {}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TftpOp {
    ReadRequest,
    WriteRequest,
    Data,
    Ack,
    Error,
    Oack,
}

impl Into<u16> for TftpOp {
    fn into(self) -> u16 {
        match self {
            TftpOp::ReadRequest => 1,
            TftpOp::WriteRequest => 2,
            TftpOp::Data => 3,
            TftpOp::Ack => 4,
            TftpOp::Error => 5,
            TftpOp::Oack => 6,
        }
    }
}

impl TryFrom<u16> for TftpOp {
    type Error = InvalidTftpOp;

    fn try_from(value: u16) -> std::result::Result<Self, InvalidTftpOp> {
        match value {
            1 => Ok(Self::ReadRequest),
            2 => Ok(Self::WriteRequest),
            3 => Ok(Self::Data),
            4 => Ok(Self::Ack),
            5 => Ok(Self::Error),
            6 => Ok(Self::Oack),
            unknown => Err(InvalidTftpOp(unknown)),
        }
    }
}

#[derive(Debug)]
pub struct InvalidTftpMode(String);

impl std::fmt::Display for InvalidTftpMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "invalid tftp mode '{}'", self.0)
    }
}

impl std::error::Error for InvalidTftpMode {}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TftpMode {
    NetAscii,
    Octet,
    Mail,
}

impl FromStr for TftpMode {
    type Err = InvalidTftpMode;

    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
        match s.to_lowercase().as_str() {
            "netascii" => Ok(Self::NetAscii),
            "octet" => Ok(Self::Octet),
            "mail" => Ok(Self::Mail),
            _ => Err(InvalidTftpMode(s.to_string())),
        }
    }
}

#[derive(Debug)]
pub enum TftpPacket {
    Request(TftpRequestPacket),
    Data(TftpDataPacket),
    Ack(TftpAckPacket),
    OAck(TftpOAckPacket),
    Error(TftpErrorPacket),
}

impl TftpPacket {
    pub fn write<W: Write>(&self, writer: W) -> Result<()> {
        match self {
            TftpPacket::Request(p) => p.write(writer),
            TftpPacket::Data(p) => p.write(writer),
            TftpPacket::Ack(p) => p.write(writer),
            TftpPacket::OAck(p) => p.write(writer),
            TftpPacket::Error(p) => p.write(writer),
        }
    }
}

#[derive(Debug)]
pub struct TftpRequestPacket {
    pub filename: String,
    pub mode: TftpMode,
    pub tsize: Option<u64>,
    pub blksize: Option<u64>,
}

impl TftpRequestPacket {
    pub fn write<W: Write>(&self, mut writer: W) -> Result<()> {
        todo!()
    }
}

#[derive(Debug)]
pub struct TftpDataPacket {
    pub block: u16,
    pub data: Vec<u8>,
}

impl TftpDataPacket {
    pub fn write<W: Write>(&self, mut writer: W) -> Result<()> {
        wire::write_u16(&mut writer, TftpOp::Data.into())?;
        wire::write_u16(&mut writer, self.block)?;
        wire::write(&mut writer, &self.data)?;
        Ok(())
    }
}

#[derive(Debug)]
pub struct TftpAckPacket {
    pub block: u16,
}

impl TftpAckPacket {
    pub fn write<W: Write>(&self, mut writer: W) -> Result<()> {
        wire::write_u16(&mut writer, TftpOp::Data.into())?;
        wire::write_u16(&mut writer, self.block)?;
        Ok(())
    }
}

#[derive(Debug)]
pub struct TftpOAckPacket {
    pub tsize: Option<u64>,
    pub blksize: Option<u64>,
}

impl TftpOAckPacket {
    pub fn write<W: Write>(&self, mut writer: W) -> Result<()> {
        wire::write_u16(&mut writer, TftpOp::Oack.into())?;

        // Only include options that were requested by the client
        if let Some(blksize_val) = self.blksize {
            wire::write(&mut writer, b"blksize")?;
            wire::write_u8(&mut writer, 0)?; // null terminator
            let blksize_str = blksize_val.to_string();
            wire::write(&mut writer, blksize_str.as_bytes())?;
            wire::write_u8(&mut writer, 0)?; // null terminator
        }

        if let Some(tsize_val) = self.tsize {
            wire::write(&mut writer, b"tsize")?;
            wire::write_u8(&mut writer, 0)?; // null terminator
            let tsize_str = tsize_val.to_string();
            wire::write(&mut writer, tsize_str.as_bytes())?;
            wire::write_u8(&mut writer, 0)?; // null terminator
        }

        Ok(())
    }
}

#[derive(Debug)]
pub struct TftpErrorPacket {
    pub code: u16,
    pub message: String,
}

impl TftpErrorPacket {
    pub fn write<W: Write>(&self, mut writer: W) -> Result<()> {
        wire::write_u16(&mut writer, TftpOp::Error.into())?;
        wire::write_u16(&mut writer, self.code)?;
        wire::write_null_terminated_string(&mut writer, &self.message)?;
        Ok(())
    }
}

pub fn parse_packet(buf: &[u8]) -> Result<TftpPacket> {
    let mut cursor = Cursor::new(buf);
    let op = TftpOp::try_from(wire::read_u16(&mut cursor)?).unwrap();

    match op {
        TftpOp::ReadRequest => {
            let filename = wire::read_null_terminated_string(&mut cursor)?;
            let mode = wire::read_null_terminated_string(&mut cursor)?
                .parse::<TftpMode>()
                .unwrap();

            let mut tsize = None;
            let mut blksize = None;

            while let Ok(opt_name) = wire::read_null_terminated_string(&mut cursor) {
                if opt_name.is_empty() {
                    break;
                }
                let opt_data = wire::read_null_terminated_string(&mut cursor)?;
                match opt_name.as_str() {
                    "tsize" => tsize = Some(opt_data.parse::<u64>().unwrap()),
                    "blksize" => blksize = Some(opt_data.parse::<u64>().unwrap()),
                    _ => eprintln!("unknown tftp request option '{opt_name}'"),
                }
            }

            Ok(TftpPacket::Request(TftpRequestPacket {
                filename,
                mode,
                tsize,
                blksize,
            }))
        }
        TftpOp::WriteRequest => unimplemented!(),
        TftpOp::Data => {
            let block = wire::read_u16(&mut cursor)?;
            let mut data = Vec::new();
            cursor.read_to_end(&mut data)?;
            Ok(TftpPacket::Data(TftpDataPacket { block, data }))
        }
        TftpOp::Ack => {
            let block = wire::read_u16(&mut cursor)?;
            Ok(TftpPacket::Ack(TftpAckPacket { block }))
        }
        TftpOp::Error => {
            let code = wire::read_u16(&mut cursor)?;
            let message = wire::read_null_terminated_string(&mut cursor)?;
            Ok(TftpPacket::Error(TftpErrorPacket { code, message }))
        }
        TftpOp::Oack => {
            let mut tsize = None;
            let mut blksize = None;

            while let Ok(opt_name) = wire::read_null_terminated_string(&mut cursor) {
                if opt_name.is_empty() {
                    break;
                }
                let opt_data = wire::read_null_terminated_string(&mut cursor)?;
                match opt_name.as_str() {
                    "tsize" => tsize = Some(opt_data.parse::<u64>().unwrap()),
                    "blksize" => blksize = Some(opt_data.parse::<u64>().unwrap()),
                    _ => eprintln!("unknown tftp ack option '{opt_name}'"),
                }
            }
            Ok(TftpPacket::OAck(TftpOAckPacket { tsize, blksize }))
        }
    }
}

pub fn serve(dir: impl AsRef<Path>) -> Result<()> {
    let dir = dir.as_ref();
    let socket = UdpSocket::bind(format!("0.0.0.0:{PORT}"))?;

    // TODO: this needs to be done per addr
    let mut last_blksize = 512u64;
    let mut current_file = PathBuf::default();

    loop {
        let mut buf = [0u8; 1500];
        let (n, addr) = socket.recv_from(&mut buf)?;
        let packet = parse_packet(&buf[..n]).unwrap();

        let response = match packet {
            TftpPacket::Request(req) => {
                println!(
                    "Request options: tsize={:?}, blksize={:?}",
                    req.tsize, req.blksize
                );

                let filepath = dir.join(req.filename);
                current_file = filepath.clone();
                let meta = std::fs::metadata(&filepath).unwrap();
                let actual_file_size = meta.len();

                // Only send OACK if client sent options
                if req.tsize.is_some() || req.blksize.is_some() {
                    if let Some(blksize) = req.blksize {
                        last_blksize = blksize;
                    }

                    let tsize_response = if req.tsize.is_some() {
                        Some(actual_file_size)
                    } else {
                        None
                    };

                    Some(TftpPacket::OAck(TftpOAckPacket {
                        tsize: req.tsize,
                        blksize: req.blksize,
                    }))
                } else {
                    // No options, send first data block directly
                    let contents = std::fs::read(&filepath).unwrap();
                    let block_size = 512;
                    let first_block = if contents.len() > block_size {
                        contents[..block_size].to_vec()
                    } else {
                        contents
                    };

                    Some(TftpPacket::Data(TftpDataPacket {
                        block: 1,
                        data: first_block,
                    }))
                }
            }
            TftpPacket::Data(dat) => unimplemented!(),
            TftpPacket::Ack(ack) => {
                println!("Received ACK packet: block {}", ack.block);

                let contents = std::fs::read(&current_file).unwrap();
                let next_block = ack.block + 1;
                let start_offset = (next_block - 1) as u64 * last_blksize;
                let end_offset = next_block as u64 * last_blksize;
                let prev_start_offset = (next_block.saturating_sub(2)) as u64 * last_blksize;
                let prev_remain = contents.len() - prev_start_offset as usize;
                if prev_remain as u64 >= last_blksize || ack.block == 0 {
                    let end = std::cmp::min(end_offset as usize, contents.len());
                    let block_data = contents[start_offset as usize..end].to_vec();
                    println!("sending tftp data packet with {} bytes", block_data.len());
                    Some(TftpPacket::Data(TftpDataPacket {
                        block: next_block,
                        data: block_data,
                    }))
                } else {
                    None
                }
            }
            TftpPacket::OAck(ack) => todo!(),
            TftpPacket::Error(err) => {
                println!(
                    "Received ERROR packet: code {}, message: {}",
                    err.code, err.message
                );
                None
            }
        };

        if let Some(response) = response {
            let mut writer = Cursor::new(&mut buf[..]);
            println!("Sending to {addr}: {response:#?}");
            response.write(&mut writer).unwrap();
            let (response, _) = writer.split();
            println!("Sending {} bytes to {addr}: {response:#?}", response.len());
            socket.send_to(&response, addr).unwrap();
        }
    }
}