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

use super::handle::CloseReceiver;
use crate::{connection, endpoint::handle::CloseSender};
use alloc::sync::Arc;
use core::{
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Poll, Waker},
};

/// Held by library. Used to receive close attempts and track close state.
#[derive(Debug)]
pub(crate) struct CloseHandle {
    /// Used to determine if the application has interest in closing the endpoint
    first_waker: Option<Waker>,
    /// A channel which is used to receive connection close attempts
    close_receiver: CloseReceiver,
    /// Track the endpoint open state
    endpoint_state: EndpointState,
}

impl CloseHandle {
    pub fn new(close_receiver: CloseReceiver, endpoint_state: EndpointState) -> Self {
        Self {
            first_waker: None,
            close_receiver,
            endpoint_state,
        }
    }

    /// Returns `Poll::Ready` if there is interest in closing the endpoint.
    pub fn poll_interest(&mut self) -> Poll<()> {
        if self.first_waker.is_some() {
            Poll::Ready(())
        } else {
            match self.close_receiver.try_next() {
                Ok(Some(waker)) => {
                    self.first_waker = Some(waker);
                    Poll::Ready(())
                }
                _ => Poll::Pending,
            }
        }
    }

    /// Marks that the endpoint has finished processing and accepting connections and is
    /// ready to be closed.
    pub fn close(&mut self) {
        self.endpoint_state.close();

        if let Some(waker) = self.first_waker.take() {
            waker.wake();
        }
        while let Ok(Some(waker)) = self.close_receiver.try_next() {
            waker.wake_by_ref();
        }
    }
}

/// Track if the endpoint is still open and handling connections
#[derive(Clone, Debug)]
pub(crate) struct EndpointState(Arc<AtomicBool>);

impl Default for EndpointState {
    fn default() -> Self {
        Self(Arc::new(AtomicBool::new(true)))
    }
}

impl EndpointState {
    /// Return `true` if the endpoint is still handling or accepting new connections
    fn is_open(&self) -> bool {
        self.0.load(Ordering::SeqCst)
    }

    /// Marks that the endpoint has finished processing and accepting connections and is
    /// ready to be closed.
    fn close(&mut self) {
        self.0.store(false, Ordering::SeqCst);
    }
}

#[derive(Clone, Debug)]
pub(crate) struct Closer {
    request_sent: bool,
    close_sender: CloseSender,
    endpoint_state: EndpointState,
}

impl Closer {
    /// Creates a Close attempt
    pub fn new(close_sender: CloseSender, endpoint_state: EndpointState) -> Self {
        Self {
            request_sent: false,
            close_sender,
            endpoint_state,
        }
    }

    pub(crate) fn poll_close(
        &mut self,
        context: &mut Context,
    ) -> Poll<Result<(), connection::Error>> {
        if !self.endpoint_state.is_open() {
            return Poll::Ready(Ok(()));
        }

        if !self.request_sent {
            match self.close_sender.poll_ready(context) {
                Poll::Ready(Ok(())) => {
                    // send a waker to the endpoint, which is woken once the endpoint has closed
                    match self.close_sender.try_send(context.waker().clone()) {
                        Ok(_) => {
                            self.request_sent = true;
                        }
                        Err(err) if err.is_full() => {
                            // yield and wake up the task since the opener misreported its ready state
                            context.waker().wake_by_ref();
                        }
                        Err(_) => {
                            // the endpoint is closed so return
                            return Poll::Ready(Ok(()));
                        }
                    }

                    return Poll::Pending;
                }
                Poll::Ready(Err(_)) => {
                    // the endpoint is closed so return
                    return Poll::Ready(Ok(()));
                }
                Poll::Pending => {
                    // pending capacity so try again
                }
            }
        }

        Poll::Pending
    }
}


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

//! Configuration parameters for `Endpoint`s

use crate::{connection, stream};
use s2n_quic_core::{
    crypto::tls, datagram, endpoint, event, packet, path, random, recovery::congestion_controller,
    stateless_reset,
};

/// Configuration parameters for a QUIC endpoint
pub trait Config: 'static + Send + Sized + core::fmt::Debug {
    /// The type of the TLS endpoint which is utilized
    type TLSEndpoint: tls::Endpoint;
    type CongestionControllerEndpoint: congestion_controller::Endpoint;
    /// The connections type
    type Connection: connection::Trait<Config = Self>;
    /// The type of lock that synchronizes connection state between threads
    type ConnectionLock: connection::Lock<Self::Connection>;
    /// The connection ID format
    type ConnectionIdFormat: connection::id::Format;
    /// The stateless reset token generator
    type StatelessResetTokenGenerator: stateless_reset::token::Generator;
    /// The random data generator
    type RandomGenerator: random::Generator;
    /// The validation token format
    type TokenFormat: s2n_quic_core::token::Format;
    /// The endpoint limits
    type EndpointLimits: endpoint::Limiter;
    /// The connection limits
    type ConnectionLimits: connection::limits::Limiter;
    /// The type of stream
    type StreamManager: stream::Manager;
    /// The connection close formatter
    type ConnectionCloseFormatter: connection::close::Formatter;
    /// The event subscriber
    type EventSubscriber: event::Subscriber;
    /// The type by which paths are identified
    type PathHandle: path::Handle;
    /// The path migration validator for the endpoint
    type PathMigrationValidator: path::migration::Validator;
    /// The packet_interceptor implementation for the endpoint
    type PacketInterceptor: packet::interceptor::Interceptor;
    /// The datagram implementation for the endpoint
    type DatagramEndpoint: datagram::Endpoint;

    /// The type of the local endpoint
    const ENDPOINT_TYPE: endpoint::Type;

    /// Returns the context for the endpoint configuration
    fn context(&mut self) -> Context<Self>;
}

#[derive(Debug)]
pub struct Context<'a, Cfg: Config> {
    /// The congestion controller endpoint associated with the endpoint config
    pub congestion_controller: &'a mut Cfg::CongestionControllerEndpoint,

    /// The connection id format associated with the endpoint config
    pub connection_id_format: &'a mut Cfg::ConnectionIdFormat,

    /// The stateless reset token generator associated with the endpoint config
    pub stateless_reset_token_generator: &'a mut Cfg::StatelessResetTokenGenerator,

    /// The random data generator associated with the endpoint config
    pub random_generator: &'a mut Cfg::RandomGenerator,

    /// The TLS endpoint associated with the endpoint config
    pub tls: &'a mut Cfg::TLSEndpoint,

    /// The endpoint limits
    pub endpoint_limits: &'a mut Cfg::EndpointLimits,

    /// Token generator / validator
    pub token: &'a mut Cfg::TokenFormat,

    /// The connection limits
    pub connection_limits: &'a mut Cfg::ConnectionLimits,

    pub connection_close_formatter: &'a mut Cfg::ConnectionCloseFormatter,

    pub event_subscriber: &'a mut Cfg::EventSubscriber,

    pub path_migration: &'a mut Cfg::PathMigrationValidator,

    pub packet_interceptor: &'a mut Cfg::PacketInterceptor,

    pub datagram: &'a mut Cfg::DatagramEndpoint,
}


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

use crate::{
    connection::{self, Connection},
    endpoint::handle::ConnectorSender,
};
use core::{
    fmt,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use futures_channel::oneshot;
use s2n_quic_core::{application::ServerName, inet::SocketAddress, path::RemoteAddress};

/// Held by connection Attempt future. Used to receive the actual connection.
pub(crate) type ConnectionReceiver = oneshot::Receiver<Result<Connection, connection::Error>>;

/// Held within the library connection_container. Used to send the actual connection once
/// its been created.
pub(crate) type ConnectionSender = oneshot::Sender<Result<Connection, connection::Error>>;

#[derive(Clone, Debug)]
pub struct Connect {
    pub(crate) remote_address: RemoteAddress,
    pub(crate) server_name: Option<ServerName>,
}

impl fmt::Display for Connect {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if f.alternate() {
            if let Some(hostname) = self.server_name.as_deref() {
                write!(f, "{hostname} at {}", &*self.remote_address)
            } else {
                write!(f, "{}", &*self.remote_address)
            }
        } else if let Some(hostname) = self.server_name.as_deref() {
            write!(f, "{hostname}")
        } else {
            write!(f, "{}", &*self.remote_address)
        }
    }
}

impl Connect {
    /// Creates a connection attempt with the specified remote address
    pub fn new<Addr: Into<SocketAddress>>(addr: Addr) -> Self {
        Self {
            remote_address: addr.into().into(),
            server_name: None,
        }
    }

    /// Specifies the server name to use for the connection
    #[must_use]
    pub fn with_server_name<Name: Into<ServerName>>(self, server_name: Name) -> Self {
        Self {
            server_name: Some(server_name.into()),
            ..self
        }
    }
}

/// Make it easy for applications to create a connection attempt without importing the `Connect` struct
impl<T: Into<SocketAddress>> From<T> for Connect {
    fn from(addr: T) -> Self {
        Self::new(addr)
    }
}

#[derive(Debug)]
#[allow(dead_code)]
pub(crate) struct Request {
    pub connect: Connect,
    pub sender: ConnectionSender,
}

#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Attempt {
    state: AttemptState,
}

impl Attempt {
    /// Creates a connection attempt
    ///
    /// The flow is currently implemented as follows:
    ///
    /// * The applications provides a `Connect` struct containing information for the remote endpoint
    /// * The attempt creates a oneshot channel and creates a `Request` with the sender and `Connect` struct
    /// * The attempt returns a `Self` while holding on to the oneshot receiver
    /// * The application polls the `Attempt` until either a successful `Connection` or `connection::Error` is
    ///   received over the oneshot receiver.
    pub(crate) fn new(opener: &ConnectorSender, connect: Connect) -> Self {
        // open a oneshot channel to receive the connection or error after the endpoint attempted the handshake
        let (response, receiver) = oneshot::channel();
        // The request includes both the connection info and response onshot channel
        let request = Request {
            connect,
            sender: response,
        };
        Self {
            state: AttemptState::Connect(request, opener.clone(), receiver),
        }
    }
}

enum AttemptState {
    /// The attempt is currently waiting for capacity in the `ConnectorSender` to make the `Request`
    Connect(Request, ConnectorSender, ConnectionReceiver),
    /// The attempt is currently waiting for a response back from the endpoint on the `ConnectionReceiver`
    Waiting(ConnectionReceiver),
    /// This is an intermediate state and should not persist across calls to `poll`
    Unreachable,
}

impl Future for Attempt {
    type Output = Result<Connection, connection::Error>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            match core::mem::replace(&mut self.state, AttemptState::Unreachable) {
                AttemptState::Connect(request, mut opener, response) => {
                    match opener.poll_ready(cx) {
                        Poll::Ready(Ok(())) => {
                            match opener.try_send(request) {
                                Ok(_) => {
                                    // transition to the waiting state
                                    self.state = AttemptState::Waiting(response);
                                    continue;
                                }
                                Err(err) if err.is_full() => {
                                    // reset to the original state
                                    self.state =
                                        AttemptState::Connect(err.into_inner(), opener, response);

                                    // yield and wake up the task since the opener misreported its ready state
                                    cx.waker().wake_by_ref();
                                }
                                Err(_) => {
                                    // The endpoint has closed
                                    return Err(connection::Error::unspecified()).into();
                                }
                            }
                        }
                        Poll::Ready(Err(_)) => {
                            // The endpoint has closed
                            return Err(connection::Error::unspecified()).into();
                        }
                        Poll::Pending => {
                            // reset to the original state
                            self.state = AttemptState::Connect(request, opener, response);
                        }
                    }

                    return Poll::Pending;
                }
                AttemptState::Waiting(mut response) => {
                    return match Pin::new(&mut response).poll(cx) {
                        Poll::Ready(Ok(res)) => Poll::Ready(res),
                        Poll::Ready(Err(_)) => {
                            // The endpoint has closed
                            Err(connection::Error::unspecified()).into()
                        }
                        Poll::Pending => {
                            self.state = AttemptState::Waiting(response);
                            Poll::Pending
                        }
                    };
                }
                AttemptState::Unreachable => {
                    unreachable!(
                        "Unreachable is an immediate state and should not exist across polls"
                    );
                }
            }
        }
    }
}


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

//! Allows to accept connections

use crate::{
    connection,
    connection::Connection,
    endpoint::{close, close::CloseHandle, connect},
};
use core::{
    pin::Pin,
    task::{Context, Poll, Waker},
};
use futures_channel::mpsc;
use futures_core::Stream;

/// Held by application. Used to accept new connections.
pub(crate) type AcceptorReceiver = mpsc::UnboundedReceiver<Connection>;
/// Held by library. Used to notify the application of newly-accepted connections.
pub(crate) type AcceptorSender = mpsc::UnboundedSender<Connection>;

/// Held by library. Used to receive connection attempts from the application.
pub(crate) type ConnectorReceiver = mpsc::Receiver<connect::Request>;
/// Held by application. Used to submit connection attempts to the library.
pub(crate) type ConnectorSender = mpsc::Sender<connect::Request>;

/// Held by library. Used to receive close attempts from the application.
pub(crate) type CloseReceiver = mpsc::Receiver<Waker>;
/// Held by the application. Used to submit connection close attempts to the library.
pub(crate) type CloseSender = mpsc::Sender<Waker>;

/// The [`Handle`] allows applications to accept and open QUIC connections on an `Endpoint`.
#[derive(Debug)]
pub(crate) struct Handle {
    pub acceptor: Acceptor,
    pub connector: Connector,
}

impl Handle {
    /// Creates a new `Handle` with a limit opening connection limit.
    pub(crate) fn new(
        max_opening_connections: usize,
    ) -> (Self, AcceptorSender, ConnectorReceiver, CloseHandle) {
        let (acceptor_sender, acceptor_receiver) = mpsc::unbounded();
        let (connector_sender, connector_receiver) = mpsc::channel(max_opening_connections);

        let (close_sender, close_receiver) = mpsc::channel(max_opening_connections);

        let endpoint_state = close::EndpointState::default();
        let closer = close::Closer::new(close_sender, endpoint_state.clone());
        let handle = Self {
            acceptor: Acceptor {
                acceptor: acceptor_receiver,
            },
            connector: Connector {
                connector: connector_sender,
                closer,
            },
        };
        (
            handle,
            acceptor_sender,
            connector_receiver,
            CloseHandle::new(close_receiver, endpoint_state),
        )
    }
}

#[derive(Debug)]
pub struct Acceptor {
    acceptor: AcceptorReceiver,
}

impl Acceptor {
    /// Polls for incoming connections and returns them.
    ///
    /// The method will return
    /// - `Poll::Ready(Some(connection))` if a connection was accepted.
    /// - `Poll::Ready(None)` if the acceptor is closed.
    /// - `Poll::Pending` if no new connection was accepted yet.
    ///   In this case the caller must retry polling as soon as a client
    ///   establishes a connection.
    ///   In order to notify the application of this condition,
    ///   the method will save the [`core::task::Waker`] which is provided as part of the
    ///   [`Context`] parameter, and notify it as soon as retrying
    ///   the method will yield a different result.
    pub fn poll_accept(&mut self, context: &mut Context) -> Poll<Option<Connection>> {
        match Stream::poll_next(Pin::new(&mut self.acceptor), context) {
            Poll::Ready(Some(connection)) => Poll::Ready(Some(connection)),
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Pending => Poll::Pending,
        }
    }
}

#[derive(Clone, Debug)]
pub struct Connector {
    connector: ConnectorSender,
    closer: close::Closer,
}

impl Connector {
    /// Attempts to establish a connection to an endpoint and returns a future to be awaited
    pub fn connect(&self, connect: connect::Connect) -> connect::Attempt {
        connect::Attempt::new(&self.connector, connect)
    }

    /// Polls to close the endpoint
    pub fn poll_close(&mut self, context: &mut Context) -> Poll<Result<(), connection::Error>> {
        self.closer.poll_close(context)
    }
}


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

use crate::{
    connection::{
        self,
        id::{ConnectionInfo, Generator as _},
        limits::{ConnectionInfo as LimitsInfo, Limiter as _},
        Trait as _,
    },
    endpoint,
    recovery::congestion_controller::{self, Endpoint as _},
    space::PacketSpaceManager,
};
use core::convert::TryInto;
use s2n_codec::DecoderBufferMut;
use s2n_quic_core::{
    crypto::{tls, tls::Endpoint as TLSEndpoint, CryptoSuite, InitialKey},
    datagram::{Endpoint, PreConnectionInfo},
    event::{self, supervisor, ConnectionPublisher, IntoEvent, Subscriber as _},
    inet::{datagram, DatagramInfo},
    packet::initial::ProtectedInitial,
    path::Handle as _,
    stateless_reset::token::Generator as _,
    transport::{self, parameters::ServerTransportParameters},
};

impl<Config: endpoint::Config> endpoint::Endpoint<Config> {
    pub(super) fn handle_initial_packet(
        &mut self,
        header: &datagram::Header<Config::PathHandle>,
        datagram: &DatagramInfo,
        packet: ProtectedInitial,
        remaining: DecoderBufferMut,
        retry_token_dcid: Option<connection::InitialId>,
    ) -> Result<(), connection::Error> {
        debug_assert!(
            Config::ENDPOINT_TYPE.is_server(),
            "only servers can accept new initial connections"
        );

        //= https://www.rfc-editor.org/rfc/rfc9000#section-14.1
        //# A client MUST expand the payload of all UDP datagrams carrying
        //# Initial packets to at least the smallest allowed maximum datagram
        //# size of 1200 bytes

        //= https://www.rfc-editor.org/rfc/rfc9000#section-14.1
        //# A server MUST discard an Initial packet that is carried in a UDP
        //# datagram with a payload that is smaller than the smallest allowed
        //# maximum datagram size of 1200 bytes.

        //= https://www.rfc-editor.org/rfc/rfc9001#section-9.3
        //# First, the packet
        //# containing a ClientHello MUST be padded to a minimum size.
        if datagram.payload_len < 1200 {
            return Err(transport::Error::PROTOCOL_VIOLATION
                .with_reason("packet too small")
                .into());
        }

        let remote_address = header.path.remote_address();

        // The first connection ID to persist and use for routing incoming packets
        let initial_connection_id;
        // The randomly generated destination connection ID that was sent from the client
        let original_destination_connection_id;

        if let Some(retry_dcid) = retry_token_dcid {
            original_destination_connection_id = retry_dcid;
            // This initial packet was in response to a Retry, so the destination connection ID
            // on the packet was generated by this server. We can use this destination connection
            // ID as the initial_connection_id.
            initial_connection_id = datagram.destination_connection_id;
        } else {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
            //# When an Initial packet is sent by a client that has not previously
            //# received an Initial or Retry packet from the server, the client
            //# populates the Destination Connection ID field with an unpredictable
            //# value.  This Destination Connection ID MUST be at least 8 bytes in
            //# length.
            original_destination_connection_id =
                datagram.destination_connection_id.try_into().map_err(|_| {
                    transport::Error::PROTOCOL_VIOLATION
                        .with_reason("destination connection id too short")
                })?;
            // The destination connection ID on the packet was randomly generated by the client
            // so we'll generate a new initial_connection_id.
            let connection_info = ConnectionInfo::new(&remote_address);
            initial_connection_id = self
                .config
                .context()
                .connection_id_format
                .generate(&connection_info);
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2
        //# Endpoints that receive a version 1 long header
        //# with a value larger than 20 MUST drop the packet.
        let source_connection_id = packet
            .source_connection_id()
            .try_into()
            .map_err(transport::Error::from)?;

        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.2
        //# Changing the Destination Connection ID field also results in
        //# a change to the keys used to protect the Initial packet.
        let (initial_key, initial_header_key) =
            <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialKey::new_server(
                datagram.destination_connection_id.as_bytes(),
            );

        let largest_packet_number = Default::default();
        let packet = packet.unprotect(&initial_header_key, largest_packet_number)?;
        let packet = packet.decrypt(&initial_key)?;

        // TODO handle token with stateless retry

        let internal_connection_id = self.connection_id_generator.generate_id();

        let initial_connection_id_expiration_time = self
            .config
            .context()
            .connection_id_format
            .lifetime()
            .map(|duration| datagram.timestamp + duration);

        let stateless_reset_token = self
            .config
            .context()
            .stateless_reset_token_generator
            .generate(initial_connection_id.as_bytes());

        let local_id_registry = self.connection_id_mapper.create_local_id_registry(
            internal_connection_id,
            &initial_connection_id,
            initial_connection_id_expiration_time,
            stateless_reset_token,
        );

        let peer_id_registry = self
            .connection_id_mapper
            .create_server_peer_id_registry(internal_connection_id, source_connection_id);

        let wakeup_handle = self
            .wakeup_queue
            .create_wakeup_handle(internal_connection_id);

        let mut transport_parameters = ServerTransportParameters::default();

        let limits = self
            .config
            .context()
            .connection_limits
            .on_connection(&LimitsInfo::new(&remote_address));

        transport_parameters.load_limits(&limits);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
        //# A server includes the Destination Connection ID field from the first
        //# Initial packet it received from the client in the
        //# original_destination_connection_id transport parameter; if the server
        //# sent a Retry packet, this refers to the first Initial packet received
        //# before sending the Retry packet.  If it sends a Retry packet, a
        //# server also includes the Source Connection ID field from the Retry
        //# packet in the retry_source_connection_id transport parameter.
        transport_parameters.original_destination_connection_id = Some(
            original_destination_connection_id
                .try_into()
                .expect("connection ID already validated"),
        );
        if retry_token_dcid.is_some() {
            transport_parameters.retry_source_connection_id = Some(
                datagram
                    .destination_connection_id
                    .try_into()
                    .expect("failed to convert source connection id"),
            );
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
        //# Each endpoint includes the value of the Source Connection ID field
        //# from the first Initial packet it sent in the
        //# initial_source_connection_id transport parameter
        transport_parameters.initial_source_connection_id = Some(
            initial_connection_id
                .as_bytes()
                .try_into()
                .expect("connection ID already validated"),
        );

        //= https://www.rfc-editor.org/rfc/rfc9000#section-18.2
        //# active_connection_id_limit (0x0e):  This is an integer value
        //#    specifying the maximum number of connection IDs from the peer that
        //#    an endpoint is willing to store.  This value includes the
        //#    connection ID received during the handshake, that received in the
        //#    preferred_address transport parameter, and those received in
        //#    NEW_CONNECTION_ID frames.
        transport_parameters.active_connection_id_limit = s2n_quic_core::varint::VarInt::from(
            connection::peer_id_registry::ACTIVE_CONNECTION_ID_LIMIT,
        )
        .try_into()
        .unwrap();

        let endpoint_context = self.config.context();

        transport_parameters.max_datagram_frame_size = endpoint_context
            .datagram
            .max_datagram_frame_size(&PreConnectionInfo::new())
            .try_into()
            .expect("Failed to convert max_datagram_frame_size");

        let tls_session = endpoint_context
            .tls
            .new_server_session(&transport_parameters);

        let path_info = congestion_controller::PathInfo::new(&remote_address);
        let congestion_controller = endpoint_context
            .congestion_controller
            .new_congestion_controller(path_info);

        let quic_version = packet.version;

        let meta = event::builder::ConnectionMeta {
            endpoint_type: Config::ENDPOINT_TYPE,
            id: internal_connection_id.into(),
            timestamp: datagram.timestamp,
        };

        let supervisor_context = supervisor::Context::new(
            self.connections.handshake_connections(),
            self.connections.len(),
            &remote_address,
            true,
        );

        let mut event_context = endpoint_context.event_subscriber.create_connection_context(
            &meta.clone().into_event(),
            &event::builder::ConnectionInfo {}.into_event(),
        );

        let mut publisher = event::ConnectionPublisherSubscriber::new(
            meta,
            quic_version,
            endpoint_context.event_subscriber,
            &mut event_context,
        );

        let space_manager = PacketSpaceManager::new(
            original_destination_connection_id,
            tls_session,
            initial_key,
            initial_header_key,
            datagram.timestamp,
            &mut publisher,
        );

        let max_mtu = self.max_mtu;
        let connection_parameters = connection::Parameters {
            internal_connection_id,
            local_id_registry,
            peer_id_registry,
            space_manager,
            wakeup_handle,
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
            //# A server MUST set the Destination Connection ID it
            //# uses for sending packets based on the first received Initial packet.
            peer_connection_id: source_connection_id,
            local_connection_id: initial_connection_id,
            path_handle: header.path,
            congestion_controller,
            timestamp: datagram.timestamp,
            quic_version,
            limits,
            max_mtu,
            event_context,
            supervisor_context: &supervisor_context,
            event_subscriber: endpoint_context.event_subscriber,
            datagram_endpoint: endpoint_context.datagram,
        };

        let mut connection = <Config as endpoint::Config>::Connection::new(connection_parameters)?;

        let endpoint_context = self.config.context();
        let handle_first_packet =
            move |connection: &mut <Config as endpoint::Config>::Connection| {
                let path_id = connection.on_datagram_received(
                    &header.path,
                    datagram,
                    endpoint_context.congestion_controller,
                    endpoint_context.path_migration,
                    max_mtu,
                    endpoint_context.event_subscriber,
                );

                let path_id = path_id.map_err(|err| {
                    connection.with_event_publisher(
                        datagram.timestamp,
                        None,
                        endpoint_context.event_subscriber,
                        |publisher, _path| {
                            publisher.on_datagram_dropped(event::builder::DatagramDropped {
                                len: datagram.payload_len as u16,
                                reason: err,
                            });
                        },
                    );
                    debug_assert!(
                        false,
                        "on_datagram_received should not fail for a newly created connection"
                    );
                    connection::Error::unspecified()
                })?;

                connection
                    .handle_cleartext_initial_packet(
                        datagram,
                        path_id,
                        packet,
                        endpoint_context.random_generator,
                        endpoint_context.event_subscriber,
                        endpoint_context.packet_interceptor,
                        endpoint_context.datagram,
                    )
                    .map_err(|err| {
                        use connection::ProcessingError;
                        match err {
                            ProcessingError::CryptoError(err) => err.into(),
                            ProcessingError::ConnectionError(err) => err,
                            // this is the first packet this connection has received
                            // so getting this error would be incorrect
                            ProcessingError::DuplicatePacket => {
                                debug_assert!(false, "got duplicate packet error on first packet");
                                transport::Error::INTERNAL_ERROR.into()
                            }
                            // this error is only raised by the client
                            ProcessingError::NonEmptyRetryToken => {
                                debug_assert!(false, "got non empty retry token error on server");
                                transport::Error::INTERNAL_ERROR.into()
                            }
                            // this error is only emitted when processing a Retry packet
                            ProcessingError::RetryScidEqualsDcid => {
                                debug_assert!(
                                    false,
                                    "got a Retry packet error while processing an Initial packet"
                                );
                                transport::Error::INTERNAL_ERROR.into()
                            }
                        }
                    })?;

                connection.handle_remaining_packets(
                    &header.path,
                    datagram,
                    path_id,
                    endpoint_context.connection_id_format,
                    remaining,
                    endpoint_context.random_generator,
                    endpoint_context.event_subscriber,
                    endpoint_context.packet_interceptor,
                    endpoint_context.datagram,
                )?;

                Ok(())
            };

        if let Err(error) = handle_first_packet(&mut connection) {
            let endpoint_context = self.config.context();

            connection.with_event_publisher(
                datagram.timestamp,
                None,
                endpoint_context.event_subscriber,
                |publisher, _path| {
                    use s2n_quic_core::event::builder::ConnectionClosed;
                    publisher.on_connection_closed(ConnectionClosed { error });
                },
            );

            return Err(error);
        }

        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.3
        //= type=TODO
        //= tracking-issue=299
        //# If the
        //# ClientHello spans multiple Initial packets, such servers would need
        //# to buffer the first received fragments, which could consume excessive
        //# resources if the client's address has not yet been validated.  To
        //# avoid this, servers MAY use the Retry feature (see Section 8.1 of
        //# [QUIC-TRANSPORT]) to only buffer partial ClientHello messages from
        //# clients with a validated address.

        let result = self
            .connection_id_mapper
            .try_insert_initial_id(original_destination_connection_id, internal_connection_id);

        debug_assert!(
            result.is_ok(),
            "Initial ID {original_destination_connection_id:?} was already in the map"
        );

        // Only persist the connection if everything is good.
        // Otherwise the connection will automatically get dropped. This
        // will also clean up all state which was already allocated for
        // the connection
        self.connections
            .insert_server_connection(connection, internal_connection_id);

        Ok(())
    }
}


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

//! This module defines a QUIC endpoint

use crate::{
    connection::{
        self,
        limits::{ConnectionInfo as LimitsInfo, Limiter as _},
        ConnectionContainer, ConnectionContainerIterationResult, ConnectionIdMapper,
        InternalConnectionId, InternalConnectionIdGenerator, ProcessingError, Trait as _,
    },
    endpoint,
    endpoint::close::CloseHandle,
    recovery::congestion_controller::{self, Endpoint as _},
    space::PacketSpaceManager,
    wakeup_queue::WakeupQueue,
};
use alloc::collections::VecDeque;
use core::{
    convert::TryInto,
    task::{self, Poll},
};
use s2n_codec::{DecoderBuffer, DecoderBufferMut};
use s2n_quic_core::{
    connection::{
        id::{ConnectionInfo, Generator},
        InitialId, LocalId, PeerId,
    },
    crypto::{tls, tls::Endpoint as _, CryptoSuite, InitialKey},
    datagram::{Endpoint as DatagramEndpoint, PreConnectionInfo},
    endpoint::{limits::Outcome, Limiter as _},
    event::{
        self, supervisor, ConnectionPublisher, EndpointPublisher as _, IntoEvent, Subscriber as _,
    },
    inet::{datagram, DatagramInfo},
    io::{rx, tx},
    packet::{initial::ProtectedInitial, interceptor::Interceptor, ProtectedPacket},
    path,
    path::{Handle as _, MaxMtu},
    random::Generator as _,
    stateless_reset::token::{Generator as _, LEN as StatelessResetTokenLen},
    time::{Clock, Timestamp},
    token::{self, Format},
    transport::parameters::ClientTransportParameters,
};

pub mod close;
mod config;
pub mod connect;
pub mod handle;
mod initial;
mod packet_buffer;
mod retry;
mod stateless_reset;
mod version;

// exports
pub use config::{Config, Context};
pub use packet_buffer::Buffer as PacketBuffer;
pub use s2n_quic_core::endpoint::*;

const DEFAULT_MAX_PEERS: usize = 1024;

/// A QUIC `Endpoint`
pub struct Endpoint<Cfg: Config> {
    /// Configuration parameters for the endpoint
    config: Cfg,
    /// Contains all active connections
    connections: ConnectionContainer<Cfg::Connection, Cfg::ConnectionLock>,
    /// Creates internal IDs for new connections
    connection_id_generator: InternalConnectionIdGenerator,
    /// Maps from external to internal connection IDs
    connection_id_mapper: ConnectionIdMapper,
    /// Allows to wakeup the endpoint task which might be blocked on waiting for packets
    /// from application tasks (which e.g. enqueued new data to send).
    wakeup_queue: WakeupQueue<InternalConnectionId>,
    /// Used to receive close attempts and track close state.
    close_handle: CloseHandle,
    /// This queue contains wakeups we retrieved from the [`Self::wakeup_queue`] earlier.
    /// This is not a local variable in order to reuse the allocated queue capacity in between
    /// [`Endpoint`] interactions.
    dequeued_wakeups: VecDeque<InternalConnectionId>,
    version_negotiator: version::Negotiator<Cfg>,
    retry_dispatch: retry::Dispatch<Cfg::PathHandle>,
    stateless_reset_dispatch: stateless_reset::Dispatch<Cfg::PathHandle>,
    close_packet_buffer: packet_buffer::Buffer,
    /// The largest maximum transmission unit (MTU) that can be sent on a path
    max_mtu: MaxMtu,
}

impl<Cfg: Config> s2n_quic_core::endpoint::Endpoint for Endpoint<Cfg> {
    type PathHandle = Cfg::PathHandle;
    type Subscriber = Cfg::EventSubscriber;

    const ENDPOINT_TYPE: s2n_quic_core::endpoint::Type = Cfg::ENDPOINT_TYPE;

    fn receive<Rx, C>(&mut self, queue: &mut Rx, clock: &C)
    where
        Rx: rx::Queue<Handle = Cfg::PathHandle>,
        C: Clock,
    {
        let mut now: Option<Timestamp> = None;

        queue.for_each(|header, payload| {
            let timestamp = match now {
                Some(time) => time,
                None => {
                    now = Some(clock.get_time());
                    now.expect("value should be set")
                }
            };

            self.receive_datagram(&header, payload, timestamp)
        });
    }

    fn transmit<Tx, C>(&mut self, queue: &mut Tx, clock: &C)
    where
        Tx: tx::Queue<Handle = Self::PathHandle>,
        C: Clock,
    {
        self.on_timeout(clock.get_time());

        // the queue doesn't have the capacity for transmissions so don't try to push
        if !queue.has_capacity() {
            return;
        }

        // Iterate over all connections which want to transmit data
        let endpoint_context = self.config.context();

        let timestamp = clock.get_time();

        self.connections.iterate_transmission_list(|connection| {
            // if we no longer have capacity, then put the connection at the front of the queue for
            // next time
            if !queue.has_capacity() {
                return ConnectionContainerIterationResult::BreakAndInsertAtFront;
            }

            // ignore the transmission error and just query the queue capacity instead
            let _ = connection.on_transmit(
                queue,
                timestamp,
                endpoint_context.event_subscriber,
                endpoint_context.packet_interceptor,
            );

            // flush the TX queue between connections
            queue.flush();

            ConnectionContainerIterationResult::Continue
        });

        if queue.has_capacity() {
            let mut publisher = event::EndpointPublisherSubscriber::new(
                event::builder::EndpointMeta {
                    endpoint_type: Cfg::ENDPOINT_TYPE,
                    timestamp,
                },
                None,
                endpoint_context.event_subscriber,
            );
            self.version_negotiator.on_transmit(queue, &mut publisher);
            self.retry_dispatch.on_transmit(queue, &mut publisher);
            self.stateless_reset_dispatch
                .on_transmit(queue, &mut publisher);
        }
    }

    fn poll_wakeups<C: Clock>(
        &mut self,
        cx: &mut task::Context<'_>,
        clock: &C,
    ) -> Poll<Result<usize, s2n_quic_core::endpoint::CloseError>> {
        if self.close_handle.poll_interest().is_ready() // poll for close interest
            && self.connections.is_empty() // wait for all connections to close gracefully
            && self.connections.is_open()
        {
            // transition to close state
            self.close_handle.close();

            // stop accepting new connections and prepare to close the endpoint
            self.connections.close();
        }

        // Drop the endpoint if there is no more progress to be made.
        if !self.connections.is_open() {
            return Poll::Ready(Err(s2n_quic_core::endpoint::CloseError));
        }

        self.wakeup_queue
            .poll_pending_wakeups(&mut self.dequeued_wakeups, cx);

        let mut now: Option<Timestamp> = None;
        let mut wakeup_count = self.dequeued_wakeups.len();
        let close_packet_buffer = &mut self.close_packet_buffer;
        let endpoint_context = self.config.context();

        for internal_id in self.dequeued_wakeups.drain(..) {
            self.connections.with_connection(internal_id, |conn| {
                let timestamp = match now {
                    Some(now) => now,
                    _ => {
                        let time = clock.get_time();
                        now = Some(time);
                        time
                    }
                };

                if let Err(error) = conn.on_wakeup(
                    timestamp,
                    endpoint_context.event_subscriber,
                    endpoint_context.datagram,
                ) {
                    conn.close(
                        error,
                        endpoint_context.connection_close_formatter,
                        close_packet_buffer,
                        timestamp,
                        endpoint_context.event_subscriber,
                        endpoint_context.packet_interceptor,
                    );
                }
            });
        }

        // try to open connection requests from the application
        if Cfg::ENDPOINT_TYPE.is_client() {
            loop {
                match self.connections.poll_connection_request(cx) {
                    Poll::Pending => break,
                    Poll::Ready(Some(request)) => {
                        wakeup_count += 1;

                        let time = clock.get_time();
                        if let Err(err) = self.create_client_connection(request, time) {
                            // TODO report that the connection was not successfully created
                            // TODO emit event
                            dbg!(err);
                        }
                    }
                    Poll::Ready(None) => {
                        // the client handle has been dropped so break from loop
                        break;
                    }
                }
            }
        }

        if wakeup_count > 0 {
            Poll::Ready(Ok(wakeup_count))
        } else {
            Poll::Pending
        }
    }

    #[inline]
    fn timeout(&self) -> Option<Timestamp> {
        self.connections.next_expiration()
    }

    #[inline]
    fn set_max_mtu(&mut self, max_mtu: MaxMtu) {
        self.max_mtu = max_mtu
    }

    #[inline]
    fn subscriber(&mut self) -> &mut Self::Subscriber {
        self.config.context().event_subscriber
    }
}

impl<Cfg: Config> Endpoint<Cfg> {
    /// Creates a new QUIC server endpoint using the given configuration
    pub fn new_server(config: Cfg) -> (Self, handle::Acceptor) {
        assert!(
            Cfg::ENDPOINT_TYPE.is_server(),
            "only server endpoints can be created with server configurations"
        );
        let (endpoint, handle) = Self::new(config);
        (endpoint, handle.acceptor)
    }

    /// Creates a new QUIC client endpoint using the given configuration
    pub fn new_client(config: Cfg) -> (Self, handle::Connector) {
        assert!(
            Cfg::ENDPOINT_TYPE.is_client(),
            "only client endpoints can be created with client configurations"
        );
        let (endpoint, handle) = Self::new(config);
        (endpoint, handle.connector)
    }

    fn new(mut config: Cfg) -> (Self, handle::Handle) {
        // TODO make this limit configurable
        let max_opening_connections = 1000;
        let (handle, acceptor_sender, connector_receiver, close_handle) =
            handle::Handle::new(max_opening_connections);

        let connection_id_mapper =
            ConnectionIdMapper::new(config.context().random_generator, Cfg::ENDPOINT_TYPE);

        let endpoint = Self {
            config,
            connections: ConnectionContainer::new(acceptor_sender, connector_receiver),
            connection_id_generator: InternalConnectionIdGenerator::new(),
            connection_id_mapper,
            wakeup_queue: WakeupQueue::new(),
            close_handle,
            dequeued_wakeups: VecDeque::new(),
            version_negotiator: version::Negotiator::default(),
            retry_dispatch: retry::Dispatch::default(),
            stateless_reset_dispatch: stateless_reset::Dispatch::default(),
            close_packet_buffer: Default::default(),
            max_mtu: Default::default(),
        };

        (endpoint, handle)
    }

    /// Determine the next step when a peer attempts a connection
    fn connection_allowed(
        &mut self,
        header: &datagram::Header<Cfg::PathHandle>,
        packet: &ProtectedInitial,
        payload_len: usize,
        timestamp: Timestamp,
    ) -> Option<()> {
        if !self.connections.can_accept() {
            return None;
        }

        let remote_address = header.path.remote_address();

        let attempt = s2n_quic_core::endpoint::limits::ConnectionAttempt::new(
            self.connections.handshake_connections(),
            self.connections.len(),
            &remote_address,
            timestamp.into_event(),
        );

        let context = self.config.context();
        let outcome = context.endpoint_limits.on_connection_attempt(&attempt);
        let mut publisher = event::EndpointPublisherSubscriber::new(
            event::builder::EndpointMeta {
                endpoint_type: Cfg::ENDPOINT_TYPE,
                timestamp,
            },
            None,
            context.event_subscriber,
        );

        match outcome {
            Outcome::Allow { .. } => Some(()),
            Outcome::Retry { .. } => {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.2
                //# A server can also use a Retry packet to defer the state and
                //# processing costs of connection establishment.  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.

                let connection_info = ConnectionInfo::new(&remote_address);

                let local_connection_id = context.connection_id_format.generate(&connection_info);

                self.retry_dispatch.queue::<
                    _,
                    <<<Cfg as Config>::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::RetryKey,
                >(
                    header.path,
                    packet,
                    local_connection_id,
                    context.random_generator,
                    context.token
                );

                None
            }
            Outcome::Close { .. } => {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.2
                //= type=TODO
                //= tracking-issue=270
                //# If a server refuses to accept a new connection, it SHOULD send an
                //# Initial packet containing a CONNECTION_CLOSE frame with error code
                //# CONNECTION_REFUSED.

                publisher.on_endpoint_datagram_dropped(event::builder::EndpointDatagramDropped {
                    len: payload_len as u16,
                    reason: event::builder::DatagramDropReason::RejectedConnectionAttempt,
                });

                None
            }
            Outcome::Drop { .. } => {
                publisher.on_endpoint_datagram_dropped(event::builder::EndpointDatagramDropped {
                    len: payload_len as u16,
                    reason: event::builder::DatagramDropReason::RejectedConnectionAttempt,
                });
                None
            }
            _ => {
                publisher.on_endpoint_datagram_dropped(event::builder::EndpointDatagramDropped {
                    len: payload_len as u16,
                    reason: event::builder::DatagramDropReason::RejectedConnectionAttempt,
                });
                // Outcome is non_exhaustive so drop on things we don't understand
                None
            }
        }
    }

    /// Ingests a single datagram
    fn receive_datagram(
        &mut self,
        header: &datagram::Header<Cfg::PathHandle>,
        payload: &mut [u8],
        timestamp: Timestamp,
    ) {
        let endpoint_context = self.config.context();

        let remote_address = header.path.remote_address();

        // Try to decode the first packet in the datagram
        let payload_len = payload.len();
        let buffer = DecoderBufferMut::new(payload);

        let buffer = {
            let subject = event::builder::Subject::Endpoint {}.into_event();
            let remote_address = header.path.remote_address();
            let local_address = header.path.local_address();

            let datagram = s2n_quic_core::packet::interceptor::Datagram {
                remote_address: remote_address.into_event(),
                local_address: local_address.into_event(),
                timestamp,
            };

            endpoint_context
                .packet_interceptor
                .intercept_rx_datagram(&subject, &datagram, buffer)
        };

        let connection_info = ConnectionInfo::new(&remote_address);
        let (packet, remaining) = if let Ok((packet, remaining)) = ProtectedPacket::decode(
            buffer,
            &connection_info,
            endpoint_context.connection_id_format,
        ) {
            (packet, remaining)
        } else {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.2
            //# Servers MUST drop incoming packets under all other circumstances.

            //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
            //# However, endpoints MUST treat any packet ending in a
            //# valid stateless reset token as a Stateless Reset, as other QUIC
            //# versions might allow the use of a long header.

            // The packet may be a stateless reset, check before returning.
            let internal_connection_id = self.close_on_matching_stateless_reset(payload, timestamp);

            if internal_connection_id.is_none() {
                // The packet didn't contain a valid stateless token
                let mut publisher = event::EndpointPublisherSubscriber::new(
                    event::builder::EndpointMeta {
                        endpoint_type: Cfg::ENDPOINT_TYPE,
                        timestamp,
                    },
                    None,
                    self.config.context().event_subscriber,
                );
                publisher.on_endpoint_datagram_dropped(event::builder::EndpointDatagramDropped {
                    len: payload_len as u16,
                    reason: event::builder::DatagramDropReason::DecodingFailed,
                });
            }

            return;
        };

        let mut publisher = event::EndpointPublisherSubscriber::new(
            event::builder::EndpointMeta {
                endpoint_type: Cfg::ENDPOINT_TYPE,
                timestamp,
            },
            packet.version(),
            endpoint_context.event_subscriber,
        );

        // Ensure the version is supported. This check occurs before the destination
        // connection ID is parsed since future versions of QUIC could have different
        // length requirements for connection IDs.
        if self
            .version_negotiator
            .on_packet(&header.path, payload_len, &packet, &mut publisher)
            .is_err()
        {
            publisher.on_endpoint_datagram_dropped(event::builder::EndpointDatagramDropped {
                len: payload_len as u16,
                reason: event::builder::DatagramDropReason::UnsupportedVersion,
            });
            return;
        }

        let destination_connection_id =
            match connection::LocalId::try_from_bytes(packet.destination_connection_id()) {
                Some(connection_id) => connection_id,
                None => {
                    // Ignore the datagram
                    publisher.on_endpoint_datagram_dropped(
                        event::builder::EndpointDatagramDropped {
                            len: payload_len as u16,
                            reason:
                                event::builder::DatagramDropReason::InvalidDestinationConnectionId,
                        },
                    );
                    return;
                }
            };

        let source_connection_id = packet
            .source_connection_id()
            .and_then(PeerId::try_from_bytes);

        let datagram = &DatagramInfo {
            timestamp,
            payload_len,
            ecn: header.ecn,
            destination_connection_id,
            source_connection_id,
        };

        // TODO validate the connection ID before looking up the connection in the map
        let close_packet_buffer = &mut self.close_packet_buffer;

        // Try to lookup the internal connection ID and dispatch the packet
        // to the Connection
        if let Some(internal_id) = self
            .connection_id_mapper
            .lookup_internal_connection_id(&datagram.destination_connection_id)
        {
            let mut check_for_stateless_reset = false;
            let max_mtu = self.max_mtu;

            let _ = self.connections.with_connection(internal_id, |conn| {
                // The path `Id` needs to be passed around instead of the path to get around `&mut self` and
                // `&mut self.path_manager` being borrowed at the same time
                let path_id = conn
                    .on_datagram_received(
                        &header.path,
                        datagram,
                        endpoint_context.congestion_controller,
                        endpoint_context.path_migration,
                        max_mtu,
                        endpoint_context.event_subscriber,
                    )
                    .map_err(|datagram_drop_reason| {
                        // An error received at this point was caused by a datagram that has not
                        // been authenticated yet, and thus the connection should not be closed.
                        conn.with_event_publisher(
                            datagram.timestamp,
                            None,
                            endpoint_context.event_subscriber,
                            |publisher, _path| {
                                publisher.on_datagram_dropped(event::builder::DatagramDropped {
                                    len: datagram.payload_len as u16,
                                    reason: datagram_drop_reason,
                                });
                            },
                        );
                    })?;

                //= https://www.rfc-editor.org/rfc/rfc9000#section-10.2.1
                //# An endpoint
                //# that is closing is not required to process any received frame.

                if let Err(err) = conn.handle_packet(
                    datagram,
                    path_id,
                    packet,
                    endpoint_context.random_generator,
                    endpoint_context.event_subscriber,
                    endpoint_context.packet_interceptor,
                    endpoint_context.datagram,
                ) {
                    match err {
                        ProcessingError::DuplicatePacket => {
                            // We discard duplicate packets
                        }
                        ProcessingError::NonEmptyRetryToken => {
                            //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
                            //# Initial packets sent by the server MUST set the Token Length field
                            //# to 0; clients that receive an Initial packet with a non-zero Token
                            //# Length field MUST either discard the packet or generate a
                            //# connection error of type PROTOCOL_VIOLATION.
                            //
                            // We discard server initials with non empty retry tokens instead of closing
                            // the connection to prevent an attacker that can spoof initial packets
                            // from gaining the ability to close a connection by setting a retry token.
                        }
                        ProcessingError::RetryScidEqualsDcid => {
                            //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.1
                            //# A client MUST
                            //# discard a Retry packet that contains a Source Connection ID field
                            //# that is identical to the Destination Connection ID field of its
                            //# Initial packet.
                        }
                        ProcessingError::ConnectionError(err) => {
                            conn.close(
                                err,
                                endpoint_context.connection_close_formatter,
                                close_packet_buffer,
                                datagram.timestamp,
                                endpoint_context.event_subscriber,
                                endpoint_context.packet_interceptor,
                            );
                            return Err(());
                        }
                        ProcessingError::CryptoError(_) => {
                            // CryptoErrors returned as a result of a packet failing decryption
                            // will be silently discarded, but are a potential indication of a
                            // stateless reset from the peer

                            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.1
                            //# Due to packet reordering or loss, a client might receive packets for
                            //# a connection that are encrypted with a key it has not yet computed.
                            //# The client MAY drop these packets, or it MAY buffer them in
                            //# anticipation of later packets that allow it to compute the key.
                            //
                            // Packets that fail decryption are discarded rather than buffered.

                            //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
                            //# Endpoints MAY skip this check if any packet from a datagram is
                            //# successfully processed.  However, the comparison MUST be performed
                            //# when the first packet in an incoming datagram either cannot be
                            //# associated with a connection, or cannot be decrypted.
                            check_for_stateless_reset = true;
                        }
                    }
                }

                if let Err(err) = conn.handle_remaining_packets(
                    &header.path,
                    datagram,
                    path_id,
                    endpoint_context.connection_id_format,
                    remaining,
                    endpoint_context.random_generator,
                    endpoint_context.event_subscriber,
                    endpoint_context.packet_interceptor,
                    endpoint_context.datagram,
                ) {
                    conn.close(
                        err,
                        endpoint_context.connection_close_formatter,
                        close_packet_buffer,
                        datagram.timestamp,
                        endpoint_context.event_subscriber,
                        endpoint_context.packet_interceptor,
                    );
                    return Err(());
                }

                Ok(())
            });

            if check_for_stateless_reset {
                self.close_on_matching_stateless_reset(payload, timestamp);
            }

            return;
        }

        match (Cfg::ENDPOINT_TYPE, packet) {
            (s2n_quic_core::endpoint::Type::Server, ProtectedPacket::Initial(packet)) => {
                let source_connection_id =
                    match connection::PeerId::try_from_bytes(packet.source_connection_id()) {
                        Some(connection_id) => connection_id,
                        None => {
                            publisher.on_endpoint_datagram_dropped(
                           event::builder::EndpointDatagramDropped {
                               len: payload_len as u16,
                               reason:
                                   event::builder::DatagramDropReason::InvalidSourceConnectionId,
                           },
                       );
                            return;
                        }
                    };

                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
                //= type=TODO
                //= tracking-issue=140
                //# Additionally, an endpoint MAY consider the peer address validated if
                //# the peer uses a connection ID chosen by the endpoint and the
                //# connection ID contains at least 64 bits of entropy

                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.2
                //# In response to processing an Initial packet containing a token that
                //# was provided in a Retry packet, a server cannot send another Retry
                //# packet; it can only refuse the connection or permit it to proceed.
                let retry_token_dcid = if !packet.token().is_empty() {
                    let mut context = token::Context::new(
                        &remote_address,
                        &source_connection_id,
                        endpoint_context.random_generator,
                    );

                    let outcome = endpoint_context
                        .token
                        .validate_token(&mut context, packet.token());

                    if outcome.is_none() {
                        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
                        //= type=TODO
                        //= tracking-issue=344
                        //# If the token is invalid, then the
                        //# server SHOULD proceed as if the client did not have a validated
                        //# address, including potentially sending a Retry packet.

                        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.2
                        //= type=TODO
                        //= tracking-issue=344
                        //# Instead, the
                        //# server SHOULD immediately close (Section 10.2) the connection with an
                        //# INVALID_TOKEN error.
                        publisher.on_endpoint_datagram_dropped(
                            event::builder::EndpointDatagramDropped {
                                len: payload_len as u16,
                                reason: event::builder::DatagramDropReason::InvalidRetryToken,
                            },
                        );

                        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
                        //# Servers MAY
                        //# discard any Initial packet that does not carry the expected token.
                        return;
                    }

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.3
                    //# If the validation succeeds, the server SHOULD then allow
                    //# the handshake to proceed.
                    outcome
                } else {
                    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.2
                    //# Upon receiving the client's Initial packet, the server can request
                    //# address validation by sending a Retry packet (Section 17.2.5)
                    //# containing a token.
                    if self
                        .connection_allowed(header, &packet, payload_len, timestamp)
                        .is_none()
                    {
                        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.1
                        //# A server MUST NOT send more than one Retry
                        //# packet in response to a single UDP datagram.
                        return;
                    }

                    None
                };

                if let Err(err) = self.handle_initial_packet(
                    header,
                    datagram,
                    packet,
                    remaining,
                    retry_token_dcid,
                ) {
                    // TODO send a minimal connection close frame
                    let mut publisher = event::EndpointPublisherSubscriber::new(
                        event::builder::EndpointMeta {
                            endpoint_type: Cfg::ENDPOINT_TYPE,
                            timestamp,
                        },
                        None,
                        self.config.context().event_subscriber,
                    );
                    publisher.on_endpoint_connection_attempt_failed(
                        event::builder::EndpointConnectionAttemptFailed { error: err },
                    );
                }
            }
            (_, packet) => {
                publisher.on_endpoint_datagram_dropped(event::builder::EndpointDatagramDropped {
                    len: payload_len as u16,
                    reason: event::builder::DatagramDropReason::UnknownDestinationConnectionId,
                });

                let is_short_header_packet = matches!(packet, ProtectedPacket::Short(_));
                //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
                //# Endpoints MAY skip this check if any packet from a datagram is
                //# successfully processed.  However, the comparison MUST be performed
                //# when the first packet in an incoming datagram either cannot be
                //# associated with a connection, or cannot be decrypted.

                //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
                //# However, endpoints MUST treat any packet ending in a
                //# valid stateless reset token as a Stateless Reset, as other QUIC
                //# versions might allow the use of a long header.
                let is_stateless_reset = self
                    .close_on_matching_stateless_reset(payload, timestamp)
                    .is_some();

                //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.2
                //# For instance, an endpoint MAY send a Stateless Reset in
                //# response to any further incoming packets.

                //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
                //# An endpoint MAY send a Stateless Reset in response to receiving a packet
                //# that it cannot associate with an active connection.

                //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
                //# Because the stateless reset token is not available
                //# until connection establishment is complete or near completion,
                //# ignoring an unknown packet with a long header might be as effective
                //# as sending a Stateless Reset.
                if !is_stateless_reset
                    && Cfg::StatelessResetTokenGenerator::ENABLED
                    && is_short_header_packet
                {
                    self.enqueue_stateless_reset(header, datagram, &destination_connection_id);
                }
            }
        }
    }

    /// Enqueues sending a stateless reset to a peer.
    ///
    /// Sending the reset was caused through the passed `datagram`.
    fn enqueue_stateless_reset(
        &mut self,
        header: &datagram::Header<Cfg::PathHandle>,
        datagram: &DatagramInfo,
        destination_connection_id: &LocalId,
    ) {
        let token = self
            .config
            .context()
            .stateless_reset_token_generator
            .generate(destination_connection_id.as_bytes());
        let max_tag_length = self.config.context().tls.max_tag_length();
        // The datagram payload length is used as the packet length since
        // a stateless reset is only sent if the first packet in a datagram is
        // a short header packet and a short header packet must be the last packet
        // in a datagram; thus the entire datagram is one packet.
        let triggering_packet_len = datagram.payload_len;
        self.stateless_reset_dispatch.queue(
            header.path,
            token,
            max_tag_length,
            triggering_packet_len,
            self.config.context().random_generator,
        );
    }

    /// Checks if the given payload contains a stateless reset token matching a known token.
    /// If there is a match, the matching connection will be closed and the `InternalConnectionId`
    /// will be returned.
    fn close_on_matching_stateless_reset(
        &mut self,
        payload: &[u8],
        timestamp: Timestamp,
    ) -> Option<InternalConnectionId> {
        let buffer = DecoderBuffer::new(payload);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
        //# The endpoint
        //# identifies a received datagram as a Stateless Reset by comparing the
        //# last 16 bytes of the datagram with all stateless reset tokens
        //# associated with the remote address on which the datagram was
        //# received.
        let token_index = payload.len().checked_sub(StatelessResetTokenLen)?;
        let buffer = buffer.skip(token_index).ok()?;
        let (token, _) = buffer.decode().ok()?;
        let endpoint_context = self.config.context();
        let internal_id = self
            .connection_id_mapper
            .remove_internal_connection_id_by_stateless_reset_token(&token)?;

        let mut publisher = event::EndpointPublisherSubscriber::new(
            event::builder::EndpointMeta {
                endpoint_type: Cfg::ENDPOINT_TYPE,
                timestamp,
            },
            None,
            endpoint_context.event_subscriber,
        );

        publisher.on_endpoint_packet_received(event::builder::EndpointPacketReceived {
            packet_header: event::builder::PacketHeader::StatelessReset {},
        });

        let close_packet_buffer = &mut self.close_packet_buffer;

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
        //# If the last 16 bytes of the datagram are identical in value to a
        //# stateless reset token, the endpoint MUST enter the draining period
        //# and not send any further packets on this connection.
        self.connections.with_connection(internal_id, |conn| {
            conn.close(
                connection::Error::stateless_reset(),
                endpoint_context.connection_close_formatter,
                close_packet_buffer,
                timestamp,
                endpoint_context.event_subscriber,
                endpoint_context.packet_interceptor,
            );
        });

        Some(internal_id)
    }

    fn on_timeout(&mut self, timestamp: Timestamp) {
        let connection_id_mapper = &mut self.connection_id_mapper;
        let close_packet_buffer = &mut self.close_packet_buffer;
        let endpoint_context = self.config.context();

        self.connections
            .iterate_timeout_list(timestamp, |conn, supervisor_context| {
                if let Err(error) = conn.on_timeout(
                    connection_id_mapper,
                    timestamp,
                    supervisor_context,
                    endpoint_context.random_generator,
                    endpoint_context.event_subscriber,
                ) {
                    conn.close(
                        error,
                        endpoint_context.connection_close_formatter,
                        close_packet_buffer,
                        timestamp,
                        endpoint_context.event_subscriber,
                        endpoint_context.packet_interceptor,
                    );
                }
            });

        // allow connections to generate a new connection id
        self.connections
            .iterate_new_connection_id_list(|connection| {
                let result = connection.on_new_connection_id(
                    endpoint_context.connection_id_format,
                    endpoint_context.stateless_reset_token_generator,
                    timestamp,
                );
                if result.is_ok() {
                    ConnectionContainerIterationResult::Continue
                } else {
                    // The provided Connection ID generator must never generate the same connection
                    // ID twice. If this happens, it is unlikely we could recover from it.
                    panic!("Generated connection ID was already in use");
                }
            });
    }

    fn create_client_connection(
        &mut self,
        request: endpoint::connect::Request,
        timestamp: Timestamp,
    ) -> Result<(), connection::Error> {
        let endpoint::connect::Request {
            connect:
                endpoint::connect::Connect {
                    remote_address,
                    server_name: hostname,
                },
            sender,
        } = request;

        let internal_connection_id = self.connection_id_generator.generate_id();
        let local_connection_id = self
            .config
            .context()
            .connection_id_format
            .generate(&ConnectionInfo::new(&remote_address));

        let local_connection_id_expiration_time = self
            .config
            .context()
            .connection_id_format
            .lifetime()
            .map(|duration| timestamp + duration);

        let local_id_registry = {
            // TODO: the client currently generates a random stateless_reset_token but doesnt
            // transmit it. Refactor `create_local_id_registry` to instead accept None for
            // stateless_reset_token.
            let stateless_reset_token = self
                .config
                .context()
                .stateless_reset_token_generator
                .generate(local_connection_id.as_bytes());
            self.connection_id_mapper.create_local_id_registry(
                internal_connection_id,
                &local_connection_id,
                local_connection_id_expiration_time,
                stateless_reset_token,
            )
        };

        let endpoint_context = self.config.context();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
        //# When an Initial packet is sent by a client that has not previously
        //# received an Initial or Retry packet from the server, the client
        //# populates the Destination Connection ID field with an unpredictable
        //# value.
        let original_destination_connection_id = {
            let mut data = [0u8; InitialId::MIN_LEN];
            endpoint_context
                .random_generator
                .public_random_fill(&mut data);
            InitialId::try_from_bytes(&data).expect("InitialId creation failed.")
        };

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
        //# Note that clients cannot use the
        //# stateless_reset_token transport parameter because their transport
        //# parameters do not have confidentiality protection.
        //
        // The original_destination_connection_id is a random value used to establish the
        // connection. Since the connection is not yet secured, the client must not set a
        // stateless_reset_token.
        let peer_id_registry = self
            .connection_id_mapper
            .create_client_peer_id_registry(internal_connection_id);

        let congestion_controller = {
            let path_info = congestion_controller::PathInfo::new(&remote_address);
            endpoint_context
                .congestion_controller
                .new_congestion_controller(path_info)
        };

        //= https://www.rfc-editor.org/rfc/rfc9000#section-15
        //# This version of the specification is identified by the number
        //# 0x00000001.
        let quic_version = 0x00000001;

        let meta = event::builder::ConnectionMeta {
            endpoint_type: Cfg::ENDPOINT_TYPE,
            id: internal_connection_id.into(),
            timestamp,
        };
        let supervisor_context = supervisor::Context::new(
            self.connections.handshake_connections(),
            self.connections.len(),
            &remote_address,
            true,
        );
        let mut event_context = endpoint_context.event_subscriber.create_connection_context(
            &meta.clone().into_event(),
            &event::builder::ConnectionInfo {}.into_event(),
        );
        let mut publisher = event::ConnectionPublisherSubscriber::new(
            meta,
            quic_version,
            endpoint_context.event_subscriber,
            &mut event_context,
        );

        let mut transport_parameters = ClientTransportParameters {
            initial_source_connection_id: Some(local_connection_id.into()),
            ..Default::default()
        };
        let limits = endpoint_context
            .connection_limits
            .on_connection(&LimitsInfo::new(&remote_address));
        transport_parameters.load_limits(&limits);

        transport_parameters.max_datagram_frame_size = endpoint_context
            .datagram
            .max_datagram_frame_size(&PreConnectionInfo::new())
            .try_into()
            .expect("Failed to convert max_datagram_frame_size");

        transport_parameters.active_connection_id_limit = s2n_quic_core::varint::VarInt::from(
            connection::peer_id_registry::ACTIVE_CONNECTION_ID_LIMIT,
        )
        .try_into()
        .unwrap();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
        //# The Destination Connection ID field from the first Initial packet
        //# sent by a client is used to determine packet protection keys for
        //# Initial packets.
        //
        // Use the randomly generated `original_destination_connection_id` to generate the packet
        // protection keys.
        let (initial_key, initial_header_key) =
            <<Cfg::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialKey::new_client(
                original_destination_connection_id.as_bytes(),
            );
        let tls_session = endpoint_context
            .tls
            // TODO should SNI be optional? rustls expects a SNI but other tls providers dont seem
            // to require this value.
            .new_client_session(
                &transport_parameters,
                hostname.expect("application should provide a valid server name"),
            );
        let space_manager = PacketSpaceManager::new(
            original_destination_connection_id,
            tls_session,
            initial_key,
            initial_header_key,
            timestamp,
            &mut publisher,
        );

        let wakeup_handle = self
            .wakeup_queue
            .create_wakeup_handle(internal_connection_id);

        let path_handle =
            <<Cfg as endpoint::Config>::PathHandle as path::Handle>::from_remote_address(
                remote_address,
            );

        let connection_parameters = connection::Parameters {
            internal_connection_id,
            local_id_registry,
            peer_id_registry,
            space_manager,
            wakeup_handle,
            peer_connection_id: original_destination_connection_id.into(),
            local_connection_id,
            path_handle,
            congestion_controller,
            timestamp,
            quic_version,
            limits,
            max_mtu: self.max_mtu,
            event_context,
            supervisor_context: &supervisor_context,
            event_subscriber: endpoint_context.event_subscriber,
            datagram_endpoint: endpoint_context.datagram,
        };
        let connection = <Cfg as crate::endpoint::Config>::Connection::new(connection_parameters)?;
        self.connections
            .insert_client_connection(connection, internal_connection_id, sender);
        Ok(())
    }
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use super::*;
    use s2n_quic_core::{endpoint, event::testing::Subscriber, path, random, stateless_reset};

    #[derive(Debug)]
    pub struct Server;

    impl Config for Server {
        type CongestionControllerEndpoint =
            crate::recovery::congestion_controller::testing::mock::Endpoint;
        type TLSEndpoint = s2n_quic_core::crypto::tls::testing::Endpoint;
        type PathHandle = path::RemoteAddress;
        type Connection = connection::Implementation<Self>;
        type ConnectionLock = std::sync::Mutex<Self::Connection>;
        type EndpointLimits = Limits;
        type ConnectionIdFormat = connection::id::testing::Format;
        type StatelessResetTokenGenerator = stateless_reset::token::testing::Generator;
        type RandomGenerator = random::testing::Generator;
        type TokenFormat = s2n_quic_core::token::testing::Format;
        type ConnectionLimits = s2n_quic_core::connection::limits::Limits;
        type StreamManager = crate::stream::DefaultStreamManager;
        type ConnectionCloseFormatter = s2n_quic_core::connection::close::Development;
        type EventSubscriber = Subscriber;
        type PathMigrationValidator = path::migration::default::Validator;
        type PacketInterceptor = s2n_quic_core::packet::interceptor::Disabled;
        type DatagramEndpoint = s2n_quic_core::datagram::Disabled;

        fn context(&mut self) -> super::Context<Self> {
            todo!()
        }

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

    #[derive(Debug)]
    pub struct Client;

    impl Config for Client {
        type CongestionControllerEndpoint =
            crate::recovery::congestion_controller::testing::mock::Endpoint;
        type TLSEndpoint = s2n_quic_core::crypto::tls::testing::Endpoint;
        type PathHandle = path::RemoteAddress;
        type Connection = connection::Implementation<Self>;
        type ConnectionLock = std::sync::Mutex<Self::Connection>;
        type EndpointLimits = Limits;
        type ConnectionIdFormat = connection::id::testing::Format;
        type StatelessResetTokenGenerator = stateless_reset::token::testing::Generator;
        type RandomGenerator = random::testing::Generator;
        type TokenFormat = s2n_quic_core::token::testing::Format;
        type ConnectionLimits = s2n_quic_core::connection::limits::Limits;
        type StreamManager = crate::stream::DefaultStreamManager;
        type ConnectionCloseFormatter = s2n_quic_core::connection::close::Development;
        type EventSubscriber = Subscriber;
        type PathMigrationValidator = path::migration::default::Validator;
        type PacketInterceptor = s2n_quic_core::packet::interceptor::Disabled;
        type DatagramEndpoint = s2n_quic_core::datagram::Disabled;

        fn context(&mut self) -> super::Context<Self> {
            todo!()
        }

        const ENDPOINT_TYPE: endpoint::Type = endpoint::Type::Client;
    }

    #[derive(Debug)]
    pub struct Limits;

    impl endpoint::Limiter for Limits {
        fn on_connection_attempt(
            &mut self,
            _attempt: &endpoint::limits::ConnectionAttempt,
        ) -> endpoint::limits::Outcome {
            endpoint::limits::Outcome::allow()
        }
    }
}


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

use bytes::{Bytes, BytesMut};
use s2n_codec::{Encoder, EncoderBuffer};
use s2n_quic_core::path::MINIMUM_MTU;

/// Allocates a large single buffer, rather than several small buffers
///
/// Used for sending packets that contain CONNECTION_CLOSE frames.
#[derive(Debug)]
pub struct Buffer {
    buffer: BytesMut,
    max_size: usize,
    count: usize,
}

// This number shouldn't be _too_ small, otherwise we're performing a bunch
// of allocations. It also shouldn't be _too_ big so we hold on to those allocations
// for an extended period of time.
const DEFAULT_PACKETS: usize = 64;

impl Default for Buffer {
    fn default() -> Self {
        Self {
            buffer: BytesMut::new(),
            max_size: MINIMUM_MTU as usize,
            count: DEFAULT_PACKETS,
        }
    }
}

impl Buffer {
    pub fn write<F: FnOnce(EncoderBuffer) -> EncoderBuffer>(
        &mut self,
        on_write: F,
    ) -> Option<Bytes> {
        let max_size = self.max_size;

        if self.buffer.capacity() < max_size {
            let len = max_size * self.count;
            let mut buffer = BytesMut::with_capacity(len);
            // extend the length of the buffer to the capacity so we can
            // take a slice of it
            //
            // We could use `bytes::UninitSlice` but EncoderBuffer uses a
            // concrete slice instead.
            unsafe {
                // Safety: the EncoderBuffer only allows writing (no reading) from
                //         uninitialized memory
                buffer.set_len(len);
            }
            self.buffer = buffer;
        }

        let buffer = EncoderBuffer::new(&mut self.buffer[..max_size]);

        let new_buff = on_write(buffer);

        let len = max_size - new_buff.remaining_capacity();

        if len == 0 {
            return None;
        }

        debug_assert!(len <= max_size);

        Some(self.buffer.split_to(len).freeze())
    }
}

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

    #[test]
    fn empty_test() {
        let mut buffer = Buffer::default();
        assert_eq!(buffer.buffer.capacity(), 0);

        assert!(
            buffer.write(|buffer| buffer).is_none(),
            "empty writes should return None"
        );

        assert!(buffer.buffer.capacity() > 0);
    }

    #[test]
    fn non_empty_test() {
        let mut buffer = Buffer::default();

        let packet = buffer
            .write(|mut buffer| {
                assert_eq!(
                    buffer.remaining_capacity(),
                    MINIMUM_MTU as usize,
                    "the provider buffer should be the MINIMUM_MTU"
                );
                buffer.encode(&1337u16);
                buffer
            })
            .expect("non-empty writes should return a packet");

        assert_eq!(packet, 1337u16.to_be_bytes()[..]);

        assert!(buffer.buffer.capacity() > 0);
        assert!(
            buffer.buffer.capacity() < MINIMUM_MTU as usize * DEFAULT_PACKETS,
            "space should be trimmed off for the returned packet"
        );
    }
}


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

use crate::endpoint;
use alloc::collections::VecDeque;
use core::ops::Range;
use s2n_quic_core::{
    connection,
    crypto::RetryKey,
    event,
    inet::ExplicitCongestionNotification,
    io::tx,
    packet,
    path::{self, MINIMUM_MTU},
    random, time, token,
};

#[derive(Debug)]
pub struct Dispatch<Path: path::Handle> {
    // TODO: Find a better datastructure capable of handling delays in transmission
    // https://github.com/aws/s2n-quic/issues/280
    transmissions: VecDeque<Transmission<Path>>,
}

impl<Path: path::Handle> Default for Dispatch<Path> {
    fn default() -> Self {
        Self::new(endpoint::DEFAULT_MAX_PEERS)
    }
}

impl<Path: path::Handle> Dispatch<Path> {
    pub fn new(max_peers: usize) -> Self {
        Self {
            transmissions: VecDeque::with_capacity(max_peers),
        }
    }

    pub fn queue<T: token::Format, C: RetryKey>(
        &mut self,
        path_handle: Path,
        packet: &packet::initial::ProtectedInitial,
        local_connection_id: connection::LocalId,
        random: &mut dyn random::Generator,
        token_format: &mut T,
    ) {
        if let Some(transmission) = Transmission::new::<_, C>(
            path_handle,
            packet,
            local_connection_id,
            random,
            token_format,
        ) {
            self.transmissions.push_back(transmission);
        }
    }

    pub fn on_transmit<Tx: tx::Queue<Handle = Path>, Pub: event::EndpointPublisher>(
        &mut self,
        queue: &mut Tx,
        publisher: &mut Pub,
    ) {
        while let Some(transmission) = self.transmissions.pop_front() {
            match queue.push(&transmission) {
                Ok(tx::Outcome { len, .. }) => {
                    publisher.on_endpoint_packet_sent(event::builder::EndpointPacketSent {
                        packet_header: event::builder::PacketHeader::Retry {
                            version: transmission.version,
                        },
                    });

                    publisher.on_endpoint_datagram_sent(event::builder::EndpointDatagramSent {
                        len: len as u16,
                        gso_offset: 0,
                    });
                }
                Err(_) => {
                    self.transmissions.push_front(transmission);
                    return;
                }
            }
        }
    }
}

pub struct Transmission<Path: path::Handle> {
    path: Path,
    packet: [u8; MINIMUM_MTU as usize],
    packet_range: Range<usize>,
    version: u32,
}

impl<Path: path::Handle> core::fmt::Debug for Transmission<Path> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Transmission")
            .field("remote_address", &self.path.remote_address())
            .field("local_address", &self.path.local_address())
            .field("packet", &&self.packet[self.packet_range.clone()])
            .finish()
    }
}

impl<Path: path::Handle> Transmission<Path> {
    pub fn new<T: token::Format, C: RetryKey>(
        path: Path,
        packet: &packet::initial::ProtectedInitial,
        local_connection_id: connection::LocalId,
        random: &mut dyn random::Generator,
        token_format: &mut T,
    ) -> Option<Self> {
        let mut packet_buf = [0u8; MINIMUM_MTU as usize];
        let packet_range = packet::retry::Retry::encode_packet::<_, C>(
            &path.remote_address(),
            packet,
            &local_connection_id,
            random,
            token_format,
            &mut packet_buf,
        )?;

        Some(Self {
            path,
            packet: packet_buf,
            packet_range,
            version: packet.version,
        })
    }
}

impl<Path: path::Handle> AsRef<[u8]> for Transmission<Path> {
    fn as_ref(&self) -> &[u8] {
        &self.packet[self.packet_range.clone()]
    }
}

impl<Path: path::Handle> tx::Message for &Transmission<Path> {
    type Handle = Path;

    #[inline]
    fn path_handle(&self) -> &Self::Handle {
        &self.path
    }

    #[inline]
    fn ecn(&mut self) -> ExplicitCongestionNotification {
        Default::default()
    }

    #[inline]
    fn delay(&mut self) -> time::Duration {
        Default::default()
    }

    #[inline]
    fn ipv6_flow_label(&mut self) -> u32 {
        0
    }

    #[inline]
    fn can_gso(&self, segment_len: usize, _segment_count: usize) -> bool {
        segment_len >= self.as_ref().len()
    }

    #[inline]
    fn write_payload(
        &mut self,
        mut buffer: tx::PayloadBuffer,
        _gso_offset: usize,
    ) -> Result<usize, tx::Error> {
        buffer.write(self.as_ref())
    }
}


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

use crate::endpoint;
use alloc::collections::VecDeque;
use s2n_quic_core::{
    event, inet::ExplicitCongestionNotification, io::tx, packet, path, path::MINIMUM_MTU, random,
    stateless_reset, time,
};

#[derive(Debug)]
pub struct Dispatch<Path: path::Handle> {
    transmissions: VecDeque<Transmission<Path>>,
}

impl<Path: path::Handle> Default for Dispatch<Path> {
    fn default() -> Self {
        Self::new(endpoint::DEFAULT_MAX_PEERS)
    }
}

impl<Path: path::Handle> Dispatch<Path> {
    pub fn new(max_peers: usize) -> Self {
        Self {
            transmissions: VecDeque::with_capacity(max_peers),
        }
    }

    pub fn queue(
        &mut self,
        path: Path,
        token: stateless_reset::Token,
        max_tag_len: usize,
        triggering_packet_len: usize,
        random_generator: &mut dyn random::Generator,
    ) {
        if let Some(transmission) = Transmission::new(
            path,
            token,
            max_tag_len,
            triggering_packet_len,
            random_generator,
        ) {
            self.transmissions.push_back(transmission);
        }
    }

    pub fn on_transmit<Tx: tx::Queue<Handle = Path>, Pub: event::EndpointPublisher>(
        &mut self,
        queue: &mut Tx,
        publisher: &mut Pub,
    ) {
        while let Some(transmission) = self.transmissions.pop_front() {
            match queue.push(&transmission) {
                Ok(tx::Outcome { len, .. }) => {
                    publisher.on_endpoint_packet_sent(event::builder::EndpointPacketSent {
                        packet_header: event::builder::PacketHeader::StatelessReset {},
                    });

                    publisher.on_endpoint_datagram_sent(event::builder::EndpointDatagramSent {
                        len: len as u16,
                        gso_offset: 0,
                    });
                }
                Err(_) => {
                    self.transmissions.push_front(transmission);
                    return;
                }
            }
        }
    }
}

pub struct Transmission<Path: path::Handle> {
    path: Path,
    packet: [u8; MINIMUM_MTU as usize],
    packet_len: usize,
}

impl<Handle: path::Handle> core::fmt::Debug for Transmission<Handle> {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Transmission")
            .field("remote_address", &self.path.remote_address())
            .field("local_address", &self.path.local_address())
            .field("packet_len", &self.packet_len)
            .field("packet", &&self.packet[0..self.packet_len])
            .finish()
    }
}

impl<Path: path::Handle> Transmission<Path> {
    pub fn new(
        path: Path,
        token: stateless_reset::Token,
        max_tag_len: usize,
        triggering_packet_len: usize,
        random_generator: &mut dyn random::Generator,
    ) -> Option<Self> {
        let mut packet_buf = [0u8; MINIMUM_MTU as usize];

        let packet_len = packet::stateless_reset::encode_packet(
            token,
            max_tag_len,
            triggering_packet_len,
            random_generator,
            &mut packet_buf,
        )?;

        Some(Self {
            path,
            packet: packet_buf,
            packet_len,
        })
    }
}

impl<Path: path::Handle> AsRef<[u8]> for Transmission<Path> {
    fn as_ref(&self) -> &[u8] {
        &self.packet[..self.packet_len]
    }
}

impl<Path: path::Handle> tx::Message for &Transmission<Path> {
    type Handle = Path;

    #[inline]
    fn path_handle(&self) -> &Self::Handle {
        &self.path
    }

    #[inline]
    fn ecn(&mut self) -> ExplicitCongestionNotification {
        Default::default()
    }

    #[inline]
    fn delay(&mut self) -> time::Duration {
        Default::default()
    }

    #[inline]
    fn ipv6_flow_label(&mut self) -> u32 {
        0
    }

    #[inline]
    fn can_gso(&self, segment_len: usize, _segment_count: usize) -> bool {
        segment_len >= self.as_ref().len()
    }

    #[inline]
    fn write_payload(
        &mut self,
        mut buffer: tx::PayloadBuffer,
        _gso_offset: usize,
    ) -> Result<usize, tx::Error> {
        buffer.write(self.as_ref())
    }
}