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

use crate::{
    decoder::{
        buffer::{DecoderBuffer, DecoderBufferResult},
        buffer_mut::{DecoderBufferMut, DecoderBufferMutResult},
    },
    unaligned::{i24, i48, u24, u48},
};
use byteorder::{ByteOrder, NetworkEndian};
use core::mem::size_of;

pub trait DecoderValue<'a>: Sized {
    fn decode(bytes: DecoderBuffer<'a>) -> DecoderBufferResult<'a, Self>;
}

pub trait DecoderValueMut<'a>: Sized {
    fn decode_mut(bytes: DecoderBufferMut<'a>) -> DecoderBufferMutResult<Self>;
}

#[macro_export]
macro_rules! decoder_value {
    (impl<$lt:lifetime $(, $generic:ident)*> $ty:ty {
        fn decode($buffer:ident: Buffer) -> Result<$ret:ty> $impl:block
    }) => {
        impl<$lt $(, $generic: $crate::DecoderValue<$lt>)*> $crate::DecoderValue<$lt> for $ty {
            #[inline]
            fn decode($buffer: $crate::DecoderBuffer<$lt>) -> $crate::DecoderBufferResult<$lt, $ret> $impl
        }

        impl<$lt $(, $generic: $crate::DecoderValueMut<$lt>)*> $crate::DecoderValueMut<$lt> for $ty {
            #[inline]
            fn decode_mut($buffer: $crate::DecoderBufferMut<$lt>) -> $crate::DecoderBufferMutResult<$lt, $ret> $impl
        }
    };
}

macro_rules! decoder_value_byte {
    ($ty:ident) => {
        decoder_value!(
            impl<'a> $ty {
                fn decode(buffer: Buffer) -> Result<Self> {
                    let (value, buffer) = buffer.decode_slice(size_of::<Self>())?;
                    let value = value.as_less_safe_slice()[0] as $ty;
                    Ok((value, buffer))
                }
            }
        );
    };
}

decoder_value_byte!(u8);
decoder_value_byte!(i8);

macro_rules! decoder_value_network_endian {
    ($call:ident, $ty:ty) => {
        decoder_value!(
            impl<'a> $ty {
                fn decode(buffer: Buffer) -> Result<Self> {
                    let (value, buffer) = buffer.decode_slice(size_of::<Self>())?;
                    let value = value.as_less_safe_slice();
                    let value = NetworkEndian::$call(value);
                    Ok((value.into(), buffer))
                }
            }
        );
    };
}

decoder_value_network_endian!(read_u16, u16);
decoder_value_network_endian!(read_i16, i16);
decoder_value_network_endian!(read_u32, u32);
decoder_value_network_endian!(read_i32, i32);
decoder_value_network_endian!(read_u64, u64);
decoder_value_network_endian!(read_i64, i64);
decoder_value_network_endian!(read_u128, u128);
decoder_value_network_endian!(read_i128, i128);
decoder_value_network_endian!(read_f32, f32);
decoder_value_network_endian!(read_f64, f64);

macro_rules! decoder_value_unaligned_integer {
    ($call:ident, $ty:ident, $bitsize:expr) => {
        decoder_value!(
            impl<'a> $ty {
                fn decode(buffer: Buffer) -> Result<Self> {
                    let (value, buffer) = buffer.decode_slice($bitsize / 8)?;
                    let value = value.as_less_safe_slice();
                    let value = NetworkEndian::$call(value);
                    Ok(($ty(value), buffer))
                }
            }
        );
    };
}

decoder_value_unaligned_integer!(read_u24, u24, 24);
decoder_value_unaligned_integer!(read_i24, i24, 24);
decoder_value_unaligned_integer!(read_u48, u48, 48);
decoder_value_unaligned_integer!(read_i48, i48, 48);

decoder_value!(
    impl<'a> DecoderBuffer<'a> {
        fn decode(buffer: Buffer) -> Result<Self> {
            let len = buffer.len();
            let (slice, buffer) = buffer.decode_slice(len)?;
            #[allow(clippy::useless_conversion)]
            let slice = slice.into();
            Ok((slice, buffer))
        }
    }
);

decoder_value!(
    impl<'a> () {
        fn decode(buffer: Buffer) -> Result<Self> {
            Ok(((), buffer))
        }
    }
);

decoder_value!(
    impl<'a, T> Option<T> {
        fn decode(buffer: Buffer) -> Result<Self> {
            if buffer.is_empty() {
                Ok((None, buffer))
            } else {
                let (value, buffer) = buffer.decode()?;
                Ok((Some(value), buffer))
            }
        }
    }
);

impl<'a> DecoderValueMut<'a> for DecoderBufferMut<'a> {
    #[inline]
    fn decode_mut(buffer: DecoderBufferMut<'a>) -> DecoderBufferMutResult<'a, Self> {
        let len = buffer.len();
        buffer.decode_slice(len)
    }
}

/// A value whose decoding implementation can be altered
/// by a provided parameter.
pub trait DecoderParameterizedValue<'a>: Sized {
    type Parameter;

    fn decode_parameterized(
        parameter: Self::Parameter,
        bytes: DecoderBuffer<'a>,
    ) -> DecoderBufferResult<'a, Self>;
}

/// A mutable value whose decoding implementation can be altered
/// by a provided parameter.
pub trait DecoderParameterizedValueMut<'a>: Sized {
    type Parameter;

    fn decode_parameterized_mut(
        parameter: Self::Parameter,
        bytes: DecoderBufferMut<'a>,
    ) -> DecoderBufferMutResult<'a, Self>;
}

#[macro_export]
macro_rules! decoder_parameterized_value {
    (impl<$lt:lifetime $(, $generic:ident)*> $ty:ty {
        fn decode($tag:ident: $tag_ty:ty, $buffer:ident: Buffer) -> Result<$ret:ty> $impl:block
    }) => {
        impl<$lt $(, $generic: $crate::DecoderValue<$lt>)*> $crate::DecoderParameterizedValue<$lt> for $ty {
            type Parameter = $tag_ty;

            #[inline]
            fn decode_parameterized($tag: Self::Parameter, $buffer: $crate::DecoderBuffer<$lt>) -> $crate::DecoderBufferResult<$lt, $ret> $impl
        }

        impl<$lt $(, $generic: $crate::DecoderValueMut<$lt>)*> $crate::DecoderParameterizedValueMut<$lt> for $ty {
            type Parameter = $tag_ty;

            #[inline]
            fn decode_parameterized_mut($tag: Self::Parameter, $buffer: $crate::DecoderBufferMut<$lt>) -> $crate::DecoderBufferMutResult<$lt, $ret> $impl
        }
    };
}


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

use crate::encoder::Encoder;

/// EncoderBuffer is a buffer for writing to a mutable slice
#[derive(Debug, Eq, PartialEq, PartialOrd, Ord)]
pub struct EncoderBuffer<'a> {
    bytes: &'a mut [u8],
    position: usize,
}

impl<'a> EncoderBuffer<'a> {
    /// Creates a new `EncoderBuffer`
    #[inline]
    pub fn new(bytes: &'a mut [u8]) -> Self {
        Self { bytes, position: 0 }
    }

    /// Sets the write cursor to a new position
    ///
    /// # Panics
    /// Panics when `position > capacity`
    #[inline]
    pub fn set_position(&mut self, position: usize) {
        debug_assert!(position <= self.capacity());
        self.position = position;
    }

    /// Advances the write cursor by offset
    ///
    /// # Panics
    /// Panics when `position > capacity`
    #[inline]
    pub fn advance_position(&mut self, offset: usize) {
        let position = self.position + offset;
        self.set_position(position)
    }

    /// Splits off the used buffer from the remaining bytes
    #[inline]
    pub fn split_off(self) -> (&'a mut [u8], &'a mut [u8]) {
        self.bytes.split_at_mut(self.position)
    }

    /// Splits the used buffer from the remaining bytes
    #[inline]
    pub fn split_mut(&mut self) -> (&mut [u8], &mut [u8]) {
        self.bytes.split_at_mut(self.position)
    }

    /// Returns the written bytes as a mutable slice
    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [u8] {
        unsafe { self.bytes.get_unchecked_mut(..self.position) }
    }

    #[inline]
    fn assert_capacity(&self, len: usize) {
        debug_assert!(
            len <= self.remaining_capacity(),
            "not enough buffer capacity. wanted: {}, available: {}",
            len,
            self.remaining_capacity()
        );
    }
}

impl<'a> Encoder for EncoderBuffer<'a> {
    #[inline]
    fn write_sized<F: FnOnce(&mut [u8])>(&mut self, len: usize, write: F) {
        self.assert_capacity(len);
        let end = self.position + len;
        let bytes = unsafe {
            // Safety: bounds already checked
            self.bytes.get_unchecked_mut(self.position..end)
        };
        write(bytes);
        self.position = end;
    }

    #[inline]
    fn write_slice(&mut self, slice: &[u8]) {
        self.write_sized(slice.len(), |dest| dest.copy_from_slice(slice));
    }

    #[inline]
    fn write_repeated(&mut self, count: usize, value: u8) {
        self.write_sized(count, |dest| {
            for byte in dest {
                *byte = value;
            }
        })
    }

    #[inline]
    fn write_zerocopy<
        T: zerocopy::AsBytes + zerocopy::FromBytes + zerocopy::Unaligned,
        F: FnOnce(&mut T),
    >(
        &mut self,
        write: F,
    ) {
        let len = core::mem::size_of::<T>();
        self.write_sized(len, |bytes| {
            let value = unsafe {
                // The `zerocopy` markers ensure this is a safe operation
                &mut *(bytes as *mut _ as *mut T)
            };
            write(value)
        })
    }

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

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


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

use crate::encoder::Encoder;

/// Estimates the `encoding_size` of an `EncoderValue`
pub struct EncoderLenEstimator {
    capacity: usize,
    len: usize,
}

impl EncoderLenEstimator {
    /// Create a new estimator with a given buffer `capacity`
    #[inline]
    pub const fn new(capacity: usize) -> Self {
        Self { capacity, len: 0 }
    }

    /// Returns true when the estimated len is greater than
    /// the allocated buffer capacity.
    #[inline]
    pub const fn overflowed(&self) -> bool {
        self.len > self.capacity
    }
}

impl Encoder for EncoderLenEstimator {
    #[inline]
    fn write_sized<F: FnOnce(&mut [u8])>(&mut self, len: usize, _write: F) {
        self.len += len;
    }

    #[inline]
    fn write_slice(&mut self, slice: &[u8]) {
        self.len += slice.len();
    }

    #[inline]
    fn write_repeated(&mut self, count: usize, _value: u8) {
        self.len += count;
    }

    #[inline]
    fn write_zerocopy<T: zerocopy::FromBytes + zerocopy::Unaligned, F: FnOnce(&mut T)>(
        &mut self,
        _write: F,
    ) {
        self.len += core::mem::size_of::<T>();
    }

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

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


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

pub mod buffer;
pub mod estimator;
pub mod value;

pub use buffer::*;
pub use estimator::*;
pub use value::*;

use core::convert::TryFrom;

pub trait Encoder: Sized {
    /// Set to `true` if the particular encoder specializes on the bytes implementation
    #[cfg(feature = "bytes")]
    const SPECIALIZES_BYTES: bool = false;

    /// Encode the given `EncoderValue` into the buffer
    #[inline]
    fn encode<T: EncoderValue>(&mut self, value: &T) {
        value.encode(self)
    }

    /// Encode the given `EncoderValue` into the buffer with a prefix of `Len`
    #[inline]
    fn encode_with_len_prefix<Len: TryFrom<usize> + EncoderValue, T: EncoderValue>(
        &mut self,
        value: &T,
    ) where
        Len::Error: core::fmt::Debug,
    {
        value.encode_with_len_prefix::<Len, Self>(self)
    }

    /// Calls `write` with a slice of `len` bytes at the current write position
    fn write_sized<F: FnOnce(&mut [u8])>(&mut self, len: usize, write: F);

    /// Copies the slice into the buffer
    fn write_slice(&mut self, slice: &[u8]);

    #[cfg(feature = "bytes")]
    #[inline]
    fn write_bytes(&mut self, bytes: bytes::Bytes) {
        self.write_slice(&bytes)
    }

    /// Writes a zerocopy value to the buffer
    fn write_zerocopy<
        T: zerocopy::AsBytes + zerocopy::FromBytes + zerocopy::Unaligned,
        F: FnOnce(&mut T),
    >(
        &mut self,
        write: F,
    );

    /// Repeatedly write a byte `value` for a given `count`
    ///
    /// ```
    /// # use s2n_codec::*;
    /// let mut buffer = vec![255; 1024];
    /// let mut encoder = EncoderBuffer::new(&mut buffer);
    /// encoder.encode(&1u8);
    /// encoder.write_repeated(4, 0);
    /// assert_eq!(&buffer[0..6], &[1, 0, 0, 0, 0, 255]);
    /// ```
    fn write_repeated(&mut self, count: usize, value: u8);

    /// Returns the total buffer capacity
    fn capacity(&self) -> usize;

    /// Returns the number of bytes written to the buffer
    fn len(&self) -> usize;

    /// Returns `true` if no bytes have been written
    #[inline]
    fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Returns the number of available bytes in the buffer
    #[inline]
    fn remaining_capacity(&self) -> usize {
        self.capacity().saturating_sub(self.len())
    }
}


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

use crate::{
    encoder::{Encoder, EncoderLenEstimator},
    i24, i48, u24, u48, DecoderBuffer, DecoderBufferMut,
};
use byteorder::{ByteOrder, NetworkEndian};
use core::{convert::TryFrom, mem::size_of};

pub trait EncoderValue: Sized {
    /// Encodes the value into the encoder
    fn encode<E: Encoder>(&self, encoder: &mut E);

    /// Encodes the value into the encoder, while potentially mutating the value itself
    #[inline]
    fn encode_mut<E: Encoder>(&mut self, encoder: &mut E) {
        self.encode(encoder)
    }

    /// Returns the encoding size with no buffer constrains
    #[inline]
    fn encoding_size(&self) -> usize {
        self.encoding_size_for_encoder(&EncoderLenEstimator::new(core::usize::MAX))
    }

    /// Returns the encoding size for the given encoder's capacity
    #[inline]
    fn encoding_size_for_encoder<E: Encoder>(&self, encoder: &E) -> usize {
        let mut estimator = EncoderLenEstimator::new(encoder.remaining_capacity());
        self.encode(&mut estimator);
        estimator.len()
    }

    /// Encodes the value into the encoder with a prefix of `Len`
    #[inline]
    fn encode_with_len_prefix<Len: TryFrom<usize> + EncoderValue, E: Encoder>(
        &self,
        encoder: &mut E,
    ) where
        Self: Sized,
        Len::Error: core::fmt::Debug,
    {
        let len = self.encoding_size_for_encoder(encoder);
        let len: Len = Len::try_from(len).expect("invalid conversion");
        len.encode(encoder);
        self.encode(encoder);
    }
}

macro_rules! encoder_value_byte {
    ($ty:ident) => {
        impl EncoderValue for $ty {
            #[inline]
            fn encode<E: Encoder>(&self, encoder: &mut E) {
                encoder.write_sized(size_of::<Self>(), |buf| {
                    buf[0] = *self as u8;
                })
            }

            #[inline]
            fn encoding_size(&self) -> usize {
                size_of::<Self>()
            }

            #[inline]
            fn encoding_size_for_encoder<E: Encoder>(&self, _encoder: &E) -> usize {
                size_of::<Self>()
            }
        }
    };
}

encoder_value_byte!(u8);
encoder_value_byte!(i8);

macro_rules! encoder_value_network_endian {
    ($call:ident, $ty:ty, $size:expr) => {
        impl EncoderValue for $ty {
            #[inline]
            fn encode<E: Encoder>(&self, encoder: &mut E) {
                encoder.write_sized($size, |buf| {
                    NetworkEndian::$call(buf, (*self).into());
                })
            }

            #[inline]
            fn encoding_size(&self) -> usize {
                $size
            }

            #[inline]
            fn encoding_size_for_encoder<E: Encoder>(&self, _encoder: &E) -> usize {
                $size
            }
        }
    };
}

encoder_value_network_endian!(write_u16, u16, size_of::<Self>());
encoder_value_network_endian!(write_i16, i16, size_of::<Self>());
encoder_value_network_endian!(write_u24, u24, 3);
encoder_value_network_endian!(write_i24, i24, 3);
encoder_value_network_endian!(write_u32, u32, size_of::<Self>());
encoder_value_network_endian!(write_i32, i32, size_of::<Self>());
encoder_value_network_endian!(write_u48, u48, 6);
encoder_value_network_endian!(write_i48, i48, 6);
encoder_value_network_endian!(write_u64, u64, size_of::<Self>());
encoder_value_network_endian!(write_i64, i64, size_of::<Self>());
encoder_value_network_endian!(write_u128, u128, size_of::<Self>());
encoder_value_network_endian!(write_i128, i128, size_of::<Self>());
encoder_value_network_endian!(write_f32, f32, size_of::<Self>());
encoder_value_network_endian!(write_f64, f64, size_of::<Self>());

macro_rules! encoder_value_slice {
    ($ty:ty, |$self:ident| $value:expr) => {
        impl EncoderValue for $ty {
            #[inline]
            fn encode<E: Encoder>(&$self, encoder: &mut E) {
                encoder.write_slice($value)
            }

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

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

encoder_value_slice!(&[u8], |self| self);
encoder_value_slice!(&mut [u8], |self| self);
encoder_value_slice!(DecoderBuffer<'_>, |self| self.as_less_safe_slice());
encoder_value_slice!(DecoderBufferMut<'_>, |self| self.as_less_safe_slice());

impl EncoderValue for &'_ [&'_ [u8]] {
    #[inline]
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        for slice in self.iter() {
            encoder.write_slice(slice)
        }
    }

    #[inline]
    fn encoding_size(&self) -> usize {
        self.iter().map(|s| s.len()).sum()
    }

    #[inline]
    fn encoding_size_for_encoder<E: Encoder>(&self, _encoder: &E) -> usize {
        self.iter().map(|s| s.len()).sum()
    }
}

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

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

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

impl<T: EncoderValue> EncoderValue for Option<T> {
    #[inline]
    fn encode<E: Encoder>(&self, buffer: &mut E) {
        if let Some(value) = self.as_ref() {
            value.encode(buffer);
        }
    }

    #[inline]
    fn encode_mut<E: Encoder>(&mut self, buffer: &mut E) {
        if let Some(value) = self.as_mut() {
            value.encode_mut(buffer);
        }
    }
}

#[cfg(feature = "bytes")]
impl EncoderValue for bytes::Bytes {
    #[inline]
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        if E::SPECIALIZES_BYTES {
            encoder.write_bytes(self.clone())
        } else {
            encoder.write_slice(self)
        }
    }

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

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

#[cfg(feature = "bytes")]
impl EncoderValue for &bytes::Bytes {
    #[inline]
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        if E::SPECIALIZES_BYTES {
            encoder.write_bytes((*self).clone())
        } else {
            encoder.write_slice(self)
        }
    }

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

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


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

pub use crate::{
    DecoderBuffer, DecoderBufferMut, DecoderError, DecoderValue, DecoderValueMut, Encoder,
    EncoderBuffer, EncoderLenEstimator, EncoderValue,
};

pub type Error = DecoderError;

#[macro_export]
macro_rules! assert_codec_round_trip_value {
    ($ty:ty, $value:expr) => {{
        $crate::ensure_codec_round_trip_value!($ty, $value).unwrap()
    }};
}

#[macro_export]
macro_rules! ensure_codec_round_trip_value {
    ($ty:ty, $value:expr) => {{
        fn execute_round_trip(expected_value: &$ty) -> Result<Vec<u8>, $crate::testing::Error> {
            let mut expected_bytes = $crate::testing::encode(expected_value)?;
            $crate::testing::ensure_decoding_matches(expected_value, &expected_bytes)?;
            $crate::testing::ensure_decoding_mut_matches(expected_value, &mut expected_bytes)?;
            Ok(expected_bytes)
        }

        execute_round_trip(&$value)
    }};
}

#[macro_export]
macro_rules! assert_codec_round_trip_value_mut {
    ($ty:ty, $value:expr) => {{
        $crate::ensure_codec_round_trip_value_mut!($ty, $value).unwrap()
    }};
}

#[macro_export]
macro_rules! ensure_codec_round_trip_value_mut {
    ($ty:ty, $value:expr) => {{
        fn execute_round_trip(expected_value: &$ty) -> Result<Vec<u8>, $crate::testing::Error> {
            let mut expected_bytes = $crate::testing::encode(expected_value)?;
            $crate::testing::ensure_decoding_mut_matches(expected_value, &mut expected_bytes)?;
            Ok(expected_bytes)
        }

        execute_round_trip(&$value)
    }};
}

macro_rules! impl_round_trip_bytes {
    ($name:ident, $round_trip:ident, $buffer:ident, $slice:ty, $as_ref:ident) => {
        #[macro_export]
        macro_rules! $name {
            ($ty: ty,$bytes: expr) => {{
                fn execute_all_round_trip(
                    buffer: $slice,
                ) -> Result<Vec<$ty>, $crate::testing::Error> {
                    let mut buffer = $crate::$buffer::new(buffer);
                    let mut values = vec![];

                    while let Ok((value, remaining)) = buffer.decode::<$ty>() {
                        $crate::$round_trip!($ty, &value)?;
                        values.push(value);
                        buffer = remaining;

                        // handle zero-sized encoded types, otherwise we loop forever
                        if buffer.is_empty() {
                            break;
                        }
                    }

                    Ok(values)
                }

                execute_all_round_trip($bytes.$as_ref())
            }};
        }
    };
}

#[macro_export]
macro_rules! assert_codec_round_trip_bytes {
    ($ty:ty, $bytes:expr) => {{
        $crate::ensure_codec_round_trip_bytes!($ty, $bytes).unwrap()
    }};
}

impl_round_trip_bytes!(
    ensure_codec_round_trip_bytes,
    ensure_codec_round_trip_value,
    DecoderBuffer,
    &[u8],
    as_ref
);

#[macro_export]
macro_rules! assert_codec_round_trip_bytes_mut {
    ($ty:ty, $bytes:expr) => {{
        $crate::ensure_codec_round_trip_bytes_mut!($ty, $bytes).unwrap()
    }};
}

impl_round_trip_bytes!(
    ensure_codec_round_trip_bytes_mut,
    ensure_codec_round_trip_value_mut,
    DecoderBufferMut,
    &mut [u8],
    as_mut
);

#[macro_export]
macro_rules! assert_codec_round_trip_sample_file {
    ($ty:ty, $path:expr) => {{
        $crate::assert_codec_round_trip_sample_file!($ty, $path, |buffer| {
            let (value, buffer) = buffer.decode::<$ty>().unwrap();
            $crate::assert_codec_round_trip_value_mut!($ty, value);
            (value, buffer)
        });
    }};
    ($ty:ty, $path:expr, | $buffer:ident | $decode:expr) => {{
        #[cfg(not(miri))]
        let mut expected = std::fs::read(concat!(env!("CARGO_MANIFEST_DIR"), "/", $path)).unwrap();
        #[cfg(miri)]
        let mut expected = include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/", $path)).to_vec();

        let mut $buffer = $crate::DecoderBufferMut::new(&mut expected);
        let mut values = vec![];

        while !$buffer.is_empty() {
            let (value, remaining) = $decode;
            values.push(value);
            $buffer = remaining;
        }

        #[cfg(not(miri))]
        insta::assert_debug_snapshot!(
            std::path::Path::new($path)
                .file_stem()
                .unwrap()
                .to_str()
                .unwrap(),
            values
        );
    }};
}

macro_rules! ensure {
    ($expr:expr, $message:expr $(,)?) => {
        if !($expr) {
            return Err($crate::testing::Error::InvariantViolation($message));
        }
    };
}

pub fn encode<T: EncoderValue>(expected_value: &T) -> Result<Vec<u8>, Error> {
    let len = expected_value.encoding_size();
    let mut buffer = vec![0; len];
    EncoderBuffer::new(&mut buffer).encode(expected_value);
    Ok(buffer)
}

pub fn ensure_encoding_matches<T: EncoderValue + PartialEq + core::fmt::Debug>(
    expected_value: &T,
    expected_bytes: &[u8],
) -> Result<Vec<u8>, Error> {
    let actual_bytes = encode(expected_value)?;
    ensure!(actual_bytes == expected_bytes, "encodings do not match");
    Ok(actual_bytes)
}

pub fn decode<'a, T: DecoderValue<'a>>(
    expected_bytes: &'a [u8],
) -> Result<(T, DecoderBuffer<'a>), Error> {
    let (actual_value, remaining) = DecoderBuffer::new(expected_bytes).decode()?;
    Ok((actual_value, remaining))
}

pub fn ensure_decoding_matches<'a, T: DecoderValue<'a> + PartialEq + core::fmt::Debug>(
    expected_value: &T,
    expected_bytes: &'a [u8],
) -> Result<(), Error> {
    let (actual_value, remaining) = decode(expected_bytes)?;
    ensure!(
        expected_value == &actual_value,
        "mut decodings do not match",
    );
    remaining.ensure_empty()?;
    Ok(())
}

pub fn decode_mut<'a, T: DecoderValueMut<'a>>(
    expected_bytes: &'a mut [u8],
) -> Result<(T, DecoderBufferMut<'a>), Error> {
    let (actual_value, remaining) = DecoderBufferMut::new(expected_bytes).decode()?;
    Ok((actual_value, remaining))
}

pub fn ensure_decoding_mut_matches<'a, T: DecoderValueMut<'a> + PartialEq + core::fmt::Debug>(
    expected_value: &T,
    expected_bytes: &'a mut [u8],
) -> Result<(), Error> {
    let (actual_value, remaining) = decode_mut(expected_bytes)?;
    ensure!(
        expected_value == &actual_value,
        "mut decodings do not match",
    );
    remaining.ensure_empty()?;
    Ok(())
}

#[cfg(test)]
mod tests {
    #[test]
    fn test_u8_round_trip_value() {
        for i in 0..core::u8::MAX {
            ensure_codec_round_trip_value!(u8, i).unwrap();
        }
    }

    #[test]
    fn test_u8_round_trip_bytes() {
        let bytes = (0..core::u8::MAX).collect::<Vec<_>>();
        ensure_codec_round_trip_bytes!(u8, &bytes).unwrap();
    }
}


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

use s2n_quic::{client::Connect, Client};
use std::{error::Error, net::SocketAddr};

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let client = Client::builder()
        .with_tls(CERT_PEM)?
        .with_io("0.0.0.0:0")?
        .start()?;

    let addr: SocketAddr = "127.0.0.1:4433".parse()?;
    let connect = Connect::new(addr).with_server_name("localhost");
    let mut connection = client.connect(connect).await?;

    // ensure the connection doesn't time out with inactivity
    connection.keep_alive(true)?;

    // open a new stream and split the receiving and sending sides
    let stream = connection.open_bidirectional_stream().await?;
    let (mut receive_stream, mut send_stream) = stream.split();

    // spawn a task that copies responses from the server to stdout
    tokio::spawn(async move {
        let mut stdout = tokio::io::stdout();
        let _ = tokio::io::copy(&mut receive_stream, &mut stdout).await;
    });

    // copy data from stdin and send it to the server
    let mut stdin = tokio::io::stdin();
    tokio::io::copy(&mut stdin, &mut send_stream).await?;

    Ok(())
}


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

use moka::sync::Cache;
use rand::{distributions::WeightedIndex, prelude::*};
use s2n_quic::{
    provider::tls::s2n_tls::{
        s2n_tls::{
            callbacks::{ConfigResolver, ConnectionFuture},
            config::Config,
            error::Error as S2nError,
        },
        ClientHelloCallback, Connection,
    },
    Server,
};
use std::{error::Error, fmt::Display, pin::Pin, sync::Arc, time::Duration};
use tokio::{fs, sync::OnceCell};

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM_PATH: &str = concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
);
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM_PATH: &str = concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
);

type Sni = String;

// A Config cache associated with as SNI (server name indication).
//
// Implements ClientHelloCallback, loading the certificates asynchronously,
// and caching the s2n_tls::config::Config for subsequent calls with the
// same SNI.
//
// An SNI, indicates which hostname the client is attempting to connect to.
// Some deployments could require configuring the s2n_tls::config::Config
// based on the SNI (certificate).
struct ConfigCache {
    cache: Cache<Sni, Arc<OnceCell<Config>>>,
}

impl ConfigCache {
    fn new() -> Self {
        ConfigCache {
            // store Config for up to 100 unique SNI
            cache: Cache::new(100),
        }
    }
}

impl ClientHelloCallback for ConfigCache {
    fn on_client_hello(
        &self,
        connection: &mut Connection,
    ) -> Result<Option<Pin<Box<dyn ConnectionFuture>>>, S2nError> {
        let sni = connection
            .server_name()
            .ok_or_else(|| S2nError::application(Box::new(CustomError)))?
            .to_string();

        let once_cell_config = self
            .cache
            .get_with(sni.clone(), || Arc::new(OnceCell::new()));
        if let Some(config) = once_cell_config.get() {
            eprintln!("Config already cached for SNI: {}", sni);
            connection.set_config(config.clone())?;
            // return `None` if the Config is already in the cache
            return Ok(None);
        }

        // simulate failure 75% of times and success 25% of the times
        let choices = [true, false];
        let weights = [3, 1];
        let dist = WeightedIndex::new(weights).unwrap();
        let mut rng = thread_rng();

        let fut = async move {
            let fut = once_cell_config.get_or_try_init(|| async {
                let simulated_network_call_failed = choices[dist.sample(&mut rng)];

                if simulated_network_call_failed {
                    eprintln!("simulated network call failed");
                    return Err(S2nError::application(Box::new(CustomError)));
                }

                eprintln!("resolving certificate for SNI: {}", sni);

                // load the cert and key file asynchronously.
                let (cert, key) = {
                    // the SNI can be used to load the appropriate cert file
                    let _sni = sni;
                    let cert = fs::read_to_string(CERT_PEM_PATH)
                        .await
                        .map_err(|_| S2nError::application(Box::new(CustomError)))?;
                    let key = fs::read_to_string(KEY_PEM_PATH)
                        .await
                        .map_err(|_| S2nError::application(Box::new(CustomError)))?;
                    (cert, key)
                };

                // sleep(async tokio task which doesn't block thread) to mimic delay
                tokio::time::sleep(Duration::from_secs(3)).await;

                let config = s2n_quic::provider::tls::s2n_tls::Server::builder()
                    .with_certificate(cert, key)?
                    .build()?
                    .into();
                Ok(config)
            });
            fut.await.map(|config| config.clone())
        };

        // return `Some(ConnectionFuture)` if the Config wasn't found in the
        // cache and we need to load it asynchronously
        Ok(Some(Box::pin(ConfigResolver::new(fut))))
    }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let tls = s2n_quic::provider::tls::s2n_tls::Server::builder()
        .with_client_hello_handler(ConfigCache::new())?
        .build()?;

    let mut server = Server::builder()
        .with_tls(tls)?
        .with_io("127.0.0.1:4433")?
        .start()?;

    while let Some(mut connection) = server.accept().await {
        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                // spawn a new task for the stream
                tokio::spawn(async move {
                    eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                    // echo any data back to the stream
                    while let Ok(Some(data)) = stream.receive().await {
                        stream.send(data).await.expect("stream should be open");
                    }
                });
            }
        });
    }

    Ok(())
}

#[derive(Debug)]
struct CustomError;

impl Display for CustomError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "custom error")?;
        Ok(())
    }
}

impl Error for CustomError {}


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

/// Example implementation of a custom congestion controller algorithm.
///
/// This example serves only to illustrate the integration points for incorporating a custom
/// congestion controller into s2n-quic, and not as an actual congestion controller implementation.
///
/// NOTE: The `CongestionController` trait is considered unstable and may be subject to change
///       in a future release.
pub mod custom_congestion_controller {
    use s2n_quic::provider::{
        congestion_controller,
        congestion_controller::{
            CongestionController, Publisher, RandomGenerator, RttEstimator, Timestamp,
        },
    };

    /// Define a congestion controller containing any state you wish to track.
    /// For this example, we track the size of the congestion window in bytes and
    /// the number of bytes in flight.
    #[derive(Debug, Clone)]
    pub struct MyCongestionController {
        congestion_window: u32,
        bytes_in_flight: u32,
    }

    /// The following is a simple implementation of the `CongestionController` trait
    /// that increases the congestion window by the number of bytes acknowledged and
    /// decreases the congestion window by half when packets are lost.
    #[allow(unused)]
    impl CongestionController for MyCongestionController {
        // A custom `PacketInfo` type may optionally be defined to include additional per-packet
        // state. This state will be stored for each in flight packet, and returned to the
        // `on_ack` and `on_packet_lost` methods.
        type PacketInfo = ();

        fn congestion_window(&self) -> u32 {
            self.congestion_window
        }

        fn bytes_in_flight(&self) -> u32 {
            self.bytes_in_flight
        }

        fn is_congestion_limited(&self) -> bool {
            self.congestion_window < self.bytes_in_flight
        }

        fn requires_fast_retransmission(&self) -> bool {
            false
        }

        fn on_packet_sent<Pub: Publisher>(
            &mut self,
            time_sent: Timestamp,
            sent_bytes: usize,
            app_limited: Option<bool>,
            rtt_estimator: &RttEstimator,
            publisher: &mut Pub,
        ) -> Self::PacketInfo {
            self.bytes_in_flight += sent_bytes as u32;
        }

        fn on_rtt_update<Pub: Publisher>(
            &mut self,
            time_sent: Timestamp,
            now: Timestamp,
            rtt_estimator: &RttEstimator,
            publisher: &mut Pub,
        ) {
            // no op
        }

        fn on_ack<Pub: Publisher>(
            &mut self,
            newest_acked_time_sent: Timestamp,
            bytes_acknowledged: usize,
            newest_acked_packet_info: Self::PacketInfo,
            rtt_estimator: &RttEstimator,
            random_generator: &mut dyn RandomGenerator,
            ack_receive_time: Timestamp,
            publisher: &mut Pub,
        ) {
            self.bytes_in_flight -= bytes_acknowledged as u32;
            self.congestion_window += bytes_acknowledged as u32;
        }

        fn on_packet_lost<Pub: Publisher>(
            &mut self,
            lost_bytes: u32,
            packet_info: Self::PacketInfo,
            persistent_congestion: bool,
            new_loss_burst: bool,
            random_generator: &mut dyn RandomGenerator,
            timestamp: Timestamp,
            publisher: &mut Pub,
        ) {
            // Loss-based congestion controllers such as New Reno or CUBIC will decrease the
            // congestion window when a packet is lost. In this simple example congestion
            // controller the congestion window is reduced for every packet; an actual congestion
            // controller should take a more nuanced approach. This reduction would typically only
            // occur once for the initial lost packet, and subsequent lost packets would not lead to
            // further reduction.
            self.bytes_in_flight -= lost_bytes;
            self.congestion_window = (self.congestion_window as f32 * 0.5) as u32;
        }

        fn on_explicit_congestion<Pub: Publisher>(
            &mut self,
            ce_count: u64,
            event_time: Timestamp,
            publisher: &mut Pub,
        ) {
            self.congestion_window = (self.congestion_window as f32 * 0.5) as u32;
        }

        fn on_mtu_update<Pub: Publisher>(&mut self, max_data_size: u16, publisher: &mut Pub) {
            // no op
        }

        fn on_packet_discarded<Pub: Publisher>(&mut self, bytes_sent: usize, publisher: &mut Pub) {
            self.bytes_in_flight -= bytes_sent as u32;
        }

        fn earliest_departure_time(&self) -> Option<Timestamp> {
            None
        }
    }

    // Define an endpoint for the custom congestion controller so it may be used as a
    // congestion controller provider to the s2n-quic server or client
    #[derive(Debug, Default)]
    pub struct MyCongestionControllerEndpoint {}

    impl congestion_controller::Endpoint for MyCongestionControllerEndpoint {
        type CongestionController = MyCongestionController;

        // This method will be called whenever a new congestion controller instance is needed.
        fn new_congestion_controller(
            &mut self,
            path_info: congestion_controller::PathInfo,
        ) -> Self::CongestionController {
            MyCongestionController {
                // Specify the initial congestion window
                congestion_window: 10 * path_info.max_datagram_size as u32,
                bytes_in_flight: 0,
            }
        }
    }
}


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

use custom_congestion_controller::custom_congestion_controller::MyCongestionControllerEndpoint;
use s2n_quic::Server;
use std::error::Error;

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Build an `s2n_quic::Server`
    let mut server = Server::builder()
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io("127.0.0.1:4433")?
        // Specify the custom congestion controller endpoint defined in `custom-congestion-controller/src/lib.rs`
        .with_congestion_controller(MyCongestionControllerEndpoint::default())?
        .start()?;

    while let Some(mut connection) = server.accept().await {
        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                // spawn a new task for the stream
                tokio::spawn(async move {
                    eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                    // echo any data back to the stream
                    while let Ok(Some(data)) = stream.receive().await {
                        stream.send(data).await.expect("stream should be open");
                    }
                });
            }
        });
    }

    Ok(())
}


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

/// Example mitigation for Slowloris-style Denial of Service attacks. For details on this attack,
/// see [QUIC Transport RFC](https://www.rfc-editor.org/rfc/rfc9000.html#name-slowloris-attacks).
pub mod slowloris {
    use s2n_quic::provider::{
        event,
        event::{events, supervisor, ConnectionMeta, Timestamp},
    };
    use std::time::Duration;

    /// The maximum number of connections that may be active concurrently before
    /// low throughput connections are closed.
    const CONNECTION_COUNT_THRESHOLD: usize = 1000;
    /// The minimum throughput a connection must sustain, in bytes per second
    const MIN_THROUGHPUT: usize = 500;

    /// Define a Connection Context containing any per-connection state you wish to track.
    /// For this example, we need to track the number of bytes transferred and the last
    /// time the transferred byte count was totalled.
    #[derive(Debug, Clone)]
    pub struct MyConnectionContext {
        transferred_bytes: usize,
        last_update: Timestamp,
    }

    /// Define a struct containing any state across all connections you wish to track.
    /// For this example, there is no additional state we need to track so the struct is empty.
    #[derive(Default)]
    pub struct MyConnectionSupervisor;

    /// Implement the `event::Subscriber` trait for your struct. The `create_connection_context`
    /// method must be implemented to initialize the Connection Context for each connection.
    /// Other methods may be implemented as needed.
    impl event::Subscriber for MyConnectionSupervisor {
        type ConnectionContext = MyConnectionContext;

        /// Initialize the Connection Context that is passed to the `supervisor_timeout` and
        /// `on_supervisor_timeout` methods, as well as each connection-related event.
        fn create_connection_context(
            &mut self,
            meta: &events::ConnectionMeta,
            _info: &events::ConnectionInfo,
        ) -> Self::ConnectionContext {
            MyConnectionContext {
                transferred_bytes: 0,
                last_update: meta.timestamp,
            }
        }

        /// Implement `supervisor_timeout` to define the period at which `on_supervisor_timeout` will
        /// be invoked. For this example, a constant of 1 second is used, but this value can be
        /// varied over time or based on the connection.
        fn supervisor_timeout(
            &mut self,
            _conn_context: &mut Self::ConnectionContext,
            _meta: &ConnectionMeta,
            _context: &supervisor::Context,
        ) -> Option<Duration> {
            Some(Duration::from_secs(1))
        }

        /// Implement `on_supervisor_timeout` to define what action should be taken on the connection
        /// when the `supervisor_timeout` expires. For this example, the connection will be closed
        /// immediately (`supervisor::Outcome::ImmediateClose`) if the number of open connections
        /// is greater than `CONNECTION_COUNT_THRESHOLD` and the the throughput of the connection
        /// since the last `supervisor_timeout` has dropped below `MIN_THROUGHPUT`.
        fn on_supervisor_timeout(
            &mut self,
            conn_context: &mut Self::ConnectionContext,
            meta: &ConnectionMeta,
            context: &supervisor::Context,
        ) -> supervisor::Outcome {
            if !context.is_handshaking && context.connection_count > CONNECTION_COUNT_THRESHOLD {
                let elapsed_time = meta.timestamp.duration_since_start()
                    - conn_context.last_update.duration_since_start();

                // Calculate throughput as bytes per second
                let throughput =
                    (conn_context.transferred_bytes as f32 / elapsed_time.as_secs_f32()) as usize;

                if throughput < MIN_THROUGHPUT {
                    // Close the connection immediately without notifying the peer
                    return supervisor::Outcome::ImmediateClose {
                        reason: "Connection throughput was below MIN_THROUGHPUT",
                    };
                }
            }

            // Update the `last_update` timestamp and reset transferred bytes
            conn_context.last_update = meta.timestamp;
            conn_context.transferred_bytes = 0;

            // Allow the connection to continue
            supervisor::Outcome::Continue
        }

        /// Implement `on_tx_stream_progress` to be notified every time forward progress is made
        /// on an outgoing stream.
        fn on_tx_stream_progress(
            &mut self,
            context: &mut Self::ConnectionContext,
            _meta: &events::ConnectionMeta,
            event: &events::TxStreamProgress,
        ) {
            context.transferred_bytes += event.bytes;
        }

        /// Implement `on_rx_progress` to be notified every time forward progress is made on an
        /// incoming stream.
        fn on_rx_stream_progress(
            &mut self,
            context: &mut Self::ConnectionContext,
            _meta: &events::ConnectionMeta,
            event: &events::RxStreamProgress,
        ) {
            context.transferred_bytes += event.bytes;
        }
    }
}


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

use dos_mitigation::slowloris;
use s2n_quic::Server;
use std::{error::Error, time::Duration};

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Limit the duration any handshake attempt may take to 5 seconds
    // By default, handshakes are limited to 10 seconds.
    let connection_limits = s2n_quic::provider::limits::Limits::new()
        .with_max_handshake_duration(Duration::from_secs(5))
        .expect("connection limits are valid");

    // Limit the number of inflight handshakes to 100.
    let endpoint_limits = s2n_quic::provider::endpoint_limits::Default::builder()
        .with_inflight_handshake_limit(100)?
        .build()?;

    // Build an `s2n_quic::Server`
    let mut server = Server::builder()
        // Provide the `connection_limits` defined above
        .with_limits(connection_limits)?
        // Provide the `endpoint_limits defined above
        .with_endpoint_limits(endpoint_limits)?
        // Provide a tuple of the `example::MyConnectionSupervisor` defined in `dos-mitigation/src/lib.rs`
        // and the default event tracing subscriber. This combination will allow for both the slowloris mitigation
        // functionality of `MyConnectionSupervisor` as well as event tracing to be utilized.
        .with_event((
            slowloris::MyConnectionSupervisor,
            s2n_quic::provider::event::tracing::Subscriber::default(),
        ))?
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io("127.0.0.1:4433")?
        .start()?;

    while let Some(mut connection) = server.accept().await {
        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                // spawn a new task for the stream
                tokio::spawn(async move {
                    eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                    // echo any data back to the stream
                    while let Ok(Some(data)) = stream.receive().await {
                        stream.send(data).await.expect("stream should be open");
                    }
                });
            }
        });
    }

    Ok(())
}


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

use s2n_quic::Server;
use std::error::Error;

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let mut server = Server::builder()
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io("127.0.0.1:4433")?
        .start()?;

    while let Some(mut connection) = server.accept().await {
        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                // spawn a new task for the stream
                tokio::spawn(async move {
                    eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                    // echo any data back to the stream
                    while let Ok(Some(data)) = stream.receive().await {
                        stream.send(data).await.expect("stream should be open");
                    }
                });
            }
        });
    }

    Ok(())
}


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

/// Example of a print subscriber which can print all events or only
/// connection related events.
pub mod print_event {
    use s2n_quic::provider::{event, event::ConnectionMeta};

    #[derive(Debug, Clone)]
    pub struct MyPrintSubscriber {
        // prints all events, including connection events
        pub print_all_events: bool,
        // prints only connection related events
        pub print_connection_events: bool,
    }

    impl event::Subscriber for MyPrintSubscriber {
        type ConnectionContext = ();

        /// Initialize the Connection Context.
        fn create_connection_context(
            &mut self,
            meta: &ConnectionMeta,
            info: &event::ConnectionInfo,
        ) -> Self::ConnectionContext {
            println!("{:?} {:?}", meta, info);
        }

        /// This event fires for all events.
        fn on_event<M: event::Meta + core::fmt::Debug, E: event::Event + core::fmt::Debug>(
            &mut self,
            meta: &M,
            event: &E,
        ) {
            if self.print_all_events {
                println!("event: {:?} {:?}", meta, event);
            }
        }

        /// This event fires only for connection-level events. Excluded are events which
        /// happen prior to connection creation, e.g. `on_version_information`,
        /// `on_endpoint_datagram_drop`.
        fn on_connection_event<E: event::Event + core::fmt::Debug>(
            &mut self,
            context: &mut Self::ConnectionContext,
            meta: &ConnectionMeta,
            event: &E,
        ) {
            if self.print_connection_events {
                println!("connection_event: {:?} {:?} {:?}", context, meta, event);
            }
        }
    }
}

/// Example of a query subscriber which can be used to store event information; which
/// can then be queried from the application.
pub mod query_event {
    use s2n_quic::provider::{
        event,
        event::{events, ConnectionMeta},
    };

    #[derive(Debug, Clone)]
    pub struct MyQueryContext {
        // Record how many data packets are received
        pub packet_sent_count: usize,
        // Flag to control the packet counter behavior
        pub count_non_data_packets: bool,
    }

    #[derive(Default)]
    pub struct MyQuerySubscriber;

    impl event::Subscriber for MyQuerySubscriber {
        type ConnectionContext = MyQueryContext;

        /// Initialize the Connection Context.
        fn create_connection_context(
            &mut self,
            _meta: &events::ConnectionMeta,
            _info: &events::ConnectionInfo,
        ) -> Self::ConnectionContext {
            MyQueryContext {
                packet_sent_count: 0,
                count_non_data_packets: true,
            }
        }

        /// This event fires for every outgoing packet that is transmitted.
        fn on_packet_sent(
            &mut self,
            context: &mut Self::ConnectionContext,
            _meta: &ConnectionMeta,
            event: &events::PacketSent,
        ) {
            match event.packet_header {
                events::PacketHeader::ZeroRtt { .. } | events::PacketHeader::OneRtt { .. } => {
                    context.packet_sent_count += 1;
                }
                _ => {
                    if context.count_non_data_packets {
                        context.packet_sent_count += 1;
                    }
                }
            }
        }
    }

    impl Drop for MyQueryContext {
        // Execute some operations on the context before the Connection is dropped.
        fn drop(&mut self) {
            println!("{:?}", self);
        }
    }
}


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

use event_framework::{print_event, query_event};
use s2n_quic::Server;
use std::error::Error;

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // It's possible to compose different Subscribers, each of which is responsible
    // for a different task.
    //
    // Note: subscriber is implemented for `(A, B)` and therefore requires the nested tuple
    let composed_event_subscriber = (
        // Our custom query subscriber allows us to store information on a custom
        // connection context and query it from the application.
        //
        // There are two types of queries:
        // - connection.query_event_context: yields a reference to the connection context
        // which can be used read data from the connection context.
        // - connection.query_event_context_mut: yields a mutable reference to the connection
        // context which can be used to read and write data to the connection context.
        query_event::MyQuerySubscriber,
        (
            // Our custom print subscriber allows us to print events to stdout.
            print_event::MyPrintSubscriber {
                print_all_events: true,
                print_connection_events: false,
            },
            // The tracing subscriber will allow applications to configure and use
            // [tracing](https://docs.rs/tracing/latest/tracing/) for instrumentation.
            s2n_quic::provider::event::tracing::Subscriber::default(),
        ),
    );
    // Build an `s2n_quic::Server`
    let mut server = Server::builder()
        .with_event(composed_event_subscriber)?
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io("127.0.0.1:4433")?
        .start()?;

    let mut connection_count = 0;
    while let Some(mut connection) = server.accept().await {
        // Change packet count behavior after processing some number of connections.
        //
        // Note: The context type, `&query_event::MyQueryContext`, provided in the query must
        // match the context on one of the subscribers. See the docs on `query_event_context_mut`
        // for a detailed explanation on how a query is executed on composed subscribers.
        //
        // The application can mutably access the connection context and modify data on
        // the context itself.
        if connection_count > 5 {
            connection.query_event_context_mut(|context: &mut query_event::MyQueryContext| {
                context.count_non_data_packets = false;
            })?;
        } else {
            connection_count += 1;
        }

        // Query the packet sent count and print it.
        //
        // Note: The context type, `&query_event::MyQueryContext`, provided in the query must
        // match the context on one of the subscribers. See the docs on `query_event_context`
        // for a detailed explanation on how a query is executed on composed subscribers.
        //
        // The application can immutably access the connection context and read data from it.
        connection.query_event_context(|context: &query_event::MyQueryContext| {
            println!(
                "{:?} data packets have been processed",
                context.packet_sent_count
            )
        })?;

        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                // spawn a new task for the stream
                tokio::spawn(async move {
                    eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                    // echo any data back to the stream
                    while let Ok(Some(data)) = stream.receive().await {
                        stream.send(data).await.expect("stream should be open");
                    }
                });
            }
        });
    }

    Ok(())
}


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

use s2n_quic::provider::event::Subscriber;

pub struct MtuEventInformer {}

pub struct MtuEventInformerContext {}

impl Subscriber for MtuEventInformer {
    type ConnectionContext = MtuEventInformerContext;

    fn create_connection_context(
        &mut self,
        _meta: &s2n_quic::provider::event::ConnectionMeta,
        _info: &s2n_quic::provider::event::ConnectionInfo,
    ) -> Self::ConnectionContext {
        MtuEventInformerContext {}
    }

    fn on_mtu_updated(
        &mut self,
        _context: &mut Self::ConnectionContext,
        _meta: &s2n_quic::provider::event::ConnectionMeta,
        event: &s2n_quic::provider::event::events::MtuUpdated,
    ) {
        eprintln!("{:?}", event);
    }
}


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

use jumbo_frame::MtuEventInformer;
use s2n_quic::{client::Connect, Client};
use std::{error::Error, net::SocketAddr};
use tokio::time::Duration;

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let address: SocketAddr = "0.0.0.0:2198".parse()?;

    // set up an io provider with jumbo mtu and larger socket buffers
    let io = s2n_quic::provider::io::Default::builder()
        .with_max_mtu(9001)?
        .with_receive_address(address)?
        .with_recv_buffer_size(12_000_000)?
        .with_send_buffer_size(12_000_000)?
        .build()?;
    let client = Client::builder()
        .with_tls(CERT_PEM)?
        .with_io(io)?
        .with_event(MtuEventInformer {})?
        .start()?;
    let addr: SocketAddr = "127.0.0.1:4433".parse()?;
    let connect = Connect::new(addr).with_server_name("localhost");
    let mut connection = client.connect(connect).await?;

    // ensure the connection doesn't time out with inactivity
    connection.keep_alive(true)?;

    // we aren't actually sending any data, but during this time the quic
    // endpoint will be probing to see if we can use a 9_001 byte mtu
    // if we were sending data that would happen concurrently with the mtu
    // probing.
    tokio::time::sleep(Duration::from_secs(3)).await;

    Ok(())
}


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

use jumbo_frame::MtuEventInformer;
use s2n_quic::Server;
use std::{error::Error, net::SocketAddr};

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let address: SocketAddr = "127.0.0.1:4433".parse()?;

    // set up an io provider with jumbo mtu and larger socket buffers
    let io = s2n_quic::provider::io::Default::builder()
        .with_receive_address(address)?
        .with_max_mtu(9_001)?
        .with_recv_buffer_size(12_000_000)?
        .with_send_buffer_size(12_000_000)?
        .build()?;

    let mut server = Server::builder()
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io(io)?
        .with_event(MtuEventInformer {})?
        .start()?;

    eprintln!("Listening for a connection");
    let connection = server.accept().await.unwrap();

    eprintln!("Connection accepted from {:?}", connection.remote_addr());

    // we aren't sending any data, but the endpoint will be probing
    // for higher mtu's during this time.
    tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;

    Ok(())
}


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

use s2n_quic::{client::Connect, Client};
use std::{error::Error, net::SocketAddr};
use tokio::io::AsyncWriteExt;

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let mut client = Client::builder()
        .with_tls(CERT_PEM)?
        .with_io("0.0.0.0:0")?
        .start()?;

    let addr: SocketAddr = "127.0.0.1:4433".parse()?;
    let connect = Connect::new(addr).with_server_name("localhost");
    let mut connection = client.connect(connect).await?;

    let mut stream = connection.open_bidirectional_stream().await?;
    stream.write_all(b"hello post-quantum server!").await?;

    let mut stdout = tokio::io::stdout();
    tokio::io::copy(&mut stream, &mut stdout).await?;

    connection.close(0u8.into());

    client.wait_idle().await?;

    Ok(())
}


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

use s2n_quic::Server;
use std::error::Error;
use tokio::io::AsyncWriteExt;

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let mut server = Server::builder()
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io("127.0.0.1:4433")?
        .start()?;

    while let Some(mut connection) = server.accept().await {
        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                // spawn a new task for the stream
                tokio::spawn(async move {
                    eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                    // respond to the client with our own message
                    if let Ok(Some(data)) = stream.receive().await {
                        eprintln!("message from the client: {:?}", data);
                        let _ = stream.write_all(b"hello post-quantum client!\n").await;
                    }
                });
            }
        });
    }

    Ok(())
}


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

use rustls::{
    cipher_suite, ClientConfig, Error, RootCertStore, ServerConfig, SupportedCipherSuite,
};
use s2n_quic::provider::{tls, tls::rustls::rustls};
use std::{io::Cursor, path::Path};
use tokio::{fs::File, io::AsyncReadExt};
use tracing::Level;

static PROTOCOL_VERSIONS: &[&rustls::SupportedProtocolVersion] = &[&rustls::version::TLS13];

pub static DEFAULT_CIPHERSUITES: &[SupportedCipherSuite] = &[
    cipher_suite::TLS13_AES_128_GCM_SHA256,
    cipher_suite::TLS13_AES_256_GCM_SHA384,
    cipher_suite::TLS13_CHACHA20_POLY1305_SHA256,
];

pub fn initialize_logger(endpoint: &str) {
    use std::sync::Once;

    static TRACING: Once = Once::new();

    // make sure this only gets initialized once (per process)
    TRACING.call_once(|| {
        // always write to the same file, and don't rotate it. This would be a
        // bad idea for a long running process, but is useful to make sure that
        // all the logs of our program end up in the same file.
        let file_appender = tracing_appender::rolling::never("logs", format!("{endpoint}.txt"));

        tracing_subscriber::fmt()
            .with_max_level(Level::DEBUG)
            // don't color the output, otherwise the text logs will have odd
            // characters
            .with_ansi(false)
            .with_writer(file_appender)
            .init();
    });
}

pub struct MtlsProvider {
    root_store: rustls::RootCertStore,
    my_cert_chain: Vec<rustls::Certificate>,
    my_private_key: rustls::PrivateKey,
}

impl tls::Provider for MtlsProvider {
    type Server = tls::rustls::Server;
    type Client = tls::rustls::Client;
    type Error = rustls::Error;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        let verifier = rustls::server::AllowAnyAuthenticatedClient::new(self.root_store);
        let mut cfg = ServerConfig::builder()
            .with_cipher_suites(DEFAULT_CIPHERSUITES)
            .with_safe_default_kx_groups()
            .with_protocol_versions(PROTOCOL_VERSIONS)?
            .with_client_cert_verifier(verifier)
            .with_single_cert(self.my_cert_chain, self.my_private_key)?;

        cfg.ignore_client_order = true;
        cfg.max_fragment_size = None;
        cfg.alpn_protocols = vec![b"h3".to_vec()];
        Ok(cfg.into())
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        let mut cfg = ClientConfig::builder()
            .with_cipher_suites(DEFAULT_CIPHERSUITES)
            .with_safe_default_kx_groups()
            .with_protocol_versions(PROTOCOL_VERSIONS)?
            .with_root_certificates(self.root_store)
            .with_single_cert(self.my_cert_chain, self.my_private_key)?;

        cfg.max_fragment_size = None;
        cfg.alpn_protocols = vec![b"h3".to_vec()];
        Ok(cfg.into())
    }
}

impl MtlsProvider {
    pub async fn new<A: AsRef<Path>, B: AsRef<Path>, C: AsRef<Path>>(
        ca_cert_pem: A,
        my_cert_pem: B,
        my_key_pem: C,
    ) -> Result<Self, Error> {
        let root_store = into_root_store(ca_cert_pem.as_ref()).await?;
        let cert_chain = into_certificate(my_cert_pem.as_ref()).await?;
        let private_key = into_private_key(my_key_pem.as_ref()).await?;
        Ok(MtlsProvider {
            root_store,
            my_cert_chain: cert_chain.into_iter().map(rustls::Certificate).collect(),
            my_private_key: rustls::PrivateKey(private_key),
        })
    }
}

async fn into_certificate(path: &Path) -> Result<Vec<Vec<u8>>, Error> {
    let mut f = File::open(path)
        .await
        .map_err(|e| Error::General(format!("Failed to load file: {}", e)))?;
    let mut buf = Vec::new();
    f.read_to_end(&mut buf)
        .await
        .map_err(|e| Error::General(format!("Failed to read file: {}", e)))?;
    let mut cursor = Cursor::new(buf);
    let certs = rustls_pemfile::certs(&mut cursor)
        .map(|certs| certs.into_iter().collect())
        .map_err(|_| Error::General("Could not read certificate".to_string()))?;
    Ok(certs)
}

async fn into_root_store(path: &Path) -> Result<RootCertStore, Error> {
    let ca_certs = into_certificate(path).await?;
    let mut cert_store = RootCertStore::empty();
    cert_store.add_parsable_certificates(ca_certs.as_slice());
    Ok(cert_store)
}

async fn into_private_key(path: &Path) -> Result<Vec<u8>, Error> {
    let mut f = File::open(path)
        .await
        .map_err(|e| Error::General(format!("Failed to load file: {}", e)))?;
    let mut buf = Vec::new();
    f.read_to_end(&mut buf)
        .await
        .map_err(|e| Error::General(format!("Failed to read file: {}", e)))?;
    let mut cursor = Cursor::new(buf);

    let parsers = [
        rustls_pemfile::rsa_private_keys,
        rustls_pemfile::pkcs8_private_keys,
    ];
    for parser in parsers.iter() {
        cursor.set_position(0);

        match parser(&mut cursor) {
            Ok(keys) if keys.is_empty() => continue,
            Ok(mut keys) if keys.len() == 1 => return Ok(rustls::PrivateKey(keys.pop().unwrap()).0),
            Ok(keys) => {
                return Err(Error::General(format!(
                    "Unexpected number of keys: {} (only 1 supported)",
                    keys.len()
                )));
            }
            // try the next parser
            Err(_) => continue,
        }
    }
    Err(Error::General(
        "could not load any valid private keys".to_string(),
    ))
}


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

use rustls_mtls::{initialize_logger, MtlsProvider};
use s2n_quic::{client::Connect, Client};
use std::{error::Error, net::SocketAddr};

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CACERT_PEM: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/certs/ca-cert.pem");
pub static MY_CERT_PEM: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/certs/client-cert.pem");
pub static MY_KEY_PEM: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/certs/client-key.pem");

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    initialize_logger("client");
    let provider = MtlsProvider::new(CACERT_PEM, MY_CERT_PEM, MY_KEY_PEM).await?;
    let client = Client::builder()
        .with_event(s2n_quic::provider::event::tracing::Subscriber::default())?
        .with_tls(provider)?
        .with_io("0.0.0.0:0")?
        .start()?;

    let addr: SocketAddr = "127.0.0.1:4433".parse()?;
    let connect = Connect::new(addr).with_server_name("localhost");
    let mut connection = client.connect(connect).await?;

    // ensure the connection doesn't time out with inactivity
    connection.keep_alive(true)?;

    // open a new stream and split the receiving and sending sides
    let stream = connection.open_bidirectional_stream().await?;
    let (mut receive_stream, mut send_stream) = stream.split();

    // spawn a task that copies responses from the server to stdout
    tokio::spawn(async move {
        let mut stdout = tokio::io::stdout();
        let _ = tokio::io::copy(&mut receive_stream, &mut stdout).await;
    });

    // copy data from stdin and send it to the server
    let mut stdin = tokio::io::stdin();
    tokio::io::copy(&mut stdin, &mut send_stream).await?;

    Ok(())
}


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

use rustls_mtls::{initialize_logger, MtlsProvider};
use s2n_quic::Server;
use std::error::Error;

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CACERT_PEM: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/certs/ca-cert.pem");
pub static MY_CERT_PEM: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/certs/server-cert.pem");
pub static MY_KEY_PEM: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/certs/server-key.pem");

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    initialize_logger("server");
    let provider = MtlsProvider::new(CACERT_PEM, MY_CERT_PEM, MY_KEY_PEM).await?;
    let mut server = Server::builder()
        .with_event(s2n_quic::provider::event::tracing::Subscriber::default())?
        .with_tls(provider)?
        .with_io("127.0.0.1:4433")?
        .start()?;

    while let Some(mut connection) = server.accept().await {
        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                // spawn a new task for the stream
                tokio::spawn(async move {
                    eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                    // echo any data back to the stream
                    while let Ok(Some(data)) = stream.receive().await {
                        stream.send(data).await.expect("stream should be open");
                    }
                });
            }
        });
    }

    Ok(())
}


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

use s2n_quic::{client::Connect, Client};
use std::{error::Error, net::SocketAddr};

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let client = Client::builder()
        .with_tls(CERT_PEM)?
        .with_io("0.0.0.0:0")?
        .start()?;

    let addr: SocketAddr = "127.0.0.1:4433".parse()?;
    let connect = Connect::new(addr).with_server_name("localhost");
    let mut connection = client.connect(connect).await?;

    // ensure the connection doesn't time out with inactivity
    connection.keep_alive(true)?;

    // open a new stream and split the receiving and sending sides
    let stream = connection.open_bidirectional_stream().await?;
    let (mut receive_stream, mut send_stream) = stream.split();

    // spawn a task that copies responses from the server to stdout
    tokio::spawn(async move {
        let mut stdout = tokio::io::stdout();
        let _ = tokio::io::copy(&mut receive_stream, &mut stdout).await;
    });

    // copy data from stdin and send it to the server
    let mut stdin = tokio::io::stdin();
    tokio::io::copy(&mut stdin, &mut send_stream).await?;

    Ok(())
}


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

#![cfg(test)]

use s2n_quic::{
    client::Connect,
    provider::{event, io},
    Client, Server,
};
use std::net::SocketAddr;
use turmoil::{lookup, Builder, Result};

/// NOTE: this certificate is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
/// NOTE: this certificate is to be used for demonstration purposes only!
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[test]
fn lossy_handshake() -> Result {
    let mut sim = Builder::new()
        .simulation_duration(core::time::Duration::from_secs(20))
        .build();

    sim.host("server", || async move {
        let io = io::turmoil::Builder::default()
            .with_address(bind_to(443))?
            .build()?;

        let mut server = Server::builder()
            .with_io(io)?
            .with_tls((CERT_PEM, KEY_PEM))?
            .with_event(events())?
            .start()?;

        while let Some(mut connection) = server.accept().await {
            tokio::spawn(async move {
                eprintln!("Connection accepted from {:?}", connection.remote_addr());

                while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                    tokio::spawn(async move {
                        eprintln!("Stream opened from {:?}", stream.connection().remote_addr());

                        // echo any data back to the stream
                        while let Ok(Some(data)) = stream.receive().await {
                            stream.send(data).await.expect("stream should be open");
                        }
                    });
                }
            });
        }

        Ok(())
    });

    sim.client("client", async move {
        let io = io::turmoil::Builder::default()
            .with_address(bind_to(1234))?
            .build()?;

        let client = Client::builder()
            .with_io(io)?
            .with_tls(CERT_PEM)?
            .with_event(events())?
            .start()?;

        // drop packets for 1 second
        drop_for(1);

        // even though we're dropping packets, the connection still goes through
        let server_addr: SocketAddr = (lookup("server"), 443).into();
        let mut connection = client
            .connect(Connect::new(server_addr).with_server_name("localhost"))
            .await?;

        // drop packets for 5 seconds
        drop_for(5);

        // even though we're dropping packets, the stream should still complete
        let mut stream = connection.open_bidirectional_stream().await?;
        stream.send(vec![1, 2, 3].into()).await?;
        stream.finish()?;

        let response = stream.receive().await?.unwrap();
        assert_eq!(&response[..], &[1, 2, 3]);

        Ok(())
    });

    sim.run()?;

    Ok(())
}

pub fn events() -> event::tracing::Provider {
    use std::sync::Once;

    static TRACING: Once = Once::new();

    // make sure this only gets initialized once
    TRACING.call_once(|| {
        use tokio::time::Instant;

        struct TokioUptime {
            epoch: Instant,
        }

        impl Default for TokioUptime {
            fn default() -> Self {
                Self {
                    epoch: Instant::now(),
                }
            }
        }

        impl tracing_subscriber::fmt::time::FormatTime for TokioUptime {
            fn format_time(
                &self,
                w: &mut tracing_subscriber::fmt::format::Writer,
            ) -> std::fmt::Result {
                write!(w, "{:?}", self.epoch.elapsed())
            }
        }

        let format = tracing_subscriber::fmt::format()
            .with_level(false) // don't include levels in formatted output
            .with_timer(TokioUptime::default())
            .with_ansi(false)
            .compact(); // Use a less verbose output format.

        tracing_subscriber::fmt()
            .with_env_filter(tracing_subscriber::EnvFilter::new("trace"))
            .event_format(format)
            .with_test_writer()
            .init();
    });

    event::tracing::Provider::default()
}

fn bind_to(port: u16) -> SocketAddr {
    (std::net::Ipv4Addr::UNSPECIFIED, port).into()
}

fn drop_for(secs: u64) {
    turmoil::partition("client", "server");
    tokio::spawn(async move {
        sleep_ms(secs * 1000).await;
        turmoil::repair("client", "server");
    });
}

async fn sleep_ms(millis: u64) {
    tokio::time::sleep(core::time::Duration::from_millis(millis)).await
}


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

use core::task::Poll;
use s2n_quic::{
    client::Connect,
    provider::datagram::default::{Endpoint, Receiver},
    Client,
};
use std::{error::Error, net::SocketAddr};

/// NOTE: this certificate/key pair is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Create a datagram provider that has recv queue capacity
    let datagram_provider = Endpoint::builder()
        .with_recv_capacity(200)?
        .build()
        .unwrap();

    // Build an `s2n_quic::Client`
    let client = Client::builder()
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io("0.0.0.0:0")?
        .with_datagram(datagram_provider)?
        .start()?;

    let addr: SocketAddr = "127.0.0.1:4433".parse()?;
    let connect = Connect::new(addr).with_server_name("localhost");
    let connection = client.connect(connect).await?;

    loop {
        let recv_result = futures::future::poll_fn(|cx| {
            // datagram_mut takes a closure which calls the requested datagram function. The type
            // of the closure parameter should be either the datagram Sender type or the
            // datagram Receiver type. The datagram_mut function will check this type against
            // its stored datagram Sender and Receiver, and if the type matches, the requested
            // function will execute. Here, that requested function is poll_recv_datagram.
            match connection.datagram_mut(|recv: &mut Receiver| recv.poll_recv_datagram(cx)) {
                // If the function is successfully called on the provider, it will return Poll<Bytes>.
                // Here we send an Ok() to wrap around the Bytes so the poll_fn doesn't complain.
                Ok(poll_value) => poll_value.map(Ok),
                // The datagram_mut function may return a query error if it can't find the type
                // referenced in the closure. Here we wrap the error in a Poll::Ready enum so the
                // poll_fn doesn't complain.
                Err(query_err) => Poll::Ready(Err(query_err)),
            }
        })
        .await;

        match recv_result {
            Ok(value) => eprintln!("RECV {:?}", value),
            Err(err) => {
                eprintln!("{:?}", err);
                return Ok(());
            }
        }
    }
}


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

use bytes::Bytes;
use s2n_quic::{
    provider::datagram::default::{Endpoint, Sender},
    Server,
};
use std::error::Error;

/// NOTE: this certificate/key pair is to be used for demonstration purposes only!
pub static CERT_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/cert.pem"
));
pub static KEY_PEM: &str = include_str!(concat!(
    env!("CARGO_MANIFEST_DIR"),
    "/../../quic/s2n-quic-core/certs/key.pem"
));

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    // Create a datagram provider that has a send queue capacity
    let datagram_provider = Endpoint::builder()
        .with_send_capacity(200)?
        .build()
        .unwrap();

    // Build an `s2n_quic::Server`
    let mut server = Server::builder()
        .with_tls((CERT_PEM, KEY_PEM))?
        .with_io("127.0.0.1:4433")?
        .with_datagram(datagram_provider)?
        .start()?;

    while let Some(connection) = server.accept().await {
        // spawn a new task for the connection
        tokio::spawn(async move {
            eprintln!("Connection accepted from {:?}", connection.remote_addr());

            loop {
                // Add datagrams to the send queue by passing in a closure that calls
                // the desired datagram send function
                let send_func = |x: &mut Sender| {
                    match x.send_datagram(Bytes::from_static(&[1, 2, 3])) {
                        Ok(_) => {
                            // The datagram was successfully inserted into the send queue
                        }
                        Err(err) => {
                            eprintln!("{}", err);
                            // An error was encountered while calling the send_datagram
                            // method. Either the peer didn't advertise support for datagrams
                            // or the send queue is at capacity.
                        }
                    }
                };

                if connection.datagram_mut(send_func).is_err() {
                    eprintln!("closed");
                    return;
                }

                tokio::time::sleep(core::time::Duration::from_secs(1)).await;
            }
        });
    }

    Ok(())
}


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

use core::task::{Context, Poll};
use std::collections::HashSet;

pub trait Checkpoints {
    fn park(&mut self, id: u64) -> Poll<()>;
    fn unpark(&mut self, id: u64, cx: &mut Context);
}

impl Checkpoints for HashSet<u64> {
    fn park(&mut self, id: u64) -> Poll<()> {
        if self.remove(&id) {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }

    fn unpark(&mut self, id: u64, cx: &mut Context) {
        self.insert(id);
        // notify the task to make more progress
        cx.waker().wake_by_ref();
    }
}


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

use crate::{
    driver::timer,
    operation as ops,
    scenario::{self, Scenario},
    timer::Timer,
    Checkpoints, Result, Trace,
};
use std::{
    collections::BTreeMap,
    future::Future,
    net::SocketAddr,
    sync::Arc,
    task::{Context, Poll},
};

mod thread;

pub trait Client<'a> {
    type Connect: Future<Output = Result<Self::Connection>> + Unpin;
    type Connection: Connection;

    fn connect(
        &mut self,
        addr: SocketAddr,
        hostname: &str,
        server_conn_id: u64,
        ops: &'a Arc<scenario::Connection>,
    ) -> Self::Connect;
}

pub trait Connection: crate::driver::timer::Provider {
    fn id(&self) -> u64;

    fn poll<T: Trace, Ch: Checkpoints>(
        &mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        now: crate::driver::timer::Timestamp,
        cx: &mut Context,
    ) -> Poll<Result<()>>;
}

pub struct Driver<'a, C: Client<'a>> {
    client: C,
    thread: thread::Thread<'a, C>,
    addresses: &'a AddressMap,
}

impl<'a, C: Client<'a>> Driver<'a, C> {
    pub fn new(client: C, scenario: &'a scenario::Client, addresses: &'a AddressMap) -> Self {
        Self {
            client,
            thread: thread::Thread::new(scenario, &scenario.scenario),
            addresses,
        }
    }

    pub async fn run<T: Trace, Ch: Checkpoints, Ti: Timer>(
        mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        timer: &mut Ti,
    ) -> Result<C> {
        futures::future::poll_fn(|cx| self.poll_with_timer(trace, checkpoints, timer, cx)).await?;
        Ok(self.client)
    }

    pub fn poll_with_timer<T: Trace, Ch: Checkpoints, Ti: Timer>(
        &mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        timer: &mut Ti,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        let now = timer.now();
        let res = self.poll(trace, checkpoints, now, cx);

        if let Some(target) = timer::Provider::next_expiration(&self) {
            // update the timer with the next expiration
            let _ = timer.poll(target, cx);
        };

        res
    }

    pub fn poll<T: Trace, Ch: Checkpoints>(
        &mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        now: timer::Timestamp,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        self.thread.poll(
            &mut self.client,
            self.addresses,
            trace,
            checkpoints,
            now,
            cx,
        )
    }
}

impl<'a, C: Client<'a>> timer::Provider for Driver<'a, C> {
    fn timers<Q>(&self, query: &mut Q) -> timer::Result
    where
        Q: timer::Query,
    {
        self.thread.timers(query)
    }
}

pub struct AddressMap {
    routers: Vec<BTreeMap<u64, SocketAddr>>,
    servers: Vec<SocketAddr>,
    hostnames: Vec<String>,
}

pub trait Resolver {
    fn server(&mut self, server_id: u64) -> Result<String>;
    fn router(&mut self, router_id: u64, server_id: u64) -> Result<String>;
}

async fn resolve_routes<R: Resolver>(
    id: u64,
    ops: &[ops::Client],
    routes: &mut BTreeMap<u64, SocketAddr>,
    resolver: &mut R,
) -> Result<()> {
    let mut pending: Vec<_> = ops.iter().collect();

    while let Some(op) = pending.pop() {
        match op {
            ops::Client::Connect {
                server_id,
                router_id,
                ..
            } if Some(id) == *router_id => {
                let host = resolver.router(id, *server_id)?;
                let mut addr = tokio::net::lookup_host(host).await?;
                let addr = addr.next().ok_or("invalid address")?;
                routes.insert(*server_id, addr);
            }
            ops::Client::Scope { threads } => {
                for thread in threads {
                    pending.extend(thread);
                }
            }
            _ => {}
        }
    }

    Ok(())
}

impl AddressMap {
    pub async fn new<R: Resolver>(
        scenario: &Scenario,
        client_id: u64,
        resolver: &mut R,
    ) -> Result<Self> {
        let client = &scenario.clients[client_id as usize];

        let mut servers = vec![];
        for id in 0..scenario.servers.len() {
            let host = resolver.server(id as u64)?;
            let mut addr = tokio::net::lookup_host(host).await?;
            let addr = addr.next().ok_or("invalid address")?;
            servers.push(addr);
        }

        let mut routers = vec![];
        for router_id in 0..scenario.routers.len() {
            let mut routes = BTreeMap::new();

            resolve_routes(router_id as u64, &client.scenario, &mut routes, resolver).await?;

            routers.push(routes);
        }

        let hostname_len = scenario
            .servers
            .iter()
            .map(|server| server.connections.len())
            .max()
            .unwrap_or(0);

        let mut hostnames = vec![];
        for idx in 0..hostname_len {
            hostnames.push(format!("{}.{}.net", idx, scenario.id));
        }

        Ok(Self {
            servers,
            routers,
            hostnames,
        })
    }

    pub fn server(&self, server_id: u64) -> SocketAddr {
        self.servers[server_id as usize]
    }

    pub fn router(&self, router_id: u64, server_id: u64) -> SocketAddr {
        *self.routers[router_id as usize]
            .get(&server_id)
            .expect("missing server route")
    }

    pub fn hostname(&self, connection_id: u64) -> &str {
        &self.hostnames[connection_id as usize]
    }
}


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

use crate::Result;
use core::task::{Context, Poll};

pub trait Connection {
    fn id(&self) -> u64;
    fn poll_open_bidirectional_stream(&mut self, id: u64, cx: &mut Context) -> Poll<Result<()>>;
    fn poll_open_send_stream(&mut self, id: u64, cx: &mut Context) -> Poll<Result<()>>;
    fn poll_accept_stream(&mut self, cx: &mut Context) -> Poll<Result<Option<u64>>>;
    fn poll_send(
        &mut self,
        owner: Owner,
        id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>>;
    fn poll_receive(
        &mut self,
        owner: Owner,
        id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>>;
    fn poll_send_finish(&mut self, owner: Owner, id: u64, cx: &mut Context) -> Poll<Result<()>>;
    fn poll_receive_finish(&mut self, owner: Owner, id: u64, cx: &mut Context) -> Poll<Result<()>>;
    fn poll_progress(&mut self, cx: &mut Context) -> Poll<Result<()>> {
        let _ = cx;
        Ok(()).into()
    }
    fn poll_finish(&mut self, cx: &mut Context) -> Poll<Result<()>> {
        let _ = cx;
        Ok(()).into()
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Owner {
    Local,
    Remote,
}

impl<T> core::ops::Index<Owner> for [T; 2] {
    type Output = T;

    fn index(&self, index: Owner) -> &Self::Output {
        match index {
            Owner::Local => &self[0],
            Owner::Remote => &self[1],
        }
    }
}

impl<T> core::ops::IndexMut<Owner> for [T; 2] {
    fn index_mut(&mut self, index: Owner) -> &mut Self::Output {
        match index {
            Owner::Local => &mut self[0],
            Owner::Remote => &mut self[1],
        }
    }
}

impl core::ops::Not for Owner {
    type Output = Self;

    fn not(self) -> Self::Output {
        match self {
            Self::Local => Self::Remote,
            Self::Remote => Self::Local,
        }
    }
}