crete/src/tftp/packet.rs

use std::io::{BufRead, BufReader, BufWriter, Cursor, Read, Write};
use std::net::SocketAddr;

use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
use num_enum::{IntoPrimitive, TryFromPrimitive};
use tokio::net::UdpSocket;

use crate::tftp::Error;

// 2 bytes op code + 2 bytes block number + max 512 bytes data
pub(crate) const PACKET_SIZE: usize = 2 + 2 + 512;

fn decode_null_terminated<R>(reader: &mut BufReader<R>) -> Result<String, Error>
where
    R: Read + ?Sized,
{
    let mut text = Vec::new();
    reader.read_until(0, &mut text)?;

    // Make sure we ended with a null terminator
    match text.pop() {
        Some(0) => {}
        None | Some(_) => return Err(Error::UnterminatedString),
    }

    Ok(String::from_utf8(text)?)
}

fn encode_null_terminated<W>(writer: &mut BufWriter<W>, text: impl AsRef<str>) -> Result<(), Error>
where
    W: Write + ?Sized,
{
    let text = text.as_ref().as_bytes();
    writer.write_all(text)?;

    // Add null terminator
    writer.write_u8(0)?;

    Ok(())
}

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

impl Mode {
    fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
    where
        R: Read + ?Sized,
    {
        let mut mode = decode_null_terminated(reader)?;
        mode.make_ascii_lowercase();
        let mode = mode;

        Ok(match mode.as_str() {
            "netascii" => Self::NetAscii,
            "octet" => Self::Octet,
            "mail" => Self::Mail,
            _ => return Err(Error::InvalidMode(mode)),
        })
    }

    fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
    where
        W: Write + ?Sized,
    {
        let mode = match &self {
            Mode::NetAscii => "netascii",
            Mode::Octet => "octet",
            Mode::Mail => "mail",
        };

        encode_null_terminated(writer, mode)?;

        Ok(())
    }
}

#[derive(Debug, PartialEq, Eq)]
pub struct Request {
    filename: String,
    mode: Mode,
}

impl Request {
    pub fn filename(&self) -> &str {
        &self.filename
    }

    fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
    where
        R: Read + ?Sized,
    {
        let filename = decode_null_terminated(reader)?;
        let mode = Mode::decode(reader)?;

        Ok(Self { filename, mode })
    }

    fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
    where
        W: Write + ?Sized,
    {
        encode_null_terminated(writer, &self.filename)?;
        self.mode.encode(writer)?;

        Ok(())
    }

    pub fn mode(&self) -> Mode {
        self.mode
    }
}

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

impl Data {
    fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
    where
        R: Read + ?Sized,
    {
        let block = reader.read_u16::<NetworkEndian>()?;

        let mut data = Vec::new();
        reader.read_to_end(&mut data)?;

        Ok(Self { data, block })
    }

    fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
    where
        W: Write + ?Sized,
    {
        writer.write_u16::<NetworkEndian>(self.block)?;
        writer.write_all(&self.data)?;

        Ok(())
    }
}

#[derive(Debug)]
pub struct Ack(u16);

impl Ack {
    pub fn block(&self) -> u16 {
        self.0
    }

    fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
    where
        R: Read + ?Sized,
    {
        let number = reader.read_u16::<NetworkEndian>()?;

        Ok(Self(number))
    }

    fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
    where
        W: Write + ?Sized,
    {
        writer.write_u16::<NetworkEndian>(self.0)?;

        Ok(())
    }
}

#[derive(Debug, TryFromPrimitive, IntoPrimitive, Clone, Copy)]
#[repr(u16)]
pub enum ErrorCode {
    NotDefined,
    FileNotFound,
    AccessViolation,
    DiskFull,
    Illegal,
    UnknownTransferId,
    FileAlreadyExists,
    NoSuchUser,
}

impl ErrorCode {
    fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
    where
        R: Read + ?Sized,
    {
        let code = reader.read_u16::<NetworkEndian>()?;
        Self::try_from_primitive(code).map_err(|err| Error::InvalidErrorCode(err.number))
    }

    fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
    where
        W: Write + ?Sized,
    {
        let code: u16 = (*self).into();

        writer.write_u16::<NetworkEndian>(code)?;

        Ok(())
    }
}

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

impl ErrorMessage {
    fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
    where
        R: Read + ?Sized,
    {
        let code = ErrorCode::decode(reader)?;
        let message = decode_null_terminated(reader)?;

        Ok(Self { code, message })
    }

    fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
    where
        W: Write + ?Sized,
    {
        self.code.encode(writer)?;

        encode_null_terminated(writer, &self.message)?;

        Ok(())
    }
}

#[derive(Debug)]
pub enum Packet {
    Read(Request),
    Write(Request),
    Data(Data),
    Ack(Ack),
    Error(ErrorMessage),
}

impl Packet {
    pub fn data(block: u16, data: &[u8]) -> Self {
        Self::Data(Data {
            block,
            data: data.into(),
        })
    }

    pub fn ack(block: u16) -> Self {
        Self::Ack(Ack(block))
    }

    pub fn error(code: ErrorCode, message: impl Into<String>) -> Self {
        Self::Error(ErrorMessage {
            code,
            message: message.into(),
        })
    }

    fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
    where
        R: Read + ?Sized,
    {
        let op_code = reader.read_u16::<NetworkEndian>()?;

        Ok(match op_code {
            1 => Self::Read(Request::decode(reader)?),
            2 => Self::Write(Request::decode(reader)?),
            3 => Self::Data(Data::decode(reader)?),
            4 => Self::Ack(Ack::decode(reader)?),
            5 => Self::Error(ErrorMessage::decode(reader)?),
            _ => return Err(Error::InvalidOpCode(op_code)),
        })
    }

    fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
    where
        W: Write + ?Sized,
    {
        match &self {
            Self::Read(inner) => {
                writer.write_u16::<NetworkEndian>(1)?;
                inner.encode(writer)?;
            }
            Self::Write(inner) => {
                writer.write_u16::<NetworkEndian>(2)?;
                inner.encode(writer)?;
            }
            Self::Data(inner) => {
                writer.write_u16::<NetworkEndian>(3)?;
                inner.encode(writer)?;
            }
            Self::Ack(inner) => {
                writer.write_u16::<NetworkEndian>(4)?;
                inner.encode(writer)?;
            }
            Self::Error(inner) => {
                writer.write_u16::<NetworkEndian>(5)?;
                inner.encode(writer)?;
            }
        };

        Ok(())
    }

    pub async fn recv_from(socket: &UdpSocket) -> Result<(Self, SocketAddr), Error> {
        let mut buf = Vec::with_capacity(PACKET_SIZE);
        let (_, src) = socket.recv_buf_from(&mut buf).await?;

        let packet = buf.try_into()?;

        Ok((packet, src))
    }

    pub async fn send_to(&self, socket: &UdpSocket, src: &SocketAddr) -> Result<(), Error> {
        let buf: Vec<_> = self.try_into()?;
        socket.send_to(&buf, src).await?;

        Ok(())
    }
}

impl TryFrom<&Packet> for Vec<u8> {
    type Error = Error;

    fn try_from(packet: &Packet) -> Result<Self, Self::Error> {
        let mut buf = Vec::with_capacity(PACKET_SIZE);
        let mut writer = BufWriter::new(Cursor::new(&mut buf));
        packet.encode(&mut writer)?;

        // Normally the writer would get dropped after we return, but since it has a mutable
        // borrow in buf we have to drop it first.
        drop(writer);

        Ok(buf)
    }
}

impl TryFrom<Vec<u8>> for Packet {
    type Error = Error;

    fn try_from(buf: Vec<u8>) -> Result<Self, <Packet as TryFrom<Vec<u8>>>::Error> {
        let mut reader = BufReader::new(Cursor::new(buf));

        Packet::decode(&mut reader)
    }
}