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

#![cfg_attr(kani, allow(unused_variables))] // kani complains about things not being used

use crate::{
    application::ServerName,
    crypto::{
        header_crypto::{LONG_HEADER_MASK, SHORT_HEADER_MASK},
        tls, CryptoSuite, HeaderKey, Key,
    },
    endpoint, transport,
};
use alloc::sync::Arc;
use bytes::Bytes;
use core::{
    fmt,
    sync::atomic::{AtomicBool, Ordering},
    task::{Poll, Waker},
};
use futures_test::task::new_count_waker;
use s2n_codec::EncoderValue;
use std::{collections::VecDeque, fmt::Debug};

pub mod certificates {
    macro_rules! pem {
        ($name:ident, $path:expr) => {
            pub static $name: &str =
                include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/certs/", $path));
        };
    }
    macro_rules! der {
        ($name:ident, $path:expr) => {
            pub static $name: &[u8] =
                include_bytes!(concat!(env!("CARGO_MANIFEST_DIR"), "/certs/", $path));
        };
    }

    pem!(KEY_PEM, "key.pem");
    pem!(CERT_PEM, "cert.pem");
    der!(KEY_DER, "key.der");
    der!(CERT_DER, "cert.der");
    pem!(UNTRUSTED_CERT_PEM, "untrusted_cert.pem");
    pem!(UNTRUSTED_KEY_PEM, "untrusted_key.pem");

    // Certificates and Keys used specifically for testing the mTLS (mutual TLS) feature.
    //
    // The MTLS_CA_CERT is the root CA cert and used to sign both MTLS_SERVER_CERT and
    // MTLS_CLIENT_CERT.
    pem!(MTLS_CA_CERT, "mtls/ca-cert.pem");
    pem!(MTLS_SERVER_CERT, "mtls/server-cert.pem");
    pem!(MTLS_SERVER_KEY, "mtls/server-key.pem");
    pem!(MTLS_CLIENT_CERT, "mtls/client-cert.pem");
    pem!(MTLS_CLIENT_KEY, "mtls/client-key.pem");
}

#[derive(Debug)]
pub struct Endpoint;

impl super::Endpoint for Endpoint {
    type Session = Session;

    fn new_server_session<Params: EncoderValue>(
        &mut self,
        _transport_parameters: &Params,
    ) -> Self::Session {
        Session
    }

    fn new_client_session<Params: EncoderValue>(
        &mut self,
        _transport_parameters: &Params,
        _server_name: ServerName,
    ) -> Self::Session {
        Session
    }

    fn max_tag_length(&self) -> usize {
        16
    }
}

#[derive(Debug)]
pub struct Session;

impl super::Session for Session {
    fn poll<C: tls::Context<Self>>(
        &mut self,
        _context: &mut C,
    ) -> Poll<Result<(), transport::Error>> {
        todo!("implement dummy handshake")
    }
}

impl CryptoSuite for Session {
    type HandshakeKey = crate::crypto::key::testing::Key;
    type HandshakeHeaderKey = crate::crypto::key::testing::HeaderKey;
    type InitialKey = crate::crypto::key::testing::Key;
    type InitialHeaderKey = crate::crypto::key::testing::HeaderKey;
    type ZeroRttKey = crate::crypto::key::testing::Key;
    type ZeroRttHeaderKey = crate::crypto::key::testing::HeaderKey;
    type OneRttKey = crate::crypto::key::testing::Key;
    type OneRttHeaderKey = crate::crypto::key::testing::HeaderKey;
    type RetryKey = crate::crypto::key::testing::Key;
}

#[derive(Debug)]
pub struct TlsEndpoint<S: tls::Session, State: Debug> {
    pub session: S,
    pub context: Context<S, State>,
}

impl<S: tls::Session, State: Debug> TlsEndpoint<S, State> {
    fn new(session: S, context: Context<S, State>) -> Self {
        Self { session, context }
    }
}

/// A pair of TLS sessions and contexts being driven to completion
#[derive(Debug)]
pub struct Pair<S: tls::Session, C: tls::Session> {
    pub server: TlsEndpoint<S, ServerState>,
    pub client: TlsEndpoint<C, ClientState>,
    pub server_name: ServerName,
}

const TEST_SERVER_TRANSPORT_PARAMS: &[u8] = &[1, 2, 3];
const TEST_CLIENT_TRANSPORT_PARAMS: &[u8] = &[3, 2, 1];

impl<S: tls::Session, C: tls::Session> Pair<S, C> {
    pub fn new<SE, CE>(
        server_endpoint: &mut SE,
        client_endpoint: &mut CE,
        server_name: ServerName,
    ) -> Self
    where
        SE: tls::Endpoint<Session = S>,
        CE: tls::Endpoint<Session = C>,
    {
        use crate::crypto::InitialKey;

        let server = server_endpoint.new_server_session(&TEST_SERVER_TRANSPORT_PARAMS);
        let mut server_context =
            Context::new(endpoint::Type::Server, ServerState::WaitingClientHello);
        server_context.initial.crypto = Some(S::InitialKey::new_server(server_name.as_bytes()));

        let client =
            client_endpoint.new_client_session(&TEST_CLIENT_TRANSPORT_PARAMS, server_name.clone());
        let mut client_context = Context::new(endpoint::Type::Client, ClientState::ClientHelloSent);
        client_context.initial.crypto = Some(C::InitialKey::new_client(server_name.as_bytes()));

        Self {
            server: TlsEndpoint::new(server, server_context),
            client: TlsEndpoint::new(client, client_context),
            server_name,
        }
    }

    /// Returns true if `poll` should be called
    pub fn is_handshaking(&self) -> bool {
        !(self.server.context.handshake_complete && self.client.context.handshake_complete)
    }

    /// Continues progress of the handshake
    pub fn poll(
        &mut self,
        client_hello_cb_done: Option<Arc<AtomicBool>>,
    ) -> Result<(), transport::Error> {
        match self.client.session.poll(&mut self.client.context) {
            Poll::Ready(res) => res?,
            Poll::Pending => (),
        }
        match self.server.session.poll(&mut self.server.context) {
            Poll::Ready(res) => res?,
            Poll::Pending => (),
        }
        self.client.context.transfer(&mut self.server.context);

        eprintln!("1/2 RTT");
        if let Some(client_hello_cb_done) = client_hello_cb_done {
            // If the server is processing the async client hello callback, then return early
            // and poll it until it completes
            if !client_hello_cb_done.load(Ordering::SeqCst)
                && matches!(self.server.context.state, ServerState::ClientHelloRead)
            {
                return Ok(());
            };
        }
        self.check_progress();

        Ok(())
    }

    fn check_progress(&mut self) {
        match (&self.client.context.state, &self.server.context.state) {
            (ClientState::ClientHelloSent, ServerState::WaitingClientHello) => {
                assert!(
                    !self.server.context.initial.rx.is_empty(),
                    "client should send ClientHello"
                );
            }
            (ClientState::WaitingServerHello, ServerState::ClientHelloRead) => {
                assert!(
                    self.server.context.handshake.crypto.is_some(),
                    "server should have handshake keys after sending the ServerHello"
                );
                assert!(
                    self.server.context.application.crypto.is_some(),
                    "server should have application keys after sending a ServerFinished"
                );
                assert!(!self.server.context.handshake_complete);
                assert!(!self.client.context.handshake_complete);
            }
            (ClientState::ServerHelloRead, ServerState::WaitingClientFinish) => {
                assert!(
                    self.client.context.handshake.crypto.is_some(),
                    "client should have handshake keys after reading the ServerHello"
                );
                assert!(
                    self.client.context.application.crypto.is_some(),
                    "client should have application keys after reading the ServerFinished"
                );
                assert!(
                    self.client.context.handshake_complete,
                    "client should complete the handshake"
                );
            }
            (ClientState::ClientFinishSent, ServerState::ClientFinishRead) => {
                assert!(
                    self.server.context.handshake_complete,
                    "server should finish after reading the ClientFinished"
                );
                // Transition to HandshakeComplete after this case
            }
            _ => unreachable!("handshake made too many iterations"),
        };

        self.client.context.state.transition();
        self.server.context.state.transition();
    }

    /// Finished the test
    pub fn finish(&self) {
        self.client.context.finish(&self.server.context);

        assert_eq!(
            self.client.context.transport_parameters.as_ref().unwrap(),
            TEST_SERVER_TRANSPORT_PARAMS,
            "client did not receive the server transport parameters"
        );
        assert_eq!(
            self.server.context.transport_parameters.as_ref().unwrap(),
            TEST_CLIENT_TRANSPORT_PARAMS,
            "server did not receive the client transport parameters"
        );
        assert_eq!(
            self.client
                .context
                .server_name
                .as_ref()
                .expect("missing SNI on client"),
            &self.server_name[..]
        );
        assert_eq!(
            self.server
                .context
                .server_name
                .as_ref()
                .expect("missing ServerName on server"),
            &self.server_name[..]
        );
        assert!(
            matches!(self.client.context.state, ClientState::HandshakeComplete),
            "client state did not complete: current state: {:?}",
            self.client.context.state,
        );
        assert!(
            matches!(self.server.context.state, ServerState::HandshakeComplete),
            "server state did not complete: current state: {:?}",
            self.server.context.state,
        );

        // TODO check 0-rtt keys
    }
}

#[derive(Debug)]
pub enum ClientState {
    ClientHelloSent,
    WaitingServerHello,
    ServerHelloRead,
    ClientFinishSent,
    HandshakeComplete,
}

impl ClientState {
    fn transition(&mut self) {
        let new_state = match &self {
            ClientState::ClientHelloSent => Self::WaitingServerHello,
            ClientState::WaitingServerHello => Self::ServerHelloRead,
            ClientState::ServerHelloRead => Self::ClientFinishSent,
            ClientState::ClientFinishSent => Self::HandshakeComplete,
            ClientState::HandshakeComplete => unreachable!("handshake made too many iterations"),
        };

        *self = new_state;
    }
}

#[derive(Debug)]
pub enum ServerState {
    WaitingClientHello,
    ClientHelloRead, // send server hello, handshake keys, 1-rtt keys
    WaitingClientFinish,
    ClientFinishRead,
    HandshakeComplete,
}

impl ServerState {
    fn transition(&mut self) {
        let new_state = match &self {
            ServerState::WaitingClientHello => Self::ClientHelloRead,
            ServerState::ClientHelloRead => Self::WaitingClientFinish,
            ServerState::WaitingClientFinish => Self::ClientFinishRead,
            ServerState::ClientFinishRead => Self::HandshakeComplete,
            ServerState::HandshakeComplete => unreachable!("handshake made too many iterations"),
        };

        *self = new_state;
    }
}

/// Harness to ensure a TLS implementation adheres to the session contract
pub struct Context<C: CryptoSuite, State: Debug> {
    pub initial: Space<C::InitialKey, C::InitialHeaderKey>,
    pub handshake: Space<C::HandshakeKey, C::HandshakeHeaderKey>,
    pub application: Space<C::OneRttKey, C::OneRttHeaderKey>,
    pub zero_rtt_crypto: Option<(C::ZeroRttKey, C::ZeroRttHeaderKey)>,
    pub handshake_complete: bool,
    pub server_name: Option<Bytes>,
    pub application_protocol: Option<Bytes>,
    pub transport_parameters: Option<Bytes>,
    endpoint: endpoint::Type,
    pub state: State,
    waker: Waker,
}

impl<C: CryptoSuite, State: Debug> fmt::Debug for Context<C, State> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Context")
            .field("initial", &self.initial)
            .field("handshake", &self.handshake)
            .field("application", &self.application)
            .field("zero_rtt_crypto", &self.zero_rtt_crypto.is_some())
            .field("handshake_complete", &self.handshake_complete)
            .field("sni", &self.server_name)
            .field("application_protocol", &self.application_protocol)
            .field("transport_parameters", &self.transport_parameters)
            .field("endpoint", &self.endpoint)
            .finish()
    }
}

impl<C: CryptoSuite, State: Debug> Context<C, State> {
    fn new(endpoint: endpoint::Type, state: State) -> Self {
        let (waker, _wake_counter) = new_count_waker();
        Self {
            initial: Space::default(),
            handshake: Space::default(),
            application: Space::default(),
            zero_rtt_crypto: None,
            handshake_complete: false,
            server_name: None,
            application_protocol: None,
            transport_parameters: None,
            endpoint,
            state,
            waker,
        }
    }

    /// Transfers incoming and outgoing buffers between two contexts
    pub fn transfer<O: CryptoSuite, OS: Debug>(&mut self, other: &mut Context<O, OS>) {
        self.initial.transfer(&mut other.initial);
        self.handshake.transfer(&mut other.handshake);
        self.application.transfer(&mut other.application);
    }

    /// Finishes the test and asserts consistency
    pub fn finish<O: CryptoSuite, OS: Debug>(&self, other: &Context<O, OS>) {
        self.assert_done();
        other.assert_done();

        assert_eq!(
            self.server_name, other.server_name,
            "sni is not consistent between endpoints"
        );
        assert_eq!(
            self.application_protocol, other.application_protocol,
            "application_protocol is not consistent between endpoints"
        );

        assert_eq!(
            self.zero_rtt_crypto.is_some(),
            other.zero_rtt_crypto.is_some(),
            "0-rtt keys are not consistent between endpoints"
        );

        self.initial.finish(&other.initial);
        self.handshake.finish(&other.handshake);
        self.application.finish(&other.application);
    }

    fn assert_done(&self) {
        assert!(self.initial.crypto.is_some(), "missing initial crypto");
        assert!(self.handshake.crypto.is_some(), "missing handshake crypto");
        assert!(
            self.application.crypto.is_some(),
            "missing application crypto"
        );
        assert!(self.handshake_complete);
        assert!(self.application_protocol.is_some());
        assert!(self.transport_parameters.is_some());
    }

    fn on_application_params(&mut self, params: tls::ApplicationParameters) {
        self.transport_parameters = Some(Bytes::copy_from_slice(params.transport_parameters));
    }

    fn log(&self, event: &str) {
        eprintln!(
            "{:?}: {:?}: {}: {}",
            self.endpoint,               // Client or Server
            self.state,                  // ClientState or ServerState
            core::any::type_name::<C>(), // rustls or s2n-tls
            event,
        );
    }
}

pub struct Space<K: Key, Hk: HeaderKey> {
    pub crypto: Option<(K, Hk)>,
    pub rx: VecDeque<Bytes>,
    pub tx: VecDeque<Bytes>,
}

impl<K: Key, Hk: HeaderKey> Default for Space<K, Hk> {
    fn default() -> Self {
        Self {
            crypto: None,
            rx: VecDeque::new(),
            tx: VecDeque::new(),
        }
    }
}

impl<K: Key, Hk: HeaderKey> fmt::Debug for Space<K, Hk> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Space")
            .field("crypto", &self.crypto.is_some())
            .field("rx", &self.rx)
            .field("tx", &self.tx)
            .finish()
    }
}

impl<K: Key, Hk: HeaderKey> Space<K, Hk> {
    pub fn transfer<O: Key, Ohk: HeaderKey>(&mut self, other: &mut Space<O, Ohk>) {
        self.rx.extend(other.tx.drain(..));
        other.rx.extend(self.tx.drain(..));
    }

    fn tx(&mut self, bytes: Bytes) {
        self.tx.push_back(bytes)
    }

    fn rx(&mut self, max_len: Option<usize>) -> Option<Bytes> {
        loop {
            let mut chunk = self.rx.pop_front()?;

            if chunk.is_empty() {
                continue;
            }

            let max_len = max_len.unwrap_or(usize::MAX);

            if chunk.len() > max_len {
                self.rx.push_front(chunk.split_off(max_len));
            }

            return Some(chunk);
        }
    }

    fn finish<O: Key, Ohk: HeaderKey>(&self, other: &Space<O, Ohk>) {
        let (crypto_a, crypto_a_hk) = self.crypto.as_ref().expect("missing crypto");
        let (crypto_b, crypto_b_hk) = other.crypto.as_ref().expect("missing crypto");

        // ensure payloads can be encrypted and decrypted in both directions
        seal_open(crypto_a, crypto_b);
        seal_open(crypto_b, crypto_a);

        // ensure packet protection keys can be applied in both directions
        protect_unprotect(crypto_a_hk, crypto_b_hk, LONG_HEADER_MASK);
        protect_unprotect(crypto_b_hk, crypto_a_hk, LONG_HEADER_MASK);
        protect_unprotect(crypto_a_hk, crypto_b_hk, SHORT_HEADER_MASK);
        protect_unprotect(crypto_b_hk, crypto_a_hk, SHORT_HEADER_MASK);
    }
}

fn seal_open<S: Key, O: Key>(sealer: &S, opener: &O) {
    let packet_number = 123;
    let header = &[1, 2, 3, 4, 5, 6];

    let cleartext_payload = (0u16..1200).map(|i| i as u8).collect::<Vec<_>>();

    let mut encrypted_payload = cleartext_payload.clone();
    encrypted_payload.resize(cleartext_payload.len() + sealer.tag_len(), 0);

    let sealer_name = core::any::type_name::<S>();
    let opener_name = core::any::type_name::<O>();

    sealer
        .encrypt(packet_number, header, &mut encrypted_payload)
        .unwrap_or_else(|err| {
            panic!("encryption error; opener={opener_name}, sealer={sealer_name} - {err:?}")
        });
    opener
        .decrypt(packet_number, header, &mut encrypted_payload)
        .unwrap_or_else(|err| {
            panic!("decryption error; opener={opener_name}, sealer={sealer_name} - {err:?}")
        });

    assert_eq!(
        cleartext_payload[..],
        encrypted_payload[..cleartext_payload.len()]
    );
}

fn protect_unprotect<P: HeaderKey, U: HeaderKey>(protect: &P, unprotect: &U, tag_mask: u8) {
    let sample = [1u8; 1000];

    let mut protected_mask =
        protect.sealing_header_protection_mask(&sample[..protect.sealing_sample_len()]);
    let mut unprotected_mask =
        unprotect.opening_header_protection_mask(&sample[..unprotect.opening_sample_len()]);

    let protect_name = core::any::type_name::<P>();
    let unprotect_name = core::any::type_name::<U>();

    // we only care about certain bits being the same so mask out the others
    protected_mask[0] &= tag_mask;
    unprotected_mask[0] &= tag_mask;

    assert_eq!(
        &protected_mask, &unprotected_mask,
        "{protect_name} -> {unprotect_name}"
    );
}

impl<C: CryptoSuite, State: Debug> tls::Context<C> for Context<C, State> {
    fn on_handshake_keys(
        &mut self,
        key: C::HandshakeKey,
        header_key: C::HandshakeHeaderKey,
    ) -> Result<(), transport::Error> {
        assert!(
            self.handshake.crypto.is_none(),
            "handshake keys emitted multiple times"
        );
        self.log("handshake keys");
        self.handshake.crypto = Some((key, header_key));
        Ok(())
    }

    fn on_zero_rtt_keys(
        &mut self,
        key: C::ZeroRttKey,
        header_key: C::ZeroRttHeaderKey,
        params: tls::ApplicationParameters,
    ) -> Result<(), transport::Error> {
        assert!(
            self.zero_rtt_crypto.is_none(),
            "0-rtt keys emitted multiple times"
        );
        self.log("0-rtt keys");
        self.zero_rtt_crypto = Some((key, header_key));
        self.on_application_params(params);
        Ok(())
    }

    fn on_one_rtt_keys(
        &mut self,
        key: C::OneRttKey,
        header_key: C::OneRttHeaderKey,
        params: tls::ApplicationParameters,
    ) -> Result<(), transport::Error> {
        assert!(
            self.application.crypto.is_none(),
            "1-rtt keys emitted multiple times"
        );
        self.log("1-rtt keys");
        self.application.crypto = Some((key, header_key));
        self.on_application_params(params);

        Ok(())
    }

    fn on_server_name(
        &mut self,
        server_name: crate::application::ServerName,
    ) -> Result<(), transport::Error> {
        self.log("server name");
        self.server_name = Some(server_name.into_bytes());
        Ok(())
    }

    fn on_application_protocol(
        &mut self,
        application_protocol: Bytes,
    ) -> Result<(), transport::Error> {
        self.log("application protocol");
        self.application_protocol = Some(application_protocol);
        Ok(())
    }

    fn on_handshake_complete(&mut self) -> Result<(), transport::Error> {
        assert!(
            !self.handshake_complete,
            "handshake complete called multiple times"
        );
        assert!(
            !self.application_protocol.as_ref().unwrap().is_empty(),
            "application_protocol is empty at handshake complete"
        );
        self.handshake_complete = true;
        self.log("handshake complete");
        Ok(())
    }

    fn receive_initial(&mut self, max_len: Option<usize>) -> Option<Bytes> {
        self.log("rx initial");
        self.initial.rx(max_len)
    }

    fn receive_handshake(&mut self, max_len: Option<usize>) -> Option<Bytes> {
        self.log("rx handshake");
        self.handshake.rx(max_len)
    }

    fn receive_application(&mut self, max_len: Option<usize>) -> Option<Bytes> {
        self.log("rx application");
        self.application.rx(max_len)
    }

    fn can_send_initial(&self) -> bool {
        true
    }

    fn send_initial(&mut self, transmission: Bytes) {
        self.log("tx initial");
        self.initial.tx(transmission)
    }

    fn can_send_handshake(&self) -> bool {
        self.handshake.crypto.is_some()
    }

    fn send_handshake(&mut self, transmission: Bytes) {
        assert!(
            self.can_send_handshake(),
            "handshake keys need to be emitted before buffering handshake crypto"
        );
        self.log("tx handshake");
        self.handshake.tx(transmission)
    }

    fn can_send_application(&self) -> bool {
        self.application.crypto.is_some()
    }

    fn send_application(&mut self, transmission: Bytes) {
        assert!(
            self.can_send_application(),
            "1-rtt keys need to be emitted before buffering application crypto"
        );
        self.log("tx application");
        self.application.tx(transmission)
    }

    fn waker(&self) -> &Waker {
        &self.waker
    }
}


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

// s2n-quic's default implementation of the datagram component

use crate::{
    connection,
    datagram::{ConnectionInfo, Packet, PreConnectionInfo, ReceiveContext},
    transport::parameters::MaxDatagramFrameSize,
};
use alloc::collections::VecDeque;
use bytes::Bytes;
use core::{
    fmt,
    task::{Context, Poll, Waker},
};

/// Handles configuring unreliable datagram support.
///
/// The datagram endpoint allows users to configure their unreliable datagram
/// behavior. It contains two types, the Sender and Receiver, which are necessary
/// for configuring sending and receiving behavior separately. The default Sender
/// and Receiver behavior can be swapped out by implementing the respective [`Sender`](s2n-quic-core::datagram::traits::Sender) and
/// [`Receiver`](s2n-quic-core::datagram::traits::Receiver) traits.
///
#[derive(Debug, Default)]
pub struct Endpoint {
    send_queue_capacity: usize,
    recv_queue_capacity: usize,
}

impl Endpoint {
    /// Creates a builder for the default datagram endpoint
    pub fn builder() -> EndpointBuilder {
        EndpointBuilder::default()
    }
}

/// A builder for the default datagram endpoint
#[derive(Debug, Default)]
pub struct EndpointBuilder {
    send_queue_capacity: usize,
    recv_queue_capacity: usize,
}

#[non_exhaustive]
#[derive(Debug)]
pub enum BuilderError {
    ZeroCapacity,
}

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

impl fmt::Display for BuilderError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::ZeroCapacity { .. } => {
                write!(f, "Cannot create a queue with zero capacity")
            }
        }
    }
}

/// Builder for the datagram endpoint
impl EndpointBuilder {
    pub fn with_send_capacity(mut self, capacity: usize) -> Result<Self, BuilderError> {
        if capacity == 0 {
            return Err(BuilderError::ZeroCapacity);
        }
        self.send_queue_capacity = capacity;
        Ok(self)
    }

    pub fn with_recv_capacity(mut self, capacity: usize) -> Result<Self, BuilderError> {
        if capacity == 0 {
            return Err(BuilderError::ZeroCapacity);
        }
        self.recv_queue_capacity = capacity;
        Ok(self)
    }

    pub fn build(self) -> Result<Endpoint, core::convert::Infallible> {
        Ok(Endpoint {
            send_queue_capacity: self.send_queue_capacity,
            recv_queue_capacity: self.recv_queue_capacity,
        })
    }
}

impl super::Endpoint for Endpoint {
    type Sender = Sender;
    type Receiver = Receiver;

    fn create_connection(&mut self, info: &ConnectionInfo) -> (Self::Sender, Self::Receiver) {
        (
            Sender::builder()
                .with_capacity(self.send_queue_capacity)
                .with_connection_info(info)
                .build()
                .unwrap(),
            Receiver::builder()
                .with_capacity(self.recv_queue_capacity)
                .with_max_datagram_frame_size(MaxDatagramFrameSize::RECOMMENDED)
                .build()
                .unwrap(),
        )
    }

    fn max_datagram_frame_size(&self, _info: &PreConnectionInfo) -> u64 {
        MaxDatagramFrameSize::RECOMMENDED
    }
}

/// Handles receiving unreliable datagrams.
///
/// Stores the queue of datagrams received from the peer.
/// Old datagrams will be popped off the queue in favor of new datagrams if the
/// queue capacity is reached.
pub struct Receiver {
    queue: VecDeque<Bytes>,
    capacity: usize,
    waker: Option<Waker>,
    max_datagram_frame_size: u64,
    error: Option<connection::Error>,
}

impl Receiver {
    /// Creates a builder for the default datagram receiver
    fn builder() -> ReceiverBuilder {
        ReceiverBuilder::default()
    }

    /// Returns a datagram if there are any on the queue
    pub fn recv_datagram(&mut self) -> Option<Bytes> {
        self.queue.pop_front()
    }

    /// Dequeues a datagram received from the peer.
    ///
    /// # Return value
    ///
    /// - `Poll::Pending` if there are no datagrams to be received on the queue. In this case,
    ///   the caller should retry receiving after the [`Waker`](core::task::Waker) on the provided
    ///   [`Context`](core::task::Context) is notified.
    /// - `Poll::Ready(Datagram)` if there exists a datagram to be received.
    /// - `Poll::Ready(DatagramError)` if a connection error occurred and no more datagrams will be received.
    pub fn poll_recv_datagram(&mut self, cx: &mut Context) -> Poll<Result<Bytes, DatagramError>> {
        if let Some(datagram) = self.queue.pop_front() {
            Poll::Ready(Ok(datagram))
        // If there was some connection-level error we don't take the waker
        // and instead error as there will never be any datagrams to receive.
        } else if let Some(err) = self.error {
            Poll::Ready(Err(DatagramError::ConnectionError { error: err }))
        } else {
            self.waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }
}

impl super::Receiver for Receiver {
    fn on_datagram(&mut self, _: &ReceiveContext, datagram: &[u8]) {
        if datagram.len() as u64 > self.max_datagram_frame_size {
            return;
        }
        // The oldest datagram on the queue is popped off if the queue is full.
        // Configure this behavior by implementing a custom Receiver for datagrams.
        if self.queue.len() == self.capacity {
            self.queue.pop_front();
        }

        self.queue
            .push_back(bytes::Bytes::copy_from_slice(datagram));
        // Since a datagram was appended to the queue, wake the waker to inform
        // the user that it can receive datagrams now.
        if let Some(w) = self.waker.take() {
            w.wake();
        }
    }

    fn on_connection_error(&mut self, error: connection::Error) {
        self.error = Some(error);
        if let Some(w) = self.waker.take() {
            w.wake();
        }
    }
}

// A builder for the default datagram receiver
///
/// Use to configure a datagram recv queue size and how large of
/// a datagram we can receive.
#[derive(Debug)]
struct ReceiverBuilder {
    queue_capacity: usize,
    max_datagram_frame_size: u64,
}

impl Default for ReceiverBuilder {
    fn default() -> Self {
        Self {
            queue_capacity: 200,
            max_datagram_frame_size: MaxDatagramFrameSize::RECOMMENDED,
        }
    }
}

impl ReceiverBuilder {
    /// Sets the capacity of the datagram receiver queue
    pub fn with_capacity(mut self, capacity: usize) -> Self {
        self.queue_capacity = capacity;
        self
    }

    pub fn with_max_datagram_frame_size(mut self, size: u64) -> Self {
        self.max_datagram_frame_size = size;
        self
    }

    /// Builds the datagram receiver
    pub fn build(self) -> Result<Receiver, core::convert::Infallible> {
        Ok(Receiver {
            queue: VecDeque::with_capacity(self.queue_capacity),
            capacity: self.queue_capacity,
            waker: None,
            max_datagram_frame_size: self.max_datagram_frame_size,
            error: None,
        })
    }
}

/// A struct to handle sending unreliable datagrams.
///
/// The Sender struct contains the queue of unreliable datagrams to be sent.
/// During transmission time, we alternate between sending datagrams and sending stream
/// data. This is to ensure there is a balance between the amount of reliable
/// and unreliable data getting sent.
///
/// Datagrams are written to the packet in the order they are added to the queue.
/// A datagram that is too large to fit in the packet will be dropped, unless the
/// packet already contains written datagrams. This attempts to prevent
/// the case where all datagrams are dropped because only a small amount of packet
/// space remains.
///
/// Note that there is currently no expiration date for datagrams to live on the queue.
/// Implement the [`Sender`](s2n-quic-core::datagram::traits::Sender) trait if
/// this behavior is necessary for your use-case.
///
#[derive(Debug)]
pub struct Sender {
    queue: VecDeque<Datagram>,
    capacity: usize,
    min_packet_space: usize,
    max_packet_space: usize,
    smoothed_packet_size: f64,
    waker: Option<Waker>,
    max_datagram_payload: u64,
    error: Option<connection::Error>,
}

#[non_exhaustive]
#[derive(Debug, PartialEq)]
pub struct Datagram {
    pub data: Bytes,
}

#[non_exhaustive]
#[derive(Debug, PartialEq)]
pub enum DatagramError {
    #[non_exhaustive]
    QueueAtCapacity,
    #[non_exhaustive]
    ExceedsPeerTransportLimits,
    #[non_exhaustive]
    ConnectionError { error: connection::Error },
}

impl fmt::Display for DatagramError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::QueueAtCapacity { .. } => {
                write!(f, "Queue does not have room for more datagrams.")
            }
            Self::ExceedsPeerTransportLimits { .. } => {
                write!(
                    f,
                    "Datagram size is larger than peer's transport parameters."
                )
            }
            Self::ConnectionError { .. } => {
                write!(f, "Connection-level error occurred.")
            }
        }
    }
}

impl Sender {
    /// Creates a builder for the default datagram sender
    fn builder() -> SenderBuilder {
        SenderBuilder::default()
    }

    /// Enqueues a datagram for sending it towards the peer.
    ///
    /// # Return value
    ///
    /// - `Poll::Pending` if the datagram's send buffer capacity is currently exhausted
    ///   and the datagram was not added to the queue. In this case, the caller should
    ///   retry sending after the [`Waker`](core::task::Waker) on the provided
    ///   [`Context`](core::task::Context) is notified.
    /// - `Poll::Ready(Ok(()))` if the datagram was enqueued for sending.
    /// - `Poll::Ready(Err(DatagramError))` if an error occurred while trying
    ///   to send the datagram.
    pub fn poll_send_datagram(
        &mut self,
        data: &mut bytes::Bytes,
        cx: &mut Context,
    ) -> Poll<Result<(), DatagramError>> {
        if data.len() as u64 > self.max_datagram_payload {
            return Poll::Ready(Err(DatagramError::ExceedsPeerTransportLimits));
        }

        // If there was some connection-level error the user is not allowed to add
        // datagrams to the queue as they will never be sent.
        if let Some(err) = self.error {
            return Poll::Ready(Err(DatagramError::ConnectionError { error: err }));
        }

        if self.queue.len() == self.capacity {
            self.waker = Some(cx.waker().clone());
            return Poll::Pending;
        }

        let datagram = Datagram {
            data: core::mem::replace(data, bytes::Bytes::new()),
        };
        self.queue.push_back(datagram);
        Poll::Ready(Ok(()))
    }

    /// Adds datagrams on the queue to be sent
    ///
    /// If the datagram queue is at capacity the oldest datagram will be popped
    /// off the queue and returned to make space for the newest datagram.
    ///
    /// # Return value
    /// - `Ok(None)` if the datagram was enqueued for sending
    /// - `Ok(Some(Bytes))` if the queue is at capacity this will be the oldest datagram on the queue
    /// - `Err(DatagramError)` if some error occurred
    pub fn send_datagram_forced(
        &mut self,
        data: bytes::Bytes,
    ) -> Result<Option<Bytes>, DatagramError> {
        if data.len() as u64 > self.max_datagram_payload {
            return Err(DatagramError::ExceedsPeerTransportLimits);
        }

        // If there was some connection-level error the user is not allowed to add
        // datagrams to the queue as they will never be sent.
        if let Some(err) = self.error {
            return Err(DatagramError::ConnectionError { error: err });
        }

        // Pop oldest datagram off the queue if it is at capacity
        let mut oldest = None;
        if self.queue.len() == self.capacity {
            oldest = self.queue.pop_front();
        }

        let datagram = Datagram { data };
        self.queue.push_back(datagram);

        match oldest {
            Some(datagram) => Ok(Some(datagram.data)),
            None => Ok(None),
        }
    }

    /// Adds datagrams on the queue to be sent
    ///
    /// If the queue is full the newest datagram is not added and an error is returned.
    ///
    /// # Return value
    /// - `Ok()` if the datagram was enqueued for sending
    /// - `Err(DatagramError)` if some error occurred
    pub fn send_datagram(&mut self, data: bytes::Bytes) -> Result<(), DatagramError> {
        if data.len() as u64 > self.max_datagram_payload {
            return Err(DatagramError::ExceedsPeerTransportLimits);
        }

        // If there was some connection-level error the user is not allowed to add
        // datagrams to the queue as they will never be sent.
        if let Some(err) = self.error {
            return Err(DatagramError::ConnectionError { error: err });
        }

        if self.queue.len() == self.capacity {
            return Err(DatagramError::QueueAtCapacity);
        }

        let datagram = Datagram { data };
        self.queue.push_back(datagram);
        Ok(())
    }

    /// Filter through the datagrams in the send queue and only keep those that
    /// match a predicate
    pub fn retain_datagrams<F>(&mut self, f: F)
    where
        F: FnMut(&Datagram) -> bool,
    {
        self.queue.retain(f);
    }

    fn record_capacity_stats(&mut self, capacity: usize) {
        if capacity < self.min_packet_space || self.min_packet_space == 0 {
            self.min_packet_space = capacity;
        }
        if capacity > self.max_packet_space {
            self.max_packet_space = capacity;
        }

        // https://www.rfc-editor.org/rfc/rfc9002#section-5.3
        self.smoothed_packet_size =
            7.0 / 8.0 * self.smoothed_packet_size + 1.0 / 8.0 * capacity as f64;
    }

    /// Returns the largest packet space for datagrams seen during this connection.
    ///
    /// Should be used to determine an appropriate datagram size that can be sent in
    /// this connection.
    pub fn max_packet_space(&self) -> usize {
        self.max_packet_space
    }

    /// Returns the smallest packet space for datagrams seen during this connection.
    ///
    /// Should be used to determine an appropriate datagram size that can be sent in
    /// this connection.
    pub fn min_packet_space(&self) -> usize {
        self.min_packet_space
    }

    /// Returns a smoothed calculation of the size of packet space for datagrams seen during this connection.
    ///
    /// Should be used to determine an appropriate datagram size that can be sent in
    /// this connection.
    pub fn smoothed_packet_space(&self) -> usize {
        self.smoothed_packet_size as usize
    }
}

impl super::Sender for Sender {
    fn on_transmit<P: Packet>(&mut self, packet: &mut P) {
        // Cede space to stream data when datagrams are not prioritized
        if packet.has_pending_streams() && !packet.datagrams_prioritized() {
            return;
        }
        self.record_capacity_stats(packet.remaining_capacity());
        let mut has_written = false;
        while packet.remaining_capacity() > 0 {
            if let Some(datagram) = self.queue.pop_front() {
                // Ensure there is enough space in the packet to send a datagram
                if packet.remaining_capacity() >= datagram.data.len() {
                    match packet.write_datagram(&datagram.data) {
                        Ok(()) => has_written = true,
                        Err(_error) => {
                            continue;
                        }
                    }
                    // Since a datagram was popped off the queue, wake the
                    // stored waker if we have one to let the application know
                    // that there is space on the queue for more datagrams.
                    if let Some(w) = self.waker.take() {
                        w.wake();
                    }
                } else {
                    // This check keeps us from popping all the datagrams off the
                    // queue when packet space remaining is smaller than the datagram.
                    if has_written {
                        self.queue.push_front(datagram);
                        return;
                    }
                }
            } else {
                // If there are no datagrams on the queue we return
                return;
            }
        }
    }

    #[inline]
    fn has_transmission_interest(&self) -> bool {
        !self.queue.is_empty()
    }

    fn on_connection_error(&mut self, error: connection::Error) {
        self.error = Some(error);
        if let Some(w) = self.waker.take() {
            w.wake();
        }
    }
}

/// A builder for the default datagram sender
///
/// Use to configure a datagram send queue size
#[derive(Debug)]
struct SenderBuilder {
    queue_capacity: usize,
    max_datagram_payload: u64,
}

impl Default for SenderBuilder {
    fn default() -> Self {
        Self {
            queue_capacity: 200,
            max_datagram_payload: 0,
        }
    }
}

impl SenderBuilder {
    /// Sets the capacity of the datagram sender queue
    pub fn with_capacity(mut self, capacity: usize) -> Self {
        self.queue_capacity = capacity;
        self
    }

    /// Gives the default sender relevant connection info
    pub fn with_connection_info(mut self, connection_info: &ConnectionInfo) -> Self {
        self.max_datagram_payload = connection_info.max_datagram_payload;
        self
    }

    /// Builds the datagram sender into a provider
    pub fn build(self) -> Result<Sender, core::convert::Infallible> {
        Ok(Sender {
            queue: VecDeque::with_capacity(self.queue_capacity),
            capacity: self.queue_capacity,
            max_datagram_payload: self.max_datagram_payload,
            max_packet_space: 0,
            min_packet_space: 0,
            smoothed_packet_size: 0.0,
            waker: None,
            error: None,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::datagram::WriteError;
    use core::task::{Context, Poll};
    use futures_test::task::{new_count_waker, noop_waker};

    #[test]
    fn send_datagram_forced() {
        let conn_info = ConnectionInfo {
            max_datagram_payload: 100,
            waker: noop_waker(),
        };
        // Create a default sender queue that only holds two elements
        let mut default_sender = Sender::builder()
            .with_capacity(2)
            .with_connection_info(&conn_info)
            .build()
            .unwrap();
        let datagram_0 = bytes::Bytes::from_static(&[1, 2, 3]);
        let datagram_1 = bytes::Bytes::from_static(&[4, 5, 6]);
        let datagram_2 = bytes::Bytes::from_static(&[7, 8, 9]);
        assert_eq!(default_sender.send_datagram_forced(datagram_0), Ok(None));
        assert_eq!(default_sender.send_datagram_forced(datagram_1), Ok(None));
        // Queue has reached capacity so oldest datagram is returned
        let result = default_sender.send_datagram_forced(datagram_2);
        assert_eq!(result, Ok(Some(bytes::Bytes::from_static(&[1, 2, 3]))));

        // Oldest datagram has been bumped off the queue and the newest two datagrams
        // are there
        let second = default_sender.queue.pop_front().unwrap();
        assert_eq!(second.data[..], [4, 5, 6]);
        let third = default_sender.queue.pop_front().unwrap();
        assert_eq!(third.data[..], [7, 8, 9]);
        assert!(default_sender.queue.is_empty());

        // Connection-level error means new datagrams are not added to the queue
        let conn_err = connection::Error::closed(crate::endpoint::Location::Remote);
        default_sender.error = Some(conn_err);
        assert_eq!(
            default_sender.send_datagram_forced(bytes::Bytes::from_static(&[7, 8, 9])),
            Err(DatagramError::ConnectionError { error: conn_err })
        );
        assert!(default_sender.queue.is_empty());
    }

    #[test]
    fn send_datagram() {
        let conn_info = ConnectionInfo {
            max_datagram_payload: 100,
            waker: noop_waker(),
        };
        // Create a default sender queue that only holds two elements
        let mut default_sender = Sender::builder()
            .with_capacity(2)
            .with_connection_info(&conn_info)
            .build()
            .unwrap();
        let datagram_0 = bytes::Bytes::from_static(&[1, 2, 3]);
        let datagram_1 = bytes::Bytes::from_static(&[4, 5, 6]);
        let datagram_2 = bytes::Bytes::from_static(&[7, 8, 9]);
        assert_eq!(default_sender.send_datagram(datagram_0), Ok(()));
        assert_eq!(default_sender.send_datagram(datagram_1), Ok(()));
        // Attempting to send a third datagram will result in an error, since the queue
        // is at capacity
        assert_eq!(
            default_sender.send_datagram(datagram_2),
            Err(DatagramError::QueueAtCapacity)
        );

        // Check that the first two datagrams are still there
        let first = default_sender.queue.pop_front().unwrap();
        assert_eq!(first.data[..], [1, 2, 3]);
        let second = default_sender.queue.pop_front().unwrap();
        assert_eq!(second.data[..], [4, 5, 6]);
        assert!(default_sender.queue.is_empty());

        // Connection-level error means new datagrams are not added to the queue
        let conn_err = connection::Error::closed(crate::endpoint::Location::Remote);
        default_sender.error = Some(conn_err);
        assert_eq!(
            default_sender.send_datagram(bytes::Bytes::from_static(&[7, 8, 9])),
            Err(DatagramError::ConnectionError { error: conn_err })
        );
        assert!(default_sender.queue.is_empty());
    }

    #[test]
    fn poll_send_datagram() {
        let conn_info = ConnectionInfo::new(100, noop_waker());
        let mut default_sender = Sender::builder()
            .with_capacity(2)
            .with_connection_info(&conn_info)
            .build()
            .unwrap();
        let mut datagram_0 = bytes::Bytes::from_static(&[1, 2, 3]);
        let mut datagram_1 = bytes::Bytes::from_static(&[4, 5, 6]);
        let mut datagram_2 = bytes::Bytes::from_static(&[7, 8, 9]);

        let (waker, wake_count) = new_count_waker();
        let mut cx = Context::from_waker(&waker);

        assert_eq!(
            default_sender.poll_send_datagram(&mut datagram_0, &mut cx),
            Poll::Ready(Ok(()))
        );

        assert_eq!(
            default_sender.poll_send_datagram(&mut datagram_1, &mut cx),
            Poll::Ready(Ok(()))
        );

        // Waker has not been set up yet
        assert!(default_sender.waker.is_none());

        // Queue is at capacity
        assert_eq!(
            default_sender.poll_send_datagram(&mut datagram_2, &mut cx),
            Poll::Pending
        );

        // Since queue is at capacity default_sender is now storing a waker that will
        // alert when the queue has more space
        assert!(default_sender.waker.is_some());

        let mut packet = MockPacket {
            remaining_capacity: 10,
            has_pending_streams: false,
            datagrams_prioritized: false,
        };
        crate::datagram::Sender::on_transmit(&mut default_sender, &mut packet);

        // Waker was called
        assert_eq!(wake_count, 1);

        // Now datagrams can be added to the queue as there is space
        let mut datagram_3 = bytes::Bytes::from_static(&[10, 11, 12]);
        assert_eq!(
            default_sender.poll_send_datagram(&mut datagram_3, &mut cx),
            Poll::Ready(Ok(()))
        );

        // Check that all datagrams we expect are on the queue
        let datagram = default_sender.queue.pop_front().unwrap();
        assert_eq!(datagram.data[..], [10, 11, 12]);
        assert!(default_sender.queue.is_empty());

        // Connection-level error means new datagrams are not added to the queue
        let conn_err = connection::Error::closed(crate::endpoint::Location::Remote);
        default_sender.error = Some(conn_err);
        assert_eq!(
            default_sender.poll_send_datagram(&mut bytes::Bytes::from_static(&[7, 8, 9]), &mut cx),
            Poll::Ready(Err(DatagramError::ConnectionError { error: conn_err }))
        );
        assert!(default_sender.queue.is_empty());
    }

    #[test]
    fn retain_datagrams() {
        let conn_info = ConnectionInfo {
            max_datagram_payload: 100,
            waker: noop_waker(),
        };
        let mut default_sender = Sender::builder()
            .with_capacity(3)
            .with_connection_info(&conn_info)
            .build()
            .unwrap();
        let datagram_0 = bytes::Bytes::from_static(&[1, 2, 3]);
        let datagram_1 = bytes::Bytes::from_static(&[4, 5, 6]);
        let datagram_2 = bytes::Bytes::from_static(&[7, 8, 9]);
        assert_eq!(default_sender.send_datagram_forced(datagram_0), Ok(None));
        assert_eq!(default_sender.send_datagram_forced(datagram_1), Ok(None));
        assert_eq!(default_sender.send_datagram_forced(datagram_2), Ok(None));

        // Keep only the third datagram
        default_sender.retain_datagrams(|datagram| datagram.data[..] == [7, 8, 9]);
        let first = default_sender.queue.pop_front().unwrap();
        assert_eq!(first.data[..], [7, 8, 9]);
        assert!(default_sender.queue.is_empty());
    }

    #[test]
    fn record_capacity_stats() {
        // Here we test that record_capacity_stats() is working as expected. We use
        // a precalculated const for the smoothed_packet_space value.
        const SMOOTHED_PACKET_SPACE: usize = 102;

        let mut default_sender = Sender::builder().build().unwrap();
        default_sender.record_capacity_stats(100);
        default_sender.record_capacity_stats(100);
        default_sender.record_capacity_stats(200);
        default_sender.record_capacity_stats(500);

        assert_eq!(default_sender.max_packet_space(), 500);
        assert_eq!(default_sender.min_packet_space(), 100);
        assert_eq!(
            default_sender.smoothed_packet_space(),
            SMOOTHED_PACKET_SPACE
        );
    }

    #[test]
    // Check that our default on_transmit function doesn't continue to pop datagrams
    // off the send queue if the remaining packet space is too small to send datagrams.
    fn has_written_test() {
        let conn_info = ConnectionInfo::new(100, noop_waker());
        let mut default_sender = Sender::builder()
            .with_connection_info(&conn_info)
            .build()
            .unwrap();
        let datagram_0 = bytes::Bytes::from_static(&[1, 2, 3]);
        let datagram_1 = bytes::Bytes::from_static(&[4, 5, 6]);
        assert_eq!(default_sender.send_datagram_forced(datagram_0), Ok(None));
        assert_eq!(default_sender.send_datagram_forced(datagram_1), Ok(None));

        // Packet size is just enough to write the first datagram with some
        // room left over, but not enough to write the second.
        let mut packet = MockPacket {
            remaining_capacity: 5,
            has_pending_streams: false,
            datagrams_prioritized: false,
        };
        crate::datagram::Sender::on_transmit(&mut default_sender, &mut packet);

        // Packet still has capacity to write datagrams
        assert!(packet.remaining_capacity > 0);
        // Send queue is not completely depleted
        assert!(!default_sender.queue.is_empty());
    }

    fn fake_receive_context() -> crate::datagram::ReceiveContext<'static> {
        crate::datagram::ReceiveContext {
            path: crate::event::api::Path {
                local_addr: crate::event::api::SocketAddress::IpV4 {
                    ip: &[0; 4],
                    port: 0,
                },
                local_cid: crate::event::api::ConnectionId { bytes: &[] },
                remote_addr: crate::event::api::SocketAddress::IpV4 {
                    ip: &[0; 4],
                    port: 0,
                },
                remote_cid: crate::event::api::ConnectionId { bytes: &[] },
                id: 0,
                is_active: true,
            },
        }
    }

    #[test]
    fn on_datagram() {
        // Create a receiver with limited capacity
        let mut receiver = Receiver::builder()
            .with_capacity(2)
            .with_max_datagram_frame_size(5)
            .build()
            .unwrap();

        let datagram_0 = vec![1, 2, 3];
        let datagram_1 = vec![4, 5, 6];
        let datagram_2 = vec![7, 8, 9];
        let ctx = fake_receive_context();
        crate::datagram::Receiver::on_datagram(&mut receiver, &ctx, &datagram_0);
        crate::datagram::Receiver::on_datagram(&mut receiver, &ctx, &datagram_1);
        // Datagram queue will be forced to drop a datagram to receive the newest one
        crate::datagram::Receiver::on_datagram(&mut receiver, &ctx, &datagram_2);

        // Oldest datagram has been dropped
        assert_eq!(receiver.queue.pop_front().unwrap(), datagram_1);
        assert_eq!(receiver.queue.pop_front().unwrap(), datagram_2);
        assert!(receiver.queue.pop_front().is_none());

        // Datagram sent by peer is larger than max_datagram_frame_size
        let datagram_3 = vec![10, 11, 12, 13, 14, 15];
        crate::datagram::Receiver::on_datagram(&mut receiver, &ctx, &datagram_3);
        // Queue is empty as datagram was not accepted
        assert!(receiver.queue.pop_front().is_none());
    }

    #[test]
    fn recv_datagram() {
        let mut receiver = Receiver::builder().build().unwrap();

        // Calling receive with no datagrams on the queue will result in a None
        assert!(receiver.recv_datagram().is_none());

        // Append a datagram to the receive queue
        receiver
            .queue
            .push_back(bytes::Bytes::from_static(&[1, 2, 3]));

        // Now the user can receive a datagram
        assert_eq!(
            receiver.recv_datagram(),
            Some(bytes::Bytes::from_static(&[1, 2, 3]),)
        );
    }

    #[test]
    fn poll_recv_datagram() {
        // Create a receiver
        let mut receiver = Receiver::builder().build().unwrap();

        let (waker, wake_count) = new_count_waker();
        let mut cx = Context::from_waker(&waker);

        // There are no datagrams on the queue to be received
        assert_eq!(receiver.poll_recv_datagram(&mut cx), Poll::Pending);

        // Adding some datagrams to the queue will wake up the stored waker
        let datagram_0 = vec![1, 2, 3];
        let datagram_1 = vec![4, 5, 6];
        let datagram_2 = vec![7, 8, 9];
        let ctx = fake_receive_context();
        crate::datagram::Receiver::on_datagram(&mut receiver, &ctx, &datagram_0);
        crate::datagram::Receiver::on_datagram(&mut receiver, &ctx, &datagram_1);
        crate::datagram::Receiver::on_datagram(&mut receiver, &ctx, &datagram_2);

        // Waker was called
        assert_eq!(wake_count, 1);

        assert_eq!(
            receiver.poll_recv_datagram(&mut cx),
            Poll::Ready(Ok(bytes::Bytes::from_static(&[1, 2, 3])))
        );

        // Mock a connection error
        let connection_error = connection::Error::closed(crate::endpoint::Location::Remote);
        receiver.error = Some(connection_error);

        // Continue to recv bytes until we run out and return the connection error
        assert_eq!(
            receiver.poll_recv_datagram(&mut cx),
            Poll::Ready(Ok(bytes::Bytes::from_static(&[4, 5, 6])))
        );

        assert_eq!(
            receiver.poll_recv_datagram(&mut cx),
            Poll::Ready(Ok(bytes::Bytes::from_static(&[7, 8, 9])))
        );

        assert_eq!(
            receiver.poll_recv_datagram(&mut cx),
            Poll::Ready(Err(DatagramError::ConnectionError {
                error: connection_error
            }))
        );
    }

    // The MockPacket mocks writing datagrams to a packet, but is not
    // a fully functional mock. It is used to test the logic in the
    // on_transmit function.
    struct MockPacket {
        has_pending_streams: bool,
        datagrams_prioritized: bool,
        remaining_capacity: usize,
    }

    impl crate::datagram::Packet for MockPacket {
        fn remaining_capacity(&self) -> usize {
            self.remaining_capacity
        }

        fn write_datagram(&mut self, data: &[u8]) -> Result<(), WriteError> {
            self.write_datagram_vectored(&[data])
        }

        fn write_datagram_vectored(&mut self, data: &[&[u8]]) -> Result<(), WriteError> {
            let data_len = data.iter().map(|d| d.len()).sum::<usize>();
            if data_len > self.remaining_capacity {
                return Err(WriteError::ExceedsPacketCapacity);
            }
            self.remaining_capacity -= data_len;
            Ok(())
        }

        fn has_pending_streams(&self) -> bool {
            self.has_pending_streams
        }

        fn datagrams_prioritized(&self) -> bool {
            self.datagrams_prioritized
        }
    }
}


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

use crate::datagram::{
    ConnectionInfo, Endpoint, Packet, PreConnectionInfo, ReceiveContext, Receiver, Sender,
};

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

impl Endpoint for Disabled {
    type Sender = DisabledSender;
    type Receiver = DisabledReceiver;

    fn create_connection(&mut self, _info: &ConnectionInfo) -> (Self::Sender, Self::Receiver) {
        (DisabledSender(()), DisabledReceiver(()))
    }

    fn max_datagram_frame_size(&self, _info: &PreConnectionInfo) -> u64 {
        0
    }
}

pub struct DisabledSender(());
pub struct DisabledReceiver(());

impl Sender for DisabledSender {
    fn on_transmit<P: Packet>(&mut self, _packet: &mut P) {}

    #[inline]
    fn has_transmission_interest(&self) -> bool {
        false
    }

    fn on_connection_error(&mut self, _error: crate::connection::Error) {}
}

impl Receiver for DisabledReceiver {
    fn on_datagram(&mut self, _: &ReceiveContext<'_>, _: &[u8]) {}

    fn on_connection_error(&mut self, _error: crate::connection::Error) {}
}


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

#[cfg(feature = "alloc")]
pub mod default;
pub mod disabled;
pub mod traits;
pub use disabled::*;
pub use traits::*;


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

use core::task::Waker;

use crate::connection;

/// The datagram endpoint trait provides a way to implement custom unreliable datagram
/// sending and receiving logic. The Sender type should be implemented for custom
/// sending behavior, and the Receiver type should be implemented for custom
/// receiving behavior.
pub trait Endpoint: 'static + Send {
    type Sender: Sender;
    type Receiver: Receiver;

    fn create_connection(&mut self, info: &ConnectionInfo) -> (Self::Sender, Self::Receiver);

    /// Returns the maximum datagram frame size the provider is willing to accept
    fn max_datagram_frame_size(&self, info: &PreConnectionInfo) -> u64;
}

/// Information about the accepted connection for which the Sender/Receiver are being created.
#[non_exhaustive]
#[derive(Debug)]
pub struct ConnectionInfo {
    /// The peer's limit on the size of datagrams they will accept. Datagrams larger than this will
    /// result in an error.
    pub max_datagram_payload: u64,

    /// The `waker` associated with this connection. When woken, the connection will check the
    /// interest in sending ([`Sender::has_transmission_interest`]), and send packets if necessary.
    ///
    /// This is useful for applications that wish to enqueue packets with `Sender` without calling
    /// `datagram_mut`, perhaps because they don't have an available handle to the connection when
    /// enqueuing packets, or wish to avoid incurring the lock/unlock required by `datagram_mut`.
    pub waker: Waker,
}

impl ConnectionInfo {
    #[doc(hidden)]
    pub fn new(max_datagram_payload: u64, waker: Waker) -> Self {
        ConnectionInfo {
            max_datagram_payload,
            waker,
        }
    }
}

/// PreConnectionInfo will contain information needed to determine whether
/// or not a provider will accept datagrams.
#[non_exhaustive]
pub struct PreConnectionInfo(());

impl PreConnectionInfo {
    #[doc(hidden)]
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        PreConnectionInfo(())
    }
}

/// ReceiveContext contains information about the connection.
#[non_exhaustive]
#[derive(Debug)]
pub struct ReceiveContext<'a> {
    /// This is the current connection path this datagram was received on.
    pub path: crate::event::api::Path<'a>,
}

impl<'a> ReceiveContext<'a> {
    #[doc(hidden)]
    pub fn new(path: crate::event::api::Path<'a>) -> Self {
        ReceiveContext { path }
    }
}

/// Allows users to configure the behavior of receiving datagrams.
pub trait Receiver: 'static + Send {
    /// A callback that gives users direct access to datagrams as they are read off a packet
    fn on_datagram(&mut self, context: &ReceiveContext<'_>, datagram: &[u8]);

    /// A callback used to notify the application in the case of a connection error
    fn on_connection_error(&mut self, error: connection::Error);
}

/// Allows users to configure the behavior of sending datagrams.
pub trait Sender: 'static + Send {
    /// A callback that allows users to write datagrams directly to the packet
    fn on_transmit<P: Packet>(&mut self, packet: &mut P);

    /// A callback that checks if a user has datagrams ready to send
    ///
    /// Use method to trigger the on_transmit callback
    fn has_transmission_interest(&self) -> bool;

    /// A callback used to notify the application in the case of a connection error
    fn on_connection_error(&mut self, error: connection::Error);
}

/// A packet will be available during the on_transmit callback. Use the methods
/// defined here to interrogate the packet struct and write datagrams to the packet.
pub trait Packet {
    /// Returns the remaining space in the packet left to write datagrams
    fn remaining_capacity(&self) -> usize;

    /// Writes a single datagram to a packet. This function should be called
    /// per datagram.
    fn write_datagram(&mut self, data: &[u8]) -> Result<(), WriteError>;

    /// Writes a single datagram to a packet from a collection of buffers. This function should be
    /// called per datagram.
    fn write_datagram_vectored(&mut self, data: &[&[u8]]) -> Result<(), WriteError>;

    /// Returns whether or not there is reliable data waiting to be sent.
    ///
    /// Use method to decide whether or not to cede the packet space to the stream data.
    fn has_pending_streams(&self) -> bool;

    /// Returns whether or not datagrams are prioritized in this packet or not.
    ///
    /// Datagrams get prioritized every other packet, which gives the application the best
    /// chance to send a large datagram.
    fn datagrams_prioritized(&self) -> bool;
}

#[non_exhaustive]
#[derive(Debug)]
pub enum WriteError {
    ExceedsPacketCapacity,
    ExceedsPeerTransportLimits,
}


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

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

/// Outcome describes how the library should proceed on a connection attempt. The implementor will
/// use information from the ConnectionAttempt object to determine how the library should handle
/// the connection attempt
#[non_exhaustive]
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Outcome {
    /// Allow the connection to continue
    ///
    /// Use `Outcome::allow()` to construct this variant
    #[non_exhaustive]
    Allow,

    /// Defer the connection by sending a Retry packet
    ///
    /// Use `Outcome::retry()` to construct this variant
    #[non_exhaustive]
    Retry,

    /// Silently drop the connection attempt
    ///
    /// Use `Outcome::drop()` to construct this variant
    #[non_exhaustive]
    Drop,

    /// Cleanly close the connection
    ///
    /// Use `Outcome::close()` to construct this variant
    #[non_exhaustive]
    Close,
}

impl Outcome {
    /// Allow the connection to continue
    pub fn allow() -> Self {
        Self::Allow
    }

    /// Defer the connection by sending a Retry packet
    pub fn retry() -> Self {
        Self::Retry
    }

    /// Silently drop the connection attempt
    pub fn drop() -> Self {
        Self::Drop
    }

    /// Cleanly close the connection
    pub fn close() -> Self {
        Self::Close
    }
}

/// A ConnectionAttempt holds information about the state of endpoint receiving a connect, along
/// with information about the connection. This can be used to make decisions about the Outcome of
/// an attempted connection
#[non_exhaustive]
#[derive(Debug)]
pub struct ConnectionAttempt<'a> {
    /// Number of handshakes the have begun but not completed
    pub inflight_handshakes: usize,

    /// Number of open connections
    pub connection_count: usize,

    /// The unverified address of the connecting peer
    /// This address comes from the datagram
    pub remote_address: SocketAddress<'a>,
    pub timestamp: Timestamp,
}

impl<'a> ConnectionAttempt<'a> {
    #[doc(hidden)]
    pub fn new(
        inflight_handshakes: usize,
        connection_count: usize,
        remote_address: &'a inet::SocketAddress,
        timestamp: Timestamp,
    ) -> Self {
        Self {
            inflight_handshakes,
            connection_count,
            remote_address: remote_address.into_event(),
            timestamp,
        }
    }
}

pub trait Limiter: 'static + Send {
    /// This trait is used to determine the outcome of connection attempts on an endpoint. The
    /// implementor returns an Outcome based on the ConnectionAttempt, or other information that the
    /// implementor may have.
    ///
    /// ```rust
    /// # mod s2n_quic { pub mod provider { pub mod endpoint_limits { pub use s2n_quic_core::endpoint::limits::*; } } }
    /// use s2n_quic::provider::endpoint_limits::{Limiter, ConnectionAttempt, Outcome};
    ///
    /// struct MyEndpointLimits {
    ///    handshake_limit: usize,
    ///    delay: core::time::Duration,
    /// }
    ///
    /// impl Limiter for MyEndpointLimits {
    ///    fn on_connection_attempt(&mut self, info: &ConnectionAttempt) -> Outcome {
    ///        if info.inflight_handshakes > self.handshake_limit {
    ///            Outcome::retry()
    ///        } else {
    ///            Outcome::allow()
    ///        }
    ///    }
    /// }
    /// ```
    fn on_connection_attempt(&mut self, info: &ConnectionAttempt) -> Outcome;
}


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

use crate::{
    io::{rx, tx},
    path::{self, MaxMtu},
    time::{Clock, Timestamp},
};
use core::{
    fmt,
    future::Future,
    task::{Context, Poll},
};

pub mod limits;
pub use limits::Limiter;

/// Enumerates endpoint types
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Type {
    /// The endpoint is a client
    Client,
    /// The endpoint is a server
    Server,
}

impl Type {
    /// Returns true if the given endpoint is a QUIC client
    pub fn is_client(self) -> bool {
        self == Self::Client
    }

    /// Returns true if the given endpoint is a QUIC server
    pub fn is_server(self) -> bool {
        self == Self::Server
    }

    /// Returns the [`Type`] of the peer.
    /// - If called on `Client` this will return `Server`
    /// - If called on `Server` this will return `Client`
    #[must_use]
    pub fn peer_type(self) -> Self {
        match self {
            Self::Client => Self::Server,
            Self::Server => Self::Client,
        }
    }
}

/// Enumerates endpoint locations
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Location {
    /// The local endpoint
    Local,
    /// The remote endpoint
    Remote,
}

impl fmt::Display for Location {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Local => write!(f, "the local endpoint"),
            Self::Remote => write!(f, "the remote endpoint"),
        }
    }
}

impl Location {
    /// Returns true if the given endpoint is local
    pub fn is_local(self) -> bool {
        self == Self::Local
    }

    /// Returns true if the given endpoint is remote
    pub fn is_remote(self) -> bool {
        self == Self::Remote
    }

    /// Returns the [`Location`] of the peer.
    /// - If called on `Local` this will return `Remote`
    /// - If called on `Remote` this will return `Local`
    #[must_use]
    pub fn peer_type(self) -> Self {
        match self {
            Self::Local => Self::Remote,
            Self::Remote => Self::Local,
        }
    }
}

/// The main interface for a QUIC endpoint
pub trait Endpoint: 'static + Send + Sized {
    type PathHandle: path::Handle;
    type Subscriber: crate::event::Subscriber;

    const ENDPOINT_TYPE: Type;

    /// Receives and processes datagrams for the Rx queue
    fn receive<Rx, C>(&mut self, rx: &mut Rx, clock: &C)
    where
        Rx: rx::Queue<Handle = Self::PathHandle>,
        C: Clock;

    /// Transmits outgoing datagrams into the Tx queue
    fn transmit<Tx, C>(&mut self, tx: &mut Tx, clock: &C)
    where
        Tx: tx::Queue<Handle = Self::PathHandle>,
        C: Clock;

    /// Returns a future which polls for application-space wakeups
    ///
    /// When successful, the number of wakeups is returned.
    fn wakeups<'a, C: Clock>(&'a mut self, clock: &'a C) -> Wakeups<'a, Self, C> {
        Wakeups {
            endpoint: self,
            clock,
        }
    }

    /// Polls for any application-space wakeups
    ///
    /// When successful, the number of wakeups is returned.
    fn poll_wakeups<C: Clock>(
        &mut self,
        cx: &mut Context<'_>,
        clock: &C,
    ) -> Poll<Result<usize, CloseError>>;

    /// Returns the latest Timestamp at which `transmit` should be called
    fn timeout(&self) -> Option<Timestamp>;

    /// Sets the largest maximum transmission unit (MTU) that can be sent on a path
    fn set_max_mtu(&mut self, max_mtu: MaxMtu);

    /// Returns the endpoint's event subscriber
    fn subscriber(&mut self) -> &mut Self::Subscriber;
}

/// A future which polls an endpoint for application-space wakeups
pub struct Wakeups<'a, E: Endpoint, C: Clock> {
    endpoint: &'a mut E,
    clock: &'a C,
}

impl<'a, E: Endpoint, C: Clock> Future for Wakeups<'a, E, C> {
    type Output = Result<usize, CloseError>;

    fn poll(mut self: core::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let clock = self.clock;
        self.endpoint.poll_wakeups(cx, clock)
    }
}

/// Indicates the endpoint is no longer processing connections.
#[derive(Clone, Copy, Debug)]
pub struct CloseError;