internet 0.0.5

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

use crate::udp::{
    AdditionalPayloadChecksum, EchoRequest, EchoResponse, EndOfOptionsList, Fragmentation, Kind,
    MaximumDatagramSize, MaximumReassembledDatagramSize, NoOperation, Option, Timestamp,
};
use crate::{Buf, BufError, BufResult};

/// UDP 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 kind_byte = unsafe { self.buffer.get_u8_unchecked(offset) };
        Kind::try_from(kind_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) })
    }

    /// Reads the next option from the given offset.
    pub fn read_option_at(&self, offset: usize) -> BufResult<Option> {
        if offset >= self.buffer.length() {
            return Err(BufError::OutOfBounds);
        }

        let kind_byte = unsafe { self.buffer.get_u8_unchecked(offset) };
        let kind = Kind::try_from(kind_byte).map_err(|_| BufError::UnexpectedValue)?;

        match kind {
            Kind::EndOfOptionList => Ok(Option::EndOfOptionsList(EndOfOptionsList)),
            Kind::NoOperation => Ok(Option::NoOperation(NoOperation)),
            Kind::AdditionalPayloadChecksum => {
                if offset + 8 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != 8 {
                    return Err(BufError::UnexpectedValue);
                }
                let crc32c = unsafe { self.buffer.get_u32_be_unchecked(offset + 4) };
                Ok(Option::AdditionalPayloadChecksum(
                    AdditionalPayloadChecksum { crc32c },
                ))
            }
            Kind::Fragmentation => {
                if offset + 2 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                Ok(Option::Fragmentation(Fragmentation { len }))
            }
            Kind::MaximumDatagramSize => {
                if offset + 4 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != 4 {
                    return Err(BufError::UnexpectedValue);
                }
                let mds = unsafe { self.buffer.get_u16_be_unchecked(offset + 2) };
                Ok(Option::MaximumDatagramSize(MaximumDatagramSize { mds }))
            }
            Kind::MaximumReassembledDatagramSize => {
                if offset + 4 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != 4 {
                    return Err(BufError::UnexpectedValue);
                }
                let mrds = unsafe { self.buffer.get_u16_be_unchecked(offset + 2) };
                Ok(Option::MaximumReassembledDatagramSize(
                    MaximumReassembledDatagramSize { mrds },
                ))
            }
            Kind::Request => {
                if offset + 4 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != 4 {
                    return Err(BufError::UnexpectedValue);
                }
                let mss = unsafe { self.buffer.get_u16_be_unchecked(offset + 2) };
                Ok(Option::Request(EchoRequest { mss }))
            }
            Kind::Response => {
                if offset + 2 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                Ok(Option::Response(EchoResponse { len }))
            }
            Kind::Timestamp => {
                if offset + 12 > self.buffer.length() {
                    return Err(BufError::OutOfBounds);
                }
                let len = unsafe { self.buffer.get_u8_unchecked(offset + 1) };
                if len != 12 {
                    return Err(BufError::UnexpectedValue);
                }
                let ts_val = unsafe { self.buffer.get_u32_be_unchecked(offset + 2) };
                let ts_ecr = unsafe { self.buffer.get_u32_be_unchecked(offset + 6) };
                Ok(Option::Timestamp(Timestamp { ts_val, ts_ecr }))
            }
        }
    }

    /// 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 kind_byte = unsafe { self.buffer.get_u8_unchecked(offset) };
        let kind = Kind::try_from(kind_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;
                Ok(len)
            }
        }
    }
}

/// Iterator over UDP 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) -> core::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::EndOfOptionsList(_)) => {
                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::OptionsMapping;
    use crate::udp::Kind;
    use crate::udp::Option;

    #[test]
    fn empty_options() {
        let buffer: &[u8] = &[];
        let mapping = OptionsMapping::new(buffer);

        assert!(mapping.is_empty());
        assert_eq!(mapping.len(), 0);
        assert_eq!(mapping.iter().count(), 0);
    }

    #[test]
    fn single_nop() {
        let buffer = &[Kind::NoOperation as u8];
        let mapping = OptionsMapping::new(&buffer[..]);

        assert!(!mapping.is_empty());
        assert_eq!(mapping.len(), 1);

        let mut iter = mapping.iter();
        if let Some(Ok(Option::NoOperation(_))) = iter.next() {
        } else {
            panic!("Expected NoOperation");
        }
        assert!(iter.next().is_none());
    }

    #[test]
    fn eool_terminates() {
        let buffer = &[
            Kind::NoOperation as u8,
            Kind::EndOfOptionList as u8,
            Kind::NoOperation as u8, // Should be ignored
        ];
        let mapping = OptionsMapping::new(&buffer[..]);

        let options: Vec<_> = mapping.iter().filter_map(|r| r.ok()).collect();
        assert_eq!(options.len(), 2);
    }

    #[test]
    fn read_kind_at() {
        let buffer = &[Kind::NoOperation as u8, Kind::Timestamp as u8];
        let mapping = OptionsMapping::new(&buffer[..]);

        assert_eq!(mapping.read_kind_at(0).unwrap(), Kind::NoOperation);
        assert_eq!(mapping.read_kind_at(1).unwrap(), Kind::Timestamp);
        assert!(mapping.read_kind_at(2).is_err());
    }

    #[test]
    fn option_size() {
        let buffer = &[
            Kind::NoOperation as u8,
            Kind::Timestamp as u8,
            12, // len
            0,
            0,
            0,
            1, // ts_val
            0,
            0,
            0,
            2, // ts_ecr
        ];
        let mapping = OptionsMapping::new(&buffer[..]);

        assert_eq!(mapping.option_size_at(0).unwrap(), 1);
        assert_eq!(mapping.option_size_at(1).unwrap(), 12);
    }
}