// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    event::IntoEvent,
    packet::number::{
        derive_truncation_range, packet_number_space::PacketNumberSpace,
        truncated_packet_number::TruncatedPacketNumber,
    },
    varint::VarInt,
};
use core::{
    cmp::Ordering,
    fmt,
    hash::{Hash, Hasher},
    mem::size_of,
    num::NonZeroU64,
};

#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

const PACKET_SPACE_BITLEN: usize = 2;
const PACKET_SPACE_SHIFT: usize = (size_of::<PacketNumber>() * 8) - PACKET_SPACE_BITLEN;
const PACKET_NUMBER_MASK: u64 = core::u64::MAX >> PACKET_SPACE_BITLEN;

/// Contains a fully-decoded packet number in a given space
///
/// Internally the packet number is represented as a [`NonZeroU64`]
/// to ensure optimal memory layout.
///
/// The lower 62 bits are used to store the actual packet number value.
/// The upper 2 bits are used to store the packet number space. Because
/// there are only 3 spaces, the zero state is never used, which is why
/// [`NonZeroU64`] can be used instead of `u64`.
#[derive(Clone, Copy, Eq)]
#[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
pub struct PacketNumber(NonZeroU64);

impl IntoEvent<u64> for PacketNumber {
    #[inline]
    fn into_event(self) -> u64 {
        self.as_u64()
    }
}

impl Default for PacketNumber {
    fn default() -> Self {
        Self::from_varint(Default::default(), PacketNumberSpace::Initial)
    }
}

impl Hash for PacketNumber {
    #[inline]
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.hash(state)
    }
}

impl PartialEq for PacketNumber {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other) == Ordering::Equal
    }
}

impl PartialOrd for PacketNumber {
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for PacketNumber {
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        if cfg!(debug_assertions) {
            self.space().assert_eq(other.space());
        }
        self.0.cmp(&other.0)
    }
}

impl fmt::Debug for PacketNumber {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_tuple("PacketNumber")
            .field(&self.space())
            .field(&self.as_u64())
            .finish()
    }
}

impl fmt::Display for PacketNumber {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.as_u64().fmt(f)
    }
}

impl PacketNumber {
    /// Creates a PacketNumber for a given VarInt and PacketNumberSpace
    #[inline]
    pub(crate) const fn from_varint(value: VarInt, space: PacketNumberSpace) -> Self {
        let tag = space.as_tag() as u64;
        let pn = (tag << PACKET_SPACE_SHIFT) | value.as_u64();
        let pn = unsafe {
            // Safety: packet number space tag is never 0
            NonZeroU64::new_unchecked(pn)
        };
        Self(pn)
    }

    /// Returns the `PacketNumberSpace` for the given `PacketNumber`
    #[inline]
    pub fn space(self) -> PacketNumberSpace {
        let tag = self.0.get() >> PACKET_SPACE_SHIFT;
        PacketNumberSpace::from_tag(tag as u8)
    }

    /// Converts the `PacketNumber` into a `VarInt` value.
    ///
    /// Note: Even though some scenarios require this function, it should be
    /// avoided in most cases, as it removes the corresponding `PacketNumberSpace`
    /// and allows math operations to be performed, which can easily result in
    /// protocol errors.
    #[allow(clippy::wrong_self_convention)] // Don't use `self` here to make conversion explicit
    pub const fn as_varint(packet_number: Self) -> VarInt {
        // Safety: when converting to a u64, we remove the top 2 bits which
        //         will force the value to fit into a VarInt.
        unsafe { VarInt::new_unchecked(packet_number.as_u64()) }
    }

    /// Truncates the `PacketNumber` into a `TruncatedPacketNumber` based on
    /// the largest acknowledged packet number
    #[inline]
    pub fn truncate(
        self,
        largest_acknowledged_packet_number: Self,
    ) -> Option<TruncatedPacketNumber> {
        Some(
            derive_truncation_range(largest_acknowledged_packet_number, self)?
                .truncate_packet_number(Self::as_varint(self)),
        )
    }

    /// Compute the next packet number in the space. If the packet number has
    /// exceeded the maximum value allowed `None` will be returned.
    #[inline]
    pub fn next(self) -> Option<Self> {
        let value = Self::as_varint(self).checked_add(VarInt::from_u8(1))?;
        let space = self.space();
        Some(Self::from_varint(value, space))
    }

    /// Compute the prev packet number in the space. If the packet number has
    /// underflowed `None` will be returned.
    #[inline]
    pub fn prev(self) -> Option<Self> {
        let value = Self::as_varint(self).checked_sub(VarInt::from_u8(1))?;
        let space = self.space();
        Some(Self::from_varint(value, space))
    }

    /// Create a nonce for crypto from the packet number value
    ///
    /// Note: This should not be used by anything other than crypto-related
    /// functionality.
    #[inline]
    pub const fn as_crypto_nonce(self) -> u64 {
        self.as_u64()
    }

    /// Returns the value with the top 2 bits removed
    #[inline]
    pub const fn as_u64(self) -> u64 {
        self.0.get() & PACKET_NUMBER_MASK
    }

    /// Computes the distance between this packet number and the given packet number,
    /// returning None if overflow occurred.
    #[inline]
    pub fn checked_distance(self, rhs: PacketNumber) -> Option<u64> {
        self.space().assert_eq(rhs.space());
        Self::as_u64(self).checked_sub(Self::as_u64(rhs))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Make sure the assumptions around packet number space tags holds true
    #[test]
    fn packet_number_space_assumptions_test() {
        assert!(PacketNumberSpace::Initial.as_tag() != 0);
        assert!(PacketNumberSpace::Handshake.as_tag() != 0);
        assert!(PacketNumberSpace::ApplicationData.as_tag() != 0);
    }

    #[test]
    fn round_trip_test() {
        let spaces = [
            PacketNumberSpace::Initial,
            PacketNumberSpace::Handshake,
            PacketNumberSpace::ApplicationData,
        ];

        let values = [
            VarInt::from_u8(0),
            VarInt::from_u8(1),
            VarInt::from_u8(2),
            VarInt::from_u8(core::u8::MAX / 2),
            VarInt::from_u8(core::u8::MAX - 1),
            VarInt::from_u8(core::u8::MAX),
            VarInt::from_u16(core::u16::MAX / 2),
            VarInt::from_u16(core::u16::MAX - 1),
            VarInt::from_u16(core::u16::MAX),
            VarInt::from_u32(core::u32::MAX / 2),
            VarInt::from_u32(core::u32::MAX - 1),
            VarInt::from_u32(core::u32::MAX),
            VarInt::MAX,
        ];

        for space in spaces.iter().cloned() {
            for value in values.iter().cloned() {
                let pn = PacketNumber::from_varint(value, space);
                assert_eq!(pn.space(), space, "{:#064b}", pn.0);
                assert_eq!(PacketNumber::as_varint(pn), value, "{:#064b}", pn.0);
            }
        }
    }
    #[test]
    #[should_panic]
    fn wrong_packet_number_space() {
        PacketNumberSpace::ApplicationData
            .new_packet_number(VarInt::from_u8(0))
            .checked_distance(PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(0)));
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    packet::number::{
        packet_number_space::PacketNumberSpace, truncated_packet_number::TruncatedPacketNumber,
        PACKET_NUMBER_LEN_MASK,
    },
    varint::VarInt,
};
use s2n_codec::{u24, DecoderBuffer, DecoderBufferResult};

/// A fully-decoded and unprotected packet number length
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct PacketNumberLen {
    pub(crate) space: PacketNumberSpace,
    pub(crate) value: PacketNumberLenValue,
}

impl PacketNumberLen {
    pub const MAX_LEN: usize = U32_SIZE;

    /// Returns the max `PacketNumberLen` value for the given `PacketNumberSpace`
    pub const fn max(space: PacketNumberSpace) -> Self {
        Self {
            value: PacketNumberLenValue::U32,
            space,
        }
    }

    /// Returns the `PacketNumberSpace` for the given `PacketNumberLen`
    #[inline]
    pub const fn space(self) -> PacketNumberSpace {
        self.space
    }

    /// Decodes a `TruncatedPacketNumber` with the given `PacketNumberLen`
    #[inline]
    pub fn decode_truncated_packet_number(
        self,
        buffer: DecoderBuffer,
    ) -> DecoderBufferResult<TruncatedPacketNumber> {
        self.value
            .decode_truncated_packet_number(buffer, self.space)
    }

    /// Returns a packet tag mask for the given `PacketNumberLen`.
    #[inline]
    pub fn into_packet_tag_mask(self) -> u8 {
        self.value.into_packet_tag_mask()
    }

    /// Returns the bytesize required for encoding the given `PacketNumberLen`
    #[inline]
    pub fn bytesize(self) -> usize {
        self.value.bytesize()
    }

    /// Returns the bitsize required for encoding the given `PacketNumberLen`
    #[inline]
    pub fn bitsize(self) -> usize {
        self.value.bitsize()
    }

    #[inline]
    pub(crate) fn truncate_packet_number(self, value: VarInt) -> TruncatedPacketNumber {
        self.value.truncate_packet_number(value, self.space)
    }

    #[inline]
    pub(crate) fn from_packet_tag(tag: u8, space: PacketNumberSpace) -> Self {
        Self {
            value: PacketNumberLenValue::from_packet_tag(tag),
            space,
        }
    }

    #[inline]
    pub(crate) fn from_varint(value: VarInt, space: PacketNumberSpace) -> Option<Self> {
        Some(Self {
            value: PacketNumberLenValue::from_varint(value)?,
            space,
        })
    }
}

const U8_TAG: u8 = 0; // (8 / 8) - 1;
const U16_TAG: u8 = (16 / 8) - 1;
const U24_TAG: u8 = (24 / 8) - 1;
const U32_TAG: u8 = (32 / 8) - 1;

const U32_SIZE: usize = 32 / 8;

const U8_MAX: u64 = (1 << 8) - 1;
const U16_MAX: u64 = (1 << 16) - 1;
const U24_MAX: u64 = (1 << 24) - 1;
const U32_MAX: u64 = (1 << 32) - 1;

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub(crate) enum PacketNumberLenValue {
    U8,
    U16,
    U24,
    U32,
}

impl PacketNumberLenValue {
    #[inline]
    pub fn decode_truncated_packet_number(
        self,
        buffer: DecoderBuffer,
        space: PacketNumberSpace,
    ) -> DecoderBufferResult<TruncatedPacketNumber> {
        match self {
            Self::U8 => TruncatedPacketNumber::decode::<u8>(buffer, space),
            Self::U16 => TruncatedPacketNumber::decode::<u16>(buffer, space),
            Self::U24 => TruncatedPacketNumber::decode::<u24>(buffer, space),
            Self::U32 => TruncatedPacketNumber::decode::<u32>(buffer, space),
        }
    }

    #[inline]
    pub(crate) fn truncate_packet_number(
        self,
        value: VarInt,
        space: PacketNumberSpace,
    ) -> TruncatedPacketNumber {
        match self {
            Self::U8 => TruncatedPacketNumber::new(*value as u8, space),
            Self::U16 => TruncatedPacketNumber::new(*value as u16, space),
            Self::U24 => TruncatedPacketNumber::new(u24::new_truncated(*value as u32), space),
            Self::U32 => TruncatedPacketNumber::new(*value as u32, space),
        }
    }

    #[inline]
    pub fn into_packet_tag_mask(self) -> u8 {
        self as u8
    }

    #[inline]
    pub fn bytesize(self) -> usize {
        self as usize + 1
    }

    #[inline]
    pub fn bitsize(self) -> usize {
        self.bytesize() * 8
    }

    #[inline]
    pub fn from_packet_tag(tag: u8) -> Self {
        match tag & PACKET_NUMBER_LEN_MASK {
            U8_TAG => Self::U8,
            U16_TAG => Self::U16,
            U24_TAG => Self::U24,
            U32_TAG => Self::U32,
            _ => unreachable!("the mask only allows for 4 valid values"),
        }
    }

    #[inline]
    pub fn from_varint(value: VarInt) -> Option<Self> {
        #[allow(clippy::match_overlapping_arm)]
        match *value {
            0..=U8_MAX => Some(Self::U8),
            0..=U16_MAX => Some(Self::U16),
            0..=U24_MAX => Some(Self::U24),
            0..=U32_MAX => Some(Self::U32),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// the code relies on the variants to be in ascending order
    #[test]
    fn ordering_test() {
        assert!(PacketNumberLenValue::U8 < PacketNumberLenValue::U16);
        assert!(PacketNumberLenValue::U16 < PacketNumberLenValue::U24);
        assert!(PacketNumberLenValue::U24 < PacketNumberLenValue::U32);
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{event::IntoEvent, packet::number::PacketNumber};
use core::ops::RangeInclusive;

/// An inclusive range of `PacketNumber`s
#[derive(Clone, Copy, Eq, PartialEq)]
pub struct PacketNumberRange {
    start: PacketNumber,
    end: PacketNumber,
    exhausted: bool,
}

impl PacketNumberRange {
    /// Creates a new packet number range.
    #[inline]
    pub fn new(start: PacketNumber, end: PacketNumber) -> Self {
        assert!(start <= end, "start must be less than or equal to end");
        Self {
            start,
            end,
            exhausted: false,
        }
    }

    /// Returns true if the range contains the given packet number
    #[inline]
    pub fn contains(&self, packet_number: PacketNumber) -> bool {
        self.start <= packet_number && packet_number <= self.end
    }

    /// Returns the lower bound of the range (inclusive).
    #[inline]
    pub fn start(&self) -> PacketNumber {
        self.start
    }

    /// Returns the upper bound of the range (inclusive).
    #[inline]
    pub fn end(&self) -> PacketNumber {
        self.end
    }
}

impl IntoEvent<RangeInclusive<u64>> for PacketNumberRange {
    #[inline]
    fn into_event(self) -> RangeInclusive<u64> {
        self.start().into_event()..=self.end().into_event()
    }
}

impl Iterator for PacketNumberRange {
    type Item = PacketNumber;

    #[inline]
    fn next(&mut self) -> Option<Self::Item> {
        if !self.exhausted && self.start <= self.end {
            let current = self.start;
            if let Some(next) = current.next() {
                self.start = next;
            } else {
                // PacketNumber range has been exceeded
                self.exhausted = true;
            }
            Some(current)
        } else {
            self.exhausted = true;
            None
        }
    }
}

impl DoubleEndedIterator for PacketNumberRange {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        if !self.exhausted && self.start <= self.end {
            let current = self.end;
            if let Some(prev) = current.prev() {
                self.end = prev;
                self.exhausted = self.start > self.end;
            } else {
                // PacketNumber range has been exceeded
                self.exhausted = true;
            }
            Some(current)
        } else {
            self.exhausted = true;
            None
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        packet::number::{PacketNumberRange, PacketNumberSpace},
        varint::VarInt,
    };

    #[test]
    fn iterator() {
        let mut counter = 1;
        let start = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(1));
        let end = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(10));

        let range = PacketNumberRange::new(start, end);
        assert_eq!(start, range.start());
        assert_eq!(end, range.end());

        for packet_number in range {
            assert_eq!(counter, packet_number.as_u64());
            counter += 1;
        }

        assert_eq!(counter, 11);
    }

    #[test]
    fn double_ended_iterator() {
        let mut counter = 10;
        let start = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(1));
        let end = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(10));

        let mut range = PacketNumberRange::new(start, end);
        assert_eq!(start, range.start());
        assert_eq!(end, range.end());

        while let Some(packet_number) = range.next_back() {
            assert_eq!(counter, packet_number.as_u64());
            counter -= 1;
        }

        assert_eq!(counter, 0);
    }

    #[test]
    fn double_ended_iterator_zero() {
        let mut items = vec![];
        let start = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(0));
        let end = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(9));

        let mut range = PacketNumberRange::new(start, end);
        assert_eq!(start, range.start());
        assert_eq!(end, range.end());

        while let Some(packet_number) = range.next_back() {
            items.push(packet_number.as_u64());
        }

        items.reverse();

        for (idx, value) in items.iter().enumerate() {
            assert_eq!(idx as u64, *value);
        }
    }

    #[test]
    fn start_equals_end() {
        let start = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(1));
        let end = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(1));

        let range = PacketNumberRange::new(start, end);

        assert_eq!(1, range.count());
        assert_eq!(start, range.last().unwrap());
    }

    #[test]
    #[should_panic]
    fn start_greater_than_end() {
        let start = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(1));
        let end = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(10));
        PacketNumberRange::new(end, start);
    }

    #[test]
    fn end_is_max_packet_number() {
        let start = PacketNumberSpace::Handshake
            .new_packet_number(VarInt::new((u64::max_value() >> 2) - 1).unwrap());
        let end = PacketNumberSpace::Handshake
            .new_packet_number(VarInt::new(u64::max_value() >> 2).unwrap());

        assert_eq!(2, PacketNumberRange::new(start, end).count());
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    packet::number::{packet_number::PacketNumber, packet_number_len::PacketNumberLen},
    varint::VarInt,
};

#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

/// Contains all of the available packet spaces for QUIC packets
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
#[repr(u8)]
pub enum PacketNumberSpace {
    // This MUST start with 1 to enable optimized memory layout
    #[default]
    Initial = 1,
    Handshake = 2,
    ApplicationData = 3,
}

impl PacketNumberSpace {
    /// Returns `true` if the `PacketNumberSpace` is set to `Initial`
    #[inline]
    pub fn is_initial(self) -> bool {
        matches!(self, Self::Initial)
    }

    /// Returns `true` if the `PacketNumberSpace` is set to `Handshake`
    #[inline]
    pub fn is_handshake(self) -> bool {
        matches!(self, Self::Handshake)
    }

    /// Returns `true` if the `PacketNumberSpace` is set to `ApplicationData`
    #[inline]
    pub fn is_application_data(self) -> bool {
        matches!(self, Self::ApplicationData)
    }

    /// Create a new `PacketNumber` for the given `PacketNumberSpace`
    #[inline]
    pub const fn new_packet_number(self, value: VarInt) -> PacketNumber {
        PacketNumber::from_varint(value, self)
    }

    /// Create a new `PacketNumberLen` for the given `PacketNumberSpace` with a packet `tag`
    #[inline]
    pub fn new_packet_number_len(self, tag: u8) -> PacketNumberLen {
        PacketNumberLen::from_packet_tag(tag, self)
    }

    /// Asserts the `PacketNumberSpace` is equal
    #[inline(always)]
    pub(crate) fn assert_eq(self, other: Self) {
        debug_assert_eq!(
            self, other,
            "PacketNumbers cannot be compared across packet spaces"
        );
    }

    /// Returns the tag representation of PacketNumberSpace
    #[inline]
    pub(crate) const fn as_tag(self) -> u8 {
        self as u8
    }

    /// Creates a PacketNumberSpace from a tag representation
    ///
    /// # Safety
    ///
    /// Callers must ensure tag is less than `3`
    #[inline]
    pub(crate) fn from_tag(tag: u8) -> Self {
        match tag {
            1 => Self::Initial,
            2 => Self::Handshake,
            3 => Self::ApplicationData,
            _ if cfg!(debug_assertions) => panic!("invalid tag for PacketNumberSpace"),
            _ => Self::Initial,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn round_trip_test() {
        let spaces = [
            PacketNumberSpace::Initial,
            PacketNumberSpace::Handshake,
            PacketNumberSpace::ApplicationData,
        ];

        for space in spaces.iter().cloned() {
            assert_eq!(PacketNumberSpace::from_tag(space.as_tag()), space);
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use core::fmt;

/// A placeholder for `PacketNumber` values.
///
/// This is used to preserve the size of partially decoded
/// packets, before packet protection is removed, and fully decoded
/// packets, after packet protection is removed.
#[derive(Clone, Copy, Default, PartialEq, Eq, PartialOrd, Ord)]
pub struct ProtectedPacketNumber;

impl fmt::Debug for ProtectedPacketNumber {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "ProtectedPacketNumber")
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::packet::number::PacketNumber;
use core::mem;

#[derive(Default, Debug)]
pub struct SlidingWindow {
    /// Bitfield representing each packet number less than
    /// the right edge up to the window width.
    window: Window,
    /// The highest packet number seen so far, which is the
    /// right edge of the window.
    right_edge: Option<PacketNumber>,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum SlidingWindowError {
    Duplicate,
    TooOld,
}

/// 128-bit wide window allowing for 128 packets, plus the highest
/// packet representing the right edge to be tracked.
type Window = u128;

/// The total width of the window = the size of the 128-bit bitfield + 1 more bit
/// representing the right edge, which is always set.
const WINDOW_WIDTH: u64 = 1 + mem::size_of::<Window>() as u64 * 8;

enum WindowPosition {
    /// Left of the window, assumed to be a duplicate.
    Left,
    /// Right of the window, the value is the offset from the right edge.
    Right(u64),
    /// Equal to the highest value tracked by the window.
    RightEdge,
    /// Within the window, the value is the offset from the right edge.
    Within(u64),
    /// The window is empty.
    Empty,
}

//= https://www.rfc-editor.org/rfc/rfc4303#section-3.4.3
//# Duplicates are rejected through the use of a sliding receive window.
//# How the window is implemented is a local matter, but the following
//# text describes the functionality that the implementation must
//# exhibit.
//#
//# The "right" edge of the window represents the highest, validated
//# Sequence Number value received on this SA.  Packets that contain
//# sequence numbers lower than the "left" edge of the window are
//# rejected.  Packets falling within the window are checked against a
//# list of received packets within the window.
impl SlidingWindow {
    /// Inserts the `packet_number` into the sliding window, returning
    /// a SlidingWindowError::Duplicate if the `packet_number` has already
    /// been inserted into the sliding window or a SlidingWindowError::TooOld
    /// if the `packet_number` is beyond the capacity of the sliding window and
    /// thus cannot be determined if it is a duplicate.
    pub fn insert(&mut self, packet_number: PacketNumber) -> Result<(), SlidingWindowError> {
        match self.window_position(packet_number) {
            WindowPosition::Left => Err(SlidingWindowError::TooOld),
            WindowPosition::RightEdge => Err(SlidingWindowError::Duplicate),
            WindowPosition::Right(delta) => {
                if delta < WINDOW_WIDTH {
                    // Make room for the current right edge
                    self.window <<= 1;
                    // Set the bit for the current right edge
                    self.window |= 1;
                    // Shift by delta while taking account of the right edge
                    self.window <<= delta - 1;
                } else {
                    // The delta is too large, reset the window
                    self.window = Window::default();
                }
                self.right_edge = Some(packet_number);
                Ok(())
            }
            WindowPosition::Within(delta) => {
                let mask = 1 << (delta - 1); // Shift by the delta - 1 to account for the right edge
                let duplicate = self.window & mask != 0;
                self.window |= mask;
                if duplicate {
                    Err(SlidingWindowError::Duplicate)
                } else {
                    Ok(())
                }
            }
            WindowPosition::Empty => {
                self.right_edge = Some(packet_number);
                Ok(())
            }
        }
    }

    /// Determines if the given packet number has already been inserted or
    /// is too old to determine if it has already been inserted.
    pub fn check(&self, packet_number: PacketNumber) -> Result<(), SlidingWindowError> {
        match self.window_position(packet_number) {
            WindowPosition::Left => Err(SlidingWindowError::TooOld),
            WindowPosition::RightEdge => Err(SlidingWindowError::Duplicate),
            WindowPosition::Right(_) | WindowPosition::Empty => Ok(()),
            WindowPosition::Within(delta) => {
                let mask = 1 << (delta - 1); // Shift by the delta - 1 to account for the right edge
                if self.window & mask != 0 {
                    Err(SlidingWindowError::Duplicate)
                } else {
                    Ok(())
                }
            }
        }
    }

    /// Gets the position of the `packet_number` relative to the window.
    fn window_position(&self, packet_number: PacketNumber) -> WindowPosition {
        if let Some(right_edge) = self.right_edge {
            match right_edge.checked_distance(packet_number) {
                Some(0) => WindowPosition::RightEdge,
                Some(delta) if delta >= WINDOW_WIDTH => WindowPosition::Left,
                Some(delta) => WindowPosition::Within(delta),
                None => WindowPosition::Right(
                    packet_number
                        .checked_distance(right_edge)
                        .expect("packet_number must be greater than right_edge"),
                ),
            }
        } else {
            WindowPosition::Empty
        }
    }
}

#[cfg(test)]
mod test {

    use super::*;
    use crate::{packet::number::PacketNumberSpace, varint::VarInt};
    use bolero::{check, generator::*};
    use SlidingWindowError::*;

    /// This macro asserts that the output of inserting the given packet number,
    /// the window, and the right edge match the expected values.
    macro_rules! assert_window {
        (
            $window:expr, $to_insert:expr, $duplicate:expr, $expected_window:expr, $right_edge:expr
        ) => {{
            assert_eq!($window.check($to_insert), $duplicate);
            assert_eq!($window.insert($to_insert), $duplicate);
            assert_eq!(
                $window.window, $expected_window,
                "Expected: {:b}, Actual: {:b}",
                $expected_window, $window.window
            );
            assert_eq!($window.right_edge.unwrap(), $right_edge);
        }};
    }

    #[test]
    #[allow(clippy::cognitive_complexity)] // several operations are needed to get the desired state
    fn insert() {
        let space = PacketNumberSpace::ApplicationData;
        let mut window = SlidingWindow::default();

        let zero = space.new_packet_number(VarInt::from_u8(0));
        let one = space.new_packet_number(VarInt::from_u8(1));
        let two = space.new_packet_number(VarInt::from_u8(2));
        let three = space.new_packet_number(VarInt::from_u8(3));
        let four = space.new_packet_number(VarInt::from_u8(4));
        let five = space.new_packet_number(VarInt::from_u8(5));
        let six = space.new_packet_number(VarInt::from_u8(6));
        let seven = space.new_packet_number(VarInt::from_u8(7));
        let eight = space.new_packet_number(VarInt::from_u8(8));
        let large = space.new_packet_number(VarInt::MAX);

        assert_eq!(window.window, Window::default());
        assert_eq!(window.right_edge, None);

        assert_window!(window, zero, Ok(()), Window::default(), zero);
        assert_window!(window, zero, Err(Duplicate), Window::default(), zero);
        assert_window!(window, one, Ok(()), 0b1, one);
        assert_window!(window, one, Err(Duplicate), 0b1, one);
        assert_window!(window, two, Ok(()), 0b11, two);
        assert_window!(window, five, Ok(()), 0b11100, five);
        assert_window!(window, eight, Ok(()), 0b1110_0100, eight);
        assert_window!(window, seven, Ok(()), 0b1110_0101, eight);
        assert_window!(window, three, Ok(()), 0b1111_0101, eight);
        assert_window!(window, six, Ok(()), 0b1111_0111, eight);
        assert_window!(window, four, Ok(()), 0b1111_1111, eight);
        assert_window!(window, seven, Err(Duplicate), 0b1111_1111, eight);
        assert_window!(window, two, Err(Duplicate), 0b1111_1111, eight);
        assert_window!(window, eight, Err(Duplicate), 0b1111_1111, eight);
        assert_window!(window, large, Ok(()), Window::default(), large);
        assert_window!(window, five, Err(TooOld), Window::default(), large);
    }

    #[test]
    #[cfg_attr(miri, ignore)] // this test is too expensive for miri
    fn incremental_insert() {
        let mut window = SlidingWindow::default();
        let space = PacketNumberSpace::ApplicationData;
        for right_edge in 0..1000 {
            let pn = space.new_packet_number(VarInt::from_u32(right_edge));
            assert_eq!(window.check(pn), Ok(()));
            assert_eq!(window.insert(pn), Ok(()));
            assert_eq!(window.right_edge.unwrap(), pn);
            for dup in 0..=right_edge {
                let expected_error = if right_edge - dup < WINDOW_WIDTH as u32 {
                    Err(Duplicate)
                } else {
                    Err(TooOld)
                };
                let dup_pn = space.new_packet_number(VarInt::from_u32(dup));
                assert_eq!(window.check(dup_pn), expected_error);
                assert_eq!(window.insert(dup_pn), expected_error);
            }
        }
    }

    #[test]
    #[allow(clippy::cognitive_complexity)] // several comparisons are needed
    fn insert_at_edge() {
        let mut window = SlidingWindow::default();
        let space = PacketNumberSpace::ApplicationData;
        let zero = space.new_packet_number(VarInt::from_u8(0));
        let window_width_minus_1 = space.new_packet_number(VarInt::new(WINDOW_WIDTH - 1).unwrap());
        let window_width = window_width_minus_1.next().unwrap();

        assert_window!(window, zero, Ok(()), Window::default(), zero);
        assert_window!(
            window,
            window_width_minus_1,
            Ok(()),
            (1_u128) << 127,
            window_width_minus_1
        );
        assert_window!(
            window,
            window_width_minus_1,
            Err(Duplicate),
            (1_u128) << 127,
            window_width_minus_1
        );
        assert_window!(window, window_width, Ok(()), 0b1, window_width);

        window = SlidingWindow::default();
        assert_window!(window, zero, Ok(()), Window::default(), zero);
        assert_window!(
            window,
            window_width,
            Ok(()),
            Window::default(),
            window_width
        );
        assert_window!(
            window,
            window_width,
            Err(Duplicate),
            Window::default(),
            window_width
        );
    }

    #[test]
    fn delta_larger_than_32_bits() {
        let mut window = SlidingWindow::default();
        let space = PacketNumberSpace::ApplicationData;
        let zero = space.new_packet_number(VarInt::from_u8(0));
        let large = space.new_packet_number(VarInt::new((1 << 32) + 1).unwrap());
        assert_eq!(window.check(zero), Ok(()));
        assert_eq!(window.insert(zero), Ok(()));
        assert_eq!(window.check(large), Ok(()));
        assert_eq!(window.insert(large), Ok(()));
        assert_eq!(window.check(large), Err(Duplicate));
        assert_eq!(window.insert(large), Err(Duplicate));
        assert_eq!(window.window, 0b0);
    }

    #[test]
    #[cfg_attr(kani, kani::proof, kani::unwind(2), kani::solver(kissat))]
    #[cfg_attr(miri, ignore)] // this test is too expensive for miri
    fn insert_test() {
        // Make sure the two packet numbers are not the same
        let gen = gen::<(VarInt, VarInt)>().filter_gen(|(a, b)| a != b);

        check!()
            .with_generator(gen)
            .cloned()
            .for_each(|(pn, other_pn)| {
                let mut window = SlidingWindow::default();
                let space = PacketNumberSpace::ApplicationData;
                let packet_number = space.new_packet_number(pn);
                let other_packet_number = space.new_packet_number(other_pn);
                assert!(window.insert(packet_number).is_ok());
                assert_eq!(Err(Duplicate), window.check(packet_number));
                assert_ne!(Err(Duplicate), window.check(other_packet_number));
            });
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::*;
use crate::{
    packet::number::{PacketNumber, PacketNumberSpace},
    varint::VarInt,
};
use bolero::{check, generator::*};
use s2n_codec::{testing::encode, DecoderBuffer};

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(cadical))]
fn round_trip() {
    check!()
        .with_generator(
            gen_packet_number_space()
                .and_then(|space| (gen_packet_number(space), gen_packet_number(space))),
        )
        .cloned()
        .for_each(|(packet_number, largest_acked_packet_number)| {
            // Try to encode the packet number to send
            if let Some((mask, bytes)) =
                encode_packet_number(packet_number, largest_acked_packet_number)
            {
                // If encoding was valid, assert that the information can be decoded
                let actual_packet_number =
                    decode_packet_number(mask, bytes, largest_acked_packet_number).unwrap();
                assert_eq!(actual_packet_number, packet_number);
            }
        });
}

fn gen_packet_number_space() -> impl ValueGenerator<Output = PacketNumberSpace> {
    (0u8..=2).map_gen(|id| match id {
        0 => PacketNumberSpace::Initial,
        1 => PacketNumberSpace::Handshake,
        2 => PacketNumberSpace::ApplicationData,
        _ => unreachable!("invalid space id {:?}", id),
    })
}

fn gen_packet_number(space: PacketNumberSpace) -> impl ValueGenerator<Output = PacketNumber> {
    gen().map(move |packet_number| {
        space.new_packet_number(match VarInt::new(packet_number) {
            Ok(packet_number) => packet_number,
            Err(_) => VarInt::from_u32(packet_number as u32),
        })
    })
}

fn encode_packet_number(
    packet_number: PacketNumber,
    largest_acked_packet_number: PacketNumber,
) -> Option<(u8, Vec<u8>)> {
    let truncated_packet_number = packet_number.truncate(largest_acked_packet_number)?;

    let bytes = encode(&truncated_packet_number).unwrap();
    let mask = truncated_packet_number.len().into_packet_tag_mask();

    Some((mask, bytes))
}

fn decode_packet_number(
    packet_tag: u8,
    packet_bytes: Vec<u8>,
    largest_acked_packet_number: PacketNumber,
) -> Result<PacketNumber, String> {
    // decode the packet number len from the packet tag
    let packet_number_len = largest_acked_packet_number
        .space()
        .new_packet_number_len(packet_tag);

    // make sure the packet_tag has the same mask as the len
    assert_eq!(packet_number_len.into_packet_tag_mask(), packet_tag);
    assert_eq!(packet_number_len.bytesize(), packet_bytes.len());

    // try decoding the truncated packet number from the packet bytes
    let (truncated_packet_number, _) = packet_number_len
        .decode_truncated_packet_number(DecoderBuffer::new(&packet_bytes))
        .map_err(|err| err.to_string())?;

    // make sure the packet_number_len round trips
    assert_eq!(truncated_packet_number.len(), packet_number_len);

    // make sure the encoding matches the original bytes
    assert_eq!(packet_bytes, encode(&truncated_packet_number).unwrap());

    // try expanding the truncated packet number
    let packet_number = truncated_packet_number.expand(largest_acked_packet_number);

    // try truncating the packet number
    let actual_truncated_packet_number = packet_number
        .truncate(largest_acked_packet_number)
        .ok_or_else(|| "Could not truncate packet number".to_string())?;

    assert_eq!(actual_truncated_packet_number, truncated_packet_number);

    Ok(packet_number)
}

fn new(value: VarInt) -> PacketNumber {
    PacketNumberSpace::Initial.new_packet_number(value)
}

/// This implementation tries to closely follow the RFC pseudo code so it's
/// easier to ensure it matches.
#[allow(clippy::blocks_in_if_conditions)]
fn rfc_decoder(largest_pn: u64, truncated_pn: u64, pn_nbits: usize) -> u64 {
    macro_rules! catch {
        ($expr:expr) => {
            (|| Some($expr))().unwrap_or(false)
        };
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.3
    //= type=test
    //# DecodePacketNumber(largest_pn, truncated_pn, pn_nbits):
    //#   expected_pn  = largest_pn + 1
    //#   pn_win       = 1 << pn_nbits
    //#   pn_hwin      = pn_win / 2
    //#   pn_mask      = pn_win - 1
    //#   // The incoming packet number should be greater than
    //#   // expected_pn - pn_hwin and less than or equal to
    //#   // expected_pn + pn_hwin
    //#   //
    //#   // This means we cannot just strip the trailing bits from
    //#   // expected_pn and add the truncated_pn because that might
    //#   // yield a value outside the window.
    //#   //
    //#   // The following code calculates a candidate value and
    //#   // makes sure it's within the packet number window.
    //#   // Note the extra checks to prevent overflow and underflow.
    //#   candidate_pn = (expected_pn & ~pn_mask) | truncated_pn
    //#   if candidate_pn <= expected_pn - pn_hwin and
    //#     candidate_pn < (1 << 62) - pn_win:
    //#     return candidate_pn + pn_win
    //#   if candidate_pn > expected_pn + pn_hwin and
    //#     candidate_pn >= pn_win:
    //#     return candidate_pn - pn_win
    //#   return candidate_pn
    let expected_pn = largest_pn + 1;
    let pn_win = 1 << pn_nbits;
    let pn_hwin = pn_win / 2;
    let pn_mask = pn_win - 1;

    let candidate_pn = (expected_pn & !pn_mask) | truncated_pn;
    if catch!(
        candidate_pn <= expected_pn.checked_sub(pn_hwin)?
            && candidate_pn < (1u64 << 62).checked_sub(pn_win)?
    ) {
        return candidate_pn + pn_win;
    }

    if catch!(candidate_pn > expected_pn.checked_add(pn_hwin)? && candidate_pn >= pn_win) {
        return candidate_pn - pn_win;
    }

    candidate_pn
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
fn truncate_expand_test() {
    check!()
        .with_type()
        .cloned()
        .for_each(|(largest_pn, expected_pn)| {
            let largest_pn = new(largest_pn);
            let expected_pn = new(expected_pn);
            if let Some(truncated_pn) = expected_pn.truncate(largest_pn) {
                assert_eq!(expected_pn, truncated_pn.expand(largest_pn));
            }
        });
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
fn rfc_differential_test() {
    check!()
        .with_type()
        .cloned()
        .for_each(|(largest_pn, truncated_pn)| {
            let largest_pn = new(largest_pn);
            let space = largest_pn.space();
            let truncated_pn = TruncatedPacketNumber {
                space,
                value: truncated_pn,
            };
            let rfc_value = rfc_decoder(
                largest_pn.as_u64(),
                truncated_pn.into_u64(),
                truncated_pn.bitsize(),
            )
            .min(VarInt::MAX.as_u64());
            let actual_value = truncated_pn.expand(largest_pn).as_u64();

            assert_eq!(
                actual_value,
                rfc_value,
                "diff: {:?}",
                actual_value
                    .checked_sub(rfc_value)
                    .or_else(|| rfc_value.checked_sub(actual_value))
            );
        });
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
fn example_test() {
    macro_rules! example {
        ($largest:expr, $truncated:expr, $expected:expr) => {{
            let largest = new(VarInt::from_u32($largest));
            let truncated = TruncatedPacketNumber::new($truncated, PacketNumberSpace::Initial);
            let expected = new(VarInt::from_u32($expected));
            assert_eq!(truncated.expand(largest), expected);
        }};
    }

    example!(0xa82e1b31, 0x9b32u16, 0xa82e9b32);
}

#[test]
#[cfg_attr(miri, ignore)] // snapshot tests don't work on miri
fn size_of_snapshots() {
    use core::mem::size_of;
    use insta::assert_debug_snapshot;

    assert_debug_snapshot!("PacketNumber", size_of::<PacketNumber>());
    assert_debug_snapshot!("PacketNumberLen", size_of::<PacketNumberLen>());
    assert_debug_snapshot!("PacketNumberSpace", size_of::<PacketNumberSpace>());
    assert_debug_snapshot!("ProtectedPacketNumber", size_of::<ProtectedPacketNumber>());
    assert_debug_snapshot!("TruncatedPacketNumber", size_of::<TruncatedPacketNumber>());
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::packet::number::{
    decode_packet_number,
    packet_number::PacketNumber,
    packet_number_len::{PacketNumberLen, PacketNumberLenValue},
    packet_number_space::PacketNumberSpace,
};
use s2n_codec::{u24, DecoderBuffer, DecoderBufferResult, DecoderValue, Encoder, EncoderValue};

#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

/// A truncated packet number, which is derived from the largest acknowledged packet number
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
pub struct TruncatedPacketNumber {
    pub(crate) space: PacketNumberSpace,
    pub(crate) value: TruncatedPacketNumberValue,
}

#[allow(clippy::len_without_is_empty)] // Clippy gets confused by the const on is_empty
impl TruncatedPacketNumber {
    /// Returns the space for the given `TruncatedPacketNumber`
    #[inline]
    pub const fn space(self) -> PacketNumberSpace {
        self.space
    }

    /// Expands the `TruncatedPacketNumber` into a `PacketNumber`
    #[inline]
    pub fn expand(self, largest_acknowledged_packet_number: PacketNumber) -> PacketNumber {
        decode_packet_number(largest_acknowledged_packet_number, self)
    }

    #[inline]
    pub fn len(&self) -> PacketNumberLen {
        self.value.len(self.space)
    }

    #[inline]
    pub const fn is_empty(self) -> bool {
        false
    }

    /// Internal function to create a `TruncatedPacketNumber`.
    ///
    /// `TruncatedPacketNumber` values should be created either
    /// by using a `PacketNumberLen` to decode a buffer, or truncating
    /// a `PacketNumber`.
    #[inline]
    pub(crate) fn new<Value: Into<TruncatedPacketNumberValue>>(
        value: Value,
        space: PacketNumberSpace,
    ) -> Self {
        Self {
            value: value.into(),
            space,
        }
    }

    /// Internal function to decode a `TruncatedPacketNumber` with a given size.
    #[inline]
    pub(crate) fn decode<'a, Value: Into<TruncatedPacketNumberValue> + DecoderValue<'a>>(
        buffer: DecoderBuffer<'a>,
        space: PacketNumberSpace,
    ) -> DecoderBufferResult<'a, Self> {
        let (value, buffer) = buffer.decode()?;
        let packet_number = Self::new::<Value>(value, space);
        Ok((packet_number, buffer))
    }

    /// Convert the `TruncatedPacketNumber` into `u64`
    #[inline]
    pub(crate) fn into_u64(self) -> u64 {
        self.value.into_u64()
    }

    /// Get the bitsize for the given `TruncatedPacketNumber`
    #[inline]
    pub(crate) fn bitsize(self) -> usize {
        self.len().bitsize()
    }
}

impl EncoderValue for TruncatedPacketNumber {
    #[inline]
    fn encode<E: Encoder>(&self, buffer: &mut E) {
        buffer.encode(&self.value)
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
pub enum TruncatedPacketNumberValue {
    U8(u8),
    U16(u16),
    U24(u24),
    U32(u32),
}

impl TruncatedPacketNumberValue {
    #[inline]
    fn into_u64(self) -> u64 {
        match self {
            Self::U8(value) => value.into(),
            Self::U16(value) => value.into(),
            Self::U24(value) => value.into(),
            Self::U32(value) => value.into(),
        }
    }

    #[inline]
    fn len(self, space: PacketNumberSpace) -> PacketNumberLen {
        let value = match self {
            Self::U8(_value) => PacketNumberLenValue::U8,
            Self::U16(_value) => PacketNumberLenValue::U16,
            Self::U24(_value) => PacketNumberLenValue::U24,
            Self::U32(_value) => PacketNumberLenValue::U32,
        };

        PacketNumberLen { space, value }
    }
}

impl EncoderValue for TruncatedPacketNumberValue {
    #[inline]
    fn encode<E: Encoder>(&self, buffer: &mut E) {
        match self {
            Self::U8(value) => buffer.encode(value),
            Self::U16(value) => buffer.encode(value),
            Self::U24(value) => buffer.encode(value),
            Self::U32(value) => buffer.encode(value),
        }
    }
}

impl From<u8> for TruncatedPacketNumberValue {
    #[inline]
    fn from(value: u8) -> Self {
        Self::U8(value)
    }
}

impl From<u16> for TruncatedPacketNumberValue {
    #[inline]
    fn from(value: u16) -> Self {
        Self::U16(value)
    }
}

impl From<u24> for TruncatedPacketNumberValue {
    #[inline]
    fn from(value: u24) -> Self {
        Self::U24(value)
    }
}

impl From<u32> for TruncatedPacketNumberValue {
    #[inline]
    fn from(value: u32) -> Self {
        Self::U32(value)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    event,
    event::{
        api::{Path, SocketAddress},
        IntoEvent,
    },
    inet,
};

#[derive(Debug)]
#[non_exhaustive]
pub struct Attempt<'a> {
    /// The path that the connection is currently actively using
    pub active_path: Path<'a>,
    /// Information about the packet triggering the migration attempt
    pub packet: PacketInfo<'a>,
}

#[derive(Debug)]
pub struct AttemptBuilder<'a> {
    /// The path that the connection is currently actively using
    pub active_path: Path<'a>,
    /// Information about the packet triggering the migration attempt
    pub packet: PacketInfo<'a>,
}

impl<'a> From<AttemptBuilder<'a>> for Attempt<'a> {
    #[inline]
    fn from(builder: AttemptBuilder<'a>) -> Self {
        Self {
            active_path: builder.active_path,
            packet: builder.packet,
        }
    }
}

#[derive(Debug)]
#[non_exhaustive]
pub struct PacketInfo<'a> {
    pub remote_address: SocketAddress<'a>,
    pub local_address: SocketAddress<'a>,
}

#[derive(Debug)]
pub struct PacketInfoBuilder<'a> {
    pub remote_address: &'a inet::SocketAddress,
    pub local_address: &'a inet::SocketAddress,
}

impl<'a> From<PacketInfoBuilder<'a>> for PacketInfo<'a> {
    #[inline]
    fn from(builder: PacketInfoBuilder<'a>) -> Self {
        Self {
            remote_address: builder.remote_address.into_event(),
            local_address: builder.local_address.into_event(),
        }
    }
}

// TODO: Add an outcome that allows the connection to be closed/stateless reset https://github.com/aws/s2n-quic/issues/317

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum Outcome {
    /// Allows the path migration attempt to continue
    ///
    /// Note that path validation will still be attempted as described in
    /// [Section 8.2](https://datatracker.ietf.org/doc/html/rfc9000#section-8.2).
    Allow,

    /// Rejects a path migration attempt
    ///
    /// The connection will drop the packet that attempted to migrate and not reserve any state
    /// for the new path.
    Deny(DenyReason),
    // Additional outcomes must be handled in the path::Manager
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum DenyReason {
    // The new address uses a port that is blocked
    BlockedPort,
    // The new address uses a port in a different scope
    PortScopeChanged,
    // The new address uses an IP in a different scope
    IpScopeChanged,
    // All connection migrations are disabled
    ConnectionMigrationDisabled,
}

impl IntoEvent<event::builder::ConnectionMigrationDenied> for DenyReason {
    #[inline]
    fn into_event(self) -> event::builder::ConnectionMigrationDenied {
        let reason = match self {
            DenyReason::BlockedPort => event::builder::MigrationDenyReason::BlockedPort,
            DenyReason::PortScopeChanged => event::builder::MigrationDenyReason::PortScopeChanged,
            DenyReason::IpScopeChanged => event::builder::MigrationDenyReason::IpScopeChange,
            DenyReason::ConnectionMigrationDisabled => {
                event::builder::MigrationDenyReason::ConnectionMigrationDisabled
            }
        };
        event::builder::ConnectionMigrationDenied { reason }
    }
}

/// Validates a path migration attempt from an active path to another
pub trait Validator: 'static + Send {
    /// Called on each connection migration attempt for a connection
    fn on_migration_attempt(&mut self, attempt: &Attempt) -> Outcome;
}

pub mod default {
    use super::*;
    use crate::path::remote_port_blocked;

    #[derive(Debug, Default)]
    pub struct Validator;

    impl super::Validator for Validator {
        #[inline]
        fn on_migration_attempt(&mut self, attempt: &Attempt) -> Outcome {
            let active_addr = to_addr(&attempt.active_path.remote_addr);
            let packet_addr = to_addr(&attempt.packet.remote_address);

            // Block migrations to a port that is blocked
            if remote_port_blocked(packet_addr.port()) {
                return Outcome::Deny(DenyReason::BlockedPort);
            }

            //= https://www.rfc-editor.org/rfc/rfc9000#section-21.5.6
            //# it might be possible over time to identify
            //# specific UDP ports that are common targets of attacks or particular
            //# patterns in datagrams that are used for attacks.  Endpoints MAY
            //# choose to avoid sending datagrams to these ports or not send
            //# datagrams that match these patterns prior to validating the
            //# destination address.

            // NOTE: this may cause reachability issues if a peer or NAT use different
            //       port scopes for the same connection. Additional research may
            //       be required to determine if this countermeasure needs to be relaxed.
            if PortScope::new(active_addr.port()) != PortScope::new(packet_addr.port()) {
                return Outcome::Deny(DenyReason::PortScopeChanged);
            }

            //= https://www.rfc-editor.org/rfc/rfc9000#section-21.5.6
            //# Endpoints MAY prevent connection attempts or
            //# migration to a loopback address.  Endpoints SHOULD NOT allow
            //# connections or migration to a loopback address if the same service
            //# was previously available at a different interface or if the address
            //# was provided by a service at a non-loopback address.

            //= https://www.rfc-editor.org/rfc/rfc9000#section-21.5.6
            //# Similarly, endpoints could regard a change in address to a link-local
            //# address [RFC4291] or an address in a private-use range [RFC1918] from
            //# a global, unique-local [RFC4193], or non-private address as a
            //# potential attempt at request forgery.

            // Here, we ensure the ip scope match so peers are unable to change after
            // establishing a connection
            match (active_addr.unicast_scope(), packet_addr.unicast_scope()) {
                (Some(a), Some(b)) if a == b => Outcome::Allow,
                _ => Outcome::Deny(DenyReason::IpScopeChanged),
            }
        }
    }

    #[derive(Debug, PartialEq, Eq)]
    enum PortScope {
        System,
        User,
        Dynamic,
    }

    impl PortScope {
        #[inline]
        pub const fn new(value: u16) -> Self {
            //= https://www.rfc-editor.org/rfc/rfc6335#section-6
            //# o  the System Ports, also known as the Well Known Ports, from 0-1023
            //#    (assigned by IANA)
            //#
            //# o  the User Ports, also known as the Registered Ports, from 1024-
            //#    49151 (assigned by IANA)
            //#
            //# o  the Dynamic Ports, also known as the Private or Ephemeral Ports,
            //#    from 49152-65535 (never assigned)
            match value {
                0..=1023 => Self::System,
                1024..=49151 => Self::User,
                49152..=65535 => Self::Dynamic,
            }
        }
    }

    fn to_addr(addr: &SocketAddress) -> crate::inet::SocketAddress {
        match addr {
            SocketAddress::IpV4 { ip, port, .. } => {
                crate::inet::SocketAddressV4::new(**ip, *port).into()
            }
            SocketAddress::IpV6 { ip, port, .. } => {
                crate::inet::SocketAddressV6::new(**ip, *port).into()
            }
        }
    }
}

pub mod disabled {
    use super::*;

    #[derive(Debug, Default)]
    pub struct Validator;

    impl super::Validator for Validator {
        fn on_migration_attempt(&mut self, _attempt: &Attempt) -> Outcome {
            // deny all migration attempts
            Outcome::Deny(DenyReason::ConnectionMigrationDisabled)
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{
    event,
    inet::{SocketAddress, SocketAddressV4, SocketAddressV6},
};
use core::{
    convert::{TryFrom, TryInto},
    fmt,
    fmt::{Display, Formatter},
    num::NonZeroU16,
};

#[cfg(any(test, feature = "generator"))]
use bolero_generator::*;

pub mod migration;

//= https://www.rfc-editor.org/rfc/rfc9000#section-14
//# QUIC MUST NOT be used if the network path cannot support a
//# maximum datagram size of at least 1200 bytes.
pub const MINIMUM_MTU: u16 = 1200;

// TODO decide on better defaults
// Safety: 1500 is greater than zero
pub const DEFAULT_MAX_MTU: MaxMtu = MaxMtu(unsafe { NonZeroU16::new_unchecked(1500) });

// Length is the length in octets of this user datagram  including  this
// header and the data. (This means the minimum value of the length is
// eight.)
// See https://www.rfc-editor.org/rfc/rfc768.txt
pub const UDP_HEADER_LEN: u16 = 8;

// IPv4 header ranges from 20-60 bytes, depending on Options
pub const IPV4_MIN_HEADER_LEN: u16 = 20;
// IPv6 header is always 40 bytes, plus extensions
pub const IPV6_MIN_HEADER_LEN: u16 = 40;
#[cfg(feature = "ipv6")]
const IP_MIN_HEADER_LEN: u16 = IPV6_MIN_HEADER_LEN;
#[cfg(not(feature = "ipv6"))]
const IP_MIN_HEADER_LEN: u16 = IPV4_MIN_HEADER_LEN;

// The minimum allowed Max MTU is the minimum UDP datagram size of 1200 bytes plus
// the UDP header length and minimal IP header length
const MIN_ALLOWED_MAX_MTU: u16 = MINIMUM_MTU + UDP_HEADER_LEN + IP_MIN_HEADER_LEN;

// Initial PTO backoff multiplier is 1 indicating no additional increase to the backoff.
pub const INITIAL_PTO_BACKOFF: u32 = 1;

/// Internal Id of a path in the manager
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Eq, Ord, Hash)]
pub struct Id(u8);

impl Id {
    /// Create a new path::Id
    ///
    /// # Safety
    /// This should only be used by the path::Manager
    pub unsafe fn new(id: u8) -> Self {
        Self(id)
    }

    pub fn as_u8(&self) -> u8 {
        self.0
    }
}

impl event::IntoEvent<u64> for Id {
    #[inline]
    fn into_event(self) -> u64 {
        self.0 as u64
    }
}

#[cfg(any(test, feature = "testing"))]
impl Id {
    pub fn test_id() -> Self {
        unsafe { Id::new(0) }
    }
}

/// An interface for an object that represents a unique path between two endpoints
pub trait Handle: 'static + Copy + Send + fmt::Debug {
    /// Creates a Handle from a RemoteAddress
    fn from_remote_address(remote_addr: RemoteAddress) -> Self;

    /// Returns the remote address for the given handle
    fn remote_address(&self) -> RemoteAddress;

    /// Returns the local address for the given handle
    fn local_address(&self) -> LocalAddress;

    /// Returns `true` if the two handles are equal from a network perspective
    ///
    /// This function is used to determine if a connection has migrated to another
    /// path.
    fn eq(&self, other: &Self) -> bool;

    /// Returns `true` if the two handles are strictly equal to each other, i.e.
    /// byte-for-byte.
    fn strict_eq(&self, other: &Self) -> bool;

    /// Depending on the current value of `self`, fields from `other` may be copied to increase the
    /// fidelity of the value.
    ///
    /// This is especially useful for clients that initiate a connection only based on the remote
    /// IP and port. They likely wouldn't know the IP address of the local socket. Once a response
    /// is received from the server, the IP information will be known at this point and the handle
    /// can be updated with the new information.
    ///
    /// Implementations should try to limit the cost of updating by checking the current value to
    /// see if it needs updating.
    fn maybe_update(&mut self, other: &Self);
}

macro_rules! impl_addr {
    ($name:ident) => {
        #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Hash)]
        #[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
        pub struct $name(pub SocketAddress);

        impl From<SocketAddress> for $name {
            #[inline]
            fn from(value: SocketAddress) -> Self {
                Self(value)
            }
        }

        impl From<SocketAddressV4> for $name {
            #[inline]
            fn from(value: SocketAddressV4) -> Self {
                Self(value.into())
            }
        }

        impl From<SocketAddressV6> for $name {
            #[inline]
            fn from(value: SocketAddressV6) -> Self {
                Self(value.into())
            }
        }

        impl core::ops::Deref for $name {
            type Target = SocketAddress;

            fn deref(&self) -> &Self::Target {
                &self.0
            }
        }

        impl core::ops::DerefMut for $name {
            fn deref_mut(&mut self) -> &mut Self::Target {
                &mut self.0
            }
        }
    };
}

impl_addr!(LocalAddress);

impl_addr!(RemoteAddress);

impl Handle for RemoteAddress {
    #[inline]
    fn from_remote_address(remote_address: RemoteAddress) -> Self {
        remote_address
    }

    #[inline]
    fn remote_address(&self) -> RemoteAddress {
        *self
    }

    #[inline]
    fn local_address(&self) -> LocalAddress {
        SocketAddressV4::UNSPECIFIED.into()
    }

    #[inline]
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(&self.unmap(), &other.unmap())
    }

    #[inline]
    fn strict_eq(&self, other: &Self) -> bool {
        PartialEq::eq(self, other)
    }

    #[inline]
    fn maybe_update(&mut self, _other: &Self) {
        // nothing to update
    }
}

#[derive(Clone, Copy, Debug, Eq)]
#[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
pub struct Tuple {
    pub remote_address: RemoteAddress,
    pub local_address: LocalAddress,
}

impl PartialEq for Tuple {
    #[inline]
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(&self.remote_address, &other.remote_address)
            && PartialEq::eq(&self.local_address, &other.local_address)
    }
}

impl Handle for Tuple {
    #[inline]
    fn from_remote_address(remote_address: RemoteAddress) -> Self {
        let local_address = SocketAddressV4::UNSPECIFIED.into();
        Self {
            remote_address,
            local_address,
        }
    }

    #[inline]
    fn remote_address(&self) -> RemoteAddress {
        self.remote_address
    }

    #[inline]
    fn local_address(&self) -> LocalAddress {
        self.local_address
    }

    #[inline]
    fn eq(&self, other: &Self) -> bool {
        PartialEq::eq(&self.local_address.unmap(), &other.local_address.unmap())
            && Handle::eq(&self.remote_address, &other.remote_address)
    }

    #[inline]
    fn strict_eq(&self, other: &Self) -> bool {
        PartialEq::eq(self, other)
    }

    #[inline]
    fn maybe_update(&mut self, other: &Self) {
        // once we discover our path, update the address local address
        if self.local_address.port() == 0 {
            self.local_address = other.local_address;
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub struct MaxMtu(NonZeroU16);

impl MaxMtu {
    /// The minimum value required for path MTU
    pub const MIN: Self = Self(unsafe { NonZeroU16::new_unchecked(MIN_ALLOWED_MAX_MTU) });
}

impl Default for MaxMtu {
    fn default() -> Self {
        DEFAULT_MAX_MTU
    }
}

impl TryFrom<u16> for MaxMtu {
    type Error = MaxMtuError;

    fn try_from(value: u16) -> Result<Self, Self::Error> {
        if value < MIN_ALLOWED_MAX_MTU {
            return Err(MaxMtuError(MIN_ALLOWED_MAX_MTU.try_into().unwrap()));
        }

        Ok(MaxMtu(value.try_into().expect(
            "Value must be greater than zero according to the check above",
        )))
    }
}

impl From<MaxMtu> for usize {
    #[inline]
    fn from(value: MaxMtu) -> Self {
        value.0.get() as usize
    }
}

impl From<MaxMtu> for u16 {
    #[inline]
    fn from(value: MaxMtu) -> Self {
        value.0.get()
    }
}

#[derive(Debug)]
pub struct MaxMtuError(NonZeroU16);

impl Display for MaxMtuError {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "MaxMtu must be at least {}", self.0)
    }
}

#[cfg(feature = "std")]
impl std::error::Error for MaxMtuError {}

//= https://www.rfc-editor.org/rfc/rfc9308#section-8.1
//# Some UDP protocols are vulnerable to reflection attacks, where an
//# attacker is able to direct traffic to a third party as a denial of
//# service.  For example, these source ports are associated with
//# applications known to be vulnerable to reflection attacks, often due
//# to server misconfiguration:
//#
//# *  port 53 - DNS [RFC1034]
//#
//# *  port 123 - NTP [RFC5905]
//#
//# *  port 1900 - SSDP [SSDP]
//#
//# *  port 5353 - mDNS [RFC6762]
//#
//# *  port 11211 - memcache

/// List of ports to refuse connections from. This list must be sorted.
///
/// Based on https://quiche.googlesource.com/quiche/+/bac04054bccb2a249d4705ecc94a646404d41c1b/quiche/quic/core/quic_dispatcher.cc#498
const BLOCKED_PORTS: [u16; 11] = [
    0,   // We cannot send to port 0 so drop that source port.
    17,  // Quote of the Day, can loop with QUIC.
    19,  // Chargen, can loop with QUIC.
    53,  // DNS, vulnerable to reflection attacks.
    111, // Portmap.
    123, // NTP, vulnerable to reflection attacks.
    137, // NETBIOS Name Service,
    138, // NETBIOS Datagram Service
    161, // SNMP.
    389, // CLDAP.
    500, // IKE, can loop with QUIC.
];
const MAX_BLOCKED_PORT: u16 = BLOCKED_PORTS[BLOCKED_PORTS.len() - 1];

#[inline]
pub fn remote_port_blocked(port: u16) -> bool {
    if port > MAX_BLOCKED_PORT {
        // Early return to avoid iteration if possible
        return false;
    }

    for blocked in BLOCKED_PORTS {
        if port == blocked {
            return true;
        }
    }

    false
}

// The below ports are also vulnerable to reflection attacks, but are within
// the original ephemeral port range of 1024â€“65535, so there is a chance
// clients may be randomly assigned them. To address this, we can throttle
// connections using these ports instead of fully blocking them.

/// List of ports to throttle connections from. This list must be sorted.
///
/// Based on https://quiche.googlesource.com/quiche/+/bac04054bccb2a249d4705ecc94a646404d41c1b/quiche/quic/core/quic_dispatcher.cc#498
const THROTTLED_PORTS: [u16; 5] = [
    1900,  // SSDP, vulnerable to reflection attacks.
    3702,  // WS-Discovery, vulnerable to reflection attacks.
    5353,  // mDNS, vulnerable to reflection attacks.
    5355,  // LLMNR, vulnerable to reflection attacks.
    11211, // memcache, vulnerable to reflection attacks.
];
const MAX_THROTTLED_PORT: u16 = THROTTLED_PORTS[THROTTLED_PORTS.len() - 1];
pub const THROTTLED_PORTS_LEN: usize = THROTTLED_PORTS.len();

#[inline]
pub fn remote_port_throttled_index(port: u16) -> Option<usize> {
    for (idx, throttled_port) in THROTTLED_PORTS.iter().enumerate() {
        if *throttled_port == port {
            return Some(idx);
        }
    }
    None
}

#[inline]
pub fn remote_port_throttled(port: u16) -> bool {
    if port > MAX_THROTTLED_PORT {
        // Early return to avoid iteration if possible
        return false;
    }

    for throttled in THROTTLED_PORTS {
        if port == throttled {
            return true;
        }
    }

    false
}

#[cfg(test)]
mod tests {
    use crate::path::{
        remote_port_blocked, remote_port_throttled, BLOCKED_PORTS, MAX_BLOCKED_PORT,
        MAX_THROTTLED_PORT, THROTTLED_PORTS,
    };

    #[test]
    fn blocked_ports_is_sorted() {
        assert_eq!(Some(MAX_BLOCKED_PORT), BLOCKED_PORTS.iter().max().copied());

        let mut sorted = BLOCKED_PORTS.to_vec();
        sorted.sort_unstable();

        for i in 0..BLOCKED_PORTS.len() {
            assert_eq!(sorted[i], BLOCKED_PORTS[i]);
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn blocked_port() {
        for port in 0..u16::MAX {
            let blocked_expected = BLOCKED_PORTS.iter().copied().any(|blocked| blocked == port);
            assert_eq!(blocked_expected, remote_port_blocked(port));
        }
    }

    #[test]
    fn throttled_ports_is_sorted() {
        assert_eq!(
            Some(MAX_THROTTLED_PORT),
            THROTTLED_PORTS.iter().max().copied()
        );

        let mut sorted = THROTTLED_PORTS.to_vec();
        sorted.sort_unstable();

        for i in 0..THROTTLED_PORTS.len() {
            assert_eq!(sorted[i], THROTTLED_PORTS[i]);
        }
    }

    #[test]
    #[cfg_attr(miri, ignore)]
    fn throttled_port() {
        for port in 0..u16::MAX {
            let throttled_expected = THROTTLED_PORTS
                .iter()
                .copied()
                .any(|throttled| throttled == port);
            assert_eq!(throttled_expected, remote_port_throttled(port));
        }
    }
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use crate::{
        connection, event,
        event::{builder::SocketAddress, IntoEvent},
    };

    impl<'a> event::builder::Path<'a> {
        pub fn test() -> Self {
            Self {
                local_addr: SocketAddress::IpV4 {
                    ip: &[127, 0, 0, 1],
                    port: 0,
                },
                local_cid: connection::LocalId::TEST_ID.into_event(),
                remote_addr: SocketAddress::IpV4 {
                    ip: &[127, 0, 0, 1],
                    port: 0,
                },
                remote_cid: connection::PeerId::TEST_ID.into_event(),
                id: 0,
                is_active: false,
            }
        }
    }
}