internet 0.0.4

//! TCP Options Mapping.

use crate::tcp::{
    EndOfOptionList, Kind, MaximumSegmentSize, Md5Signature, NoOperation, Option, Sack, SackBlock,
    SackPermitted, Timestamps, WindowScale,
};
use crate::{Buf, BufError, BufResult};

/// TCP options mapping.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct OptionsMapping<T: Buf> {
    buffer: T,
}

impl<T: Buf> OptionsMapping<T> {
    /// Creates a new options mapping.
    pub fn new(buffer: T) -> Self {
        Self { buffer }
    }

    /// Consumes the mapping and returns the underlying buffer.
    pub fn into_inner(self) -> T {
        self.buffer
    }

    /// Returns a reference to the underlying buffer.
    pub fn as_inner(&self) -> &T {
        &self.buffer
    }

    /// Returns the length of the options buffer.
    pub fn len(&self) -> usize {
        self.buffer.length()
    }

    /// Returns true if the options buffer is empty.
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }

    /// Iterates over all options in the buffer.
    pub fn iter(&self) -> Options<'_, T> {
        Options {
            mapping: self,
            offset: 0,
        }
    }
}

impl<T: Buf> OptionsMapping<T> {
    /// Reads the kind of the option at the given offset.
    pub fn read_kind_at(&self, offset: usize) -> BufResult<Kind> {
        if offset >= self.buffer.length() {
            return Err(BufError::OutOfBounds);
        }
        let byte = unsafe { self.buffer.get_u8_unchecked(offset) };
        Kind::try_from(byte).map_err(|_| BufError::UnexpectedValue)
    }

    /// Reads the length of the option at the given offset.
    pub fn read_len_at(&self, offset: usize) -> BufResult<u8> {
        if offset + 1 >= self.buffer.length() {
            return Err(BufError::OutOfBounds);
        }
        Ok(unsafe { self.buffer.get_u8_unchecked(offset + 1) })
    }

    /// Returns the size of the option at the given offset.
    pub fn option_size_at(&self, offset: usize) -> BufResult<usize> {
        if offset >= self.buffer.length() {
            return Err(BufError::OutOfBounds);
        }
        let byte = unsafe { self.buffer.get_u8_unchecked(offset) };
        let kind = Kind::try_from(byte).map_err(|_| BufError::UnexpectedValue)?;
        match kind {
            Kind::EndOfOptionList | Kind::NoOperation => Ok(1),
            _ => {
                if offset + 1 >= self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) } as usize;
                if len < 2 {
                    return Err(BufError::InvalidLength);
                }
                Ok(len)
            }
        }
    }

    /// Reads the option at the given offset.
    pub fn read_option_at(&self, offset: usize) -> BufResult<Option> {
        let kind = self.read_kind_at(offset)?;
        match kind {
            Kind::EndOfOptionList => Ok(Option::EndOfOptionList(EndOfOptionList)),
            Kind::NoOperation => Ok(Option::NoOperation(NoOperation)),
            Kind::MaximumSegmentSize => {
                if offset + 4 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != MaximumSegmentSize::LEN {
                    return Err(BufError::InvalidLength);
                }
                let mss = unsafe { self.buffer.get_u16_be_unchecked(offset + 2) };
                Ok(Option::MaximumSegmentSize(MaximumSegmentSize { mss }))
            }
            Kind::WindowScale => {
                if offset + 3 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != WindowScale::LEN {
                    return Err(BufError::InvalidLength);
                }
                let shift_count = unsafe { self.buffer.get_u8_unchecked(offset + 2) };
                Ok(Option::WindowScale(WindowScale { shift_count }))
            }
            Kind::SackPermitted => {
                if offset + 2 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != SackPermitted::LEN {
                    return Err(BufError::InvalidLength);
                }
                Ok(Option::SackPermitted(SackPermitted))
            }
            Kind::Sack => {
                if offset + 2 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) } as usize;
                if len < 2 || (len - 2) % 8 != 0 || offset + len > self.buffer.length() {
                    return Err(BufError::InvalidLength);
                }
                let block_count = (len - 2) / 8;
                let mut blocks = Vec::with_capacity(block_count);
                for i in 0..block_count {
                    let left = unsafe { self.buffer.get_u32_be_unchecked(offset + 2 + i * 8) };
                    let right = unsafe { self.buffer.get_u32_be_unchecked(offset + 6 + i * 8) };
                    blocks.push(SackBlock {
                        left_edge: left,
                        right_edge: right,
                    });
                }
                Ok(Option::Sack(Sack { blocks }))
            }
            Kind::Timestamps => {
                if offset + 10 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != Timestamps::LEN {
                    return Err(BufError::InvalidLength);
                }
                let tsval = unsafe { self.buffer.get_u32_be_unchecked(offset + 2) };
                let tsecr = unsafe { self.buffer.get_u32_be_unchecked(offset + 6) };
                Ok(Option::Timestamps(Timestamps { tsval, tsecr }))
            }
            Kind::Md5Signature => {
                if offset + 18 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != Md5Signature::LEN {
                    return Err(BufError::InvalidLength);
                }
                let mut digest = [0u8; 16];
                for i in 0..16 {
                    digest[i] = unsafe { self.buffer.get_u8_unchecked(offset + 2 + i) };
                }
                Ok(Option::Md5Signature(Md5Signature { digest }))
            }
        }
    }
}

/// Iterator over TCP options.
#[derive(Debug, Clone)]
pub struct Options<'a, T: Buf> {
    mapping: &'a OptionsMapping<T>,
    offset: usize,
}

impl<'a, T: Buf> Iterator for Options<'a, T> {
    type Item = BufResult<Option>;

    fn next(&mut self) -> std::option::Option<Self::Item> {
        if self.offset >= self.mapping.len() {
            return None;
        }

        let result = self.mapping.read_option_at(self.offset);

        match &result {
            Ok(Option::EndOfOptionList(_)) => {
                self.offset = self.mapping.len();
            }
            Ok(_) => {
                if let Ok(size) = self.mapping.option_size_at(self.offset) {
                    self.offset += size;
                } else {
                    self.offset = self.mapping.len();
                }
            }
            Err(_) => {
                self.offset = self.mapping.len();
            }
        }

        Some(result)
    }
}

#[cfg(test)]
mod tests {
    use super::{Option, OptionsMapping};
    use crate::tcp::Kind;

    #[test]
    fn empty() {
        let buffer: &[u8] = &[];
        let mapping = OptionsMapping::new(buffer);
        assert!(mapping.is_empty());
    }

    #[test]
    fn eool_terminates() {
        let buffer = &[
            Kind::NoOperation as u8,
            Kind::EndOfOptionList as u8,
            Kind::NoOperation as u8,
        ];
        let mapping = OptionsMapping::new(&buffer[..]);
        let options: Vec<_> = mapping.iter().filter_map(|r| r.ok()).collect();
        assert_eq!(options.len(), 2);
    }

    #[test]
    fn read_mss() {
        let buffer = &[Kind::MaximumSegmentSize as u8, 4, 0x05, 0xB4]; // 1460
        let mapping = OptionsMapping::new(&buffer[..]);
        if let Some(Ok(Option::MaximumSegmentSize(mss))) = mapping.iter().next() {
            assert_eq!(mss.mss, 1460);
        } else {
            panic!("expected MSS");
        }
    }

    #[test]
    fn read_window_scale() {
        let buffer = &[Kind::WindowScale as u8, 3, 7];
        let mapping = OptionsMapping::new(&buffer[..]);
        if let Some(Ok(Option::WindowScale(ws))) = mapping.iter().next() {
            assert_eq!(ws.shift_count, 7);
        } else {
            panic!("expected WS");
        }
    }

    #[test]
    fn read_timestamps() {
        let buffer = [
            Kind::Timestamps as u8,
            10,
            0x00,
            0x00,
            0x00,
            0x01, // tsval
            0x00,
            0x00,
            0x00,
            0x02, // tsecr
        ];
        let mapping = OptionsMapping::new(&buffer[..]);
        if let Some(Ok(Option::Timestamps(ts))) = mapping.iter().next() {
            assert_eq!(ts.tsval, 1);
            assert_eq!(ts.tsecr, 2);
        } else {
            panic!("expected timestamps");
        }
    }

    #[test]
    fn read_sack() {
        let buffer = [
            Kind::Sack as u8,
            10,
            0x00,
            0x00,
            0x00,
            0x01, // left
            0x00,
            0x00,
            0x00,
            0x02, // right
        ];
        let mapping = OptionsMapping::new(&buffer[..]);
        if let Some(Ok(Option::Sack(sack))) = mapping.iter().next() {
            assert_eq!(sack.blocks.len(), 1);
            assert_eq!(sack.blocks[0].left_edge, 1);
            assert_eq!(sack.blocks[0].right_edge, 2);
        } else {
            panic!("expected SACK");
        }
    }
}