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

//! Provides connection id support for an endpoint

pub use s2n_quic_core::connection::id::{ConnectionInfo, Format, Generator, LocalId, Validator};

pub trait Provider: 'static {
    type Format: 'static + Format;
    type Error: core::fmt::Display;

    fn start(self) -> Result<Self::Format, Self::Error>;
}

pub use default::Provider as Default;

impl_provider_utils!();

impl<T: 'static + Format> Provider for T {
    type Format = T;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Format, Self::Error> {
        Ok(self)
    }
}

pub mod default {
    use core::{
        convert::{Infallible, TryInto},
        time::Duration,
    };
    use rand::prelude::*;
    use s2n_quic_core::connection::{
        self,
        id::{ConnectionInfo, Generator, Validator},
    };

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

    impl super::Provider for Provider {
        type Format = Format;
        type Error = Infallible;

        fn start(self) -> Result<Self::Format, Self::Error> {
            Ok(self.0)
        }
    }

    /// 16 bytes should be big enough for a randomly generated Id
    const DEFAULT_LEN: usize = 16;

    /// Randomly generated connection Id format.
    ///
    /// By default, connection Ids of length 16 bytes are generated.
    #[derive(Debug)]
    pub struct Format {
        len: usize,
        lifetime: Option<Duration>,
    }

    impl Default for Format {
        fn default() -> Self {
            Self {
                len: DEFAULT_LEN,
                lifetime: None,
            }
        }
    }

    impl Format {
        /// Creates a builder for the format
        pub fn builder() -> Builder {
            Builder::default()
        }
    }

    /// A builder for [`Format`] providers
    #[derive(Debug)]
    pub struct Builder {
        len: usize,
        lifetime: Option<Duration>,
    }

    impl Default for Builder {
        fn default() -> Self {
            Self {
                len: DEFAULT_LEN,
                lifetime: None,
            }
        }
    }

    impl Builder {
        /// Sets the length of the generated connection Id
        pub fn with_len(mut self, len: usize) -> Result<Self, connection::id::Error> {
            if !(connection::LocalId::MIN_LEN..=connection::id::MAX_LEN).contains(&len) {
                return Err(connection::id::Error::InvalidLength);
            }
            self.len = len;
            Ok(self)
        }

        /// Sets the lifetime of each generated connection Id
        pub fn with_lifetime(mut self, lifetime: Duration) -> Result<Self, connection::id::Error> {
            if !(connection::id::MIN_LIFETIME..=connection::id::MAX_LIFETIME).contains(&lifetime) {
                return Err(connection::id::Error::InvalidLifetime);
            }
            self.lifetime = Some(lifetime);
            Ok(self)
        }

        /// Builds the [`Format`] into a provider
        pub fn build(self) -> Result<Format, core::convert::Infallible> {
            Ok(Format {
                len: self.len,
                lifetime: self.lifetime,
            })
        }
    }

    impl Generator for Format {
        fn generate(&mut self, _connection_info: &ConnectionInfo) -> connection::LocalId {
            let mut id = [0u8; connection::id::MAX_LEN];
            let id = &mut id[..self.len];
            rand::thread_rng().fill_bytes(id);
            (&*id).try_into().expect("length already checked")
        }

        fn lifetime(&self) -> Option<Duration> {
            self.lifetime
        }
    }

    impl Validator for Format {
        fn validate(&self, _connection_info: &ConnectionInfo, buffer: &[u8]) -> Option<usize> {
            if buffer.len() >= self.len {
                Some(self.len)
            } else {
                None
            }
        }
    }

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

        #[test]
        fn generator_test() {
            let remote_address = &s2n_quic_core::inet::SocketAddress::default();
            let connection_info = ConnectionInfo::new(remote_address);

            for len in connection::LocalId::MIN_LEN..connection::id::MAX_LEN {
                let mut format = Format::builder().with_len(len).unwrap().build().unwrap();

                let id = format.generate(&connection_info);

                //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.2
                //= type=test
                //# An endpoint that uses this design MUST
                //# either use the same connection ID length for all connections or
                //# encode the length of the connection ID such that it can be recovered
                //# without state.
                assert_eq!(format.validate(&connection_info, id.as_ref()), Some(len));
                assert_eq!(id.len(), len);
                assert_eq!(format.lifetime(), None);
            }

            assert_eq!(
                Some(connection::id::Error::InvalidLength),
                Format::builder()
                    .with_len(connection::id::MAX_LEN + 1)
                    .err()
            );

            assert_eq!(
                Some(connection::id::Error::InvalidLength),
                Format::builder()
                    .with_len(connection::LocalId::MIN_LEN - 1)
                    .err()
            );

            let lifetime = Duration::from_secs(1000);
            let format = Format::builder()
                .with_lifetime(lifetime)
                .unwrap()
                .build()
                .unwrap();
            assert_eq!(Some(lifetime), format.lifetime());

            assert_eq!(
                Some(connection::id::Error::InvalidLifetime),
                Format::builder()
                    .with_lifetime(connection::id::MIN_LIFETIME - Duration::from_millis(1))
                    .err()
            );

            assert_eq!(
                Some(connection::id::Error::InvalidLifetime),
                Format::builder()
                    .with_lifetime(connection::id::MAX_LIFETIME + Duration::from_millis(1))
                    .err()
            );
        }
    }
}


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

//! Provides unreliable datagram support

use s2n_quic_core::datagram::Disabled;
pub use s2n_quic_core::datagram::{
    default,
    traits::{
        ConnectionInfo, Endpoint, Packet, PreConnectionInfo, ReceiveContext, Receiver, Sender,
        WriteError,
    },
};

pub trait Provider {
    type Endpoint: Endpoint;
    type Error: 'static + core::fmt::Display;

    fn start(self) -> Result<Self::Endpoint, Self::Error>;
}

impl_provider_utils!();

pub type Default = Disabled;

impl<T: 'static + Send + Endpoint> Provider for T {
    type Endpoint = T;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Endpoint, Self::Error> {
        Ok(self)
    }
}


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

//! Allows applications to limit peer's ability to open new connections

pub use s2n_quic_core::endpoint::{
    limits::{ConnectionAttempt, Outcome},
    Limiter,
};
use s2n_quic_core::{event::Timestamp, path::THROTTLED_PORTS_LEN};

pub trait Provider: 'static {
    type Limits: 'static + Limiter;
    type Error: core::fmt::Display;

    /// Starts the token provider
    fn start(self) -> Result<Self::Limits, Self::Error>;
}

use core::time::Duration;
pub use default::Limits as Default;

impl_provider_utils!();

impl<T: 'static + Limiter> Provider for T {
    type Limits = T;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Limits, Self::Error> {
        Ok(self)
    }
}

const THROTTLED_PORT_LIMIT: usize = 10;
const THROTTLE_FREQUENCY: Duration = Duration::from_secs(1);

#[derive(Default, Debug, Clone, Copy)]
struct BasicRateLimiter {
    last_throttle_reset: Option<Timestamp>,
    count: usize,
}

impl BasicRateLimiter {
    /// Throttles a connection based on a limit.
    /// Returns True if the `should_throttle` invoke count is greater than `limit` and the
    /// connection has not been throttled in the last `throttle_frequency` duration.
    ///
    /// If the throttle timer expires the count is reset and `should_throttle` returns False.
    ///
    /// Returns False if the `should_throttle` invoke count is less than `limit`.
    fn should_throttle(
        &mut self,
        limit: usize,
        throttle_frequency: Duration,
        connection_attempt: &ConnectionAttempt,
    ) -> bool {
        self.count += 1;
        let timestamp = connection_attempt.timestamp;

        if self.count > limit {
            match self.last_throttle_reset {
                // If the throttle timer is still within the throttle_frequency
                // then throttle the connection.
                Some(last_throttle_reset)
                    if timestamp.saturating_duration_since(last_throttle_reset)
                        < throttle_frequency =>
                {
                    return true;
                }
                // If the throttle timer is greater than the throttle_frequency
                // then reset the throttle count and the throttle timer.
                // Let the connection through.
                _ => {
                    self.count = 0;
                    self.last_throttle_reset = Some(timestamp);
                    return false;
                }
            };
        }

        // If this is the first time calling instantiate the throttle timer.
        if self.last_throttle_reset.is_none() {
            self.last_throttle_reset = Some(timestamp);
        }

        false
    }
}

#[cfg(test)]
mod tests {
    use super::{BasicRateLimiter, THROTTLED_PORT_LIMIT, THROTTLE_FREQUENCY};
    use core::time::Duration;
    use s2n_quic_core::{
        endpoint::limits::ConnectionAttempt,
        event::IntoEvent,
        inet::SocketAddress,
        time::{testing::Clock as MockClock, Clock},
    };

    #[test]
    fn first_throttle_reset() {
        let remote_address = SocketAddress::default();
        let mock_clock = MockClock::default();
        let info =
            ConnectionAttempt::new(0, 0, &remote_address, mock_clock.get_time().into_event());

        let mut rate_limiter = BasicRateLimiter::default();
        // The first time the throttle limit is hit the timer will be created so we expect to be
        // able to connect THROTTLED_PORT_LIMIT amount of times before the connection is throttled.
        // Note: This test should run fast enough to not let the timer reset.
        let very_long_freq = Duration::MAX;

        for request in 0..(THROTTLED_PORT_LIMIT * 3) {
            if request >= THROTTLED_PORT_LIMIT {
                assert!(rate_limiter.should_throttle(THROTTLED_PORT_LIMIT, very_long_freq, &info));
            } else {
                assert!(!rate_limiter.should_throttle(THROTTLED_PORT_LIMIT, very_long_freq, &info));
            }
        }
    }

    #[test]
    fn throttle_timer_reset() {
        let remote_address = SocketAddress::default();
        let mut mock_clock = MockClock::default();

        let mut rate_limiter = BasicRateLimiter::default();
        let short_freq = Duration::from_millis(10);
        let sleep_longer_than_short_freq = Duration::from_millis(500);

        // This test should never throttle because everytime the limit is about to get hit the
        // thread sleeps long enough for the throttle reset timer to fire.
        for request in 0..(THROTTLED_PORT_LIMIT * 3) {
            let info =
                ConnectionAttempt::new(0, 0, &remote_address, mock_clock.get_time().into_event());
            if request % THROTTLED_PORT_LIMIT == 0 {
                mock_clock.inc_by(sleep_longer_than_short_freq)
            }
            assert!(!rate_limiter.should_throttle(THROTTLED_PORT_LIMIT, short_freq, &info));
        }
    }

    #[test]
    fn throttle_constants_changed() {
        // If the constants change consider modifying the above test cases to make sure we are
        // confident that we are hitting all the correct conditions.
        //
        // For example if we increase THROTTLE_FREQUENCY to a very large period, do the above
        // tests still make sense?
        assert_eq!(THROTTLED_PORT_LIMIT, 10);
        assert_eq!(THROTTLE_FREQUENCY, Duration::from_secs(1));
    }
}

pub mod default {
    //! Default provider for the endpoint limits.

    use super::*;
    use core::convert::Infallible;

    /// Allows the endpoint limits to be built with specific values
    ///
    /// # Examples
    ///
    /// Set the maximum inflight handshakes for this endpoint.
    ///
    /// ```rust
    /// use s2n_quic::provider::endpoint_limits;
    /// # use std::error::Error;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn Error>> {
    /// let limits = endpoint_limits::Default::builder()
    ///     .with_inflight_handshake_limit(100)?
    ///     .build();
    ///
    ///     Ok(())
    /// # }
    /// ```
    #[derive(Default)]
    pub struct Builder {
        max_inflight_handshake_limit: Option<usize>,
    }

    impl Builder {
        /// Sets limit on inflight handshakes
        pub fn with_inflight_handshake_limit(mut self, limit: usize) -> Result<Self, Infallible> {
            self.max_inflight_handshake_limit = Some(limit);
            Ok(self)
        }

        /// Build the limits
        pub fn build(self) -> Result<Limits, Infallible> {
            Ok(Limits {
                max_inflight_handshake_limit: self.max_inflight_handshake_limit,
                rate_limiter: [BasicRateLimiter::default(); THROTTLED_PORTS_LEN],
            })
        }
    }

    #[derive(Clone, Copy, Debug)]
    pub struct Limits {
        /// Maximum number of handshakes to allow before Retry packets are queued
        max_inflight_handshake_limit: Option<usize>,
        rate_limiter: [BasicRateLimiter; THROTTLED_PORTS_LEN],
    }

    impl Limits {
        pub fn builder() -> Builder {
            Builder::default()
        }
    }

    /// Default implementation for the Limits
    impl super::Limiter for Limits {
        fn on_connection_attempt(&mut self, info: &ConnectionAttempt) -> Outcome {
            let remote_port = info.remote_address.port();
            if s2n_quic_core::path::remote_port_blocked(remote_port) {
                return Outcome::drop();
            }

            if let Some(port_index) = s2n_quic_core::path::remote_port_throttled_index(remote_port)
            {
                let rate_limiter = &mut self.rate_limiter[port_index];
                if rate_limiter.should_throttle(THROTTLED_PORT_LIMIT, THROTTLE_FREQUENCY, info) {
                    return Outcome::drop();
                }
            }

            if let Some(limit) = self.max_inflight_handshake_limit {
                if info.inflight_handshakes >= limit {
                    return Outcome::retry();
                }
            }

            Outcome::allow()
        }
    }

    /// Default limit values are as non-intrusive as possible
    impl std::default::Default for Limits {
        fn default() -> Self {
            Self {
                max_inflight_handshake_limit: None,
                rate_limiter: [BasicRateLimiter::default(); THROTTLED_PORTS_LEN],
            }
        }
    }

    #[test]
    fn builder_test() {
        let elp = Limits::builder()
            .with_inflight_handshake_limit(100)
            .unwrap()
            .build()
            .unwrap();
        assert_eq!(elp.max_inflight_handshake_limit, Some(100));
    }

    #[test]
    fn blocked_port_connection_attempt() {
        use s2n_quic_core::{
            event::IntoEvent,
            inet::SocketAddress,
            time::{testing::Clock as MockClock, Clock},
        };

        let mut remote_address = SocketAddress::default();
        let mut limits = Limits::builder().build().unwrap();
        let mock_clock = MockClock::default();

        for port in 0..u16::MAX {
            let blocked_expected = s2n_quic_core::path::remote_port_blocked(port);

            remote_address.set_port(port);
            let info =
                ConnectionAttempt::new(0, 0, &remote_address, mock_clock.get_time().into_event());
            let outcome = limits.on_connection_attempt(&info);

            if blocked_expected {
                assert_eq!(Outcome::drop(), outcome);
            } else {
                assert_eq!(Outcome::allow(), outcome);
            }
        }
    }
}


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

//! Provides IO support for an endpoint

use s2n_quic_core::{endpoint::Endpoint, inet::SocketAddress, path::Handle as PathHandle};
use std::io;

pub trait Provider: 'static {
    type PathHandle: PathHandle;
    type Error: 'static + core::fmt::Display;

    fn start<E: Endpoint<PathHandle = Self::PathHandle>>(
        self,
        endpoint: E,
    ) -> Result<SocketAddress, Self::Error>;
}

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

#[cfg(feature = "unstable-provider-io-turmoil")]
pub mod turmoil;

#[cfg(feature = "unstable-provider-io-xdp")]
pub mod xdp;

pub mod tokio;

pub use self::tokio as default;

pub use default::Provider as Default;

impl TryInto for u16 {
    type Error = io::Error;
    type Provider = Default;

    fn try_into(self) -> io::Result<Self::Provider> {
        Default::new(("::", self))
    }
}

macro_rules! impl_socket_addrs {
    ($ty:ty) => {
        impl TryInto for $ty {
            type Error = io::Error;
            type Provider = Default;

            fn try_into(self) -> io::Result<Self::Provider> {
                Default::new(self)
            }
        }
    };
}

impl_socket_addrs!((&str, u16));
impl_socket_addrs!((std::net::IpAddr, u16));
impl_socket_addrs!((std::net::Ipv4Addr, u16));
impl_socket_addrs!((std::net::Ipv6Addr, u16));
impl_socket_addrs!(&str);
impl_socket_addrs!(std::net::SocketAddr);
impl_socket_addrs!(std::net::SocketAddrV4);
impl_socket_addrs!(std::net::SocketAddrV6);

impl_provider_utils!();


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

//! Provides limits support for a connection

pub use s2n_quic_core::connection::limits::{ConnectionInfo, Limiter, Limits};

pub trait Provider {
    type Limits: 'static + Send + Limiter;
    type Error: 'static + core::fmt::Display;

    fn start(self) -> Result<Self::Limits, Self::Error>;
}

pub use default::Provider as Default;

impl_provider_utils!();

impl<T: 'static + Send + Limiter> Provider for T {
    type Limits = T;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Limits, Self::Error> {
        Ok(self)
    }
}

pub mod default {
    #[derive(Debug, Default)]
    pub struct Provider(());

    impl super::Provider for Provider {
        type Limits = super::Limits;
        type Error = core::convert::Infallible;

        fn start(self) -> Result<Self::Limits, Self::Error> {
            Ok(Self::Limits::default())
        }
    }
}


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

macro_rules! impl_provider_utils {
    () => {
        /// Converts a value into a [`Provider`]
        pub trait TryInto {
            type Provider: Provider;
            type Error: 'static + core::fmt::Display;

            fn try_into(self) -> Result<Self::Provider, Self::Error>;
        }

        #[doc(hidden)]
        pub trait With<T: Provider> {
            type Output;

            fn with(self, provider: T) -> Self::Output;
        }

        /// Automatically implement anything that implements [`Provider`]
        impl<T: Provider> TryInto for T {
            type Error = core::convert::Infallible;
            type Provider = T;

            fn try_into(self) -> Result<Self::Provider, Self::Error> {
                Ok(self)
            }
        }
    };
}

macro_rules! impl_providers_state {
    (
        $(#[$($attr:tt)*])* struct Providers {
            $(
                $(
                    #[doc = $($doc:tt)*]
                )*
                $field:ident: $field_ty:ident
            ),* $(,)?
        }

        $(
            #[doc = $($trait_doc:tt)*]
        )*
        trait $trait:ident {

        }
    ) => {
        #[doc(hidden)]
        $(
            #[$($attr)*]
        )*
        pub struct Providers<$($field_ty,)*> {
            $(
                $(
                    #[doc = $($doc)*]
                )*
                $field: $field_ty,
            )*
        }

        $(
            #[doc = $($trait_doc)*]
        )*
        pub trait $trait {
            $(
                #[doc(hidden)]
                type $field_ty: $field::Provider;
            )*

            #[doc(hidden)]
            fn build(self) -> Providers<$(Self::$field_ty,)*>;

            #[doc(hidden)]
            fn as_ref(&self) -> Providers<$(&Self::$field_ty,)*>;

            #[doc(hidden)]
            fn as_mut(&mut self) -> Providers<$(&mut Self::$field_ty,)*>;
        }

        #[doc(hidden)]
        impl<$($field_ty: $field::Provider,)*> $trait for Providers<$($field_ty,)*> {
            $(
                type $field_ty = $field_ty;
            )*

            fn build(self) -> Providers<$(Self::$field_ty,)*> {
                self
            }

            fn as_ref(&self) -> Providers<$(&Self::$field_ty,)*> {
                Providers {
                    $(
                        $field: &self.$field,
                    )*
                }
            }

            fn as_mut(&mut self) -> Providers<$(&mut Self::$field_ty,)*> {
                Providers {
                    $(
                        $field: &mut self.$field,
                    )*
                }
            }
        }

        /// The recommended providers for the endpoint.
        ///
        /// The implementation details are intentionally hidden and may
        /// change between releases.
        #[derive(Debug, Default)]
        pub struct DefaultProviders {
            providers: Providers<$($field::Default,)*>
        }

        #[doc(hidden)]
        impl $trait for DefaultProviders {
            $(
                type $field_ty = $field::Default;
            )*

            fn build(self) -> Providers<$(Self::$field_ty,)*> {
                self.providers
            }

            fn as_ref(&self) -> Providers<$(&Self::$field_ty,)*> {
                Providers {
                    $(
                        $field: &self.providers.$field,
                    )*
                }
            }

            fn as_mut(&mut self) -> Providers<$(&mut Self::$field_ty,)*> {
                Providers {
                    $(
                        $field: &mut self.providers.$field,
                    )*
                }
            }
        }

        impl_providers_state!(@with, $trait, {}, {$($field: $field_ty),*});
    };
    (@with, $trait:ident, { $($field:ident: $field_ty:ident),* }, {}) => {
        // done
    };
    (@with, $trait:ident, { $($prev:ident: $prev_ty:ident),* }, { $field:ident: $field_ty:ident $(, $rest:ident: $rest_ty:ident)* }) => {
        impl<Provider: $trait, New: $field::Provider> $field::With<New> for Builder<Provider> {
            type Output = Builder<Providers<$(Provider::$prev_ty, )* New $(, Provider::$rest_ty)*>>;

            fn with(self, $field: New) -> Self::Output {
                let providers = self.0.build();
                Builder(Providers {
                    $field,
                    $(
                        $prev: providers.$prev,
                    )*
                    $(
                        $rest: providers.$rest,
                    )*
                })
            }
        }

        impl_providers_state!(@with, $trait, { $($prev: $prev_ty, )* $field: $field_ty }, {$($rest: $rest_ty),*});
    }
}

macro_rules! impl_provider_method {
    ($(#[$($attr:tt)*])* $name:ident, $field:ident, $trait:ident) => {
            $(
                #[$($attr)*]
            )*
            pub fn $name<T, U>(self, $field: T) -> Result<Builder<U>, T::Error>
            where
                T: $field::TryInto,
                U: $trait,
                Self: $field::With<T::Provider, Output = Builder<U>>,
            {
                let $field = $field.try_into()?;
                let builder = $field::With::with(self, $field);
                Ok(builder)
            }

    };
}


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

pub use s2n_quic_core::packet::interceptor::{
    loss, Disabled, Havoc, Interceptor as PacketInterceptor, Loss,
};

/// Provides packet_interceptor support for an endpoint
pub trait Provider: 'static {
    type PacketInterceptor: 'static + PacketInterceptor;
    type Error: core::fmt::Display;

    fn start(self) -> Result<Self::PacketInterceptor, Self::Error>;
}

pub type Default = Disabled;

impl_provider_utils!();

impl<T: 'static + Send + PacketInterceptor> Provider for T {
    type PacketInterceptor = T;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::PacketInterceptor, Self::Error> {
        Ok(self)
    }
}


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

pub use s2n_quic_core::path::migration::{
    default::{self, Validator as Default},
    disabled, Attempt, Outcome, Validator,
};

/// Provides limits support for an endpoint
pub trait Provider {
    type Validator: 'static + Send + Validator;
    type Error: 'static + core::fmt::Display;

    fn start(self) -> Result<Self::Validator, Self::Error>;
}

impl_provider_utils!();

impl<T: 'static + Send + Validator> Provider for T {
    type Validator = T;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Validator, Self::Error> {
        Ok(self)
    }
}


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

pub use s2n_quic_core::random::Generator;

/// Provides random number generation support for an endpoint
pub trait Provider: 'static {
    type Generator: 'static + Generator;
    type Error: core::fmt::Display;

    fn start(self) -> Result<Self::Generator, Self::Error>;
}

pub use self::rand::Provider as Default;

impl_provider_utils!();

mod rand {
    use core::convert::Infallible;
    use rand::{
        prelude::*,
        rngs::{adapter::ReseedingRng, OsRng},
    };
    use rand_chacha::ChaChaCore;
    use s2n_quic_core::random;

    // Number of generated bytes after which to reseed the public and private random
    // generators. This value is based on THREAD_RNG_RESEED_THRESHOLD from rand::rngs::thread.rs
    const RESEED_THRESHOLD: u64 = 1024 * 64;

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

    impl super::Provider for Provider {
        type Generator = Generator;
        type Error = Infallible;

        fn start(self) -> Result<Self::Generator, Self::Error> {
            Ok(self.0)
        }
    }

    impl super::TryInto for Generator {
        type Provider = Provider;
        type Error = Infallible;

        fn try_into(self) -> Result<Self::Provider, Self::Error> {
            Ok(Provider(self))
        }
    }

    /// Randomly generated bits.
    #[derive(Debug)]
    pub struct Generator {
        public: ReseedingRng<ChaChaCore, OsRng>,
        private: ReseedingRng<ChaChaCore, OsRng>,
    }

    impl Default for Generator {
        fn default() -> Self {
            Self {
                public: build_rng(),
                private: build_rng(),
            }
        }
    }

    // Constructs a `ReseedingRng` with a ChaCha RNG initially seeded from the OS,
    // that will reseed from the OS after RESEED_THRESHOLD is exceeded
    fn build_rng() -> ReseedingRng<ChaChaCore, OsRng> {
        let prng = ChaChaCore::from_rng(OsRng::default())
            .unwrap_or_else(|err| panic!("could not initialize random generator: {err}"));
        ReseedingRng::new(prng, RESEED_THRESHOLD, OsRng::default())
    }

    impl random::Generator for Generator {
        fn public_random_fill(&mut self, dest: &mut [u8]) {
            self.public.fill_bytes(dest)
        }

        fn private_random_fill(&mut self, dest: &mut [u8]) {
            self.private.fill_bytes(dest)
        }
    }

    #[cfg(test)]
    mod tests {
        use s2n_quic_core::random::Generator;

        #[test]
        fn generator_test() {
            let mut generator = super::Generator::default();

            let mut dest_1 = [0; 20];
            let mut dest_2 = [0; 20];

            generator.public_random_fill(&mut dest_1);
            generator.public_random_fill(&mut dest_2);

            assert_ne!(dest_1, dest_2);

            generator.private_random_fill(&mut dest_1);
            generator.private_random_fill(&mut dest_2);

            assert_ne!(dest_1, dest_2);
        }
    }
}


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

//! Provides stateless reset token support for an endpoint

pub use s2n_quic_core::stateless_reset::token::Generator;

pub trait Provider: 'static {
    type Generator: 'static + Generator;
    type Error: core::fmt::Display;

    fn start(self) -> Result<Self::Generator, Self::Error>;
}

pub use random::Provider as Default;

impl_provider_utils!();

mod random {
    use core::convert::Infallible;
    use rand::prelude::*;
    use s2n_quic_core::{frame::new_connection_id::STATELESS_RESET_TOKEN_LEN, stateless_reset};

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

    impl super::Provider for Provider {
        type Generator = Generator;
        type Error = Infallible;

        fn start(self) -> Result<Self::Generator, Self::Error> {
            Ok(self.0)
        }
    }

    impl super::TryInto for Generator {
        type Provider = Provider;
        type Error = Infallible;

        fn try_into(self) -> Result<Self::Provider, Self::Error> {
            Ok(Provider(self))
        }
    }

    /// Randomly generated stateless reset token.
    #[derive(Debug, Default)]
    pub struct Generator {}

    impl stateless_reset::token::Generator for Generator {
        /// Since a random stateless reset token will not be recognized by the peer, this generator
        /// is not enabled and no stateless reset packet will be sent to the peer.
        const ENABLED: bool = false;
        /// This stateless reset token generator produces a random token on each call, and
        /// thus does not enable stateless reset functionality, as the token provided to the
        /// peer with a new connection ID will be different than the token sent in a stateless
        /// reset. To enable stateless reset functionality, the stateless reset token must
        /// be generated the same for a given `local_connection_id` before and after loss of state.
        fn generate(&mut self, _local_connection_id: &[u8]) -> stateless_reset::Token {
            let mut token = [0u8; STATELESS_RESET_TOKEN_LEN];
            rand::thread_rng().fill_bytes(&mut token);
            token.into()
        }
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use s2n_quic_core::{connection, stateless_reset::token::Generator as _};

        #[test]
        fn stateless_reset_token_test() {
            let mut generator = Generator::default();
            let id = connection::LocalId::try_from_bytes(b"id01").unwrap();

            let token_1 = generator.generate(id.as_bytes());
            let token_2 = generator.generate(id.as_bytes());

            assert_ne!(token_1, token_2);
        }
    }
}


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

/// Provides synchronization support for an endpoint
pub trait Provider {
    type Sync: 'static + Send;
    type Error: 'static + core::fmt::Display;

    fn start(self) -> Result<Self::Sync, Self::Error>;
}

pub use default::Provider as Default;

impl_provider_utils!();

pub mod default {
    #[derive(Debug, Default)]
    pub struct Provider;

    impl super::Provider for Provider {
        type Sync = (); // TODO
        type Error = core::convert::Infallible;

        fn start(self) -> Result<Self::Sync, Self::Error> {
            // TODO
            Ok(())
        }
    }
}


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

//! Provides TLS support for an endpoint

use cfg_if::cfg_if;
use s2n_quic_core::crypto;

pub trait Provider {
    type Server: 'static + crypto::tls::Endpoint;
    type Client: 'static + crypto::tls::Endpoint;
    type Error: 'static + core::fmt::Display;

    /// Creates a server endpoint for the given provider
    fn start_server(self) -> Result<Self::Server, Self::Error>;

    /// Creates a client endpoint for the given provider
    fn start_client(self) -> Result<Self::Client, Self::Error>;
}

impl_provider_utils!();

cfg_if! {
    if #[cfg(feature = "provider-tls-default")] {
        pub mod default {
            //! Provides the recommended implementation of TLS using platform detection
            pub use super::default_tls::*;
        }
    } else if #[cfg(feature = "provider-tls-s2n")] {
        pub use s2n_tls as default;
    } else if #[cfg(feature = "provider-tls-rustls")] {
        pub use rustls as default;
    } else {
        pub mod default {
            //! Provides the recommended implementation of TLS using platform detection
            pub use super::default_tls::*;
        }
    }
}

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

impl Provider for Default {
    type Server = default::Server;
    type Client = default::Client;
    type Error = core::convert::Infallible;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        Ok(Self::Server::default())
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        Ok(Self::Client::default())
    }
}

impl Provider for (&std::path::Path, &std::path::Path) {
    type Server = <Default as Provider>::Server;
    type Client = <Default as Provider>::Client;
    type Error = Box<dyn std::error::Error>;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        let server = default::Server::builder()
            .with_certificate(self.0, self.1)?
            .build()?;

        Ok(server)
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        // TODO support private key
        let client = default::Client::builder()
            .with_certificate(self.0)?
            .build()?;

        Ok(client)
    }
}

impl Provider for &std::path::Path {
    type Server = <Default as Provider>::Server;
    type Client = <Default as Provider>::Client;
    type Error = Box<dyn std::error::Error>;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        let empty_cert: &[u8] = &[];
        let server = default::Server::builder()
            .with_certificate(empty_cert, self)?
            .build()?;

        Ok(server)
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        let client = default::Client::builder().with_certificate(self)?.build()?;

        Ok(client)
    }
}

impl Provider for (&[u8], &[u8]) {
    type Server = <Default as Provider>::Server;
    type Client = <Default as Provider>::Client;
    type Error = Box<dyn std::error::Error>;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        let server = default::Server::builder()
            .with_certificate(self.0, self.1)?
            .build()?;

        Ok(server)
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        // TODO support private key
        let client = default::Client::builder()
            .with_certificate(self.0)?
            .build()?;

        Ok(client)
    }
}

impl Provider for &[u8] {
    type Server = <Default as Provider>::Server;
    type Client = <Default as Provider>::Client;
    type Error = Box<dyn std::error::Error>;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        let empty_cert = &[][..];
        let server = default::Server::builder()
            .with_certificate(empty_cert, self)?
            .build()?;

        Ok(server)
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        let client = default::Client::builder().with_certificate(self)?.build()?;

        Ok(client)
    }
}

impl Provider for (&str, &str) {
    type Server = <Default as Provider>::Server;
    type Client = <Default as Provider>::Client;
    type Error = Box<dyn std::error::Error>;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        let server = default::Server::builder()
            .with_certificate(self.0, self.1)?
            .build()?;

        Ok(server)
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        // TODO support private key
        let client = default::Client::builder()
            .with_certificate(self.0)?
            .build()?;

        Ok(client)
    }
}

impl Provider for &str {
    type Server = <Default as Provider>::Server;
    type Client = <Default as Provider>::Client;
    type Error = Box<dyn std::error::Error>;

    fn start_server(self) -> Result<Self::Server, Self::Error> {
        let empty_cert = "";
        let server = default::Server::builder()
            .with_certificate(empty_cert, self)?
            .build()?;

        Ok(server)
    }

    fn start_client(self) -> Result<Self::Client, Self::Error> {
        let client = default::Client::builder().with_certificate(self)?.build()?;

        Ok(client)
    }
}

#[cfg(feature = "provider-tls-default")]
mod default_tls {
    pub use s2n_quic_tls_default::*;

    // We need to implement the provider trait for whatever the default is as long as it's
    // not called out explicitly and we're not on the platform that uses it by default.
    //
    // Note: I know this looks like a mess. And it is. Hopefully in the future cargo will support
    // platform-specific default features.
    #[cfg(not(any(
        all(not(unix), feature = "s2n-quic-rustls"),
        all(unix, feature = "s2n-quic-tls")
    )))]
    mod default_provider {
        use super::*;
        use crate::provider::tls;

        impl tls::Provider for Server {
            type Server = Self;
            type Client = Client;
            type Error = core::convert::Infallible;

            fn start_server(self) -> Result<Self::Server, Self::Error> {
                Ok(self)
            }

            fn start_client(self) -> Result<Self::Client, Self::Error> {
                panic!("cannot create a client from a server");
            }
        }

        impl tls::Provider for Client {
            type Server = Server;
            type Client = Self;
            type Error = core::convert::Infallible;

            fn start_server(self) -> Result<Self::Server, Self::Error> {
                panic!("cannot create a server from a client");
            }

            fn start_client(self) -> Result<Self::Client, Self::Error> {
                Ok(self)
            }
        }
    }
}
#[cfg(not(feature = "provider-tls-default"))]
mod default_tls {
    // TODO stub out default that fails with error when started
}

#[cfg(feature = "s2n-quic-rustls")]
pub mod rustls {
    //! Provides the [rustls](https://docs.rs/rustls/) implementation of TLS

    pub use s2n_quic_rustls::*;

    impl super::Provider for Server {
        type Server = Self;
        type Client = Client;
        type Error = core::convert::Infallible;

        fn start_server(self) -> Result<Self::Server, Self::Error> {
            Ok(self)
        }

        fn start_client(self) -> Result<Self::Client, Self::Error> {
            panic!("cannot create a client from a server");
        }
    }

    impl super::Provider for Client {
        type Server = Server;
        type Client = Self;
        type Error = core::convert::Infallible;

        fn start_server(self) -> Result<Self::Server, Self::Error> {
            panic!("cannot create a server from a client");
        }

        fn start_client(self) -> Result<Self::Client, Self::Error> {
            Ok(self)
        }
    }
}

#[cfg(feature = "s2n-quic-tls")]
pub mod s2n_tls {
    //! Provides the [s2n-tls](https://github.com/aws/s2n-tls) implementation of TLS
    pub use s2n_quic_tls::*;

    impl<L: ConfigLoader> super::Provider for Server<L> {
        type Server = Self;
        type Client = Client;
        type Error = core::convert::Infallible;

        fn start_server(self) -> Result<Self::Server, Self::Error> {
            Ok(self)
        }

        fn start_client(self) -> Result<Self::Client, Self::Error> {
            panic!("cannot create a client from a server");
        }
    }

    impl<L: ConfigLoader> super::Provider for Client<L> {
        type Server = Server;
        type Client = Self;
        type Error = core::convert::Infallible;

        fn start_server(self) -> Result<Self::Server, Self::Error> {
            panic!("cannot create a server from a client");
        }

        fn start_client(self) -> Result<Self::Client, Self::Error> {
            Ok(self)
        }
    }
}


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

//! Default provider for Address Validation tokens
//!
//! The default provider will randomly generate a 256 bit key. This key will be used to sign and
//! verify tokens. The key can be rotated at a duration set by the user.
//!
//! The default provider does not support tokens delivered in a NEW_TOKEN frame.

use core::{mem::size_of, time::Duration};
use hash_hasher::HashHasher;
use s2n_codec::{DecoderBuffer, DecoderBufferMut};
use s2n_quic_core::{
    connection, event::api::SocketAddress, random, time::Timestamp, token::Source,
};
use s2n_quic_crypto::{constant_time, digest, hmac};
use std::hash::{Hash, Hasher};
use zerocopy::{AsBytes, FromBytes, Unaligned};
use zeroize::Zeroizing;

struct BaseKey {
    active_duration: Duration,

    // HMAC key for signing and verifying
    key: Option<(Timestamp, hmac::Key)>,

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //# To protect against such attacks, servers MUST ensure that
    //# replay of tokens is prevented or limited.
    duplicate_filter: cuckoofilter::CuckooFilter<HashHasher>,
}

impl BaseKey {
    pub fn new(active_duration: Duration) -> Self {
        Self {
            active_duration,
            key: None,
            duplicate_filter: cuckoofilter::CuckooFilter::with_capacity(
                cuckoofilter::DEFAULT_CAPACITY,
            ),
        }
    }

    pub fn hasher(&mut self, random: &mut dyn random::Generator) -> Option<hmac::Context> {
        let key = self.poll_key(random)?;
        Some(hmac::Context::with_key(&key))
    }

    fn poll_key(&mut self, random: &mut dyn random::Generator) -> Option<hmac::Key> {
        let now = s2n_quic_platform::time::now();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-21.3
        //# Servers SHOULD provide mitigations for this attack by limiting the
        //# usage and lifetime of address validation tokens; see Section 8.1.3.
        if let Some((expires_at, key)) = self.key.as_ref() {
            if expires_at > &now {
                // key is still valid
                return Some(key.clone());
            }
        }

        let expires_at = now.checked_add(self.active_duration)?;

        // TODO in addition to generating new key material, clear out the filter used for detecting
        // duplicates.
        let mut key_material = Zeroizing::new([0; digest::SHA256_OUTPUT_LEN]);
        random.private_random_fill(&mut key_material[..]);
        let key = hmac::Key::new(hmac::HMAC_SHA256, key_material.as_ref());

        // TODO clear the filter instead of recreating. This is pending a merge to crates.io
        // (https://github.com/axiomhq/rust-cuckoofilter/pull/52)
        self.duplicate_filter =
            cuckoofilter::CuckooFilter::with_capacity(cuckoofilter::DEFAULT_CAPACITY);

        self.key = Some((expires_at, key));

        self.key.as_ref().map(|key| key.1.clone())
    }
}

const DEFAULT_KEY_ROTATION_PERIOD: Duration = Duration::from_millis(1000);

#[derive(Debug)]
pub struct Provider {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
    //# Thus, a token SHOULD have an
    //# expiration time, which could be either an explicit expiration time or
    //# an issued timestamp that can be used to dynamically calculate the
    //# expiration time.
    /// To fulfill this SHOULD, we rotate the key periodically. This allows
    /// customers to control the token lifetime without adding bytes to the token itself.
    key_rotation_period: Duration,
}

impl Default for Provider {
    fn default() -> Self {
        Self {
            key_rotation_period: DEFAULT_KEY_ROTATION_PERIOD,
        }
    }
}

impl super::Provider for Provider {
    type Format = Format;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Format, Self::Error> {
        // The keys must remain valid for two rotation periods or they will regenerate their
        // material and validation will fail.
        let format = Format {
            key_rotation_period: self.key_rotation_period,
            current_key_rotates_at: s2n_quic_platform::time::now(),
            current_key: 0,
            keys: [
                BaseKey::new(self.key_rotation_period * 2),
                BaseKey::new(self.key_rotation_period * 2),
            ],
        };

        Ok(format)
    }
}

pub struct Format {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //= type=exception
    //= reason=We use a duplicate filter to prevent tokens from being used more than once.
    //# Servers are encouraged to allow tokens to be used only
    //# once, if possible; tokens MAY include additional information about
    //# clients to further narrow applicability or reuse.
    /// Key validity period
    key_rotation_period: Duration,

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //# Servers SHOULD ensure that
    //# tokens sent in Retry packets are only accepted for a short time.
    /// Timestamp to rotate current key
    current_key_rotates_at: s2n_quic_core::time::Timestamp,

    /// Which key is used to sign
    current_key: u8,

    /// Key used to sign keys
    keys: [BaseKey; 2],
}

impl Format {
    fn current_key(&mut self) -> u8 {
        let now = s2n_quic_platform::time::now();
        if now > self.current_key_rotates_at {
            self.current_key ^= 1;
            self.current_key_rotates_at = now + self.key_rotation_period;

            // TODO either clear the duplicate filter here, or implement in the BaseKey logic
            // https://github.com/aws/s2n-quic/issues/173
        }
        self.current_key
    }

    // Retry Tokens need to include the original destination connection id from the transport
    // parameters. This OCID is included in the tag.
    fn tag_retry_token(
        &mut self,
        token: &Token,
        context: &mut super::Context<'_>,
    ) -> Option<hmac::Tag> {
        let mut ctx = self.keys[token.header.key_id() as usize].hasher(context.random)?;

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //# Tokens
        //# sent in Retry packets SHOULD include information that allows the
        //# server to verify that the source IP address and port in client
        //# packets remain constant.
        ctx.update(&token.original_destination_connection_id);
        ctx.update(&token.nonce);
        ctx.update(context.peer_connection_id);
        match context.remote_address {
            SocketAddress::IpV4 { ip, port, .. } => {
                ctx.update(ip);
                ctx.update(&port.to_be_bytes());
            }
            SocketAddress::IpV6 { ip, port, .. } => {
                ctx.update(ip);
                ctx.update(&port.to_be_bytes());
            }
            _ => {
                // we are unable to hash the address so bail
                return None;
            }
        };

        Some(ctx.sign())
    }

    // Using the key id in the token, verify the token
    fn validate_retry_token(
        &mut self,
        context: &mut super::Context<'_>,
        token: &Token,
    ) -> Option<connection::InitialId> {
        if self.keys[token.header.key_id() as usize]
            .duplicate_filter
            .contains(token)
        {
            return None;
        }

        let tag = self.tag_retry_token(token, context)?;

        if constant_time::verify_slices_are_equal(&token.hmac, tag.as_ref()).is_ok() {
            // Only add the token once it has been validated. This will prevent the filter from
            // being filled with garbage tokens.

            // Ignore the outcome of adding a token to the filter because we always want to
            // continue the connection if the filter fails.
            let _ = self.keys[token.header.key_id() as usize]
                .duplicate_filter
                .add(token);

            return token.original_destination_connection_id();
        }

        None
    }
}

impl super::Format for Format {
    const TOKEN_LEN: usize = size_of::<Token>();

    /// The default provider does not support NEW_TOKEN frame tokens
    fn generate_new_token(
        &mut self,
        _context: &mut super::Context<'_>,
        _source_connection_id: &connection::LocalId,
        _output_buffer: &mut [u8],
    ) -> Option<()> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=418
        //# A server MAY provide clients with an address validation token during
        //# one connection that can be used on a subsequent connection.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=TODO
        //= tracking-issue=346
        //# Tokens sent in NEW_TOKEN frames MUST include information that allows
        //# the server to verify that the client IP address has not changed from
        //# when the token was issued.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=345
        //# A token issued with NEW_TOKEN MUST NOT include information that would
        //# allow values to be linked by an observer to the connection on which
        //# it was issued.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=387
        //# A server MUST ensure that every NEW_TOKEN frame it sends
        //# is unique across all clients, with the exception of those sent to
        //# repair losses of previously sent NEW_TOKEN frames.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=394
        //# A server MAY provide clients with an address validation token during
        //# one connection that can be used on a subsequent connection.

        None
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.2
    //# Requiring the server
    //# to provide a different connection ID, along with the
    //# original_destination_connection_id transport parameter defined in
    //# Section 18.2, forces the server to demonstrate that it, or an entity
    //# it cooperates with, received the original Initial packet from the
    //# client.
    fn generate_retry_token(
        &mut self,
        context: &mut super::Context<'_>,
        original_destination_connection_id: &connection::InitialId,
        output_buffer: &mut [u8],
    ) -> Option<()> {
        let buffer = DecoderBufferMut::new(output_buffer);
        let (token, _) = buffer
            .decode::<&mut Token>()
            .expect("Provided output buffer did not match TOKEN_LEN");

        let header = Header::new(Source::RetryPacket, self.current_key());

        token.header = header;
        token.original_destination_connection_id[..original_destination_connection_id.len()]
            .copy_from_slice(original_destination_connection_id.as_bytes());
        token.odcid_len = original_destination_connection_id.len() as u8;

        // ensure the other CID bytes are zeroed out
        for b in token
            .original_destination_connection_id
            .iter_mut()
            .skip(original_destination_connection_id.len())
        {
            *b = 0;
        }

        // Populate the nonce before signing
        context.random.public_random_fill(&mut token.nonce[..]);

        let tag = self.tag_retry_token(token, context)?;

        token.hmac.copy_from_slice(tag.as_ref());

        Some(())
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
    //# When a server receives an Initial packet with an address validation
    //# token, it MUST attempt to validate the token, unless it has already
    //# completed address validation.
    fn validate_token(
        &mut self,
        context: &mut super::Context<'_>,
        token: &[u8],
    ) -> Option<connection::InitialId> {
        let buffer = DecoderBuffer::new(token);
        let (token, remaining) = buffer.decode::<&Token>().ok()?;

        // Verify the provided token doesn't have any additional data
        remaining.ensure_empty().ok()?;

        if token.header.version() != TOKEN_VERSION {
            return None;
        }

        let source = token.header.token_source();

        match source {
            Source::RetryPacket => self.validate_retry_token(context, token),
            Source::NewTokenFrame => None, // Not supported in the default provider
        }
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=TODO
        //= tracking-issue=347
        //# Tokens that are provided
        //# in NEW_TOKEN frames (Section 19.7) need to be valid for longer but
        //# SHOULD NOT be accepted multiple times.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=TODO
        //= tracking-issue=388
        //# Clients that want to break continuity of identity with a server can
        //# discard tokens provided using the NEW_TOKEN frame.
    }
}

#[derive(Clone, Copy, Debug, FromBytes, AsBytes, Unaligned)]
#[repr(C)]
pub(crate) struct Header(u8);

const TOKEN_VERSION: u8 = 0x00;

const VERSION_SHIFT: u8 = 7;
const VERSION_MASK: u8 = 0x80;

const TOKEN_SOURCE_SHIFT: u8 = 6;
const TOKEN_SOURCE_MASK: u8 = 0x40;

const KEY_ID_SHIFT: u8 = 5;
const KEY_ID_MASK: u8 = 0x20;

impl Header {
    fn new(source: Source, key_id: u8) -> Header {
        let mut header: u8 = 0;
        header |= TOKEN_VERSION << VERSION_SHIFT;
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //# Information that
        //# allows the server to distinguish between tokens from Retry and
        //# NEW_TOKEN MAY be accessible to entities other than the server.
        header |= match source {
            Source::NewTokenFrame => 0 << TOKEN_SOURCE_SHIFT,
            Source::RetryPacket => 1 << TOKEN_SOURCE_SHIFT,
        };

        // The key_id can only be 0 or 1
        debug_assert!(key_id <= 1);
        header |= (key_id & 0x01) << KEY_ID_SHIFT;

        Header(header)
    }

    fn version(self) -> u8 {
        (self.0 & VERSION_MASK) >> VERSION_SHIFT
    }

    fn key_id(self) -> u8 {
        (self.0 & KEY_ID_MASK) >> KEY_ID_SHIFT
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.1
    //# A token sent in a NEW_TOKEN frame or a Retry packet MUST be
    //# constructed in a way that allows the server to identify how it was
    //# provided to a client.  These tokens are carried in the same field but
    //# require different handling from servers.
    fn token_source(self) -> Source {
        match (self.0 & TOKEN_SOURCE_MASK) >> TOKEN_SOURCE_SHIFT {
            0 => Source::NewTokenFrame,
            1 => Source::RetryPacket,
            _ => Source::NewTokenFrame,
        }
    }
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
//#   There is no need for a single well-defined format for the token
//#   because the server that generates the token also consumes it.
#[derive(Copy, Clone, Debug, FromBytes, AsBytes, Unaligned)]
#[repr(C)]
struct Token {
    header: Header,

    odcid_len: u8,
    original_destination_connection_id: [u8; 20],

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //# An address validation token MUST be difficult to guess.  Including a
    //# random value with at least 128 bits of entropy in the token would be
    //# sufficient, but this depends on the server remembering the value it
    //# sends to clients.
    nonce: [u8; 32],

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
    //# A token-based scheme allows the server to offload any state
    //# associated with validation to the client.  For this design to work,
    //# the token MUST be covered by integrity protection against
    //# modification or falsification by clients.  Without integrity
    //# protection, malicious clients could generate or guess values for
    //# tokens that would be accepted by the server.  Only the server
    //# requires access to the integrity protection key for tokens.
    hmac: [u8; 32],
}

s2n_codec::zerocopy_value_codec!(Token);

impl Hash for Token {
    /// Token hashes are taken from the hmac
    fn hash<H: Hasher>(&self, state: &mut H) {
        state.write(&self.hmac);
    }
}

impl Token {
    pub fn original_destination_connection_id(&self) -> Option<connection::InitialId> {
        let dcid = self
            .original_destination_connection_id
            .get(..self.odcid_len as usize)?;
        connection::InitialId::try_from_bytes(dcid)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use s2n_quic_core::{
        inet::SocketAddress,
        random,
        token::{Context, Format as FormatTrait, Source},
    };
    use s2n_quic_platform::time;
    use std::{net::SocketAddr, sync::Arc};

    const TEST_KEY_ROTATION_PERIOD: Duration = Duration::from_millis(1000);

    fn get_test_format() -> Format {
        Format {
            key_rotation_period: TEST_KEY_ROTATION_PERIOD,
            keys: [
                BaseKey::new(TEST_KEY_ROTATION_PERIOD * 2),
                BaseKey::new(TEST_KEY_ROTATION_PERIOD * 2),
            ],
            current_key_rotates_at: time::now(),
            current_key: 0,
        }
    }

    #[test]
    fn test_header() {
        // Test all combinations of values to create a header and verify the header returns the
        // expected values.
        for source in &[Source::NewTokenFrame, Source::RetryPacket] {
            for key_id in [0, 1] {
                let header = Header::new(*source, key_id);
                // The version should always be the constant TOKEN_VERSION
                assert_eq!(header.version(), TOKEN_VERSION);
                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.1
                //= type=test
                //# A token sent in a NEW_TOKEN frames or a Retry packet MUST be
                //# constructed in a way that allows the server to identify how it was
                //# provided to a client.

                assert_eq!(header.token_source(), *source);
                assert_eq!(header.key_id(), key_id);
            }
        }
    }

    #[test]
    fn test_valid_retry_tokens() {
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock.clone());

        let mut format = get_test_format();
        let first_conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let second_conn_id = connection::PeerId::try_from_bytes(&[1, 3, 5, 7, 9]).unwrap();
        let orig_conn_id =
            connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut first_token = [0; Format::TOKEN_LEN];
        let mut second_token = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &first_conn_id, &mut random);

        // Generate two tokens for different connections
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut first_token)
            .unwrap();

        context = Context::new(&addr, &second_conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut second_token)
            .unwrap();

        clock.adjust_by(TEST_KEY_ROTATION_PERIOD);
        context = Context::new(&addr, &first_conn_id, &mut random);
        assert_eq!(
            format.validate_token(&mut context, &first_token),
            Some(orig_conn_id)
        );
        context = Context::new(&addr, &second_conn_id, &mut random);
        assert_eq!(
            format.validate_token(&mut context, &second_token),
            Some(orig_conn_id)
        );
        context = Context::new(&addr, &first_conn_id, &mut random);
        assert_eq!(format.validate_token(&mut context, &second_token), None);
    }

    #[test]
    fn test_retry_ip_port_validation() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# Tokens
        //# sent in Retry packets SHOULD include information that allows the
        //# server to verify that the source IP address and port in client
        //# packets remain constant.
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let orig_conn_id =
            connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();

        let mut token = [0; Format::TOKEN_LEN];
        let ip_address = "127.0.0.1:443";
        let addr: SocketAddr = ip_address.parse().unwrap();
        let correct_address: SocketAddress = addr.into();
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&correct_address, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut token)
            .unwrap();

        let ip_address = "127.0.0.2:443";
        let addr: SocketAddr = ip_address.parse().unwrap();
        let incorrect_address: SocketAddress = addr.into();
        context = Context::new(&incorrect_address, &conn_id, &mut random);
        assert_eq!(format.validate_token(&mut context, &token), None);

        let ip_address = "127.0.0.1:444";
        let addr: SocketAddr = ip_address.parse().unwrap();
        let incorrect_port: SocketAddress = addr.into();
        context = Context::new(&incorrect_port, &conn_id, &mut random);
        assert_eq!(format.validate_token(&mut context, &token), None);

        // Verify the token is still valid after the failed attempts
        context = Context::new(&correct_address, &conn_id, &mut random);
        assert!(format.validate_token(&mut context, &token).is_some());
    }

    #[test]
    fn test_key_rotation() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
        //= type=test
        //# Thus, a token SHOULD have an
        //# expiration time, which could be either an explicit expiration time or
        //# an issued timestamp that can be used to dynamically calculate the
        //# expiration time.
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock.clone());

        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let orig_conn_id = connection::InitialId::TEST_ID;
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut buf)
            .unwrap();

        // Validation should succeed because the signing key is still valid, even
        // though it has been rotated from the current signing key
        clock.adjust_by(TEST_KEY_ROTATION_PERIOD);
        assert!(format.validate_token(&mut context, &buf).is_some());

        // Validation should fail because the key used for signing has been regenerated
        clock.adjust_by(TEST_KEY_ROTATION_PERIOD);
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    #[test]
    fn test_expired_retry_token() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# Servers SHOULD ensure that
        //# tokens sent in Retry packets are only accepted for a short time.
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock.clone());

        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let orig_conn_id = connection::InitialId::TEST_ID;
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut buf)
            .unwrap();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-21.3
        //= type=test
        //# Servers SHOULD provide mitigations for this attack by limiting the
        //# usage and lifetime of address validation tokens; see Section 8.1.3.
        // Validation should fail because multiple rotation periods have elapsed
        clock.adjust_by(TEST_KEY_ROTATION_PERIOD * 2);
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    #[test]
    fn test_retry_validation_default_format() {
        let clock = Arc::new(time::testing::MockClock::new());
        time::testing::set_local_clock(clock);

        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let odcid = connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &odcid, &mut buf)
            .unwrap();

        assert_eq!(format.validate_token(&mut context, &buf), Some(odcid));

        let wrong_conn_id = connection::PeerId::try_from_bytes(&[0, 1, 2]).unwrap();
        context = Context::new(&addr, &wrong_conn_id, &mut random);
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    #[test]
    fn test_duplicate_token_detection() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# To protect against such attacks, servers MUST ensure that
        //# replay of tokens is prevented or limited.
        let mut format = get_test_format();
        let conn_id = connection::PeerId::TEST_ID;
        let odcid = connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut buf = [0; Format::TOKEN_LEN];
        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        format
            .generate_retry_token(&mut context, &odcid, &mut buf)
            .unwrap();

        assert_eq!(format.validate_token(&mut context, &buf), Some(odcid));

        // Second attempt with the same token should fail because the token is a duplicate
        assert!(format.validate_token(&mut context, &buf).is_none());
    }

    #[test]
    fn test_token_modification_detection() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# For this design to work,
        //# the token MUST be covered by integrity protection against
        //# modification or falsification by clients.
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let orig_conn_id =
            connection::InitialId::try_from_bytes(&[0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
        let addr = SocketAddress::default();
        let mut token = [0; Format::TOKEN_LEN];

        let mut random = random::testing::Generator(5);
        let mut context = Context::new(&addr, &conn_id, &mut random);
        // Generate two tokens for different connections
        format
            .generate_retry_token(&mut context, &orig_conn_id, &mut token)
            .unwrap();

        for i in 0..Format::TOKEN_LEN {
            random = random::testing::Generator(5);
            context = Context::new(&addr, &conn_id, &mut random);
            token[i] = !token[i];
            assert!(format.validate_token(&mut context, &token).is_none());
            token[i] = !token[i];
        }
    }

    #[test]
    fn test_token_length_check() {
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let addr = SocketAddress::default();

        bolero::check!().for_each(move |token| {
            let mut random = random::testing::Generator(5);
            let mut context = Context::new(&addr, &conn_id, &mut random);
            assert!(format.validate_token(&mut context, token).is_none())
        });
    }

    #[test]
    fn test_token_falsification_detection() {
        let mut format = get_test_format();
        let conn_id = connection::PeerId::try_from_bytes(&[2, 4, 6, 8, 10]).unwrap();
        let addr = SocketAddress::default();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# For this design to work,
        //# the token MUST be covered by integrity protection against
        //# modification or falsification by clients.
        let generator = bolero::generator::gen::<Vec<u8>>()
            .with()
            .len(Format::TOKEN_LEN);
        bolero::check!()
            .with_generator(generator)
            .for_each(move |token| {
                let mut random = random::testing::Generator(5);
                let mut context = Context::new(&addr, &conn_id, &mut random);
                assert!(format.validate_token(&mut context, token).is_none())
            });
    }
}


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

//! Provides address validation token functionality to a QUIC endpoint.

pub use s2n_quic_core::token::{Context, Format};

pub trait Provider: 'static {
    type Format: 'static + Format;
    type Error: 'static + core::fmt::Display;

    /// Starts the token provider
    fn start(self) -> Result<Self::Format, Self::Error>;
}

pub mod default;

pub use default::Provider as Default;

impl_provider_utils!();


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

use crate::provider::event::{ConnectionInfo, ConnectionMeta};

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

impl super::Provider for Provider {
    type Subscriber = Subscriber;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Subscriber, Self::Error> {
        Ok(Subscriber)
    }
}

pub struct Subscriber;

impl super::Subscriber for Subscriber {
    type ConnectionContext = ();

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


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

//! Allows endpoints to subscribe to connection-level and endpoint-level events

use cfg_if::cfg_if;
pub use s2n_quic_core::{
    endpoint::Location,
    event::{
        api as events,
        api::{ConnectionInfo, ConnectionMeta},
        supervisor, Event, Meta, Subscriber, Timestamp,
    },
    query,
};

/// Provides event handling support for an endpoint
pub trait Provider {
    type Subscriber: 'static + Subscriber;
    type Error: 'static + core::fmt::Display;

    fn start(self) -> Result<Self::Subscriber, Self::Error>;
}

/// Provides an implementation to disable all events
pub mod disabled;

/// This module contains event integration with [`tracing`](https://docs.rs/tracing)
#[cfg(any(feature = "provider-event-tracing", test))]
pub mod tracing;

cfg_if! {
    if #[cfg(any(feature = "provider-event-tracing", test))] {
        pub use self::tracing as default;
    } else {
        // Events are disabled by default.
        pub use self::disabled as default;
    }
}

pub use default::Provider as Default;

impl<S> Provider for S
where
    S: 'static + Subscriber,
{
    type Subscriber = S;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<S, Self::Error> {
        Ok(self)
    }
}

impl_provider_utils!();


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

pub use s2n_quic_core::event::tracing::Subscriber;

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

impl super::Provider for Provider {
    type Subscriber = Subscriber;
    type Error = core::convert::Infallible;

    fn start(self) -> Result<Self::Subscriber, Self::Error> {
        Ok(Subscriber::default())
    }
}


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

use core::time::Duration;
use s2n_quic_core::{endpoint::Endpoint, inet::SocketAddress};
pub use s2n_quic_platform::io::testing;
use std::io;

pub use self::testing::*;

pub type Error = Box<dyn 'static + std::error::Error>;
pub type Result<T = (), E = Error> = core::result::Result<T, E>;

impl super::Provider for Io {
    type PathHandle = PathHandle;
    type Error = io::Error;

    fn start<E: Endpoint<PathHandle = Self::PathHandle>>(
        self,
        endpoint: E,
    ) -> Result<SocketAddress, Self::Error> {
        let (_join_handle, local_addr) = Io::start(self, endpoint)?;
        Ok(local_addr)
    }
}

/// Runs a single test with the given network
///
/// Returns the total runtime of the test
pub fn test<N: Network, F: FnOnce(&Handle) -> Result<O>, O>(network: N, f: F) -> Result<Duration> {
    test_seed(network, 123456789, f)
}

/// Runs a single test with the given network and seed value
///
/// Returns the total runtime of the test
pub fn test_seed<N: Network, F: FnOnce(&Handle) -> Result<O>, O>(
    network: N,
    seed: u64,
    f: F,
) -> Result<Duration> {
    let mut executor = Executor::new(network, seed);
    let handle = executor.handle().clone();

    executor.enter(|| f(&handle))?;

    executor.run();

    // return the total runtime of the test
    let now = executor.enter(time::now);

    let now = unsafe { now.as_duration() };

    Ok(now)
}


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

//! Provides an implementation of the [`io::Provider`](crate::provider::io::Provider)
//! using the [`Tokio runtime`](https://docs.rs/tokio/latest/tokio/runtime/index.html)

use s2n_quic_core::{endpoint::Endpoint, inet::SocketAddress};
use s2n_quic_platform::io::tokio;
use std::io;

pub use self::tokio::{Builder, Io as Provider};

impl super::Provider for Provider {
    type PathHandle = tokio::PathHandle;
    type Error = io::Error;

    fn start<E: Endpoint<PathHandle = Self::PathHandle>>(
        self,
        endpoint: E,
    ) -> Result<SocketAddress, Self::Error> {
        let (_join_handle, local_addr) = Provider::start(self, endpoint)?;
        Ok(local_addr)
    }
}


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

//! Provides an implementation of the [`io::Provider`](crate::provider::io::Provider)
//! using the [`Turmoil network simulator`](https://docs.rs/turmoil).

use s2n_quic_core::{endpoint::Endpoint, inet::SocketAddress};
use s2n_quic_platform::io::turmoil;
use std::io;

pub use self::turmoil::{Builder, Io as Provider};

impl super::Provider for Provider {
    type PathHandle = turmoil::PathHandle;
    type Error = io::Error;

    fn start<E: Endpoint<PathHandle = Self::PathHandle>>(
        self,
        endpoint: E,
    ) -> Result<SocketAddress, Self::Error> {
        let (_join_handle, local_addr) = Provider::start(self, endpoint)?;
        Ok(local_addr)
    }
}


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

//! Provides an implementation of the [`io::Provider`](crate::provider::io::Provider)
//! using [AF_XDP](https://www.kernel.org/doc/html/latest/networking/af_xdp.html) sockets.

/// Export the platform items
pub use s2n_quic_platform::io::xdp::*;

impl<Rx, Tx> super::Provider for Provider<Rx, Tx>
where
    Rx: 'static + s2n_quic_core::io::rx::Rx<PathHandle = PathHandle> + Send,
    Tx: 'static + s2n_quic_core::io::tx::Tx<PathHandle = PathHandle> + Send,
{
    type PathHandle = PathHandle;
    type Error = std::io::Error;

    fn start<E: s2n_quic_core::endpoint::Endpoint<PathHandle = Self::PathHandle>>(
        self,
        endpoint: E,
    ) -> Result<s2n_quic_core::inet::SocketAddress, Self::Error> {
        let (_join_handle, local_addr) = Provider::start(self, endpoint)?;
        Ok(local_addr)
    }
}


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

use crate::{
    provider::*,
    server::{DefaultProviders, Server, ServerProviders},
};

/// A builder for configuring [`Server`] providers
#[derive(Debug)]
pub struct Builder<Providers>(pub(crate) Providers);

impl Default for Builder<DefaultProviders> {
    fn default() -> Self {
        Self(Default::default())
    }
}

impl<Providers: ServerProviders> Builder<Providers> {
    impl_provider_method!(
        /// Sets the connection ID provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// Uses the default connection ID provider with the default configuration.
        ///
        /// ```rust,no_run
        /// # use std::{error::Error, time::Duration};
        /// use s2n_quic::{Server, provider::connection_id};
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_connection_id(connection_id::Default::default())?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        ///
        /// Sets a custom connection ID provider for the server
        ///
        /// ```rust,no_run
        /// # use std::{error::Error, time::Duration};
        /// use s2n_quic::{Server, provider::connection_id};
        /// use rand::prelude::*;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// struct MyConnectionIdFormat;
        ///
        /// impl connection_id::Generator for MyConnectionIdFormat {
        ///     fn generate(&mut self, _conn_info: &connection_id::ConnectionInfo) -> connection_id::LocalId {
        ///         let mut id = [0u8; 16];
        ///         rand::thread_rng().fill_bytes(&mut id);
        ///         connection_id::LocalId::try_from_bytes(&id[..]).unwrap()
        ///     }
        /// }
        ///
        /// impl connection_id::Validator for MyConnectionIdFormat {
        ///     fn validate(&self, _conn_info: &connection_id::ConnectionInfo, packet: &[u8]) -> Option<usize> {
        ///         // this connection id format is always 16 bytes
        ///         Some(16)
        ///     }
        /// }
        ///
        /// let server = Server::builder()
        ///     .with_connection_id(MyConnectionIdFormat)?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_connection_id,
        connection_id,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the stateless reset token provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// Sets a custom stateless reset token provider for the server
        ///
        /// ```rust,ignore
        /// # use std::error::Error;
        /// use s2n_quic::Server;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_stateless_reset_token(MyStatelessResetTokenGenerator::new())?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_stateless_reset_token,
        stateless_reset_token,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the limits provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// Sets the max idle time, while inheriting the remaining default limits
        ///
        /// ```rust,no_run
        /// # use std::{error::Error, time::Duration};
        /// use s2n_quic::{Server, provider::limits};
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_limits(limits::Default::default())?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_limits,
        limits,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the endpoint limits provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// Sets the max inflight handshakes for an endpoint, while inheriting the remaining default limits
        ///
        /// ```rust,no_run
        /// # use std::{error::Error, time::Duration};
        /// use s2n_quic::{Server, provider::endpoint_limits};
        ///
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_endpoint_limits(endpoint_limits::Default::default())?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_endpoint_limits,
        endpoint_limits,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the event provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// Sets a custom event subscriber for the server
        ///
        /// ```rust,ignore
        /// # use std::{error::Error, time::Duration};
        /// use s2n_quic::Server;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_event(MyEventLogger::new())?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_event,
        event,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the token provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// Sets a custom token provider for the server
        ///
        /// ```rust,ignore
        /// # use std::{error::Error, time::Duration};
        /// use s2n_quic::Server;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_address_token(MyTokenProvider::new())?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_address_token,
        address_token,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the IO provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// Starts listening on [`127.0.0.1:443`](https://127.0.0.1)
        ///
        /// ```rust,no_run
        /// # use std::error::Error;
        /// use s2n_quic::Server;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_io("127.0.0.1:443")?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        ///
        /// Configures a socket with the provided `Builder`
        ///
        /// ```rust,no_run
        /// # use std::error::Error;
        /// use s2n_quic::{Server, provider::io::tokio::Builder as IoBuilder};
        /// use std::net::ToSocketAddrs;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let addr = "127.0.0.1:443".to_socket_addrs()?.next().unwrap();
        ///
        /// let io = IoBuilder::default()
        ///     .with_receive_address(addr)?
        ///     .build()?;
        ///
        /// let server = Server::builder()
        ///     .with_io(io)?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_io,
        io,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the TLS provider for the [`Server`]
        ///
        /// # Examples
        ///
        /// The default TLS provider and configuration will be used with the
        /// path to the private key.
        ///
        /// ```rust,no_run
        /// # use std::{error::Error, path::Path};
        /// # use s2n_quic::Server;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let server = Server::builder()
        ///     .with_tls((Path::new("./certs/cert.pem"), Path::new("./certs/key.pem")))?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        ///
        /// Sets the TLS provider to a TLS server builder
        ///
        /// ```rust,ignore
        /// # use std::{error::Error, path::Path};
        /// # use s2n_quic::Server;
        /// #
        /// # #[tokio::main]
        /// # async fn main() -> Result<(), Box<dyn Error>> {
        /// let tls = s2n::tls::Server::builder()
        ///     .with_certificate(Path::new("./certs/cert.pem"), Path::new("./certs/key.pem"))?
        ///     .with_security_policy(s2n::tls::security_policy::S2N_20190802)?;
        ///
        /// let server = Server::builder()
        ///     .with_tls(tls)?
        ///     .start()?;
        /// #
        /// #    Ok(())
        /// # }
        /// ```
        with_tls,
        tls,
        ServerProviders
    );

    #[cfg(any(test, feature = "unstable-provider-packet-interceptor"))]
    impl_provider_method!(
        /// Sets the packet interceptor provider for the [`Server`]
        with_packet_interceptor,
        packet_interceptor,
        ServerProviders
    );

    #[cfg(any(test, feature = "unstable-provider-random"))]
    impl_provider_method!(
        /// Sets the random provider for the [`Server`]
        with_random,
        random,
        ServerProviders
    );

    #[cfg(feature = "unstable-provider-datagram")]
    impl_provider_method!(
        /// Sets the datagram provider for the [`Server`]
        with_datagram,
        datagram,
        ServerProviders
    );

    impl_provider_method!(
        /// Sets the congestion controller provider for the [`Server`]
        with_congestion_controller,
        congestion_controller,
        ServerProviders
    );

    /// Starts the [`Server`] with the configured providers
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use std::{error::Error, path::Path};
    /// # use s2n_quic::Server;
    /// #
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn Error>> {
    /// let server = Server::builder()
    ///     .with_tls((Path::new("./certs/cert.pem"), Path::new("./certs/key.pem")))?
    ///     .with_io("127.0.0.1:443")?
    ///     .start()?;
    /// #
    /// #    Ok(())
    /// # }
    /// ```
    pub fn start(self) -> Result<Server, StartError> {
        self.0.build().start()
    }
}


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

use super::*;
use core::marker::PhantomData;
use s2n_quic_core::{connection::id::Generator, crypto, path};
use s2n_quic_transport::{connection, endpoint, stream};

impl_providers_state! {
    #[derive(Debug, Default)]
    struct Providers {
        congestion_controller: CongestionController,
        connection_close_formatter: ConnectionCloseFormatter,
        connection_id: ConnectionID,
        packet_interceptor: PacketInterceptor,
        stateless_reset_token: StatelessResetToken,
        random: Random,
        endpoint_limits: EndpointLimits,
        event: Event,
        limits: Limits,
        io: IO,
        path_migration: PathMigration,
        sync: Sync,
        tls: Tls,
        address_token: AddressToken,
        datagram: Datagram,
    }

    /// Opaque trait containing all of the configured providers
    trait ServerProviders {}
}

impl<
        CongestionController: congestion_controller::Provider,
        ConnectionCloseFormatter: connection_close_formatter::Provider,
        ConnectionID: connection_id::Provider,
        PacketInterceptor: packet_interceptor::Provider,
        StatelessResetToken: stateless_reset_token::Provider,
        Random: random::Provider,
        EndpointLimits: endpoint_limits::Provider,
        Event: event::Provider,
        Limits: limits::Provider,
        IO: io::Provider,
        PathMigration: path_migration::Provider,
        Sync: sync::Provider,
        Tls: tls::Provider,
        AddressToken: address_token::Provider,
        Datagram: datagram::Provider,
    >
    Providers<
        CongestionController,
        ConnectionCloseFormatter,
        ConnectionID,
        PacketInterceptor,
        StatelessResetToken,
        Random,
        EndpointLimits,
        Event,
        Limits,
        IO,
        PathMigration,
        Sync,
        Tls,
        AddressToken,
        Datagram,
    >
{
    pub fn start(self) -> Result<Server, StartError> {
        let Self {
            congestion_controller,
            connection_close_formatter,
            connection_id,
            packet_interceptor,
            stateless_reset_token,
            random,
            endpoint_limits,
            event,
            limits,
            address_token,
            io,
            path_migration,
            sync,
            tls,
            datagram,
        } = self;

        let congestion_controller = congestion_controller.start().map_err(StartError::new)?;
        let connection_close_formatter = connection_close_formatter
            .start()
            .map_err(StartError::new)?;
        let connection_id = connection_id.start().map_err(StartError::new)?;
        let packet_interceptor = packet_interceptor.start().map_err(StartError::new)?;
        let stateless_reset_token = stateless_reset_token.start().map_err(StartError::new)?;
        let random = random.start().map_err(StartError::new)?;
        let endpoint_limits = endpoint_limits.start().map_err(StartError::new)?;
        let limits = limits.start().map_err(StartError::new)?;
        let event = event.start().map_err(StartError::new)?;
        let address_token = address_token.start().map_err(StartError::new)?;
        let sync = sync.start().map_err(StartError::new)?;
        let path_migration = path_migration.start().map_err(StartError::new)?;
        let tls = tls.start_server().map_err(StartError::new)?;
        let datagram = datagram.start().map_err(StartError::new)?;

        // Validate providers
        // TODO: Add more validation https://github.com/aws/s2n-quic/issues/285
        let valid_lifetime = |lifetime| {
            (connection::id::MIN_LIFETIME..=connection::id::MAX_LIFETIME).contains(&lifetime)
        };
        if connection_id
            .lifetime()
            .map_or(false, |lifetime| !valid_lifetime(lifetime))
        {
            return Err(StartError::new(connection::id::Error::InvalidLifetime));
        };

        let endpoint_config = EndpointConfig {
            congestion_controller,
            connection_close_formatter,
            connection_id,
            packet_interceptor,
            stateless_reset_token,
            random,
            endpoint_limits,
            event,
            limits,
            sync,
            tls,
            address_token,
            path_handle: PhantomData,
            path_migration,
            datagram,
        };

        let (endpoint, acceptor) = endpoint::Endpoint::new_server(endpoint_config);

        // Start the IO last
        let local_addr = io.start(endpoint).map_err(StartError::new)?;

        Ok(Server {
            acceptor,
            local_addr,
        })
    }
}

#[allow(dead_code)] // don't warn on unused providers for now
struct EndpointConfig<
    CongestionController,
    ConnectionCloseFormatter,
    ConnectionID,
    PacketInterceptor,
    PathHandle,
    PathMigration,
    StatelessResetToken,
    Random,
    EndpointLimits,
    Event,
    Limits,
    Sync,
    Tls,
    AddressToken,
    Datagram,
> {
    congestion_controller: CongestionController,
    connection_close_formatter: ConnectionCloseFormatter,
    connection_id: ConnectionID,
    packet_interceptor: PacketInterceptor,
    stateless_reset_token: StatelessResetToken,
    random: Random,
    endpoint_limits: EndpointLimits,
    event: Event,
    limits: Limits,
    sync: Sync,
    tls: Tls,
    address_token: AddressToken,
    path_handle: PhantomData<PathHandle>,
    path_migration: PathMigration,
    datagram: Datagram,
}

impl<
        CongestionController: congestion_controller::Endpoint,
        ConnectionCloseFormatter: connection_close_formatter::Formatter,
        ConnectionID: connection::id::Format,
        PacketInterceptor: packet_interceptor::PacketInterceptor,
        PathMigration: path_migration::Validator,
        PathHandle: path::Handle,
        StatelessResetToken: stateless_reset_token::Generator,
        Random: s2n_quic_core::random::Generator,
        EndpointLimits: s2n_quic_core::endpoint::Limiter,
        Event: s2n_quic_core::event::Subscriber,
        Limits: s2n_quic_core::connection::limits::Limiter,
        Sync,
        Tls: crypto::tls::Endpoint,
        AddressToken: address_token::Format,
        Datagram: s2n_quic_core::datagram::Endpoint,
    > core::fmt::Debug
    for EndpointConfig<
        CongestionController,
        ConnectionCloseFormatter,
        ConnectionID,
        PacketInterceptor,
        PathHandle,
        PathMigration,
        StatelessResetToken,
        Random,
        EndpointLimits,
        Event,
        Limits,
        Sync,
        Tls,
        AddressToken,
        Datagram,
    >
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ServerConfig").finish()
    }
}

impl<
        CongestionController: congestion_controller::Endpoint,
        ConnectionCloseFormatter: connection_close_formatter::Formatter,
        ConnectionID: connection::id::Format,
        PacketInterceptor: packet_interceptor::PacketInterceptor,
        PathHandle: path::Handle,
        PathMigration: path_migration::Validator,
        StatelessResetToken: stateless_reset_token::Generator,
        Random: s2n_quic_core::random::Generator,
        EndpointLimits: s2n_quic_core::endpoint::Limiter,
        Event: s2n_quic_core::event::Subscriber,
        Limits: s2n_quic_core::connection::limits::Limiter,
        Sync: 'static + Send,
        Tls: crypto::tls::Endpoint,
        AddressToken: address_token::Format,
        Datagram: s2n_quic_core::datagram::Endpoint,
    > endpoint::Config
    for EndpointConfig<
        CongestionController,
        ConnectionCloseFormatter,
        ConnectionID,
        PacketInterceptor,
        PathHandle,
        PathMigration,
        StatelessResetToken,
        Random,
        EndpointLimits,
        Event,
        Limits,
        Sync,
        Tls,
        AddressToken,
        Datagram,
    >
{
    type ConnectionIdFormat = ConnectionID;
    type ConnectionCloseFormatter = ConnectionCloseFormatter;
    type PathHandle = PathHandle;
    type StatelessResetTokenGenerator = StatelessResetToken;
    type RandomGenerator = Random;
    type Connection = connection::Implementation<Self>;
    // TODO allow users to specify another lock type
    type ConnectionLock = std::sync::Mutex<Self::Connection>;
    type CongestionControllerEndpoint = CongestionController;
    type EndpointLimits = EndpointLimits;
    type EventSubscriber = Event;
    type TLSEndpoint = Tls;
    type TokenFormat = AddressToken;
    type ConnectionLimits = Limits;
    type StreamManager = stream::DefaultStreamManager;
    type PathMigrationValidator = PathMigration;
    type PacketInterceptor = PacketInterceptor;
    type DatagramEndpoint = Datagram;

    const ENDPOINT_TYPE: endpoint::Type = endpoint::Type::Server;

    fn context(&mut self) -> endpoint::Context<Self> {
        endpoint::Context {
            congestion_controller: &mut self.congestion_controller,
            connection_close_formatter: &mut self.connection_close_formatter,
            connection_id_format: &mut self.connection_id,
            packet_interceptor: &mut self.packet_interceptor,
            stateless_reset_token_generator: &mut self.stateless_reset_token,
            random_generator: &mut self.random,
            tls: &mut self.tls,
            endpoint_limits: &mut self.endpoint_limits,
            token: &mut self.address_token,
            connection_limits: &mut self.limits,
            event_subscriber: &mut self.event,
            path_migration: &mut self.path_migration,
            datagram: &mut self.datagram,
        }
    }
}