internet 0.1.0

//! UDP Options.
//!
//! As defined in [RFC 9868].
//!
//! [RFC 9868]: https://datatracker.ietf.org/doc/html/rfc9868

use crate::{Buf, BufError, BufMut, BufResult, Codec, Cursor};

/// UDP Options.
///
/// Contains the raw bytes of UDP options. Implements `Iterator` to sequentially
/// decode individual options.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Options {
    length: u8,
    data: [u8; 40],
    position: usize,
}

impl TryFrom<&[u8]> for Options {
    type Error = BufError;

    fn try_from(slice: &[u8]) -> Result<Self, Self::Error> {
        if slice.len() > 40 {
            return Err(BufError::UnexpectedValue);
        }
        let mut data = [0u8; 40];
        data[..slice.len()].copy_from_slice(slice);
        Ok(Self {
            length: slice.len() as u8,
            data,
            position: 0,
        })
    }
}

impl<const N: usize> TryFrom<&[u8; N]> for Options {
    type Error = BufError;

    fn try_from(slice: &[u8; N]) -> Result<Self, Self::Error> {
        if slice.len() > 40 {
            return Err(BufError::UnexpectedValue);
        }
        let mut data = [0u8; 40];
        data[..slice.len()].copy_from_slice(slice);
        Ok(Self {
            length: slice.len() as u8,
            data,
            position: 0,
        })
    }
}

impl Codec<usize> for Options {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: usize) -> BufResult<()> {
        writer.write_slice(&self.data[..self.length as usize])
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, len: usize) -> BufResult<Self> {
        if len > 40 {
            return Err(BufError::UnexpectedValue);
        }
        let mut data = [0u8; 40];
        if len > 0 {
            reader.read_into(&mut data[..len])?;
        }
        Ok(Self {
            length: len as u8,
            data,
            position: 0,
        })
    }
}

impl Iterator for Options {
    type Item = Result<self::Option, BufError>;

    fn next(&mut self) -> core::option::Option<Self::Item> {
        if self.position >= self.length as usize {
            return None;
        }

        let slice = &self.data[self.position..self.length as usize];
        let mut cursor = Cursor::new(slice);
        let result = self::Option::decode(&mut cursor, ());
        self.position += cursor.position();

        Some(result)
    }
}

/// A UDP option.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Option {
    /// End of Options List.
    EndOfOptionsList(EndOfOptionsList),
    /// No Operation.
    NoOperation(NoOperation),
    /// Additional Payload Checksum.
    AdditionalPayloadChecksum(AdditionalPayloadChecksum),
    /// Fragmentation.
    Fragmentation(Fragmentation),
    /// Maximum Datagram Size.
    MaximumDatagramSize(MaximumDatagramSize),
    /// Maximum Reassembled Datagram Size.
    MaximumReassembledDatagramSize(MaximumReassembledDatagramSize),
    /// Request.
    Request(EchoRequest),
    /// Response.
    Response(EchoResponse),
    /// Timestamp.
    Timestamp(Timestamp),
}

impl Codec for Option {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        match self {
            Option::EndOfOptionsList(opt) => opt.encode(writer, ()),
            Option::NoOperation(opt) => opt.encode(writer, ()),
            Option::AdditionalPayloadChecksum(opt) => opt.encode(writer, ()),
            Option::Fragmentation(opt) => opt.encode(writer, ()),
            Option::MaximumDatagramSize(opt) => opt.encode(writer, ()),
            Option::MaximumReassembledDatagramSize(opt) => opt.encode(writer, ()),
            Option::Request(opt) => opt.encode(writer, ()),
            Option::Response(opt) => opt.encode(writer, ()),
            Option::Timestamp(opt) => opt.encode(writer, ()),
        }
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let kind = Kind::try_from(reader.peek_u8()?)?;
        match kind {
            Kind::EndOfOptionList => Ok(Option::EndOfOptionsList(EndOfOptionsList::decode(
                reader,
                (),
            )?)),
            Kind::NoOperation => Ok(Option::NoOperation(NoOperation::decode(reader, ())?)),
            Kind::AdditionalPayloadChecksum => Ok(Option::AdditionalPayloadChecksum(
                AdditionalPayloadChecksum::decode(reader, ())?,
            )),
            Kind::Fragmentation => Ok(Option::Fragmentation(Fragmentation::decode(reader, ())?)),
            Kind::MaximumDatagramSize => Ok(Option::MaximumDatagramSize(
                MaximumDatagramSize::decode(reader, ())?,
            )),
            Kind::MaximumReassembledDatagramSize => Ok(Option::MaximumReassembledDatagramSize(
                MaximumReassembledDatagramSize::decode(reader, ())?,
            )),
            Kind::Request => Ok(Option::Request(EchoRequest::decode(reader, ())?)),
            Kind::Response => Ok(Option::Response(EchoResponse::decode(reader, ())?)),
            Kind::Timestamp => Ok(Option::Timestamp(Timestamp::decode(reader, ())?)),
        }
    }
}

/// A UDP option kind.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum Kind {
    /// End of Option List.
    EndOfOptionList = 0,
    /// No Operation.
    NoOperation = 1,
    /// Additional Payload Checksum.
    AdditionalPayloadChecksum = 2,
    /// Fragmentation.
    Fragmentation = 3,
    /// Maximum Datagram Size.
    MaximumDatagramSize = 4,
    /// Maximum Reassembled Datagram Size.
    MaximumReassembledDatagramSize = 5,
    /// Request.
    Request = 6,
    /// Response.
    Response = 7,
    /// Timestamp.
    Timestamp = 8,
}

impl TryFrom<u8> for Kind {
    type Error = BufError;

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(Kind::EndOfOptionList),
            1 => Ok(Kind::NoOperation),
            2 => Ok(Kind::AdditionalPayloadChecksum),
            3 => Ok(Kind::Fragmentation),
            4 => Ok(Kind::MaximumDatagramSize),
            5 => Ok(Kind::MaximumReassembledDatagramSize),
            6 => Ok(Kind::Request),
            7 => Ok(Kind::Response),
            8 => Ok(Kind::Timestamp),
            _ => Err(BufError::UnexpectedValue),
        }
    }
}

impl Codec for Kind {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        (*self as u8).encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let byte = u8::decode(reader, ())?;
        Self::try_from(byte).map_err(|_| BufError::UnexpectedValue)
    }
}

/// End Of Options List option.
///
/// As defined in [Section 11.1].
///
/// [Section 11.1]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.1
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct EndOfOptionsList;

impl EndOfOptionsList {
    /// The option kind.
    pub const KIND: Kind = Kind::EndOfOptionList;
}

impl Codec for EndOfOptionsList {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        Ok(Self)
    }
}

/// No Operation option.
///
/// As defined in [Section 11.2].
///
/// [Section 11.2]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.2
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct NoOperation;

impl NoOperation {
    /// The option kind.
    pub const KIND: Kind = Kind::NoOperation;
}

impl Codec for NoOperation {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        Ok(Self)
    }
}

/// Additional Payload Checksum option.
///
/// As defined in [Section 11.3].
///
/// [Section 11.3]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.3
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AdditionalPayloadChecksum {
    /// CRC32C checksum.
    pub crc32c: u32,
}

impl AdditionalPayloadChecksum {
    /// The option kind.
    pub const KIND: Kind = Kind::AdditionalPayloadChecksum;
    /// The option length.
    pub const LEN: u8 = 8;
}

impl Codec for AdditionalPayloadChecksum {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        0u16.encode(writer, ())?;
        self.crc32c.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        if u8::decode(reader, ())? != Self::LEN {
            return Err(BufError::UnexpectedValue);
        }
        let _reserved = u16::decode(reader, ())?;
        let crc32c = u32::decode(reader, ())?;
        Ok(Self { crc32c })
    }
}

/// Fragmentation option.
///
/// As defined in [Section 11.4].
///
/// [Section 11.4]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.4
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Fragmentation {
    /// Length of the fragmentation data.
    pub len: u8,
}

impl Fragmentation {
    /// The option kind.
    pub const KIND: Kind = Kind::Fragmentation;
}

impl Codec for Fragmentation {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        self.len.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        let len = u8::decode(reader, ())?;
        Ok(Self { len })
    }
}

/// Maximum Datagram Size option.
///
/// As defined in [Section 11.5].
///
/// [Section 11.5]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.5
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MaximumDatagramSize {
    /// Maximum datagram size.
    pub mds: u16,
}

impl MaximumDatagramSize {
    /// The option kind.
    pub const KIND: Kind = Kind::MaximumDatagramSize;
    /// The option length.
    pub const LEN: u8 = 4;
}

impl Codec for MaximumDatagramSize {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        self.mds.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        if u8::decode(reader, ())? != Self::LEN {
            return Err(BufError::UnexpectedValue);
        }
        let mds = u16::decode(reader, ())?;
        Ok(Self { mds })
    }
}

/// Maximum Reassembled Datagram Size option.
///
/// As defined in [Section 11.6].
///
/// [Section 11.6]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.6
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MaximumReassembledDatagramSize {
    /// Maximum reassembled datagram size.
    pub mrds: u16,
}

impl MaximumReassembledDatagramSize {
    /// The option kind.
    pub const KIND: Kind = Kind::MaximumReassembledDatagramSize;
    /// The option length.
    pub const LEN: u8 = 4;
}

impl Codec for MaximumReassembledDatagramSize {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        self.mrds.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        if u8::decode(reader, ())? != Self::LEN {
            return Err(BufError::UnexpectedValue);
        }
        let mrds = u16::decode(reader, ())?;
        Ok(Self { mrds })
    }
}

/// Echo Request option.
///
/// As defined in [Section 11.7].
///
/// [Section 11.7]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.7
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct EchoRequest {
    /// Maximum segment size.
    pub mss: u16,
}

impl EchoRequest {
    /// The option kind.
    pub const KIND: Kind = Kind::Request;
    /// The option length.
    pub const LEN: u8 = 4;
}

impl Codec for EchoRequest {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        self.mss.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        if u8::decode(reader, ())? != Self::LEN {
            return Err(BufError::UnexpectedValue);
        }
        let mss = u16::decode(reader, ())?;
        Ok(Self { mss })
    }
}

/// Echo Response option.
///
/// As defined in [Section 11.7].
///
/// [Section 11.7]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.7
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct EchoResponse {
    /// Length of the response data.
    pub len: u8,
}

impl EchoResponse {
    /// The option kind.
    pub const KIND: Kind = Kind::Response;
}

impl Codec for EchoResponse {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        self.len.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        let len = u8::decode(reader, ())?;
        Ok(Self { len })
    }
}

/// Timestamp option.
///
/// As defined in [Section 11.8].
///
/// [Section 11.8]: https://datatracker.ietf.org/doc/html/rfc9868#section-11.8
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Timestamp {
    /// Timestamp value.
    pub ts_val: u32,
    /// Timestamp echo reply.
    pub ts_ecr: u32,
}

impl Timestamp {
    /// The option kind.
    pub const KIND: Kind = Kind::Timestamp;
    /// The option length.
    pub const LEN: u8 = 12;
}

impl Codec for Timestamp {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        0u16.encode(writer, ())?;
        self.ts_val.encode(writer, ())?;
        self.ts_ecr.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        if u8::decode(reader, ())? != Self::LEN {
            return Err(BufError::UnexpectedValue);
        }
        let _reserved = u16::decode(reader, ())?;
        let ts_val = u32::decode(reader, ())?;
        let ts_ecr = u32::decode(reader, ())?;
        Ok(Self { ts_val, ts_ecr })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{Codec, Cursor};
    use core::fmt::Debug;

    fn codec_roundtrip<T: Codec<C> + Debug + Eq, C: Copy>(
        etalon_struct: T,
        etalon_bytes: &[u8],
        context: C,
    ) {
        let mut encoded_bytes = vec![];
        {
            let writer = &mut Cursor::new(&mut encoded_bytes);
            etalon_struct.encode(writer, context).unwrap();
        }
        assert_eq!(etalon_bytes, &encoded_bytes);

        let decoded_struct = {
            let reader = &mut Cursor::new(&encoded_bytes);
            T::decode(reader, context).unwrap()
        };
        assert_eq!(etalon_struct, decoded_struct);
    }

    #[test]
    fn end_of_options_list() {
        let etalon_bytes = &[0x00];
        let etalon_struct = Option::EndOfOptionsList(EndOfOptionsList);
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn no_operation() {
        let etalon_bytes = &[0x01];
        let etalon_struct = Option::NoOperation(NoOperation);
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn additional_payload_checksum() {
        let etalon_bytes = &[0x02, 0x08, 0x00, 0x00, 0x12, 0x34, 0x56, 0x78];
        let etalon_struct =
            Option::AdditionalPayloadChecksum(AdditionalPayloadChecksum { crc32c: 0x12345678 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn fragmentation() {
        let etalon_bytes = &[0x03, 0x0a];
        let etalon_struct = Option::Fragmentation(Fragmentation { len: 10 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn maximum_datagram_size() {
        let etalon_bytes = &[0x04, 0x04, 0x05, 0xdc];
        let etalon_struct = Option::MaximumDatagramSize(MaximumDatagramSize { mds: 1500 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn maximum_reassembled_datagram_size() {
        let etalon_bytes = &[0x05, 0x04, 0xff, 0xff];
        let etalon_struct =
            Option::MaximumReassembledDatagramSize(MaximumReassembledDatagramSize { mrds: 65535 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn echo_request() {
        let etalon_bytes = &[0x06, 0x04, 0x05, 0xb4];
        let etalon_struct = Option::Request(EchoRequest { mss: 1460 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn echo_response() {
        let etalon_bytes = &[0x07, 0x0a];
        let etalon_struct = Option::Response(EchoResponse { len: 10 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn timestamp() {
        let etalon_bytes = &[
            0x08, 0x0c, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x02,
        ];
        let etalon_struct = Option::Timestamp(Timestamp {
            ts_val: 1,
            ts_ecr: 2,
        });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn options_struct_roundtrip() {
        let etalon_bytes = &[0x01, 0x04, 0x04, 0x05, 0xdc, 0x00];
        let etalon_struct = Options::try_from(etalon_bytes).unwrap();
        codec_roundtrip(etalon_struct, etalon_bytes, etalon_bytes.len());
    }

    #[test]
    fn options_iterator() {
        let bytes = &[0x01, 0x04, 0x04, 0x05, 0xdc, 0x00];
        let mut options = Options::try_from(bytes).unwrap();

        let opt1 = options.next().unwrap().unwrap();
        assert!(matches!(opt1, Option::NoOperation(_)));

        let opt2 = options.next().unwrap().unwrap();
        assert!(matches!(opt2, Option::MaximumDatagramSize(mds) if mds.mds == 1500));

        let opt3 = options.next().unwrap().unwrap();
        assert!(matches!(opt3, Option::EndOfOptionsList(_)));

        assert!(options.next().is_none());
    }
}