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

use core::{
    convert::{TryFrom, TryInto},
    fmt,
    ops::Deref,
};
use s2n_codec::{decoder_value, Encoder, EncoderValue};

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

use crate::event::IntoEvent;

#[cfg(test)]
mod tests;

//= https://www.rfc-editor.org/rfc/rfc9000#section-16
//# QUIC packets and frames commonly use a variable-length encoding for
//# non-negative integer values.  This encoding ensures that smaller
//# integer values need fewer bytes to encode.

//# The QUIC variable-length integer encoding reserves the two most
//# significant bits of the first byte to encode the base 2 logarithm of
//# the integer encoding length in bytes.  The integer value is encoded
//# on the remaining bits, in network byte order.

//= https://www.rfc-editor.org/rfc/rfc9000#section-16
//# This means that integers are encoded on 1, 2, 4, or 8 bytes and can
//# encode 6-, 14-, 30-, or 62-bit values, respectively.  Table 4
//# summarizes the encoding properties.
//#
//#        +======+========+=============+=======================+
//#        | 2MSB | Length | Usable Bits | Range                 |
//#        +======+========+=============+=======================+
//#        | 00   | 1      | 6           | 0-63                  |
//#        +------+--------+-------------+-----------------------+
//#        | 01   | 2      | 14          | 0-16383               |
//#        +------+--------+-------------+-----------------------+
//#        | 10   | 4      | 30          | 0-1073741823          |
//#        +------+--------+-------------+-----------------------+
//#        | 11   | 8      | 62          | 0-4611686018427387903 |
//#        +------+--------+-------------+-----------------------+

pub const MAX_VARINT_VALUE: u64 = 4_611_686_018_427_387_903;

#[derive(Debug)]
pub struct VarIntError;

impl fmt::Display for VarIntError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "varint range exceeded")
    }
}

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

// https://godbolt.org/z/ToTvPD
#[inline(always)]
fn read_table(x: u64) -> (u64, usize, u64) {
    debug_assert!(x <= MAX_VARINT_VALUE);

    macro_rules! table {
        ($(($two_bit:expr, $length:expr, $usable_bits:expr, $max_value:expr);)*) => {{
            let mut two_bit = 0;
            let leading_zeros = x.leading_zeros();
            $(
                two_bit += if leading_zeros < (64 - $usable_bits) {
                    1
                } else {
                    0
                };
            )*

            let len = 1 << two_bit;
            let usable_bits = len * 8 - 2;

            debug_assert_eq!(len as usize, encoding_size(x));

            (two_bit, len as usize, usable_bits)
        }};
    }

    table! {
        (0b00, 1, 6 , 63);
        (0b01, 2, 14, 16_383);
        (0b10, 4, 30, 1_073_741_823);
    }
}

#[inline(always)]
fn encoding_size(x: u64) -> usize {
    debug_assert!(x <= MAX_VARINT_VALUE);

    macro_rules! table {
        ($(($two_bit:expr, $length:expr, $usable_bits:expr, $max_value:expr);)*) => {{
            let leading_zeros = x.leading_zeros();
            let mut len = 1;

            $(
                if leading_zeros < (64 - $usable_bits) {
                    len = $length * 2;
                };
            )*

            len
        }};
    }

    table! {
        (0b00, 1, 6 , 63);
        (0b01, 2, 14, 16_383);
        (0b10, 4, 30, 1_073_741_823);
    }
}

// === API ===

#[derive(Clone, Copy, Debug, Default, Eq, Hash, PartialEq, PartialOrd, Ord)]
#[cfg_attr(any(feature = "generator", test), derive(TypeGenerator))]
pub struct VarInt(#[cfg_attr(any(feature = "generator", test), generator(Self::GENERATOR))] u64);

impl fmt::Display for VarInt {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl VarInt {
    pub const MAX: Self = Self(MAX_VARINT_VALUE);

    pub const ZERO: Self = Self(0);

    #[cfg(any(feature = "generator", test))]
    const GENERATOR: core::ops::RangeInclusive<u64> = 0..=MAX_VARINT_VALUE;

    pub fn new(v: u64) -> Result<Self, VarIntError> {
        if v > MAX_VARINT_VALUE {
            return Err(VarIntError);
        }
        Ok(Self(v))
    }

    /// Returns a `VarInt` without validating the value is less than VarInt::MAX
    ///
    /// # Safety
    ///
    /// Callers need to ensure the value is less than or equal to VarInt::MAX
    pub const unsafe fn new_unchecked(value: u64) -> Self {
        Self(value)
    }

    pub const fn from_u8(v: u8) -> Self {
        Self(v as u64)
    }

    pub const fn from_u16(v: u16) -> Self {
        Self(v as u64)
    }

    pub const fn from_u32(v: u32) -> Self {
        Self(v as u64)
    }

    pub const fn as_u64(self) -> u64 {
        self.0
    }

    #[inline]
    pub fn checked_add(self, value: Self) -> Option<Self> {
        Self::new(self.0.checked_add(value.0)?).ok()
    }

    #[inline]
    pub fn checked_add_usize(self, value: usize) -> Option<Self> {
        let value = value.try_into().ok()?;
        self.checked_add(value)
    }

    #[inline]
    #[must_use]
    pub fn saturating_add(self, value: Self) -> Self {
        Self::new(self.0.saturating_add(value.0)).unwrap_or(Self::MAX)
    }

    #[inline]
    pub fn checked_sub(self, value: Self) -> Option<Self> {
        Some(Self(self.0.checked_sub(value.0)?))
    }

    #[inline]
    #[must_use]
    pub fn saturating_sub(self, value: Self) -> Self {
        Self(self.0.saturating_sub(value.0))
    }

    #[inline]
    pub fn checked_mul(self, value: Self) -> Option<Self> {
        Self::new(self.0.checked_mul(value.0)?).ok()
    }

    #[inline]
    #[must_use]
    pub fn saturating_mul(self, value: Self) -> Self {
        Self::new(self.0.saturating_mul(value.0)).unwrap_or(Self::MAX)
    }

    #[inline]
    pub fn checked_div(self, value: Self) -> Option<Self> {
        Some(Self(self.0.checked_div(value.0)?))
    }

    /// Re-encodes a replacement value where `self` was used as a placeholder.
    #[inline]
    pub fn encode_updated<E: Encoder>(self, replacement: Self, encoder: &mut E) {
        debug_assert!(
            self.encoding_table_entry().1 >= replacement.encoding_table_entry().1,
            "the replacement encoding_size should not be greater than the previous value"
        );

        replacement.encode_with_table_entry(self.encoding_table_entry(), encoder);
    }

    #[inline]
    fn encode_with_table_entry<E: Encoder>(
        self,
        (two_bit, len, usable_bits): (u64, usize, u64),
        encoder: &mut E,
    ) {
        encoder.write_sized(len, |buffer| {
            let bytes = (two_bit << usable_bits | self.0).to_be_bytes();

            unsafe {
                // Safety: the encoder will have checked the buffer size
                //         before passing it so we don't need to pay the bounds
                //         check cost twice.

                // This code looks a little scary so here's some comments describing what
                // is happening.

                // We bitwise-and the two bit value to ensure the compiler can prove the
                // `unreachable` is actually unreachable.
                match two_bit & 0b11 {
                    0b00 => {
                        // If the two bit value is 0b00 it means we only have 1 byte to
                        // encode so we copy the last byte from our big endian encoded value
                        // into the first byte of the buffer
                        debug_assert_eq!(buffer.len(), 1);
                        *buffer.get_unchecked_mut(0) = *bytes.get_unchecked(7);
                    }
                    0b01 => {
                        // If the two bit value is 0b01 it means we have a 2 byte value to
                        // encode so we copy the last 2 bytes from our big endian encoded value
                        // into the first 2 bytes of the buffer
                        debug_assert_eq!(buffer.len(), 2);
                        buffer
                            .get_unchecked_mut(..2)
                            .copy_from_slice(bytes.get_unchecked(6..));
                    }
                    0b10 => {
                        // If the two bit value is 0b10 it means we have a 4 byte value to
                        // encode so we copy the last 4 bytes from our big endian encoded value
                        // into the first 4 bytes of the buffer
                        debug_assert_eq!(buffer.len(), 4);
                        buffer
                            .get_unchecked_mut(..4)
                            .copy_from_slice(bytes.get_unchecked(4..));
                    }
                    0b11 => {
                        // If the two bit value is 0b11 it means we have a 8 byte value to
                        // encode so we copy all of the bytes into the buffer
                        debug_assert_eq!(buffer.len(), 8);
                        buffer
                            .get_unchecked_mut(..8)
                            .copy_from_slice(bytes.get_unchecked(..8));
                    }
                    _ => unreachable!(),
                }
            }
        })
    }

    #[inline]
    fn encoding_table_entry(self) -> (u64, usize, u64) {
        read_table(self.0)
    }
}

impl EncoderValue for VarInt {
    #[inline]
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        self.encode_with_table_entry(self.encoding_table_entry(), encoder);
    }

    #[inline]
    fn encoding_size(&self) -> usize {
        encoding_size(self.0)
    }

    #[inline]
    fn encoding_size_for_encoder<E: Encoder>(&self, _encoder: &E) -> usize {
        encoding_size(self.0)
    }
}

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

decoder_value!(
    impl<'a> VarInt {
        fn decode(buffer: Buffer) -> Result<Self> {
            let header = buffer.peek_byte(0)?;

            Ok(match (header >> 6) & 0b11 {
                0b00 => {
                    let value = header & (2u8.pow(6) - 1);
                    let buffer = buffer.skip(1)?;
                    (Self(value.into()), buffer)
                }
                0b01 => {
                    let (value, buffer) = buffer.decode::<u16>()?;
                    let value = value & (2u16.pow(14) - 1);
                    (Self(value.into()), buffer)
                }
                0b10 => {
                    let (value, buffer) = buffer.decode::<u32>()?;
                    let value = value & (2u32.pow(30) - 1);
                    (Self(value.into()), buffer)
                }
                0b11 => {
                    let (value, buffer) = buffer.decode::<u64>()?;
                    let value = value & (2u64.pow(62) - 1);
                    (Self(value), buffer)
                }
                _ => unreachable!(),
            })
        }
    }
);

impl AsRef<u64> for VarInt {
    #[inline]
    fn as_ref(&self) -> &u64 {
        &self.0
    }
}

impl Deref for VarInt {
    type Target = u64;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

macro_rules! impl_from_lesser {
    ($ty:ty) => {
        impl From<$ty> for VarInt {
            #[inline]
            fn from(value: $ty) -> Self {
                Self(value.into())
            }
        }
    };
}

impl_from_lesser!(u8);
impl_from_lesser!(u16);
impl_from_lesser!(u32);

impl From<VarInt> for u64 {
    #[inline]
    fn from(v: VarInt) -> u64 {
        v.0
    }
}

impl TryFrom<usize> for VarInt {
    type Error = VarIntError;

    #[inline]
    fn try_from(value: usize) -> Result<Self, Self::Error> {
        Self::new(value as u64)
    }
}

impl TryInto<usize> for VarInt {
    type Error = <usize as TryFrom<u64>>::Error;

    #[inline]
    fn try_into(self) -> Result<usize, Self::Error> {
        self.0.try_into()
    }
}

impl TryFrom<u64> for VarInt {
    type Error = VarIntError;

    #[inline]
    fn try_from(value: u64) -> Result<Self, Self::Error> {
        Self::new(value)
    }
}

impl TryFrom<u128> for VarInt {
    type Error = VarIntError;

    #[inline]
    fn try_from(value: u128) -> Result<Self, Self::Error> {
        if value > MAX_VARINT_VALUE as u128 {
            Err(VarIntError)
        } else {
            Ok(Self(value as u64))
        }
    }
}

impl core::ops::Add for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn add(self, rhs: Self) -> Self {
        if cfg!(debug_assertions) {
            self.checked_add(rhs).expect("VarInt overflow occurred")
        } else {
            Self(self.0 + rhs.0)
        }
    }
}

impl core::ops::Add<usize> for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn add(self, rhs: usize) -> Self {
        if cfg!(debug_assertions) {
            self.checked_add(VarInt::new(rhs as u64).expect("VarInt overflow occurred"))
                .expect("VarInt overflow occurred")
        } else {
            Self(self.0 + rhs as u64)
        }
    }
}

impl core::ops::AddAssign<Self> for VarInt {
    #[inline]
    #[track_caller]
    fn add_assign(&mut self, rhs: Self) {
        if cfg!(debug_assertions) {
            *self = self.checked_add(rhs).expect("VarInt overflow occurred")
        } else {
            self.0 += rhs.0
        }
    }
}

impl core::ops::AddAssign<usize> for VarInt {
    #[inline]
    #[track_caller]
    fn add_assign(&mut self, rhs: usize) {
        if cfg!(debug_assertions) {
            *self = self
                .checked_add(VarInt::new(rhs as u64).expect("VarInt overflow occurred"))
                .expect("VarInt overflow occurred")
        } else {
            self.0 += rhs as u64
        }
    }
}

impl core::ops::Sub for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn sub(self, rhs: Self) -> Self {
        // Bounds check is inherited from u64
        Self(self.0 - rhs.0)
    }
}

impl core::ops::Sub<usize> for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn sub(self, rhs: usize) -> Self {
        // Bounds check is inherited from u64
        Self(self.0 - rhs as u64)
    }
}

impl core::ops::SubAssign<Self> for VarInt {
    #[inline]
    #[track_caller]
    fn sub_assign(&mut self, rhs: Self) {
        // Bounds check is inherited from u64
        self.0 -= rhs.0
    }
}

impl core::ops::SubAssign<usize> for VarInt {
    #[inline]
    #[track_caller]
    fn sub_assign(&mut self, rhs: usize) {
        // Bounds check is inherited from u64
        self.0 -= rhs as u64
    }
}

impl core::ops::Mul for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn mul(self, rhs: Self) -> Self {
        if cfg!(debug_assertions) {
            self.checked_mul(rhs).expect("VarInt overflow occurred")
        } else {
            Self(self.0 * rhs.0)
        }
    }
}

impl core::ops::Mul<usize> for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn mul(self, rhs: usize) -> Self {
        if cfg!(debug_assertions) {
            self.checked_mul(VarInt::new(rhs as u64).expect("VarInt overflow occurred"))
                .expect("VarInt overflow occurred")
        } else {
            Self(self.0 * rhs as u64)
        }
    }
}

impl core::ops::MulAssign<Self> for VarInt {
    #[inline]
    #[track_caller]
    fn mul_assign(&mut self, rhs: Self) {
        if cfg!(debug_assertions) {
            *self = self.checked_mul(rhs).expect("VarInt overflow occurred")
        } else {
            self.0 *= rhs.0
        }
    }
}

impl core::ops::MulAssign<usize> for VarInt {
    #[inline]
    #[track_caller]
    fn mul_assign(&mut self, rhs: usize) {
        if cfg!(debug_assertions) {
            *self = self
                .checked_mul(VarInt::new(rhs as u64).expect("VarInt overflow occurred"))
                .expect("VarInt overflow occurred")
        } else {
            self.0 *= rhs as u64
        }
    }
}

impl core::ops::Div for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn div(self, rhs: Self) -> Self {
        // Bounds check is inherited from u64
        Self(self.0 / rhs.0)
    }
}

impl core::ops::Div<usize> for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn div(self, rhs: usize) -> Self {
        // Bounds check is inherited from u64
        Self(self.0 / rhs as u64)
    }
}

impl core::ops::DivAssign<Self> for VarInt {
    #[inline]
    #[track_caller]
    fn div_assign(&mut self, rhs: Self) {
        // Bounds check is inherited from u64
        self.0 /= rhs.0
    }
}

impl core::ops::DivAssign<usize> for VarInt {
    #[inline]
    #[track_caller]
    fn div_assign(&mut self, rhs: usize) {
        // Bounds check is inherited from u64
        self.0 /= rhs as u64
    }
}

impl core::ops::Rem for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn rem(self, rhs: Self) -> Self {
        // Bounds check is inherited from u64
        Self(self.0.rem(rhs.0))
    }
}

impl core::ops::Rem<usize> for VarInt {
    type Output = Self;

    #[inline]
    #[track_caller]
    fn rem(self, rhs: usize) -> Self {
        // Bounds check is inherited from u64
        Self(self.0.rem(rhs as u64))
    }
}

impl core::ops::RemAssign<Self> for VarInt {
    #[inline]
    #[track_caller]
    fn rem_assign(&mut self, rhs: Self) {
        // Bounds check is inherited from u64
        self.0 %= rhs.0
    }
}

impl core::ops::RemAssign<usize> for VarInt {
    #[inline]
    #[track_caller]
    fn rem_assign(&mut self, rhs: usize) {
        // Bounds check is inherited from u64
        self.0 %= rhs as u64
    }
}

impl PartialEq<u64> for VarInt {
    #[inline]
    fn eq(&self, other: &u64) -> bool {
        self.0.eq(other)
    }
}

impl PartialEq<usize> for VarInt {
    #[inline]
    fn eq(&self, other: &usize) -> bool {
        self.0.eq(&(*other as u64))
    }
}

impl PartialOrd<u64> for VarInt {
    #[inline]
    fn partial_cmp(&self, other: &u64) -> Option<core::cmp::Ordering> {
        self.0.partial_cmp(other)
    }
}

impl PartialOrd<usize> for VarInt {
    #[inline]
    fn partial_cmp(&self, other: &usize) -> Option<core::cmp::Ordering> {
        self.0.partial_cmp(&(*other as u64))
    }
}


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

use super::*;
use crate::varint::VarInt;
use bolero::check;
use core::mem::size_of;
use s2n_codec::{
    assert_codec_round_trip_bytes, assert_codec_round_trip_value, DecoderBuffer, EncoderBuffer,
};

#[test]
#[cfg_attr(miri, ignore)] // This test is too expensive for miri to complete in a reasonable amount of time
fn round_trip_bytes_test() {
    check!().for_each(|input| {
        assert_codec_round_trip_bytes!(VarInt, input);
    });
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(10), kani::solver(cadical))]
fn round_trip_values_test() {
    check!().with_type().cloned().for_each(|v| {
        if let Ok(v) = VarInt::new(v) {
            assert_codec_round_trip_value!(VarInt, v);
        } else {
            assert!(v > MAX_VARINT_VALUE);
        }
    })
}

#[test]
#[cfg_attr(miri, ignore)] // snapshot tests don't work on miri
fn table_snapshot_test() {
    use insta::assert_debug_snapshot;
    assert_debug_snapshot!("max_value", MAX_VARINT_VALUE);

    // These values are derived from the "usable bits" column in the table: V and V-1
    for i in [0, 1, 5, 6, 13, 14, 29, 30, 61] {
        assert_debug_snapshot!(format!("table_2_pow_{}_", i), read_table(2u64.pow(i)));
    }
}

//= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.1
//# For example, the eight-byte sequence 0xc2197c5eff14e88c decodes to
//# the decimal value 151,288,809,941,952,652; the four-byte sequence
//# 0x9d7f3e7d decodes to 494,878,333; the two-byte sequence 0x7bbd
//# decodes to 15,293; and the single byte 0x25 decodes to 37 (as does
//# the two-byte sequence 0x4025).

macro_rules! sequence_test {
    ($name:ident($input:expr, $expected:expr)) => {
        #[test]
        fn $name() {
            use s2n_codec::assert_codec_round_trip_value;

            let input = $input;
            let expected = VarInt::new($expected).unwrap();
            let actual_bytes = assert_codec_round_trip_value!(VarInt, expected);
            assert_eq!(&input[..], &actual_bytes[..]);
        }
    };
}

sequence_test!(eight_byte_sequence_example(
    [0xc2, 0x19, 0x7c, 0x5e, 0xff, 0x14, 0xe8, 0x8c],
    151_288_809_941_952_652
));

sequence_test!(four_byte_sequence_example(
    [0x9d, 0x7f, 0x3e, 0x7d],
    494_878_333
));

sequence_test!(two_byte_sequence_example([0x7b, 0xbd], 15293));

sequence_test!(one_byte_sequence_example([0x25], 37));

//= https://www.rfc-editor.org/rfc/rfc9000#section-16
//= type=test
// # +======+========+=============+=======================+
// # | 2MSB | Length | Usable Bits | Range                 |
// # +======+========+=============+=======================+
// # | 00   | 1      | 6           | 0-63                  |
// # +------+--------+-------------+-----------------------+
// # | 01   | 2      | 14          | 0-16383               |
// # +------+--------+-------------+-----------------------+
// # | 10   | 4      | 30          | 0-1073741823          |
// # +------+--------+-------------+-----------------------+
// # | 11   | 8      | 62          | 0-4611686018427387903 |
// # +------+--------+-------------+-----------------------+
#[test]
#[cfg_attr(miri, ignore)]
#[cfg_attr(kani, kani::proof, kani::unwind(3), kani::solver(cadical))]
fn one_byte_sequence_test() {
    bolero::check!()
        .with_type()
        .cloned()
        .for_each(|mut first_byte: u8| {
            first_byte &= 0x3f;
            let byte_sequence = [first_byte];
            let expected = VarInt::new(first_byte as u64).unwrap();
            let actual_bytes = assert_codec_round_trip_value!(VarInt, expected);
            assert_eq!(&byte_sequence[..], &actual_bytes[..]);
        });
}

#[test]
#[cfg_attr(miri, ignore)]
#[cfg_attr(kani, kani::proof, kani::unwind(4), kani::solver(cadical))]
fn two_byte_sequence_test() {
    // the s2n-quic implementation always chooses the smallest encoding possible.
    // this means if we wish to test two-byte sequences, we need to encode a number
    // that is > 63.
    let g = gen::<(u8, u8)>().filter_gen(|(_, b)| *b > 63);

    bolero::check!()
        .with_generator(g)
        .cloned()
        .for_each(|(mut first_byte, second_byte)| {
            first_byte = (first_byte & 0x3f) | 0x40;
            let byte_sequence = [first_byte, second_byte];
            let expected_val = u16::from_be_bytes([first_byte & 0x3f, second_byte]);
            assert!(expected_val <= 16383);
            let expected = VarInt::new(expected_val as u64).unwrap();
            let actual_bytes = assert_codec_round_trip_value!(VarInt, expected);
            assert_eq!(&byte_sequence[..], &actual_bytes[..]);
        });
}

#[test]
#[cfg_attr(miri, ignore)]
#[cfg_attr(kani, kani::proof, kani::unwind(10), kani::solver(cadical))]
fn four_byte_sequence_test() {
    // The s2n-quic implementation always chooses the smallest encoding possible.
    // This means if we wish to test the four-byte sequences, we need to encode a number
    // that is > 16383 or 0b0011 1111 1111 1111.
    let g = gen::<(u8, u8, u8)>().filter_gen(|(_, _, c)| *c > 0x3f);

    bolero::check!().with_generator(g).cloned().for_each(
        |(mut first_byte, second_byte, third_byte)| {
            first_byte = (first_byte & 0x3f) | 0x80;
            let byte_sequence = [first_byte, second_byte, third_byte, 0xff];
            let expected_val: u32 =
                u32::from_be_bytes([first_byte & 0x3f, second_byte, third_byte, 0xff]);
            assert!(expected_val <= 1073741823);
            let expected = VarInt::new(expected_val as u64).unwrap();
            let actual_bytes = assert_codec_round_trip_value!(VarInt, expected);
            assert_eq!(&byte_sequence[..], &actual_bytes[..]);
        },
    );
}

#[test]
#[cfg_attr(miri, ignore)]
#[cfg_attr(kani, kani::proof, kani::unwind(10), kani::solver(cadical))]
fn eight_byte_sequence_test() {
    // The s2n-quic implementation always chooses the smallest encoding possible.
    // This means if we wish to test eight-byte sequences, we need to encode a number
    // that is > 1073741823 or 0b0011 1111 1111 1111 1111 1111 1111 1111
    let g = gen::<(u8, u8, u8, u8)>().filter_gen(|(_, _, _, d)| *d > 0x3f);

    bolero::check!().with_generator(g).cloned().for_each(
        |(mut first_byte, second_byte, third_byte, fourth_byte)| {
            first_byte = (first_byte & 0x3f) | 0xc0;
            let byte_sequence = [
                first_byte,
                second_byte,
                third_byte,
                fourth_byte,
                0xff,
                0xff,
                0xff,
                0xff,
            ];
            let expected_val: u64 = u64::from_be_bytes([
                first_byte & 0x3f,
                second_byte,
                third_byte,
                fourth_byte,
                0xff,
                0xff,
                0xff,
                0xff,
            ]);
            assert!(expected_val <= 4611686018427387903);
            let expected = VarInt::new(expected_val).unwrap();
            let actual_bytes = assert_codec_round_trip_value!(VarInt, expected);
            assert_eq!(&byte_sequence[..], &actual_bytes[..]);
        },
    );
}

fn test_update(initial: VarInt, expected: VarInt, encoder: &mut EncoderBuffer) {
    encoder.set_position(0);
    initial.encode_updated(expected, encoder);
    let decoder = DecoderBuffer::new(encoder.as_mut_slice());
    let (actual, _) = decoder.decode::<VarInt>().unwrap();
    assert_eq!(expected, actual);
}

#[test]
fn encode_updated_test() {
    let mut buffer = [0u8; size_of::<VarInt>()];
    let mut encoder = EncoderBuffer::new(&mut buffer);
    let initial = VarInt::from_u16(1 << 14);
    encoder.encode(&initial);

    test_update(initial, VarInt::from_u32(0), &mut encoder);
    test_update(initial, VarInt::from_u32(1 << 14), &mut encoder);
    test_update(initial, VarInt::from_u32(1 << 29), &mut encoder);
}

#[test]
#[should_panic]
fn encode_updated_invalid_test() {
    let mut buffer = [0u8; size_of::<VarInt>()];
    let mut encoder = EncoderBuffer::new(&mut buffer);
    let initial = VarInt::from_u16(1 << 14);
    encoder.encode(&initial);

    test_update(initial, VarInt::from_u32(1 << 30), &mut encoder);
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
fn table_differential_test() {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-16
    //= type=test
    //# This means that integers are encoded on 1, 2, 4, or 8 bytes and can
    //# encode 6-, 14-, 30-, or 62-bit values, respectively.  Table 4
    //# summarizes the encoding properties.
    //#
    //#        +======+========+=============+=======================+
    //#        | 2MSB | Length | Usable Bits | Range                 |
    //#        +======+========+=============+=======================+
    //#        | 00   | 1      | 6           | 0-63                  |
    //#        +------+--------+-------------+-----------------------+
    //#        | 01   | 2      | 14          | 0-16383               |
    //#        +------+--------+-------------+-----------------------+
    //#        | 10   | 4      | 30          | 0-1073741823          |
    //#        +------+--------+-------------+-----------------------+
    //#        | 11   | 8      | 62          | 0-4611686018427387903 |
    //#        +------+--------+-------------+-----------------------+

    check!()
        .with_generator(0..=MAX_VARINT_VALUE)
        .cloned()
        .for_each(|v| {
            let actual = read_table(v);

            // make sure the encoding_size function matches the table
            assert_eq!(actual.1, encoding_size(v));

            // write the table from the RFC in a non-optimized format and make sure the
            // actual results match
            #[allow(clippy::match_overlapping_arm)]
            let expected = match v {
                0..=63 => (0b00, 1, 6),
                0..=16383 => (0b01, 2, 14),
                0..=1073741823 => (0b10, 4, 30),
                0..=4611686018427387903 => (0b11, 8, 62),
                _ => unreachable!(),
            };

            assert_eq!(actual, expected);
        })
}

#[test]
#[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
fn checked_ops_test() {
    check!().with_type().cloned().for_each(|(a, b)| {
        if let (Ok(a_v), Ok(b_v)) = (VarInt::new(a), VarInt::new(b)) {
            // compares checked operations against the underlying u64 operations
            macro_rules! checked {
                ($checked:ident, $op:tt) => {{
                    if let Some(res) = a_v.$checked(b_v) {
                        // make sure the behavior matches that of u64
                        assert_eq!(Some(res.as_u64()), a.$checked(b));

                        // make sure the normal operation works as well
                        assert_eq!(res, a_v $op b_v);
                    } else {
                        // make sure the behavior matches that of u64
                        let expected = a.$checked(b).and_then(|v| VarInt::new(v).ok());
                        assert!(expected.is_none());
                    }
                }};
            }

            checked!(checked_add, +);
            checked!(checked_sub, -);

            // kani currently spins for a really long time to handle these checks
            #[cfg(any(not(kani), kani_slow))]
            checked!(checked_mul, *);
            #[cfg(any(not(kani), kani_slow))]
            checked!(checked_div, /);

            // compares saturating operations against the underlying u64 operations
            macro_rules! saturating {
                ($name:ident) => {{
                    let actual = a_v.$name(b_v);
                    let expected = VarInt::new(a.$name(b)).unwrap_or(VarInt::MAX);
                    assert_eq!(actual, expected);
                }};
            }

            saturating!(saturating_add);

            // kani currently spins for a really long time to handle these checks
            #[cfg(any(not(kani), kani_slow))]
            saturating!(saturating_sub);
            #[cfg(any(not(kani), kani_slow))]
            saturating!(saturating_mul);
        }
    })
}


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

use s2n_codec::{DecoderBufferMut as Inner, DecoderError, DecoderValue, DecoderValueMut};

/// A BPF-aware version of [`s2n_codec::DecoderBufferMut`]
///
/// The Linux BPF verifier requires that every pointer be checked against the `end` pointer. This
/// means that it struggles with regular Rust slices that use `ptr + len` instead of `ptr + end`.
pub struct DecoderBufferMut<'a> {
    buffer: Inner<'a>,
    end: *mut u8,
}

impl<'a> DecoderBufferMut<'a> {
    /// Creates a new DecoderBufferMut.
    ///
    /// # Safety
    ///
    /// The `start` and `end` pointers must be a valid range of bytes, ideally directly coming
    /// from the BPF/XDP context argument.
    #[inline]
    pub unsafe fn new(start: *mut u8, end: *mut u8) -> Self {
        let len = end as usize - start as usize;
        let data = core::slice::from_raw_parts_mut(start as *mut u8, len);
        let buffer = Inner::new(data);
        Self { buffer, end }
    }

    /// Validates that the starting pointer is still within the bounds of the end pointer
    #[inline]
    fn new_checked(buffer: Inner<'a>, end: *mut u8) -> Result<Self, DecoderError> {
        // The Linux BPF verifier needs to prove that no pointers go beyond the "end" pointer
        if buffer.as_less_safe_slice().as_ptr() > end {
            return Err(DecoderError::UnexpectedEof(0));
        }

        Ok(Self { buffer, end })
    }

    /// Decodes a T from the buffer, if possible
    #[inline]
    pub fn decode<T: DecoderValueMut<'a>>(self) -> Result<(T, Self), DecoderError> {
        let end = self.end;
        let (v, buffer) = self.buffer.decode()?;
        let buffer = Self::new_checked(buffer, end)?;
        Ok((v, buffer))
    }

    /// Decodes a slice of bytes with the given len, if possible
    #[inline]
    pub fn decode_slice(self, len: usize) -> Result<(Self, Self), DecoderError> {
        let end = self.end;
        let (slice, buffer) = self.buffer.decode_slice(len)?;
        let slice = Self::new_checked(slice, end)?;
        let buffer = Self::new_checked(buffer, end)?;
        Ok((slice, buffer))
    }

    #[inline]
    pub fn into_less_safe_slice(self) -> &'a mut [u8] {
        self.buffer.into_less_safe_slice()
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.buffer.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.buffer.is_empty()
    }
}

/// A generic interface over a decoder buffer
pub trait Decoder<'a>: Sized {
    fn decode<T: DecoderValue<'a> + DecoderValueMut<'a>>(
        self,
    ) -> core::result::Result<(T, Self), DecoderError>;
    fn decode_slice(self, len: usize) -> core::result::Result<(Self, Self), DecoderError>;
}

impl<'a> Decoder<'a> for DecoderBufferMut<'a> {
    #[inline]
    fn decode<T: DecoderValue<'a> + DecoderValueMut<'a>>(
        self,
    ) -> core::result::Result<(T, Self), DecoderError> {
        Self::decode(self)
    }

    #[inline]
    fn decode_slice(self, len: usize) -> core::result::Result<(Self, Self), DecoderError> {
        Self::decode_slice(self, len)
    }
}

impl<'a> Decoder<'a> for s2n_codec::DecoderBuffer<'a> {
    #[inline]
    fn decode<T: DecoderValue<'a> + DecoderValueMut<'a>>(
        self,
    ) -> core::result::Result<(T, Self), DecoderError> {
        Self::decode(self)
    }

    #[inline]
    fn decode_slice(self, len: usize) -> core::result::Result<(Self, Self), DecoderError> {
        Self::decode_slice(self, len)
    }
}

impl<'a> Decoder<'a> for s2n_codec::DecoderBufferMut<'a> {
    #[inline]
    fn decode<T: DecoderValue<'a> + DecoderValueMut<'a>>(
        self,
    ) -> core::result::Result<(T, Self), DecoderError> {
        Self::decode(self)
    }

    #[inline]
    fn decode_slice(self, len: usize) -> core::result::Result<(Self, Self), DecoderError> {
        Self::decode_slice(self, len)
    }
}


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

use super::{bpf::Decoder, path};
use crate::inet::{
    datagram,
    ethernet::{self, EtherType},
    ip, ipv4, ipv6, udp,
};
use s2n_codec::DecoderError;

pub type Result<D = ()> = core::result::Result<Option<D>, DecoderError>;

pub trait EventHandler: Sized {
    #[inline(always)]
    fn decode_packet<'a, D: Decoder<'a>>(&mut self, buffer: D) -> Result<D> {
        decode_packet_with_event(buffer, self)
    }

    #[inline(always)]
    fn on_ethernet_header(&mut self, header: &ethernet::Header) -> Result {
        let _ = header;
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_ipv4_header(&mut self, header: &ipv4::Header) -> Result {
        let _ = header;
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_ipv6_header(&mut self, header: &ipv6::Header) -> Result {
        let _ = header;
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_udp_header(&mut self, header: &udp::Header) -> Result {
        let _ = header;
        Ok(Some(()))
    }
}

impl EventHandler for () {}

impl EventHandler for path::Tuple {
    #[inline(always)]
    fn on_ethernet_header(&mut self, header: &ethernet::Header) -> Result {
        self.remote_address.mac = header.source;
        self.local_address.mac = header.destination;
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_ipv4_header(&mut self, header: &ipv4::Header) -> Result {
        self.remote_address.ip = header.source.into();
        self.local_address.ip = header.destination.into();
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_ipv6_header(&mut self, header: &ipv6::Header) -> Result {
        self.remote_address.ip = header.source.into();
        self.local_address.ip = header.destination.into();
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_udp_header(&mut self, header: &udp::Header) -> Result {
        self.remote_address.port = header.source.get();
        self.local_address.port = header.destination.get();
        Ok(Some(()))
    }
}

impl<P: EventHandler> EventHandler for datagram::Header<P> {
    #[inline(always)]
    fn on_ethernet_header(&mut self, header: &ethernet::Header) -> Result {
        self.path.on_ethernet_header(header)
    }

    #[inline(always)]
    fn on_ipv4_header(&mut self, header: &ipv4::Header) -> Result {
        self.path.on_ipv4_header(header)?;
        self.ecn = header.tos().ecn();
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_ipv6_header(&mut self, header: &ipv6::Header) -> Result {
        self.path.on_ipv6_header(header)?;
        self.ecn = header.vtcfl().ecn();
        Ok(Some(()))
    }

    #[inline(always)]
    fn on_udp_header(&mut self, header: &udp::Header) -> Result {
        self.path.on_udp_header(header)
    }
}

/// Decodes a path tuple and payload from a raw packet
#[inline(always)]
pub fn decode_packet<'a, D: Decoder<'a>>(
    buffer: D,
) -> core::result::Result<Option<(datagram::Header<path::Tuple>, D)>, DecoderError> {
    let mut header = datagram::Header {
        path: path::Tuple::UNSPECIFIED,
        ecn: Default::default(),
    };
    match decode_packet_with_event(buffer, &mut header)? {
        Some(buffer) => Ok(Some((header, buffer))),
        None => Ok(None),
    }
}

/// Decodes a path tuple and payload from a raw packet
#[inline(always)]
pub fn decode_packet_with_event<'a, D: Decoder<'a>, E: EventHandler>(
    buffer: D,
    events: &mut E,
) -> Result<D> {
    let (header, buffer) = buffer.decode::<&ethernet::Header>()?;

    if events.on_ethernet_header(header)?.is_none() {
        return Ok(None);
    }

    match *header.ethertype() {
        EtherType::IPV4 => decode_ipv4(buffer, events),
        EtherType::IPV6 => decode_ipv6(buffer, events),
        // pass the packet on to the OS network stack if we don't understand it
        _ => Ok(None),
    }
}

#[inline(always)]
fn decode_ipv4<'a, D: Decoder<'a>, E: EventHandler>(buffer: D, events: &mut E) -> Result<D> {
    let (header, buffer) = buffer.decode::<&ipv4::Header>()?;

    if events.on_ipv4_header(header)?.is_none() {
        return Ok(None);
    }

    let protocol = header.protocol();

    //= https://www.rfc-editor.org/rfc/rfc791#section-3.1
    //# IHL:  4 bits
    //#
    //# Internet Header Length is the length of the internet header in 32
    //# bit words, and thus points to the beginning of the data.  Note that
    //# the minimum value for a correct header is 5.

    // subtract the fixed header size
    let count_without_header = header
        .vihl()
        .header_len()
        .checked_sub(5)
        .ok_or(DecoderError::InvariantViolation("invalid IPv4 IHL value"))?;

    // skip the options and go to the actual payload
    let options_len = count_without_header as usize * (32 / 8);
    let (_options, buffer) = buffer.decode_slice(options_len)?;

    parse_ip_protocol(protocol, buffer, events)
}

#[inline(always)]
fn decode_ipv6<'a, D: Decoder<'a>, E: EventHandler>(buffer: D, events: &mut E) -> Result<D> {
    let (header, buffer) = buffer.decode::<&ipv6::Header>()?;

    if events.on_ipv6_header(header)?.is_none() {
        return Ok(None);
    }

    let protocol = header.next_header();

    // TODO parse Hop-by-hop/Options headers, for now we'll just forward the packet on to the OS

    parse_ip_protocol(protocol, buffer, events)
}

#[inline]
fn parse_ip_protocol<'a, D: Decoder<'a>, E: EventHandler>(
    protocol: &ip::Protocol,
    buffer: D,
    events: &mut E,
) -> Result<D> {
    match *protocol {
        ip::Protocol::UDP | ip::Protocol::UDPLITE => parse_udp(buffer, events),
        // pass the packet on to the OS network stack if we don't understand it
        _ => Ok(None),
    }
}

#[inline(always)]
fn parse_udp<'a, D: Decoder<'a>, E: EventHandler>(buffer: D, events: &mut E) -> Result<D> {
    let (header, buffer) = buffer.decode::<&udp::Header>()?;

    if events.on_udp_header(header)?.is_none() {
        return Ok(None);
    }

    // NOTE: duvet doesn't know how to parse this RFC since it doesn't follow more modern formatting
    //# https://www.rfc-editor.org/rfc/rfc768
    //# 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.)
    let total_len = header.len().get();
    let payload_len = total_len
        .checked_sub(8)
        .ok_or(DecoderError::InvariantViolation("invalid UDP length"))?;
    let (udp_payload, _remaining) = buffer.decode_slice(payload_len as usize)?;

    Ok(Some(udp_payload))
}

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

    // Tests to ensure memory safety and no panics
    #[test]
    #[cfg_attr(kani, kani::proof, kani::unwind(258), kani::solver(cadical))]
    fn decode_test() {
        check!().for_each(|bytes| {
            let buffer = s2n_codec::DecoderBuffer::new(bytes);
            let _ = decode_packet(buffer);
        });
    }
}


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

use super::path;
use crate::{
    inet::{
        checksum::Checksum,
        ethernet::{self, EtherType},
        ip::{self, IpAddress},
        ipv4, ipv6, udp,
    },
    io::tx::{self, Message, PayloadBuffer},
};
use core::{hash::Hasher, mem::size_of};
use s2n_codec::{Encoder, EncoderBuffer};

/// The default TTL/Hop Limit for the packets
///
/// This value comes from the default value for Linux.
///
/// ```text
/// $ sudo sysctl net.ipv4.ip_default_ttl
/// net.ipv4.ip_default_ttl = 64
/// ```
const DEFAULT_TTL: u8 = 64;

pub struct State {
    ipv4_id_counter: u16,
    ipv4_checksum: bool,
    // stores a copy of Checksum so we don't have to probe the platform function every time
    cached_checksum: Checksum,
}

impl Default for State {
    fn default() -> Self {
        Self {
            ipv4_id_counter: 0,
            ipv4_checksum: true,
            cached_checksum: Default::default(),
        }
    }
}

impl State {
    #[inline]
    pub fn set_checksum(&mut self, enabled: bool) -> &mut Self {
        self.ipv4_checksum = enabled;
        self
    }

    #[inline]
    fn ipv4_id(&mut self) -> u16 {
        let id = self.ipv4_id_counter;
        self.ipv4_id_counter = self.ipv4_id_counter.wrapping_add(1);
        id
    }

    #[inline]
    fn ipv4_checksum(&self) -> Option<Checksum> {
        if self.ipv4_checksum {
            Some(self.cached_checksum)
        } else {
            None
        }
    }
}

#[inline]
pub fn encode_packet<M: Message<Handle = path::Tuple>>(
    buffer: &mut EncoderBuffer,
    message: &mut M,
    state: &mut State,
) -> Result<u16, tx::Error> {
    unsafe {
        assume!(
            buffer.remaining_capacity()
                > size_of::<ethernet::Header>()
                    + size_of::<ipv6::Header>().max(size_of::<ipv4::Header>())
                    + size_of::<udp::Header>(),
            "buffer too small"
        );
    }

    let path = message.path_handle();
    match (path.local_address.ip, path.remote_address.ip) {
        (IpAddress::Ipv4(local_ip), IpAddress::Ipv4(remote_ip)) => {
            buffer.encode(&ethernet::Header {
                destination: path.remote_address.mac,
                source: path.local_address.mac,
                ethertype: EtherType::IPV4,
            });

            encode_ipv4(buffer, local_ip, remote_ip, message, state)
        }
        (local_ip, remote_ip) => {
            buffer.encode(&ethernet::Header {
                destination: path.remote_address.mac,
                source: path.local_address.mac,
                ethertype: EtherType::IPV6,
            });

            // if either/both of the addresses are IPv6 then both need to be mapped
            let local_ip = local_ip.to_ipv6_mapped();
            let remote_ip = remote_ip.to_ipv6_mapped();

            encode_ipv6(buffer, local_ip, remote_ip, message, state)
        }
    }
}

#[inline]
fn encode_ipv4<M: Message<Handle = path::Tuple>>(
    buffer: &mut EncoderBuffer,
    local_ip: ipv4::IpV4Address,
    remote_ip: ipv4::IpV4Address,
    message: &mut M,
    state: &mut State,
) -> Result<u16, tx::Error> {
    const HEADER_LEN: u16 = (size_of::<ipv4::Header>() + size_of::<udp::Header>()) as _;

    let checksum = state.ipv4_checksum();

    let mut outcome = encode_payload(buffer, message, HEADER_LEN, checksum)?;

    buffer.write_zerocopy(|header: &mut ipv4::Header| {
        header.vihl_mut().set_version(4).set_header_len(5);
        header.tos_mut().set_dscp(0).set_ecn(message.ecn());
        header
            .flag_fragment_mut()
            .set_reserved(false)
            .set_dont_fragment(true)
            .set_more_fragments(false)
            .set_fragment_offset(0);
        header.id.set(state.ipv4_id());
        header.total_len_mut().set(HEADER_LEN + outcome.len);
        *header.ttl_mut() = DEFAULT_TTL;
        // set the checksum to zero for the initial pass
        header.checksum_mut().set(0);
        *header.protocol_mut() = ip::Protocol::UDP;
        *header.source_mut() = local_ip;
        *header.destination_mut() = remote_ip;

        // calculate the IPv4 header checksum
        header.update_checksum();

        // NOTE: duvet doesn't know how to parse this RFC since it doesn't follow more modern formatting
        //# https://www.rfc-editor.org/rfc/rfc768#Fields
        //# The pseudo  header  conceptually prefixed to the UDP header contains the
        //# source  address,  the destination  address,  the protocol,  and the  UDP
        //# length.   This information gives protection against misrouted datagrams.
        //# This checksum procedure is the same as is used in TCP.
        //#
        //#                  0      7 8     15 16    23 24    31
        //#                 +--------+--------+--------+--------+
        //#                 |          source address           |
        //#                 +--------+--------+--------+--------+
        //#                 |        destination address        |
        //#                 +--------+--------+--------+--------+
        //#                 |  zero  |protocol|   UDP length    |
        //#                 +--------+--------+--------+--------+
        if let Some(checksum) = outcome.checksum.as_mut() {
            // the addresses start at byte offset 12 in the header
            checksum.write(&header.as_bytes()[12..]);

            let payload_len = outcome.len + size_of::<udp::Header>() as u16;
            let payload_len = payload_len.to_be_bytes();

            let parts = [0, ip::Protocol::UDP.id, payload_len[0], payload_len[1]];

            checksum.write(&parts);
        }
    });

    encode_udp(buffer, outcome, message, state);

    Ok(outcome.len)
}

#[inline]
fn encode_ipv6<M: Message<Handle = path::Tuple>>(
    buffer: &mut EncoderBuffer,
    local_ip: ipv6::IpV6Address,
    remote_ip: ipv6::IpV6Address,
    message: &mut M,
    state: &mut State,
) -> Result<u16, tx::Error> {
    const HEADER_LEN: u16 = (size_of::<ipv6::Header>() + size_of::<udp::Header>()) as _;

    // Ipv6 checksums are required
    let checksum = Some(state.cached_checksum);

    let mut outcome = encode_payload(buffer, message, HEADER_LEN, checksum)?;

    buffer.write_zerocopy(|header: &mut ipv6::Header| {
        let payload_len = size_of::<udp::Header>() as u16 + outcome.len;

        header
            .vtcfl_mut()
            .set_version(6)
            .set_dscp(0)
            .set_ecn(message.ecn())
            .set_flow_label(message.ipv6_flow_label());
        header.payload_len_mut().set(payload_len);
        *header.next_header_mut() = ip::Protocol::UDP;
        *header.hop_limit_mut() = DEFAULT_TTL;
        *header.source_mut() = local_ip;
        *header.destination_mut() = remote_ip;

        //= https://www.rfc-editor.org/rfc/rfc2460#section-8.1
        //# Any transport or other upper-layer protocol that includes the
        //# addresses from the IP header in its checksum computation must be
        //# modified for use over IPv6, to include the 128-bit IPv6 addresses
        //# instead of 32-bit IPv4 addresses.  In particular, the following
        //# illustration shows the TCP and UDP "pseudo-header" for IPv6:
        //#
        //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //# |                                                               |
        //# +                                                               +
        //# |                                                               |
        //# +                         Source Address                        +
        //# |                                                               |
        //# +                                                               +
        //# |                                                               |
        //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //# |                                                               |
        //# +                                                               +
        //# |                                                               |
        //# +                      Destination Address                      +
        //# |                                                               |
        //# +                                                               +
        //# |                                                               |
        //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //# |                   Upper-Layer Packet Length                   |
        //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        //# |                      zero                     |  Next Header  |
        //# +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        if let Some(checksum) = outcome.checksum.as_mut() {
            // the addresses start at byte offset 8 in the header
            checksum.write(&header.as_bytes()[8..]);

            let mut parts = [0; 8];
            parts[..4].copy_from_slice(&(payload_len as u32).to_be_bytes());
            parts[7] = ip::Protocol::UDP.id;

            checksum.write(&parts);
        }
    });

    encode_udp(buffer, outcome, message, state);

    Ok(outcome.len)
}

#[inline]
fn encode_udp<M: Message<Handle = path::Tuple>>(
    buffer: &mut EncoderBuffer,
    outcome: PayloadOutcome,
    message: &mut M,
    _state: &mut State,
) {
    let path = message.path_handle();

    buffer.write_zerocopy(|header: &mut udp::Header| {
        header.source_mut().set(path.local_address.port);
        header.destination_mut().set(path.remote_address.port);
        // the length includes the UDP header
        header
            .len_mut()
            .set(size_of::<udp::Header>() as u16 + outcome.len);
        // initialize the checksum to 0
        header.checksum_mut().set(0);

        // write the checksum after we've written the header
        if let Some(mut checksum) = outcome.checksum {
            checksum.write(header.as_bytes());
            header.checksum_mut().set(checksum.finish());
        }
    });

    unsafe {
        assume!(
            buffer.remaining_capacity() >= outcome.len as usize,
            "buffer too small"
        );
    }

    // forward the buffer cursor to the end of the payload
    buffer.advance_position(outcome.len as _);
}

#[inline]
fn encode_payload<M: Message<Handle = path::Tuple>>(
    buffer: &mut EncoderBuffer,
    message: &mut M,
    header_size: u16,
    checksum: Option<Checksum>,
) -> Result<PayloadOutcome, tx::Error> {
    let header_position = buffer.len();
    buffer.advance_position(header_size as usize);

    let max_len = buffer
        .remaining_capacity()
        .min((u16::MAX - header_size) as usize);

    let mut outcome = PayloadOutcome { len: 0, checksum };

    unsafe {
        assume!(
            buffer.capacity() >= buffer.len(),
            "encoder cursors should be correct"
        );
    }
    let (_headers, payload) = buffer.split_mut();
    let payload = &mut payload[..max_len];
    {
        let payload = PayloadBuffer::new(payload);
        outcome.len = message.write_payload(payload, 0)? as u16;

        debug_assert!(outcome.len as usize <= max_len, "write exceeded max length");
    }

    if let Some(checksum) = outcome.checksum.as_mut() {
        unsafe {
            assume!(payload.len() >= outcome.len as usize);
        }
        checksum.write_padded(&payload[..outcome.len as usize]);
    }

    buffer.set_position(header_position);

    Ok(outcome)
}

#[derive(Clone, Copy, Debug, Default)]
struct PayloadOutcome {
    len: u16,
    checksum: Option<Checksum>,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{inet::ExplicitCongestionNotification, path::Handle};
    use bolero::{check, generator::*};
    use s2n_codec::DecoderBufferMut;

    #[derive(Debug, TypeGenerator)]
    pub struct Message {
        path: path::Tuple,
        ecn: ExplicitCongestionNotification,
        ipv4_id: u16,
        ipv4_checksum: bool,
        ipv6_flow_label: u32,
        payload: Vec<u8>,
    }

    impl<'a> tx::Message for &'a Message {
        type Handle = path::Tuple;

        fn path_handle(&self) -> &Self::Handle {
            &self.path
        }

        fn ecn(&mut self) -> ExplicitCongestionNotification {
            self.ecn
        }

        fn delay(&mut self) -> core::time::Duration {
            Default::default()
        }

        fn ipv6_flow_label(&mut self) -> u32 {
            self.ipv6_flow_label
        }

        fn can_gso(&self, _: usize, _: usize) -> bool {
            true
        }

        fn write_payload(
            &mut self,
            mut buffer: PayloadBuffer,
            _gso_offset: usize,
        ) -> Result<usize, tx::Error> {
            buffer.write(&self.payload)
        }
    }

    #[test]
    fn round_trip() {
        check!().with_type().for_each(|mut message: &Message| {
            let mut buffer = [0u8; 1500];
            let mut state = State {
                ipv4_id_counter: message.ipv4_id,
                ipv4_checksum: message.ipv4_checksum,
                cached_checksum: Checksum::default(),
            };

            let mut encoder = EncoderBuffer::new(&mut buffer);

            if encode_packet(&mut encoder, &mut message, &mut state).is_err() {
                return;
            }

            let (mut header, payload) =
                crate::xdp::decoder::decode_packet(DecoderBufferMut::new(&mut buffer))
                    .unwrap()
                    .unwrap();

            header.path.swap();

            assert!(Handle::eq(&header.path, &message.path));
            assert_eq!(header.ecn, message.ecn);
            assert_eq!(payload.into_less_safe_slice(), &message.payload);
        });
    }
}


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

use crate::{
    inet::{ethernet::MacAddress, ipv4, IpAddress},
    path::{self, Handle},
};

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

macro_rules! define_address {
    ($name:ident) => {
        #[derive(Clone, Copy, Debug, Eq, PartialEq)]
        #[cfg_attr(any(test, feature = "generator"), derive(TypeGenerator))]
        pub struct $name {
            pub mac: MacAddress,
            pub ip: IpAddress,
            pub port: u16,
        }

        impl $name {
            pub const UNSPECIFIED: Self = Self {
                mac: MacAddress::UNSPECIFIED,
                ip: IpAddress::Ipv4(ipv4::IpV4Address::UNSPECIFIED),
                port: 0,
            };

            #[inline]
            pub fn unmap(self) -> Self {
                Self {
                    mac: self.mac,
                    ip: self.ip.unmap(),
                    port: self.port,
                }
            }
        }

        impl From<path::$name> for $name {
            #[inline]
            fn from(addr: path::$name) -> Self {
                Self {
                    mac: MacAddress::UNSPECIFIED,
                    ip: addr.ip(),
                    port: addr.port(),
                }
            }
        }

        impl From<$name> for path::$name {
            #[inline]
            fn from(addr: $name) -> Self {
                addr.ip.with_port(addr.port).into()
            }
        }
    };
}

define_address!(RemoteAddress);
define_address!(LocalAddress);

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

impl Tuple {
    pub const UNSPECIFIED: Self = Self {
        remote_address: RemoteAddress::UNSPECIFIED,
        local_address: LocalAddress::UNSPECIFIED,
    };

    #[inline]
    pub fn swap(&mut self) {
        core::mem::swap(&mut self.remote_address.mac, &mut self.local_address.mac);
        core::mem::swap(&mut self.remote_address.ip, &mut self.local_address.ip);
        core::mem::swap(&mut self.remote_address.port, &mut self.local_address.port);
    }
}

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

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

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

    #[inline]
    fn eq(&self, other: &Self) -> bool {
        // TODO only compare everything if the other is all filled out
        PartialEq::eq(&self.local_address.unmap(), &other.local_address.unmap())
            && PartialEq::eq(&self.remote_address.unmap(), &other.remote_address.unmap())
    }

    #[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 full address
        if self.local_address.port == 0 {
            *self = *other;
        }
    }
}


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

pub use s2n_quic_core::crypto::CryptoError as Error;
pub type Result<T = (), E = Error> = core::result::Result<T, E>;

pub trait Aead {
    type Nonce;
    type Tag;

    fn encrypt(
        &self,
        nonce: &Self::Nonce,
        aad: &[u8],
        payload: &mut [u8],
        tag: &mut Self::Tag,
    ) -> Result;

    fn decrypt(
        &self,
        nonce: &Self::Nonce,
        aad: &[u8],
        payload: &mut [u8],
        tag: &Self::Tag,
    ) -> Result;
}


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

use crate::block::BatchMut;
use zeroize::Zeroize;

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub mod x86;

#[cfg(any(test, feature = "testing"))]
pub mod testing;

pub trait Encrypt {
    type Block;
    const KEY_LEN: usize;
    const ROUNDS: usize;

    /// Encrypts a batch of blocks with the AES key
    #[inline(always)]
    fn encrypt<B: BatchMut<Block = Self::Block>>(&self, block: &mut B) {
        self.encrypt_interleaved(block, |_| {})
    }

    /// Encrypts a batch of blocks with the AES key, while interleaving the provided function
    ///
    /// The provided function will be called for each AES round of the key. This can lead to
    /// improved performance if multiple operations are being done in parallel.
    fn encrypt_interleaved<B: BatchMut<Block = Self::Block>, F: FnMut(usize)>(
        &self,
        block: &mut B,
        f: F,
    );
}

pub trait Decrypt {
    type Block;
    const KEY_LEN: usize;
    const ROUNDS: usize;

    /// Decrypts a batch of blocks with the AES key
    #[inline(always)]
    fn decrypt<B: BatchMut<Block = Self::Block>>(&self, block: &mut B) {
        self.decrypt_interleaved(block, |_| {})
    }

    /// Decrypts a batch of blocks with the AES key, while interleaving the provided function
    ///
    /// The provided function will be called for each AES round of the key. This can lead to
    /// improved performance if multiple operations are being done in parallel.
    fn decrypt_interleaved<B: BatchMut<Block = Self::Block>, F: FnMut(usize)>(
        &self,
        block: &mut B,
        f: F,
    );
}

pub trait KeyRound {
    type Block;

    fn xor<B: BatchMut<Block = Self::Block>>(&self, block: &mut B);
    fn encrypt<B: BatchMut<Block = Self::Block>>(&self, block: &mut B);
    fn encrypt_finish<B: BatchMut<Block = Self::Block>>(&self, block: &mut B);
    fn decrypt<B: BatchMut<Block = Self::Block>>(&self, block: &mut B);
    fn decrypt_finish<B: BatchMut<Block = Self::Block>>(&self, block: &mut B);
}

#[cfg(any(test, feature = "testing"))]
pub const BLOCK_LEN: usize = 16;

pub mod aes128 {
    use super::*;

    pub const KEY_LEN: usize = 16;
    // https://github.com/awslabs/aws-lc/blob/aed75eb04d322d101941e1377f274484f5e4f5b8/crypto/fipsmodule/aes/asm/aesni-x86_64.pl#L4378
    // mov	\$9,$bits			# 10 rounds for 128-bit key
    pub const ROUNDS: usize = 10;

    #[derive(Zeroize)]
    pub struct Key<T>(pub T);

    impl<Blk, T> Encrypt for Key<T>
    where
        T: EncryptionKey<Block = Blk>,
    {
        type Block = Blk;
        const KEY_LEN: usize = KEY_LEN;
        const ROUNDS: usize = ROUNDS;

        #[inline(always)]
        fn encrypt_interleaved<B: BatchMut<Block = Blk>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            f: F,
        ) {
            aes128::EncryptionKey::encrypt_interleaved(&self.0, block, f)
        }
    }

    pub trait EncryptionKey {
        type Block;
        type KeyRound: KeyRound<Block = Self::Block>;

        fn keyround(&self, index: usize) -> &Self::KeyRound;

        #[inline(always)]
        fn encrypt_interleaved<B: BatchMut<Block = Self::Block>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            mut f: F,
        ) {
            // NOTE: instead of looping here, manually unroll the loop so the CPU has a large run of instructions
            //       without any branches.
            self.keyround(0).xor(block);
            self.keyround(1).encrypt(block);
            f(0);
            self.keyround(2).encrypt(block);
            f(1);
            self.keyround(3).encrypt(block);
            f(2);
            self.keyround(4).encrypt(block);
            f(3);
            self.keyround(5).encrypt(block);
            f(4);
            self.keyround(6).encrypt(block);
            f(5);
            self.keyround(7).encrypt(block);
            f(6);
            self.keyround(8).encrypt(block);
            f(7);
            self.keyround(9).encrypt(block);
            f(8);
            self.keyround(10).encrypt_finish(block);
        }
    }

    impl<Blk, T> Decrypt for Key<T>
    where
        T: DecryptionKey<Block = Blk>,
    {
        type Block = Blk;
        const KEY_LEN: usize = KEY_LEN;
        const ROUNDS: usize = ROUNDS;

        #[inline(always)]
        fn decrypt_interleaved<B: BatchMut<Block = Blk>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            f: F,
        ) {
            aes128::DecryptionKey::decrypt_interleaved(&self.0, block, f)
        }
    }

    pub trait DecryptionKey {
        type Block;
        type KeyRound: KeyRound<Block = Self::Block>;

        fn keyround(&self, index: usize) -> &Self::KeyRound;

        #[inline(always)]
        fn decrypt_interleaved<B: BatchMut<Block = Self::Block>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            mut f: F,
        ) {
            // NOTE: instead of looping here, manually unroll the loop so the CPU has a large run of instructions
            //       without any branches.
            self.keyround(10).xor(block);
            self.keyround(9).decrypt(block);
            f(0);
            self.keyround(8).decrypt(block);
            f(1);
            self.keyround(7).decrypt(block);
            f(2);
            self.keyround(6).decrypt(block);
            f(3);
            self.keyround(5).decrypt(block);
            f(4);
            self.keyround(4).decrypt(block);
            f(5);
            self.keyround(3).decrypt(block);
            f(6);
            self.keyround(2).decrypt(block);
            f(7);
            self.keyround(1).decrypt(block);
            f(8);
            self.keyround(0).decrypt_finish(block);
        }
    }
}

pub mod aes256 {
    use super::*;

    pub const KEY_LEN: usize = 32;
    // https://github.com/awslabs/aws-lc/blob/aed75eb04d322d101941e1377f274484f5e4f5b8/crypto/fipsmodule/aes/asm/aesni-x86_64.pl#L4548
    // mov	\$13,$bits			# 14 rounds for 256
    pub const ROUNDS: usize = 14;

    #[derive(Zeroize)]
    pub struct Key<T>(pub T);

    impl<Blk, T> Encrypt for Key<T>
    where
        T: EncryptionKey<Block = Blk>,
    {
        type Block = Blk;
        const KEY_LEN: usize = KEY_LEN;
        const ROUNDS: usize = ROUNDS;

        #[inline(always)]
        fn encrypt_interleaved<B: BatchMut<Block = Blk>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            f: F,
        ) {
            aes256::EncryptionKey::encrypt_interleaved(&self.0, block, f)
        }
    }

    pub trait EncryptionKey {
        type Block;
        type KeyRound: KeyRound<Block = Self::Block>;

        fn keyround(&self, index: usize) -> &Self::KeyRound;

        #[inline(always)]
        fn encrypt_interleaved<B: BatchMut<Block = Self::Block>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            mut f: F,
        ) {
            // NOTE: instead of looping here, manually unroll the loop so the CPU has a large run of instructions
            //       without any branches.
            self.keyround(0).xor(block);
            self.keyround(1).encrypt(block);
            f(0);
            self.keyround(2).encrypt(block);
            f(1);
            self.keyround(3).encrypt(block);
            f(2);
            self.keyround(4).encrypt(block);
            f(3);
            self.keyround(5).encrypt(block);
            f(4);
            self.keyround(6).encrypt(block);
            f(5);
            self.keyround(7).encrypt(block);
            f(6);
            self.keyround(8).encrypt(block);
            f(7);
            self.keyround(9).encrypt(block);
            f(8);
            self.keyround(10).encrypt(block);
            f(9);
            self.keyround(11).encrypt(block);
            f(10);
            self.keyround(12).encrypt(block);
            f(11);
            self.keyround(13).encrypt(block);
            f(12);
            self.keyround(14).encrypt_finish(block);
        }
    }

    impl<Blk, T> Decrypt for Key<T>
    where
        T: DecryptionKey<Block = Blk>,
    {
        type Block = Blk;
        const KEY_LEN: usize = KEY_LEN;
        const ROUNDS: usize = ROUNDS;

        #[inline(always)]
        fn decrypt_interleaved<B: BatchMut<Block = Blk>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            f: F,
        ) {
            aes256::DecryptionKey::decrypt_interleaved(&self.0, block, f)
        }
    }

    pub trait DecryptionKey {
        type Block;
        type KeyRound: KeyRound<Block = Self::Block>;

        fn keyround(&self, index: usize) -> &Self::KeyRound;

        #[inline(always)]
        fn decrypt_interleaved<B: BatchMut<Block = Self::Block>, F: FnMut(usize)>(
            &self,
            block: &mut B,
            mut f: F,
        ) {
            // NOTE: instead of looping here, manually unroll the loop so the CPU has a large run of instructions
            //       without any branches.
            self.keyround(14).xor(block);
            self.keyround(13).decrypt(block);
            f(0);
            self.keyround(12).decrypt(block);
            f(1);
            self.keyround(11).decrypt(block);
            f(2);
            self.keyround(10).decrypt(block);
            f(3);
            self.keyround(9).decrypt(block);
            f(4);
            self.keyround(8).decrypt(block);
            f(5);
            self.keyround(7).decrypt(block);
            f(6);
            self.keyround(6).decrypt(block);
            f(7);
            self.keyround(5).decrypt(block);
            f(8);
            self.keyround(4).decrypt(block);
            f(9);
            self.keyround(3).decrypt(block);
            f(10);
            self.keyround(2).decrypt(block);
            f(11);
            self.keyround(1).decrypt(block);
            f(12);
            self.keyround(0).decrypt_finish(block);
        }
    }
}

#[cfg(test)]
mod tests;


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

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
mod x86;

mod ring;

pub mod generic;
pub mod payload;

#[cfg(any(test, feature = "testing"))]
pub mod testing;

pub const NONCE_LEN: usize = 12;

pub use crate::ghash::TAG_LEN;

#[cfg(test)]
mod tests;


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

use cfg_if::cfg_if;

cfg_if! {
    if #[cfg(target_arch = "x86")] {
        pub use core::arch::x86::*;
        mod x86;
        pub use x86::*;
    } else if #[cfg(target_arch = "x86_64")] {
        pub use core::arch::x86_64::*;
        mod x86;
        pub use x86::*;
    }
}

pub trait Arch {
    fn is_supported() -> bool;

    unsafe fn call<F: FnOnce() -> R, R>(f: F) -> R;

    #[inline]
    fn call_supported<F: FnOnce()>(f: F) {
        if Self::is_supported() {
            unsafe {
                // Safety: this is only called if the arch is supported
                Self::call(f)
            }
        }
    }
}


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

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub mod x86;

pub const LEN: usize = 16;

pub trait Block: Copy + Zeroed {
    fn from_array(array: [u8; LEN]) -> Self;
    fn into_array(self) -> [u8; LEN];
    fn xor(self, other: Self) -> Self;
    fn ct_ensure_eq(self, b: Self) -> Result<(), ()>;
}

pub trait Batch {
    type Block: Block;

    fn for_each<F: FnMut(usize, &Self::Block)>(&self, f: F);
}

pub trait BatchMut: Batch {
    fn update<F: FnMut(usize, &mut Self::Block)>(&mut self, f: F);
}

pub trait Zeroed {
    fn zeroed() -> Self;
}

macro_rules! impl_array {
    ($n:expr, [$($idx:expr),*]) => {
        impl<B: Block> Batch for [B; $n] {
            type Block = B;

            #[inline(always)]
            fn for_each<F: FnMut(usize, &Self::Block)>(&self, mut f: F) {
                $(
                    f($idx, &self[$idx]);
                )*
            }
        }

        impl<B: Block> BatchMut for [B; $n] {
            fn update<F: FnMut(usize, &mut Self::Block)>(&mut self, mut f: F) {
                $(
                    f($idx, &mut self[$idx]);
                )*
            }
        }

        impl<B: Zeroed> Zeroed for [B; $n] {
            #[inline(always)]
            fn zeroed() -> Self {
                [
                    $({
                        let _ = $idx;
                        B::zeroed()
                    }),*
                ]
            }
        }
    };
}

impl_array!(1, [0]);
impl_array!(2, [0, 1]);
impl_array!(3, [0, 1, 2]);
impl_array!(4, [0, 1, 2, 3]);
impl_array!(5, [0, 1, 2, 3, 4]);
impl_array!(6, [0, 1, 2, 3, 4, 5]);
impl_array!(7, [0, 1, 2, 3, 4, 5, 6]);
impl_array!(8, [0, 1, 2, 3, 4, 5, 6, 7]);


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

use crate::{aead::Aead, header_key::HeaderKey, hkdf, iv, ring_aead as aead};
use core::fmt;
use s2n_quic_core::{
    assume,
    crypto::{label, CryptoError},
};
use zeroize::{Zeroize, Zeroizing};

mod negotiated;
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[path = "cipher_suite/x86.rs"]
mod platform;
mod ring;

#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
use self::ring as platform;

pub use negotiated::NegotiatedCipherSuite;

macro_rules! impl_cipher_suite {
    (
        $name:ident,
        $lower:ident,
        $digest:path,
        $cipher:path,
        $cipher_key_len:expr,
        $header_protection:path,
        $key_label:expr,
        $iv_label:expr,
        $hp_label:expr,
        $key_update_label:expr,
        $confidentiality_limit:expr,
        $integrity_limit:expr,
        $test_name:ident
    ) => {
        mod $lower {
            use super::*;

            pub const KEY_LEN: usize = $cipher_key_len;
            pub const TAG_LEN: usize = 16;
            pub const NONCE_LEN: usize = crate::aesgcm::NONCE_LEN;

            type Key = platform::$lower::Key;

            // ignore casing warnings in order to preserve the IANA name
            #[allow(non_camel_case_types, clippy::all)]
            pub struct $name {
                secret: hkdf::Prk,
                iv: iv::Iv,
                key: Key,
            }

            impl $name {
                pub fn new(secret: hkdf::Prk) -> (Self, HeaderKey) {
                    let iv = Self::new_iv(&secret);
                    let key = {
                        let secret = Self::new_key_secret(&secret);
                        Key::new(&*secret)
                    };
                    let header_key = Self::new_header_key(&secret);

                    let key = Self { secret, iv, key };

                    (key, header_key)
                }

                /// Update the cipher_suite as defined in
                /// https://www.rfc-editor.org/rfc/rfc9001#section-6
                #[inline]
                pub fn update(&self) -> Self {
                    let secret: hkdf::Prk = self
                        .secret
                        .expand(&[&$key_update_label], $digest)
                        .expect("label size verified")
                        .into();

                    let iv = Self::new_iv(&secret);
                    let key = {
                        let key = Self::new_key_secret(&secret);
                        // ask the existing key to derive the next one so it can persist any
                        // configuration
                        self.key.update(&*key)
                    };
                    Self { secret, iv, key }
                }

                #[inline]
                pub fn update_pmtu(&mut self, mtu: u16) {
                    if self.key.should_update_pmtu(mtu) {
                        let secret = Self::new_key_secret(&self.secret);
                        self.key.update_pmtu(&*secret, mtu);
                    }
                }

                fn new_key_secret(secret: &hkdf::Prk) -> Zeroizing<[u8; KEY_LEN]> {
                    let mut key = Zeroizing::new([0u8; KEY_LEN]);

                    secret
                        .expand(&[&$key_label], &$cipher)
                        .expect("label size verified")
                        .fill(&mut key.as_mut())
                        .expect("fill size verified");

                    key
                }

                fn new_iv(secret: &hkdf::Prk) -> iv::Iv {
                    iv::Iv::new(secret, &$iv_label)
                }

                fn new_header_key(secret: &hkdf::Prk) -> HeaderKey {
                    HeaderKey::new::<{ KEY_LEN }>(secret, &$hp_label, &$header_protection)
                }
            }

            impl Zeroize for $name {
                fn zeroize(&mut self) {
                    self.iv.zeroize();
                    self.key.zeroize();
                }
            }

            impl s2n_quic_core::crypto::Key for $name {
                #[inline]
                fn decrypt(
                    &self,
                    packet_number: u64,
                    header: &[u8],
                    payload: &mut [u8],
                ) -> Result<(), CryptoError> {
                    let nonce = self.iv.nonce(packet_number);

                    let payload_len = payload
                        .len()
                        .checked_sub(TAG_LEN)
                        .ok_or_else(|| CryptoError::DECRYPT_ERROR)?;

                    let (payload, tag) = payload.split_at_mut(payload_len);
                    let tag = {
                        use core::convert::TryInto;
                        let res = (&tag[..]).try_into();
                        unsafe {
                            assume!(res.is_ok());
                        }
                        res.unwrap()
                    };

                    self.key.decrypt(&nonce, header, payload, tag)?;

                    Ok(())
                }

                #[inline]
                fn encrypt(
                    &self,
                    packet_number: u64,
                    header: &[u8],
                    payload: &mut [u8],
                ) -> Result<(), CryptoError> {
                    let nonce = self.iv.nonce(packet_number);

                    let payload_len = payload
                        .len()
                        .checked_sub(TAG_LEN)
                        .ok_or_else(|| CryptoError::DECRYPT_ERROR)?;

                    let (payload, tag) = payload.split_at_mut(payload_len);
                    let tag = {
                        use core::convert::TryInto;
                        let res = tag.try_into();
                        unsafe {
                            assume!(res.is_ok());
                        }
                        res.unwrap()
                    };

                    self.key.encrypt(&nonce, header, payload, tag)?;

                    Ok(())
                }

                #[inline]
                fn tag_len(&self) -> usize {
                    $cipher.tag_len()
                }

                //= https://www.rfc-editor.org/rfc/rfc9001#section-6.6
                //# Any TLS cipher suite that is specified for use with QUIC MUST define
                //# limits on the use of the associated AEAD function that preserves
                //# margins for confidentiality and integrity.
                #[inline]
                fn aead_confidentiality_limit(&self) -> u64 {
                    $confidentiality_limit
                }

                #[inline]
                fn aead_integrity_limit(&self) -> u64 {
                    $integrity_limit
                }

                #[inline]
                fn cipher_suite(&self) -> s2n_quic_core::crypto::tls::CipherSuite {
                    s2n_quic_core::crypto::tls::CipherSuite::$name
                }
            }

            impl fmt::Debug for $name {
                fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                    f.debug_struct(stringify!($name)).finish()
                }
            }

            impl Drop for $name {
                fn drop(&mut self) {
                    self.zeroize();
                }
            }

            #[test]
            fn $test_name() {
                fn compute_vec_label(len: usize, label: &[u8]) -> Vec<u8> {
                    let mut out = vec![];
                    label::compute_label(len, label, &mut out);
                    out
                }

                insta::assert_debug_snapshot!(
                    concat!("integrity_", stringify!($test_name)),
                    $integrity_limit
                );
                insta::assert_debug_snapshot!(
                    concat!("confidentiality_", stringify!($test_name)),
                    $confidentiality_limit
                );

                assert_eq!(KEY_LEN, $cipher.key_len(), "key len mismatch");

                assert_eq!(
                    compute_vec_label($cipher.key_len(), b"quic key"),
                    $key_label,
                    "key label mismatch"
                );

                assert_eq!(
                    compute_vec_label(iv::NONCE_LEN, b"quic iv"),
                    $iv_label,
                    "iv label mismatch"
                );

                assert_eq!(
                    compute_vec_label($header_protection.key_len(), b"quic hp"),
                    $hp_label,
                    "hp label mismatch"
                );

                assert_eq!(
                    compute_vec_label(
                        $digest.hmac_algorithm().digest_algorithm().output_len,
                        b"quic ku"
                    ),
                    $key_update_label,
                    "key update label mismatch"
                );
            }
        }

        pub use $lower::$name;
    };
}

//= https://www.rfc-editor.org/rfc/rfc9001#section-6.6
//# For AEAD_AES_128_GCM and AEAD_AES_256_GCM, the confidentiality limit
//# is 2^23 encrypted packets; see Appendix B.1.
impl_cipher_suite!(
    TLS_AES_256_GCM_SHA384,
    aes256_gcm,
    hkdf::HKDF_SHA384,
    aead::AES_256_GCM,
    256 / 8, // 256-bit key
    aead::quic::AES_256,
    label::QUIC_KEY_32,
    label::QUIC_IV_12,
    label::QUIC_HP_32,
    label::QUIC_KU_48,
    u64::pow(2, 23), // Confidentiality limit
    u64::pow(2, 52), //Â Integrity limit
    tls_aes_256_gcm_sha384_test
);

//= https://www.rfc-editor.org/rfc/rfc9001#section-6.6
//# For
//# AEAD_CHACHA20_POLY1305, the confidentiality limit is greater than the
//# number of possible packets (2^62) and so can be disregarded.
impl_cipher_suite!(
    TLS_CHACHA20_POLY1305_SHA256,
    chacha20_poly1305,
    hkdf::HKDF_SHA256,
    aead::CHACHA20_POLY1305,
    256 / 8, // 256-bit key
    aead::quic::CHACHA20,
    label::QUIC_KEY_32,
    label::QUIC_IV_12,
    label::QUIC_HP_32,
    label::QUIC_KU_32,
    u64::pow(2, 62), // Confidentiality limit even though specification notes it can be disregarded
    u64::pow(2, 36), // Integrity limit
    tls_chacha20_poly1305_sha256_test
);

// See above annotation regarding AEAD_AES_128 and AEAD_AES_256
impl_cipher_suite!(
    TLS_AES_128_GCM_SHA256,
    aes128_gcm,
    hkdf::HKDF_SHA256,
    aead::AES_128_GCM,
    128 / 8, // 128-bit key
    aead::quic::AES_128,
    label::QUIC_KEY_16,
    label::QUIC_IV_12,
    label::QUIC_HP_16,
    label::QUIC_KU_32,
    u64::pow(2, 23), // Confidentiality limit
    u64::pow(2, 52), // Integrity limit
    tls_aes_128_gcm_sha256_test
);


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

use crate::{aesgcm::NONCE_LEN, block::Block};

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub mod x86;

pub trait Ctr {
    type Block: Block;

    fn new(nonce: &[u8; NONCE_LEN]) -> Self;
    fn block(&self) -> Self::Block;
    fn increment(&mut self);
}


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

use crate::block::{Batch, Block};

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
pub mod x86;

#[cfg(any(test, feature = "testing"))]
pub mod testing;

pub const TAG_LEN: usize = 16;
#[cfg(any(target_arch = "x86", target_arch = "x86_64", test))]
pub const KEY_LEN: usize = 16;

pub trait Constructor {
    type GHash: GHash;

    #[cfg(any(target_arch = "x86", target_arch = "x86_64", test))]
    fn create(&self, key: [u8; KEY_LEN]) -> Self::GHash;
}

pub trait GHash {
    type State;
    type Block: Block;

    fn start(&self, required_blocks: usize) -> Self::State;
    fn update<B: Batch<Block = Self::Block>>(&self, state: &mut Self::State, block: &B);
    fn finish(&self, state: Self::State) -> Self::Block;
}

#[cfg(test)]
mod tests;


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

use s2n_quic_core::crypto;

header_key!(HandshakeHeaderKey);
negotiated_crypto!(HandshakeKey, HandshakeHeaderKey);

impl crypto::HandshakeKey for HandshakeKey {}

impl crypto::HandshakeHeaderKey for HandshakeHeaderKey {}


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

use crate::{hkdf, ring_aead as aead};
use core::fmt;
use s2n_quic_core::crypto::{self, HeaderProtectionMask};

pub struct HeaderKey(pub(crate) aead::quic::HeaderProtectionKey);

impl crypto::HeaderKey for HeaderKey {
    #[inline]
    fn opening_header_protection_mask(&self, sample: &[u8]) -> HeaderProtectionMask {
        self.header_protection_mask(sample)
    }

    #[inline]
    fn opening_sample_len(&self) -> usize {
        self.0.algorithm().sample_len()
    }

    #[inline]
    fn sealing_header_protection_mask(&self, sample: &[u8]) -> HeaderProtectionMask {
        self.header_protection_mask(sample)
    }

    #[inline]
    fn sealing_sample_len(&self) -> usize {
        self.0.algorithm().sample_len()
    }
}

impl HeaderKey {
    pub fn new<const KEY_LEN: usize>(
        secret: &hkdf::Prk,
        label: &[u8],
        alg: &'static aead::quic::Algorithm,
    ) -> Self {
        let mut bytes = zeroize::Zeroizing::new([0u8; KEY_LEN]);

        secret
            .expand(&[label], alg)
            .expect("label size verified")
            .fill(bytes.as_mut())
            .expect("fill size verified");

        let key = aead::quic::HeaderProtectionKey::new(alg, bytes.as_ref())
            .expect("header secret length already checked");
        Self(key)
    }

    #[inline]
    fn header_protection_mask(&self, sample: &[u8]) -> HeaderProtectionMask {
        self.0
            .new_mask(sample)
            .expect("sample length already checked")
    }
}

impl fmt::Debug for HeaderKey {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("HeaderKey").finish()
    }
}

impl From<aead::quic::HeaderProtectionKey> for HeaderKey {
    fn from(key: aead::quic::HeaderProtectionKey) -> Self {
        Self(key)
    }
}

#[derive(Debug)]
pub struct HeaderKeyPair {
    pub(crate) sealer: HeaderKey,
    pub(crate) opener: HeaderKey,
}

impl crypto::HeaderKey for HeaderKeyPair {
    #[inline]
    fn opening_header_protection_mask(&self, sample: &[u8]) -> HeaderProtectionMask {
        self.opener.opening_header_protection_mask(sample)
    }

    #[inline]
    fn opening_sample_len(&self) -> usize {
        self.opener.opening_sample_len()
    }

    #[inline]
    fn sealing_header_protection_mask(&self, sample: &[u8]) -> HeaderProtectionMask {
        self.sealer.sealing_header_protection_mask(sample)
    }

    #[inline]
    fn sealing_sample_len(&self) -> usize {
        self.sealer.sealing_sample_len()
    }
}

macro_rules! header_key {
    ($name:ident) => {
        #[derive(Debug)]
        pub struct $name(crate::header_key::HeaderKeyPair);

        impl s2n_quic_core::crypto::HeaderKey for $name {
            #[inline]
            fn opening_header_protection_mask(
                &self,
                sample: &[u8],
            ) -> s2n_quic_core::crypto::HeaderProtectionMask {
                self.0.opening_header_protection_mask(sample)
            }

            #[inline]
            fn opening_sample_len(&self) -> usize {
                self.0.opening_sample_len()
            }

            #[inline]
            fn sealing_header_protection_mask(
                &self,
                sample: &[u8],
            ) -> s2n_quic_core::crypto::HeaderProtectionMask {
                self.0.sealing_header_protection_mask(sample)
            }

            #[inline]
            fn sealing_sample_len(&self) -> usize {
                self.0.sealing_sample_len()
            }
        }

        impl From<crate::header_key::HeaderKeyPair> for $name {
            fn from(key: crate::header_key::HeaderKeyPair) -> Self {
                Self(key)
            }
        }
    };
}


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

use crate::{cipher_suite::TLS_AES_128_GCM_SHA256 as CipherSuite, header_key::HeaderKeyPair, hkdf};
use s2n_quic_core::{
    crypto::{
        self,
        label::{CLIENT_IN, SERVER_IN},
        CryptoError, Key, INITIAL_SALT,
    },
    endpoint,
};

header_key!(InitialHeaderKey);

impl crypto::InitialHeaderKey for InitialHeaderKey {}

#[derive(Debug)]
pub struct InitialKey {
    sealer: CipherSuite,
    opener: CipherSuite,
}

lazy_static::lazy_static! {
    /// Compute the Initial salt once, as the seed is constant
    static ref INITIAL_SIGNING_KEY: hkdf::Salt = hkdf::Salt::new(hkdf::HKDF_SHA256, &INITIAL_SALT);
}

impl InitialKey {
    fn new(endpoint: endpoint::Type, connection_id: &[u8]) -> (Self, InitialHeaderKey) {
        let initial_secret = INITIAL_SIGNING_KEY.extract(connection_id);
        let digest = INITIAL_SIGNING_KEY.algorithm();

        let client_secret = initial_secret
            .expand(&[&CLIENT_IN], digest)
            .expect("label size verified")
            .into();

        let server_secret = initial_secret
            .expand(&[&SERVER_IN], digest)
            .expect("label size verified")
            .into();

        let (sealer, opener) = match endpoint {
            endpoint::Type::Client => (
                CipherSuite::new(client_secret),
                CipherSuite::new(server_secret),
            ),
            endpoint::Type::Server => (
                CipherSuite::new(server_secret),
                CipherSuite::new(client_secret),
            ),
        };

        let (key_sealer, header_sealer) = sealer;
        let (key_opener, header_opener) = opener;
        let key = Self {
            sealer: key_sealer,
            opener: key_opener,
        };
        let header_key = InitialHeaderKey(HeaderKeyPair {
            sealer: header_sealer,
            opener: header_opener,
        });

        (key, header_key)
    }
}

impl crypto::InitialKey for InitialKey {
    type HeaderKey = InitialHeaderKey;

    fn new_server(connection_id: &[u8]) -> (Self, Self::HeaderKey) {
        Self::new(endpoint::Type::Server, connection_id)
    }

    fn new_client(connection_id: &[u8]) -> (Self, Self::HeaderKey) {
        Self::new(endpoint::Type::Client, connection_id)
    }
}

impl Key for InitialKey {
    #[inline]
    fn decrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        self.opener.decrypt(packet_number, header, payload)
    }

    #[inline]
    fn encrypt(
        &self,
        packet_number: u64,
        header: &[u8],
        payload: &mut [u8],
    ) -> Result<(), CryptoError> {
        self.sealer.encrypt(packet_number, header, payload)
    }

    #[inline]
    fn tag_len(&self) -> usize {
        self.sealer.tag_len()
    }

    #[inline]
    fn aead_confidentiality_limit(&self) -> u64 {
        self.sealer.aead_confidentiality_limit()
    }

    #[inline]
    fn aead_integrity_limit(&self) -> u64 {
        self.opener.aead_integrity_limit()
    }

    #[inline]
    fn cipher_suite(&self) -> s2n_quic_core::crypto::tls::CipherSuite {
        self.opener.cipher_suite()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use s2n_codec::{DecoderBufferMut, EncoderBuffer};
    use s2n_quic_core::{
        connection::id::ConnectionInfo,
        crypto::{
            initial::{
                EXAMPLE_CLIENT_INITIAL_PAYLOAD, EXAMPLE_CLIENT_INITIAL_PROTECTED_PACKET,
                EXAMPLE_DCID, EXAMPLE_SERVER_INITIAL_PAYLOAD,
                EXAMPLE_SERVER_INITIAL_PROTECTED_PACKET,
            },
            InitialKey as _,
        },
        inet::SocketAddress,
        packet::{encoding::PacketEncoder, initial::CleartextInitial, ProtectedPacket},
    };

    #[test]
    fn rfc_example_server_test() {
        test_round_trip(
            &InitialKey::new_client(&EXAMPLE_DCID),
            &InitialKey::new_server(&EXAMPLE_DCID),
            &EXAMPLE_CLIENT_INITIAL_PROTECTED_PACKET,
            &EXAMPLE_CLIENT_INITIAL_PAYLOAD,
        );
    }

    #[test]
    fn rfc_example_client_test() {
        test_round_trip(
            &InitialKey::new_server(&EXAMPLE_DCID),
            &InitialKey::new_client(&EXAMPLE_DCID),
            &EXAMPLE_SERVER_INITIAL_PROTECTED_PACKET,
            &EXAMPLE_SERVER_INITIAL_PAYLOAD,
        );
    }

    fn test_round_trip(
        sealer: &(InitialKey, InitialHeaderKey),
        opener: &(InitialKey, InitialHeaderKey),
        protected_packet: &[u8],
        cleartext_payload: &[u8],
    ) {
        let (sealer_key, sealer_header_key) = sealer;
        let (opener_key, opener_header_key) = opener;
        let (version, dcid, scid, token, sealed_packet) = decrypt(
            opener_key,
            opener_header_key,
            protected_packet.to_vec(),
            cleartext_payload,
            |packet| {
                let version = packet.version;
                let dcid = packet.destination_connection_id.to_vec();
                let scid = packet.source_connection_id.to_vec();
                let token = packet.token.to_vec();

                let mut output_buffer = vec![0; protected_packet.len()];
                packet
                    .encode_packet(
                        sealer_key,
                        sealer_header_key,
                        Default::default(),
                        None,
                        EncoderBuffer::new(&mut output_buffer),
                    )
                    .unwrap();

                (version, dcid, scid, token, output_buffer)
            },
        );

        // We have to decrypt instead of assert_eq on the sealed_packet
        // because of potential encoding differences. But the resulting
        // packets should be equal
        decrypt(
            opener_key,
            opener_header_key,
            sealed_packet,
            cleartext_payload,
            |packet| {
                assert_eq!(packet.version, version);
                assert_eq!(packet.destination_connection_id, &dcid[..]);
                assert_eq!(packet.source_connection_id, &scid[..]);
                assert_eq!(packet.token, &token[..]);
            },
        );
    }

    fn decrypt<F: FnOnce(CleartextInitial) -> O, O>(
        opener_key: &InitialKey,
        opener_header_key: &InitialHeaderKey,
        mut protected_packet: Vec<u8>,
        cleartext_payload: &[u8],
        on_decrypt: F,
    ) -> O {
        let decoder = DecoderBufferMut::new(&mut protected_packet);
        let remote_address = SocketAddress::default();
        let connection_info = ConnectionInfo::new(&remote_address);
        let (packet, _) = ProtectedPacket::decode(decoder, &connection_info, &20).unwrap();

        let packet = match packet {
            ProtectedPacket::Initial(initial) => initial,
            _ => panic!("expected initial packet type"),
        };

        let packet = packet
            .unprotect(opener_header_key, Default::default())
            .unwrap();
        let packet = packet.decrypt(opener_key).unwrap();

        // trim any padding off of the end
        let actual_payload = &packet.payload.as_less_safe_slice()[..cleartext_payload.len()];
        assert_eq!(cleartext_payload, actual_payload);

        on_decrypt(packet)
    }
}


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

use crate::hkdf;
use s2n_codec::{Encoder, EncoderBuffer};
use zeroize::Zeroize;

pub use crate::ring_aead::NONCE_LEN;

pub struct Iv([u8; NONCE_LEN]);

impl Iv {
    #[inline]
    pub fn new(secret: &hkdf::Prk, label: &[u8]) -> Self {
        let mut bytes = [0u8; NONCE_LEN];

        secret
            .expand(&[label], IvLen)
            .expect("label size verified")
            .fill(&mut bytes)
            .expect("fill size verified");

        Self(bytes)
    }

    #[inline]
    pub fn nonce(&self, packet_number: u64) -> [u8; NONCE_LEN] {
        let mut nonce = [0; NONCE_LEN];
        let mut encoder = EncoderBuffer::new(&mut nonce);

        encoder.encode(&0u32);
        encoder.encode(&packet_number);

        for (a, b) in nonce.iter_mut().zip(self.0.iter()) {
            *a ^= b;
        }

        nonce
    }
}

impl Zeroize for Iv {
    #[inline]
    fn zeroize(&mut self) {
        // deref to a slice to we can take advantage of the bulk zeroization
        self.0.zeroize()
    }
}

struct IvLen;

impl hkdf::KeyType for IvLen {
    #[inline]
    fn len(&self) -> usize {
        NONCE_LEN
    }
}


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

#[macro_use]
mod negotiated;
#[macro_use]
mod header_key;

mod aead;
mod aes;
mod aesgcm;
mod arch;
mod block;
mod cipher_suite;
mod ctr;
mod ghash;
mod iv;

#[cfg(target_os = "linux")]
use aws_lc_rs as ring;

#[doc(hidden)]
pub use ring::{
    aead as ring_aead,
    aead::{Algorithm, MAX_TAG_LEN},
    constant_time, digest, hkdf,
    hkdf::Prk,
    hmac,
};

#[derive(Clone)]
pub struct SecretPair {
    pub server: Prk,
    pub client: Prk,
}

pub mod handshake;
pub mod initial;
pub mod one_rtt;
pub mod retry;
pub mod zero_rtt;

#[derive(Clone, Copy, Debug, Default)]
pub struct Suite;

impl s2n_quic_core::crypto::CryptoSuite for Suite {
    type HandshakeKey = handshake::HandshakeKey;
    type HandshakeHeaderKey = handshake::HandshakeHeaderKey;
    type InitialKey = initial::InitialKey;
    type InitialHeaderKey = initial::InitialHeaderKey;
    type OneRttKey = one_rtt::OneRttKey;
    type OneRttHeaderKey = one_rtt::OneRttHeaderKey;
    type ZeroRttKey = zero_rtt::ZeroRttKey;
    type ZeroRttHeaderKey = zero_rtt::ZeroRttHeaderKey;
    type RetryKey = retry::RetryKey;
}

#[cfg(any(test, feature = "testing"))]
pub mod testing;