internet 0.1.0

//! TCP Options encoding.
//!
//! As defined in [RFC 9293].
//!
//! [RFC 9293]: https://datatracker.ietf.org/doc/html/rfc9293

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

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

impl Default for Options {
    fn default() -> Self {
        Self {
            length: 0,
            data: [0u8; 40],
            position: 0,
        }
    }
}

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 res = self::Option::decode(&mut cursor, ());
        self.position += cursor.position();

        Some(res)
    }
}

/// TCP Option enum.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Option {
    /// End of Option List.
    EndOfOptionList(EndOfOptionList),
    /// No-Operation.
    NoOperation(NoOperation),
    /// Maximum Segment Size.
    MaximumSegmentSize(MaximumSegmentSize),
    /// Window Scale.
    WindowScale(WindowScale),
    /// SACK Permitted.
    SackPermitted(SackPermitted),
    /// Selective Acknowledgment.
    Sack(Sack),
    /// Timestamps.
    Timestamps(Timestamps),
    /// MD5 Signature.
    Md5Signature(Md5Signature),
}

impl Codec for Option {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        match self {
            Option::EndOfOptionList(opt) => opt.encode(writer, ()),
            Option::NoOperation(opt) => opt.encode(writer, ()),
            Option::MaximumSegmentSize(opt) => opt.encode(writer, ()),
            Option::WindowScale(opt) => opt.encode(writer, ()),
            Option::SackPermitted(opt) => opt.encode(writer, ()),
            Option::Sack(opt) => opt.encode(writer, ()),
            Option::Timestamps(opt) => opt.encode(writer, ()),
            Option::Md5Signature(opt) => opt.encode(writer, ()),
        }
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        let byte = reader.peek_u8()?;
        let kind = Kind::try_from(byte).map_err(|_| BufError::UnexpectedValue)?;
        match kind {
            Kind::EndOfOptionList => Ok(Option::EndOfOptionList(EndOfOptionList::decode(
                reader,
                (),
            )?)),
            Kind::NoOperation => Ok(Option::NoOperation(NoOperation::decode(reader, ())?)),
            Kind::MaximumSegmentSize => Ok(Option::MaximumSegmentSize(MaximumSegmentSize::decode(
                reader,
                (),
            )?)),
            Kind::WindowScale => Ok(Option::WindowScale(WindowScale::decode(reader, ())?)),
            Kind::SackPermitted => Ok(Option::SackPermitted(SackPermitted::decode(reader, ())?)),
            Kind::Sack => Ok(Option::Sack(Sack::decode(reader, ())?)),
            Kind::Timestamps => Ok(Option::Timestamps(Timestamps::decode(reader, ())?)),
            Kind::Md5Signature => Ok(Option::Md5Signature(Md5Signature::decode(reader, ())?)),
        }
    }
}

/// TCP 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,
    /// Maximum Segment Size.
    MaximumSegmentSize = 2,
    /// Window Scale.
    WindowScale = 3,
    /// SACK Permitted.
    SackPermitted = 4,
    /// Selective Acknowledgment.
    Sack = 5,
    /// Timestamps.
    Timestamps = 8,
    /// MD5 Signature.
    Md5Signature = 19,
}

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

    fn try_from(value: u8) -> Result<Self, Self::Error> {
        match value {
            0 => Ok(Kind::EndOfOptionList),
            1 => Ok(Kind::NoOperation),
            2 => Ok(Kind::MaximumSegmentSize),
            3 => Ok(Kind::WindowScale),
            4 => Ok(Kind::SackPermitted),
            5 => Ok(Kind::Sack),
            8 => Ok(Kind::Timestamps),
            19 => Ok(Kind::Md5Signature),
            _ => Err(()),
        }
    }
}

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 Option List.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct EndOfOptionList;

impl EndOfOptionList {
    /// Kind field.
    pub const KIND: Kind = Kind::EndOfOptionList;
}

impl Codec for EndOfOptionList {
    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.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct NoOperation;

impl NoOperation {
    /// Kind field.
    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)
    }
}

/// Maximum Segment Size.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct MaximumSegmentSize {
    /// MSS value.
    pub mss: u16,
}

impl MaximumSegmentSize {
    /// Kind field.
    pub const KIND: Kind = Kind::MaximumSegmentSize;
    /// Fixed length.
    pub const LEN: u8 = 4;
}

impl Codec for MaximumSegmentSize {
    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);
        }
        Ok(Self {
            mss: u16::decode(reader, ())?,
        })
    }
}

/// Window Scale.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct WindowScale {
    /// Shift count.
    pub shift_count: u8,
}

impl WindowScale {
    /// Kind field.
    pub const KIND: Kind = Kind::WindowScale;
    /// Fixed length.
    pub const LEN: u8 = 3;
}

impl Codec for WindowScale {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        self.shift_count.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);
        }
        Ok(Self {
            shift_count: u8::decode(reader, ())?,
        })
    }
}

/// SACK Permitted.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SackPermitted;

impl SackPermitted {
    /// Kind field.
    pub const KIND: Kind = Kind::SackPermitted;
    /// Fixed length.
    pub const LEN: u8 = 2;
}

impl Codec for SackPermitted {
    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);
        }
        if u8::decode(reader, ())? != Self::LEN {
            return Err(BufError::UnexpectedValue);
        }
        Ok(Self)
    }
}

/// SACK block (left edge, right edge).
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct SackBlock {
    /// Left edge of the block.
    pub left_edge: u32,
    /// Right edge of the block.
    pub right_edge: u32,
}

impl Codec for SackBlock {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        self.left_edge.encode(writer, ())?;
        self.right_edge.encode(writer, ())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        Ok(Self {
            left_edge: u32::decode(reader, ())?,
            right_edge: u32::decode(reader, ())?,
        })
    }
}

/// Selective Acknowledgment.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Sack {
    /// SACK blocks (1 to 4).
    pub blocks: Vec<SackBlock>,
}

impl Sack {
    /// Kind field.
    pub const KIND: Kind = Kind::Sack;
    /// Minimum length (kind + length).
    pub const MIN_LEN: u8 = 2;

    /// Calculates the encoded length.
    pub fn encoded_len(&self) -> u8 {
        (2 + self.blocks.len() * 8) as u8
    }
}

impl Codec for Sack {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        self.encoded_len().encode(writer, ())?;
        for block in &self.blocks {
            block.encode(writer, ())?;
        }
        Ok(())
    }

    fn decode<R: Buf>(reader: &mut Cursor<R>, _: ()) -> BufResult<Self> {
        if Kind::decode(reader, ())? != Self::KIND {
            return Err(BufError::UnexpectedValue);
        }
        let length = u8::decode(reader, ())?;
        if length < Self::MIN_LEN || (length - Self::MIN_LEN) % 8 != 0 {
            return Err(BufError::UnexpectedValue);
        }
        let block_count = ((length - Self::MIN_LEN) / 8) as usize;
        let mut blocks = Vec::with_capacity(block_count);
        for _ in 0..block_count {
            blocks.push(SackBlock::decode(reader, ())?);
        }
        Ok(Self { blocks })
    }
}

/// Timestamps.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Timestamps {
    /// Timestamp value.
    pub tsval: u32,
    /// Timestamp echo reply.
    pub tsecr: u32,
}

impl Timestamps {
    /// Kind field.
    pub const KIND: Kind = Kind::Timestamps;
    /// Fixed length.
    pub const LEN: u8 = 10;
}

impl Codec for Timestamps {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        self.tsval.encode(writer, ())?;
        self.tsecr.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);
        }
        Ok(Self {
            tsval: u32::decode(reader, ())?,
            tsecr: u32::decode(reader, ())?,
        })
    }
}

/// MD5 Signature.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Md5Signature {
    /// MD5 digest (always 16 bytes).
    pub digest: [u8; 16],
}

impl Md5Signature {
    /// Kind field.
    pub const KIND: Kind = Kind::Md5Signature;
    /// Fixed length.
    pub const LEN: u8 = 18;
}

impl Codec for Md5Signature {
    fn encode<W: BufMut>(&self, writer: &mut Cursor<W>, _: ()) -> BufResult<()> {
        Self::KIND.encode(writer, ())?;
        Self::LEN.encode(writer, ())?;
        self.digest.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);
        }
        Ok(Self {
            digest: reader.read_array::<16>()?,
        })
    }
}

#[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_option_list() {
        let etalon_bytes = &[0x00];
        let etalon_struct = Option::EndOfOptionList(EndOfOptionList);
        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 maximum_segment_size() {
        let etalon_bytes = &[0x02, 0x04, 0x05, 0xb4]; // MSS 1460
        let etalon_struct = Option::MaximumSegmentSize(MaximumSegmentSize { mss: 1460 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn window_scale() {
        let etalon_bytes = &[0x03, 0x03, 0x04]; // Shift count 4
        let etalon_struct = Option::WindowScale(WindowScale { shift_count: 4 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn sack_permitted() {
        let etalon_bytes = &[0x04, 0x02];
        let etalon_struct = Option::SackPermitted(SackPermitted);
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn sack() {
        let etalon_bytes = &[
            0x05, 0x0a, // Kind 5, Length 10
            0x00, 0x00, 0x00, 0x01, // Left edge 1
            0x00, 0x00, 0x00, 0x02, // Right edge 2
        ];
        let etalon_struct = Option::Sack(Sack {
            blocks: vec![SackBlock {
                left_edge: 1,
                right_edge: 2,
            }],
        });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn timestamps() {
        let etalon_bytes = &[
            0x08, 0x0a, // Kind 8, Length 10
            0x00, 0x00, 0x00, 0x01, // TSval 1
            0x00, 0x00, 0x00, 0x02, // TSecr 2
        ];
        let etalon_struct = Option::Timestamps(Timestamps { tsval: 1, tsecr: 2 });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn md5_signature() {
        let etalon_bytes = &[
            0x13, 0x12, // Kind 19, Length 18
            0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d,
            0x0e, 0x0f,
        ];
        let etalon_struct = Option::Md5Signature(Md5Signature {
            digest: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15],
        });
        codec_roundtrip(etalon_struct, etalon_bytes, ());
    }

    #[test]
    fn options_struct_roundtrip() {
        let etalon_bytes = &[0x02, 0x04, 0x05, 0xb4, 0x01, 0x00]; // MSS, NOP, EOL
        let etalon_struct = Options::try_from(etalon_bytes).unwrap();

        // Context is the length of the options field to read
        codec_roundtrip(etalon_struct, etalon_bytes, etalon_bytes.len());
    }

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

        let opt1 = options.next().unwrap().unwrap();
        assert!(matches!(opt1, Option::MaximumSegmentSize(mss) if mss.mss == 1460));

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

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

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