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

use crate::{
    connection, endpoint, path, path::Path, space::PacketSpaceManager, transmission,
    transmission::interest::Provider,
};
use core::time::Duration;
use s2n_codec::{Encoder, EncoderBuffer};
use s2n_quic_core::{
    event::{self, ConnectionPublisher as _},
    frame::ack_elicitation::AckElicitable,
    inet::ExplicitCongestionNotification,
    io::tx,
    packet::{encoding::PacketEncodingError, number::PacketNumberSpace},
    recovery::{CongestionController, MAX_BURST_PACKETS},
    time::Timestamp,
};

#[derive(Debug)]
pub struct ConnectionTransmissionContext<'a, 'sub, Config: endpoint::Config> {
    pub quic_version: u32,
    pub timestamp: Timestamp,
    pub path_id: path::Id,
    pub path_manager: &'a mut path::Manager<Config>,
    pub local_id_registry: &'a mut connection::LocalIdRegistry,
    pub outcome: &'a mut transmission::Outcome,
    pub ecn: ExplicitCongestionNotification,
    pub min_packet_len: Option<usize>,
    pub transmission_mode: transmission::Mode,
    pub publisher: &'a mut event::ConnectionPublisherSubscriber<'sub, Config::EventSubscriber>,
    pub packet_interceptor: &'a mut Config::PacketInterceptor,
}

impl<'a, 'sub, Config: endpoint::Config> ConnectionTransmissionContext<'a, 'sub, Config> {
    pub fn path(&self) -> &Path<Config> {
        &self.path_manager[self.path_id]
    }

    pub fn path_mut(&mut self) -> &mut Path<Config> {
        &mut self.path_manager[self.path_id]
    }
}

pub struct ConnectionTransmission<'a, 'sub, Config: endpoint::Config> {
    pub context: ConnectionTransmissionContext<'a, 'sub, Config>,
    pub space_manager: &'a mut PacketSpaceManager<Config>,
}

impl<'a, 'sub, Config: endpoint::Config> tx::Message for ConnectionTransmission<'a, 'sub, Config> {
    type Handle = Config::PathHandle;

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

    #[inline]
    fn ecn(&mut self) -> ExplicitCongestionNotification {
        self.context.ecn
    }

    #[inline]
    fn delay(&mut self) -> Duration {
        // TODO return delay from pacer
        Default::default()
    }

    #[inline]
    fn ipv6_flow_label(&mut self) -> u32 {
        // TODO compute flow label from connection id
        0
    }

    #[inline]
    fn can_gso(&self, segment_len: usize, segment_count: usize) -> bool {
        // MTU probes can cause delivery issues for other types of packets so they should always be
        // transmitted in single packets
        if self.context.transmission_mode.is_mtu_probing() {
            return false;
        }

        if let Some(min_packet_len) = self.context.min_packet_len {
            if segment_len < min_packet_len {
                return false;
            }
        }

        if let Some(send_quantum) = self.context.path().congestion_controller.send_quantum() {
            if segment_len * segment_count >= send_quantum {
                return false;
            }
        }

        // If a packet can be GSO'd it means it's limited to the previously written packet
        // size. We want to avoid sending several small packets and artificially clamping packets to
        // less than an MTU.
        segment_len >= self.context.path().mtu(self.context.transmission_mode)
    }

    #[inline]
    fn write_payload(
        &mut self,
        buffer: tx::PayloadBuffer,
        gso_offset: usize,
    ) -> Result<usize, tx::Error> {
        let space_manager = &mut self.space_manager;
        let buffer = unsafe {
            // the WriterContext has its own checks for buffer capacity so convert `buffer` into a
            // `&mut [u8]`
            buffer.into_mut_slice()
        };

        let mtu = self
            .context
            .path()
            .clamp_mtu(buffer.len(), self.context.transmission_mode);
        debug_assert_ne!(
            mtu, 0,
            "the amplification limit should be checked before trying to transmit"
        );

        // limit the number of retries to the MAX_BURST_PACKETS
        for _ in 0..MAX_BURST_PACKETS {
            let encoder = EncoderBuffer::new(&mut buffer[..mtu]);
            let initial_capacity = encoder.capacity();

            //= https://www.rfc-editor.org/rfc/rfc9002#section-7
            //# An endpoint MUST NOT send a packet if it would cause bytes_in_flight
            //# (see Appendix B.2) to be larger than the congestion window, unless
            //# the packet is sent on a PTO timer expiration (see Section 6.2) or
            //# when entering recovery (see Section 7.3.2).

            //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.4
            //# In addition to sending data in the packet number space for which the
            //# timer expired, the sender SHOULD send ack-eliciting packets from
            //# other packet number spaces with in-flight data, coalescing packets if
            //# possible.
            let transmission_constraint =
                if space_manager.requires_probe() && self.context.transmission_mode.is_normal() {
                    //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.4
                    //# When a PTO timer expires, a sender MUST send at least one ack-
                    //# eliciting packet in the packet number space as a probe.

                    //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.4
                    //# An endpoint SHOULD include new data in packets that are sent on PTO
                    //# expiration.  Previously sent data MAY be sent if no new data can be
                    //# sent.

                    //= https://www.rfc-editor.org/rfc/rfc9002#section-7.5
                    //# Probe packets MUST NOT be blocked by the congestion controller.
                    self.context.transmission_mode = transmission::Mode::LossRecoveryProbing;
                    transmission::Constraint::None
                } else {
                    self.context.path().transmission_constraint()
                };

            //= 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 by adding PADDING frames to the Initial packet or
            //# by coalescing the Initial packet; see Section 12.2.

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

            // If the transmission contains an Initial packet, it must be padded. However, we should
            // add padding only if necessary, after packets from all packet number spaces have been
            // coalesced. Therefore, after confirming there will be an Initial packet, we first check
            // if there will be an ApplicationData packet, since those packets come at the end of the
            // datagram. If there is no ApplicationData packet, the Handshake packet will come at the
            // end, so we check that next. Finally, if there is no ApplicationData or Handshake packet
            // to transmit, the Initial packet itself will be padded.
            let mut pn_space_to_pad = {
                let needs_padding = if Config::ENDPOINT_TYPE.is_client() {
                    // if we're the client and haven't completed the handshake, we need to pad out packets
                    // in order to give the server as many amplification credits as possible.
                    !space_manager.is_handshake_confirmed()
                } else {
                    // if we're the server, we only need to pad while we have the initial space
                    has_transmission(space_manager.initial(), transmission_constraint)
                };

                if !needs_padding {
                    // There is no Initial packet, so no padding is needed
                    None
                } else if has_transmission(space_manager.application(), transmission_constraint) {
                    Some(PacketNumberSpace::ApplicationData)
                } else if has_transmission(space_manager.handshake(), transmission_constraint) {
                    Some(PacketNumberSpace::Handshake)
                } else {
                    //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9
                    //# These packets MAY also include PADDING frames.
                    Some(PacketNumberSpace::Initial)
                }
            };

            //= https://www.rfc-editor.org/rfc/rfc9001#section-4
            //# When packets of different types need to be sent,
            //# endpoints SHOULD use coalesced packets to send them in the same UDP
            //# datagram.
            // here we query all of the spaces to try and fill the current datagram

            let is_mtu_probing = self.context.transmission_mode.is_mtu_probing();

            let encoder = if let Some((space, handshake_status)) = space_manager
                .initial_mut()
                // MTU probes are only sent in the Application Space
                .filter(|_| !is_mtu_probing)
            {
                self.context.min_packet_len = pn_space_to_pad
                    .filter(|pn_space| pn_space.is_initial())
                    .map(|_| encoder.capacity());

                match space.on_transmit(
                    &mut self.context,
                    transmission_constraint,
                    handshake_status,
                    encoder,
                ) {
                    Ok((outcome, encoder)) => {
                        if Config::ENDPOINT_TYPE.is_server()
                            && !outcome.ack_elicitation().is_ack_eliciting()
                        {
                            //= https://www.rfc-editor.org/rfc/rfc9000#section-14.1
                            //# Similarly, a
                            //# server MUST expand the payload of all UDP datagrams carrying ack-
                            //# eliciting Initial packets to at least the smallest allowed maximum
                            //# datagram size of 1200 bytes.

                            // The Initial packet was not ack eliciting so there is no need to pad
                            pn_space_to_pad = None;
                        }
                        *self.context.outcome += outcome;
                        encoder
                    }
                    Err(PacketEncodingError::PacketNumberTruncationError(encoder)) => {
                        // TODO handle this
                        encoder
                    }
                    Err(PacketEncodingError::InsufficientSpace(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                    Err(PacketEncodingError::EmptyPayload(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                    Err(PacketEncodingError::AeadLimitReached(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                }
            } else {
                encoder
            };

            let encoder = if let Some((space, handshake_status)) = space_manager
                .handshake_mut()
                // MTU probes are only sent in the Application Space
                .filter(|_| !is_mtu_probing)
            {
                self.context.min_packet_len = pn_space_to_pad
                    .filter(|pn_space| pn_space.is_handshake())
                    .map(|_| encoder.capacity());

                let encoder = match space.on_transmit(
                    &mut self.context,
                    transmission_constraint,
                    handshake_status,
                    encoder,
                ) {
                    Ok((outcome, encoder)) => {
                        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.1
                        //# a client MUST discard Initial keys when it first sends a
                        //# Handshake packet

                        if Config::ENDPOINT_TYPE.is_client() {
                            let path = &mut self.context.path_manager[self.context.path_id];
                            space_manager.discard_initial(
                                path,
                                self.context.path_id,
                                self.context.publisher,
                            );
                        }

                        *self.context.outcome += outcome;
                        encoder
                    }
                    Err(PacketEncodingError::PacketNumberTruncationError(encoder)) => {
                        // TODO handle this
                        encoder
                    }
                    Err(PacketEncodingError::InsufficientSpace(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                    Err(PacketEncodingError::EmptyPayload(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                    Err(PacketEncodingError::AeadLimitReached(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                };

                //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.2
                //# An endpoint MUST discard its handshake keys when the TLS handshake is
                //# confirmed (Section 4.1.2).
                if space_manager.is_handshake_confirmed() {
                    let path = &mut self.context.path_manager[self.context.path_id];
                    space_manager.discard_handshake(
                        path,
                        self.context.path_id,
                        self.context.publisher,
                    );
                }

                encoder
            } else {
                encoder
            };

            //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9
            //# Though an endpoint might retain older keys, new data MUST be sent at
            //# the highest currently-available encryption level.

            // This requirement is automatically support with s2n-quic's implementation. Each space
            // acts mostly independent from another and will buffer its own CRYPTO and ACK state. Other
            // frames are only allowed in the ApplicationData space, which will always be the highest
            // current-available encryption level.

            let encoder = if let Some((space, handshake_status)) = space_manager.application_mut() {
                self.context.min_packet_len = pn_space_to_pad
                    .filter(|pn_space| pn_space.is_application_data())
                    .map(|_| encoder.capacity());

                // Pad the packet when sending path validation frames so that MTU is also validated.
                let path = &self.context.path_manager[self.context.path_id];

                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
                //# An endpoint MUST expand datagrams that contain a PATH_CHALLENGE frame
                //# to at least the smallest allowed maximum datagram size of 1200 bytes.
                //
                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
                //# An endpoint MUST expand datagrams that contain a PATH_RESPONSE frame
                //# to at least the smallest allowed maximum datagram size of 1200 bytes.
                // Pad the packet when sending path validation frames so that MTU is also validated.
                //
                // The path's transmission_interest indicates if a PATH_CHALLENGE or PATH_RESPONSE
                // frame is to be written.
                //
                // We need to check is_validated because it is possible to receive a PATH_CHALLENGE on
                // an active path for Off-Path Packet Forwarding prevention. However, we would only
                // like to pad when validating the MTU.
                if !path.is_validated() && path.has_transmission_interest() {
                    self.context.min_packet_len = Some(encoder.capacity());
                }

                match space.on_transmit(
                    &mut self.context,
                    transmission_constraint,
                    handshake_status,
                    encoder,
                ) {
                    Ok((outcome, encoder)) => {
                        *self.context.outcome += outcome;
                        encoder
                    }
                    Err(PacketEncodingError::PacketNumberTruncationError(encoder)) => {
                        // TODO handle this
                        encoder
                    }
                    Err(PacketEncodingError::InsufficientSpace(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                    Err(PacketEncodingError::EmptyPayload(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                    Err(PacketEncodingError::AeadLimitReached(encoder)) => {
                        // move to the next packet space
                        encoder
                    }
                }
            } else {
                encoder
            };

            let datagram_len = initial_capacity - encoder.capacity();

            // the spaces didn't write anything so we're done
            if datagram_len == 0 {
                return Err(tx::Error::EmptyPayload);
            }

            // Emit the transmission event
            //
            // Even though the interceptor could alter the outgoing bytes, we're going to pretend
            // that it doesn't so it's closer to on-path datagram corruption.
            self.context.path_mut().on_bytes_transmitted(datagram_len);
            self.context
                .publisher
                .on_datagram_sent(event::builder::DatagramSent {
                    len: datagram_len as u16,
                    gso_offset,
                });

            let datagram_len = {
                use s2n_quic_core::{
                    event::IntoEvent,
                    packet::interceptor::{Datagram, Interceptor},
                };

                let mut encoder = EncoderBuffer::new(buffer);
                encoder.set_position(datagram_len);

                let subject = self.context.publisher.subject();

                let remote_address = self.context.path().remote_address();
                let local_address = self.context.path().local_address();

                let datagram = Datagram {
                    remote_address: remote_address.into_event(),
                    local_address: local_address.into_event(),
                    timestamp: self.context.timestamp,
                };

                self.context.packet_interceptor.intercept_tx_datagram(
                    &subject,
                    &datagram,
                    &mut encoder,
                );

                // if the packet interceptor cleared the encoder buffer, then try again
                if encoder.is_empty() {
                    continue;
                }

                encoder.len()
            };

            return Ok(datagram_len);
        }

        Err(tx::Error::EmptyPayload)
    }
}

fn has_transmission<P: transmission::interest::Provider>(
    transmission_interest_provider: Option<&P>,
    transmission_constraint: transmission::Constraint,
) -> bool {
    transmission_interest_provider.map_or(false, |provider| {
        provider.can_transmit(transmission_constraint)
    })
}


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

use super::*;
use crate::{
    connection::{
        self, connection_interests::ConnectionInterests,
        internal_connection_id::InternalConnectionId, InternalConnectionIdGenerator,
        ProcessingError,
    },
    endpoint, path, stream,
};
use bolero::{check, generator::*};
use bytes::Bytes;
use core::{
    task::{Context, Poll},
    time::Duration,
};
use s2n_quic_core::{
    application, event,
    event::builder::DatagramDropReason,
    inet::{DatagramInfo, SocketAddress},
    io::tx,
    packet::{
        handshake::ProtectedHandshake,
        initial::{CleartextInitial, ProtectedInitial},
        retry::ProtectedRetry,
        short::ProtectedShort,
        version_negotiation::ProtectedVersionNegotiation,
        zero_rtt::ProtectedZeroRtt,
    },
    path::MaxMtu,
    query,
    time::{Timer, Timestamp},
};
use std::sync::Mutex;

struct TestConnection {
    is_handshaking: bool,
    has_been_accepted: bool,
    is_closed: bool,
    interests: ConnectionInterests,
    close_timer: Timer,
}

impl Default for TestConnection {
    fn default() -> Self {
        Self {
            is_handshaking: true,
            has_been_accepted: false,
            is_closed: false,
            interests: ConnectionInterests {
                transmission: true,
                ..Default::default()
            },
            close_timer: Default::default(),
        }
    }
}

impl connection::Trait for TestConnection {
    type Config = crate::endpoint::testing::Server;

    fn new(_params: connection::Parameters<Self::Config>) -> Result<Self, connection::Error> {
        Ok(Self::default())
    }

    fn internal_connection_id(&self) -> InternalConnectionId {
        todo!()
    }

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

    fn close(
        &mut self,
        _error: connection::Error,
        _close_formatter: &<Self::Config as endpoint::Config>::ConnectionCloseFormatter,
        _packet_buffer: &mut endpoint::PacketBuffer,
        timestamp: Timestamp,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
    ) {
        assert!(!self.is_closed);
        assert!(!self.close_timer.is_armed());
        self.close_timer.set(timestamp + Duration::from_secs(1));
    }

    fn mark_as_accepted(&mut self) {
        assert!(!self.has_been_accepted);
        self.has_been_accepted = true;
        self.interests.accept = false;
    }

    fn on_new_connection_id(
        &mut self,
        _connection_id_format: &mut <Self::Config as endpoint::Config>::ConnectionIdFormat,
        _stateless_reset_token_generator: &mut <Self::Config as endpoint::Config>::StatelessResetTokenGenerator,
        _timestamp: Timestamp,
    ) -> Result<(), connection::local_id_registry::LocalIdRegistrationError> {
        Ok(())
    }

    fn on_transmit<Tx: tx::Queue>(
        &mut self,
        _queue: &mut Tx,
        _timestamp: Timestamp,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
    ) -> Result<(), crate::contexts::ConnectionOnTransmitError> {
        Ok(())
    }

    fn on_timeout(
        &mut self,
        _connection_id_mapper: &mut connection::ConnectionIdMapper,
        timestamp: Timestamp,
        _supervisor_context: &supervisor::Context,
        _random_generator: &mut <Self::Config as endpoint::Config>::RandomGenerator,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
    ) -> Result<(), connection::Error> {
        if self.close_timer.poll_expiration(timestamp).is_ready() {
            self.is_closed = true;
        }
        Ok(())
    }

    fn on_wakeup(
        &mut self,
        _timestamp: Timestamp,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _datagram: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
    ) -> Result<(), connection::Error> {
        Ok(())
    }

    fn handle_initial_packet(
        &mut self,
        _datagram: &DatagramInfo,
        _path_id: path::Id,
        _packet: ProtectedInitial,
        _random_generator: &mut <Self::Config as endpoint::Config>::RandomGenerator,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
        _datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
    ) -> Result<(), ProcessingError> {
        Ok(())
    }

    /// Is called when an unprotected initial packet had been received
    fn handle_cleartext_initial_packet(
        &mut self,
        _datagram: &DatagramInfo,
        _path_id: path::Id,
        _packet: CleartextInitial,
        _random_generator: &mut <Self::Config as endpoint::Config>::RandomGenerator,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
        _datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
    ) -> Result<(), ProcessingError> {
        Ok(())
    }

    /// Is called when a handshake packet had been received
    fn handle_handshake_packet(
        &mut self,
        _datagram: &DatagramInfo,
        _path_id: path::Id,
        _packet: ProtectedHandshake,
        _random_generator: &mut <Self::Config as endpoint::Config>::RandomGenerator,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
        _datagram_endpoint: &mut <Self::Config as endpoint::Config>::DatagramEndpoint,
    ) -> Result<(), ProcessingError> {
        Ok(())
    }

    /// Is called when a short packet had been received
    fn handle_short_packet(
        &mut self,
        _datagram: &DatagramInfo,
        _path_id: path::Id,
        _packet: ProtectedShort,
        _random_generator: &mut <Self::Config as endpoint::Config>::RandomGenerator,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
    ) -> Result<(), ProcessingError> {
        Ok(())
    }

    /// Is called when a version negotiation packet had been received
    fn handle_version_negotiation_packet(
        &mut self,
        _datagram: &DatagramInfo,
        _path_id: path::Id,
        _packet: ProtectedVersionNegotiation,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
    ) -> Result<(), ProcessingError> {
        Ok(())
    }

    /// Is called when a zero rtt packet had been received
    fn handle_zero_rtt_packet(
        &mut self,
        _datagram: &DatagramInfo,
        _path_id: path::Id,
        _packet: ProtectedZeroRtt,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
    ) -> Result<(), ProcessingError> {
        Ok(())
    }

    /// Is called when a retry packet had been received
    fn handle_retry_packet(
        &mut self,
        _datagram: &DatagramInfo,
        _path_id: path::Id,
        _packet: ProtectedRetry,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _packet_interceptor: &mut <Self::Config as endpoint::Config>::PacketInterceptor,
    ) -> Result<(), ProcessingError> {
        Ok(())
    }

    /// Notifies a connection it has received a datagram from a peer
    fn on_datagram_received(
        &mut self,
        _path: &<Self::Config as endpoint::Config>::PathHandle,
        _datagram: &DatagramInfo,
        _congestion_controller_endpoint: &mut <Self::Config as endpoint::Config>::CongestionControllerEndpoint,
        _path_migration: &mut <Self::Config as endpoint::Config>::PathMigrationValidator,
        _max_mtu: MaxMtu,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
    ) -> Result<path::Id, DatagramDropReason> {
        todo!()
    }

    /// Returns the Connections interests
    fn interests(&self) -> ConnectionInterests {
        self.interests
    }

    /// Returns the QUIC version selected for the current connection
    fn quic_version(&self) -> u32 {
        123
    }

    fn poll_stream_request(
        &mut self,
        _stream_id: stream::StreamId,
        _request: &mut stream::ops::Request,
        _context: Option<&Context>,
    ) -> Result<stream::ops::Response, stream::StreamError> {
        todo!()
    }

    fn poll_accept_stream(
        &mut self,
        _stream_type: Option<stream::StreamType>,
        _context: &Context,
    ) -> Poll<Result<Option<stream::StreamId>, connection::Error>> {
        todo!()
    }

    fn poll_open_stream(
        &mut self,
        _stream_type: stream::StreamType,
        _token: &mut connection::OpenToken,
        _context: &Context,
    ) -> Poll<Result<stream::StreamId, connection::Error>> {
        todo!()
    }

    fn application_close(&mut self, _error: Option<application::Error>) {
        // no-op
    }

    fn server_name(&self) -> Option<ServerName> {
        todo!()
    }

    fn application_protocol(&self) -> Bytes {
        todo!()
    }

    fn ping(&mut self) -> Result<(), connection::Error> {
        todo!()
    }

    fn keep_alive(&mut self, _enabled: bool) -> Result<(), connection::Error> {
        todo!()
    }

    fn local_address(&self) -> Result<SocketAddress, connection::Error> {
        todo!()
    }

    fn remote_address(&self) -> Result<SocketAddress, connection::Error> {
        Ok(SocketAddress::default())
    }

    fn error(&self) -> Option<connection::Error> {
        None
    }

    fn query_event_context(&self, _query: &mut dyn query::Query) {
        todo!()
    }

    fn query_event_context_mut(&mut self, _query: &mut dyn query::QueryMut) {
        todo!()
    }

    fn datagram_mut(&mut self, _query: &mut dyn query::QueryMut) {
        todo!()
    }

    fn with_event_publisher<F>(
        &mut self,
        _timestamp: Timestamp,
        _path_id: Option<path::Id>,
        _subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
        _f: F,
    ) where
        F: FnOnce(
            &mut event::ConnectionPublisherSubscriber<
                <Self::Config as endpoint::Config>::EventSubscriber,
            >,
            &path::Path<Self::Config>,
        ),
    {
        todo!()
    }
}

struct TestLock {
    connection: Mutex<(TestConnection, bool)>,
}

impl TestLock {
    fn poison(&self) {
        if let Ok(mut lock) = self.connection.lock() {
            lock.1 = true;
        }
    }
}

impl connection::Lock<TestConnection> for TestLock {
    type Error = ();

    fn new(connection: TestConnection) -> Self {
        Self {
            connection: std::sync::Mutex::new((connection, false)),
        }
    }

    fn read<F: FnOnce(&TestConnection) -> R, R>(&self, f: F) -> Result<R, Self::Error> {
        let lock = self.connection.lock().map_err(|_| ())?;
        let (conn, is_poisoned) = &*lock;
        if *is_poisoned {
            return Err(());
        }
        let result = f(conn);
        Ok(result)
    }

    fn write<F: FnOnce(&mut TestConnection) -> R, R>(&self, f: F) -> Result<R, Self::Error> {
        let mut lock = self.connection.lock().map_err(|_| ())?;
        let (conn, is_poisoned) = &mut *lock;
        if *is_poisoned {
            return Err(());
        }
        let result = f(conn);
        Ok(result)
    }
}

#[derive(Debug, TypeGenerator)]
enum Operation {
    Insert,
    UpdateInterests {
        index: usize,
        finalization: bool,
        closing: bool,
        accept: bool,
        transmission: bool,
        new_connection_id: bool,
        timeout: Option<u16>,
    },
    CloseApp,
    Receive,
    Timeout(u16),
    Transmit(u16),
    NewConnId(u16),
    Finalize,
    Poison(usize),
}

#[test]
fn container_test() {
    use core::time::Duration;

    check!().with_type::<Vec<Operation>>().for_each(|ops| {
        let mut id_gen = InternalConnectionIdGenerator::new();
        let mut connections = vec![];
        let (handle, acceptor, connector, _close_handle) = endpoint::handle::Handle::new(100);
        let (waker, _wake_count) = futures_test::task::new_count_waker();
        let mut now = unsafe { Timestamp::from_duration(Duration::from_secs(0)) };

        let mut handle = Some(handle);
        let mut container: ConnectionContainer<TestConnection, TestLock> =
            ConnectionContainer::new(acceptor, connector);

        for op in ops.iter() {
            match op {
                Operation::Insert => {
                    let id = id_gen.generate_id();
                    let connection = TestConnection::default();
                    container.insert_connection(connection, id);
                    connections.push(id);

                    let mut was_called = false;
                    container.with_connection(id, |_conn| {
                        was_called = true;
                    });
                    assert!(was_called);
                }
                Operation::UpdateInterests {
                    index,
                    finalization,
                    closing,
                    accept,
                    transmission,
                    new_connection_id,
                    timeout,
                } => {
                    if connections.is_empty() {
                        continue;
                    }
                    let index = index % connections.len();
                    let id = connections[index];

                    let mut was_called = false;
                    container.with_connection(id, |conn| {
                        was_called = true;

                        let i = &mut conn.interests;

                        if *finalization {
                            // in the finalization state, that should be the only interest
                            *i = ConnectionInterests {
                                finalization: true,
                                ..Default::default()
                            };
                            return;
                        }

                        i.closing = *closing;
                        if !conn.has_been_accepted {
                            i.accept = *accept;
                        }
                        if *accept {
                            conn.is_handshaking = false;
                        }
                        i.transmission = *transmission;
                        i.new_connection_id = *new_connection_id;
                        i.timeout = timeout.map(|ms| now + Duration::from_millis(ms as _));

                        // we need to express at least one interest to ensure progress
                        if !(i.transmission || i.new_connection_id || i.timeout.is_some()) {
                            i.transmission = true;
                        }
                    });

                    if *finalization {
                        connections.remove(index);
                    }

                    assert!(was_called);
                }
                Operation::CloseApp => {
                    handle = None;
                }
                Operation::Receive => {
                    if let Some(handle) = handle.as_mut() {
                        while let Poll::Ready(Some(_accepted)) = handle
                            .acceptor
                            .poll_accept(&mut Context::from_waker(&waker))
                        {
                            // TODO assert that the accepted connection expressed accept
                            // interest
                        }
                    }
                }
                Operation::Timeout(ms) => {
                    now += Duration::from_millis(*ms as _);
                    container.iterate_timeout_list(now, |conn, _context| {
                        let i = &mut conn.interests;

                        assert!(
                            i.timeout.take().unwrap() <= now,
                            "connections should only be present when timeout interest is expressed"
                        );

                        // we need to express at least one interest to ensure progress
                        if !(i.transmission || i.new_connection_id || i.timeout.is_some()) {
                            i.transmission = true;
                        }
                    });
                }
                Operation::Transmit(count) => {
                    let mut count = *count;
                    container.iterate_transmission_list(|conn| {
                        assert!(conn.interests.transmission);

                        if count == 0 {
                            ConnectionContainerIterationResult::BreakAndInsertAtFront
                        } else {
                            count -= 1;
                            ConnectionContainerIterationResult::Continue
                        }
                    })
                }
                Operation::NewConnId(count) => {
                    let mut count = *count;
                    container.iterate_new_connection_id_list(|conn| {
                        assert!(conn.interests.new_connection_id);

                        if count == 0 {
                            ConnectionContainerIterationResult::BreakAndInsertAtFront
                        } else {
                            count -= 1;
                            ConnectionContainerIterationResult::Continue
                        }
                    })
                }
                Operation::Finalize => {
                    container.finalize_done_connections();
                }
                Operation::Poison(index) => {
                    if connections.is_empty() {
                        continue;
                    }
                    let index = index % connections.len();
                    let id = connections[index];

                    let node = container.connection_map.find(&id).get().unwrap();
                    node.inner.poison();

                    let mut was_called = false;
                    container.with_connection(id, |_conn| {
                        was_called = true;
                    });
                    assert!(!was_called);
                    connections.remove(index);
                }
            }
        }

        container.finalize_done_connections();

        let mut connections = connections.drain(..);
        let mut cursor = container.connection_map.front();

        while let Some(conn) = cursor.get() {
            assert_eq!(conn.internal_connection_id, connections.next().unwrap());
            cursor.move_next();
        }

        assert!(connections.next().is_none());
    });
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use super::*;
use s2n_quic_core::{
    connection,
    connection::id::MIN_LIFETIME,
    frame::{Frame, NewConnectionId},
    packet::number::PacketNumberRange,
    random,
    stateless_reset::token::testing::*,
    time::{clock::testing as time, timer::Provider as _},
    varint::VarInt,
};

use crate::{
    connection::{
        connection_id_mapper::*,
        local_id_registry::{
            LocalIdInfo, LocalIdRegistrationError, LocalIdRegistry, EXPIRATION_BUFFER,
            MAX_ACTIVE_CONNECTION_ID_LIMIT, RTT_MULTIPLIER,
        },
        InternalConnectionIdGenerator,
    },
    contexts::testing::{MockWriteContext, OutgoingFrameBuffer},
    endpoint, transmission,
    transmission::interest::Provider,
};
use core::time::Duration;

impl LocalIdRegistry {
    fn get_connection_id_info(&self, id: &connection::LocalId) -> Option<&LocalIdInfo> {
        self.registered_ids.iter().find(|id_info| id_info.id == *id)
    }

    fn get_connection_id_info_mut(&mut self, id: &connection::LocalId) -> Option<&mut LocalIdInfo> {
        self.registered_ids
            .iter_mut()
            .find(|id_info| id_info.id == *id)
    }
}

// Helper function to easily generate a LocalId from bytes
fn id(bytes: &[u8]) -> connection::LocalId {
    connection::LocalId::try_from_bytes(bytes).unwrap()
}

// Helper function to easily create a LocalIdRegistry and Mapper
fn mapper(
    handshake_id: connection::LocalId,
    handshake_id_expiration_time: Option<Timestamp>,
    token: stateless_reset::Token,
) -> (ConnectionIdMapper, LocalIdRegistry) {
    let mut random_generator = random::testing::Generator(123);

    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let registry = mapper.create_local_id_registry(
        InternalConnectionIdGenerator::new().generate_id(),
        &handshake_id,
        handshake_id_expiration_time,
        token,
    );
    (mapper, registry)
}

// Verify that an expiration with the earliest possible time results in a valid retirement time
#[test]
fn minimum_lifetime() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");

    let expiration = time::now() + MIN_LIFETIME;

    let (_mapper, mut reg1) = mapper(ext_id_1, Some(expiration), TEST_TOKEN_1);
    reg1.set_active_connection_id_limit(3);
    assert!(reg1
        .register_connection_id(&ext_id_2, Some(expiration), TEST_TOKEN_2)
        .is_ok());
    assert_eq!(
        Some(expiration - EXPIRATION_BUFFER),
        reg1.get_connection_id_info(&ext_id_1)
            .unwrap()
            .retirement_time
    );
    assert_eq!(
        Some(expiration - EXPIRATION_BUFFER),
        reg1.get_connection_id_info(&ext_id_2)
            .unwrap()
            .retirement_time
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1
//= type=test
//# As a trivial example, this means the same connection ID
//# MUST NOT be issued more than once on the same connection.
#[test]
fn same_connection_id_must_not_be_issued_for_same_connection() {
    let ext_id = id(b"id01");
    let (_, mut reg) = mapper(ext_id, None, TEST_TOKEN_1);

    assert_eq!(
        Err(LocalIdRegistrationError::ConnectionIdInUse),
        reg.register_connection_id(&ext_id, None, TEST_TOKEN_1)
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//= type=test
//# The sequence number on
//# each newly issued connection ID MUST increase by 1.
#[test]
fn sequence_number_must_increase_by_one() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");

    let (_, mut reg) = mapper(ext_id_1, None, TEST_TOKEN_1);
    reg.set_active_connection_id_limit(3);
    reg.register_connection_id(&ext_id_2, None, TEST_TOKEN_2)
        .unwrap();

    let seq_num_1 = reg
        .get_connection_id_info(&ext_id_1)
        .unwrap()
        .sequence_number;
    let seq_num_2 = reg
        .get_connection_id_info(&ext_id_2)
        .unwrap()
        .sequence_number;

    assert_eq!(1, seq_num_2 - seq_num_1);
}

#[test]
fn connection_mapper_test() {
    let mut id_generator = InternalConnectionIdGenerator::new();
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);

    let id1 = id_generator.generate_id();
    let id2 = id_generator.generate_id();

    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");
    let ext_id_4 = id(b"id04");

    let now = time::now();
    let handshake_id_expiration_time = now + Duration::from_secs(60);

    let mut reg1 = mapper.create_local_id_registry(
        id1,
        &ext_id_1,
        Some(handshake_id_expiration_time),
        TEST_TOKEN_1,
    );
    let mut reg2 = mapper.create_local_id_registry(
        id2,
        &ext_id_3,
        Some(handshake_id_expiration_time),
        TEST_TOKEN_3,
    );

    reg1.set_active_connection_id_limit(3);
    reg2.set_active_connection_id_limit(3);

    assert_eq!(
        0,
        reg1.get_connection_id_info(&ext_id_1)
            .unwrap()
            .sequence_number
    );
    assert_eq!(
        Some(handshake_id_expiration_time - EXPIRATION_BUFFER),
        reg1.get_connection_id_info(&ext_id_1)
            .unwrap()
            .retirement_time
    );
    assert_eq!(Some(id1), mapper.lookup_internal_connection_id(&ext_id_1));
    assert_eq!(
        TEST_TOKEN_1,
        reg1.get_connection_id_info(&ext_id_1)
            .unwrap()
            .stateless_reset_token
    );

    assert_eq!(
        Err(LocalIdRegistrationError::ConnectionIdInUse),
        reg2.register_connection_id(&ext_id_1, None, TEST_TOKEN_1)
    );

    let exp_2 = now + Duration::from_secs(60);

    assert!(reg1
        .register_connection_id(&ext_id_2, Some(exp_2), TEST_TOKEN_2)
        .is_ok());
    assert_eq!(
        Some(exp_2 - EXPIRATION_BUFFER),
        reg1.get_connection_id_info(&ext_id_2)
            .unwrap()
            .retirement_time
    );
    assert!(reg2
        .register_connection_id(&ext_id_4, None, TEST_TOKEN_4)
        .is_ok());
    assert_eq!(Some(id1), mapper.lookup_internal_connection_id(&ext_id_2));
    assert_eq!(Some(id2), mapper.lookup_internal_connection_id(&ext_id_3));
    assert_eq!(Some(id2), mapper.lookup_internal_connection_id(&ext_id_4));

    // Unregister id 3 (sequence number 0)
    reg2.get_connection_id_info_mut(&ext_id_3).unwrap().status = PendingRemoval(now);
    reg2.unregister_expired_ids(now);
    assert_eq!(None, mapper.lookup_internal_connection_id(&ext_id_3));
    assert_eq!(Some(id2), mapper.lookup_internal_connection_id(&ext_id_4));

    reg2.get_connection_id_info_mut(&ext_id_4).unwrap().status =
        PendingRetirementConfirmation(Some(now));
    reg2.unregister_expired_ids(now);
    assert_eq!(None, mapper.lookup_internal_connection_id(&ext_id_4));

    // Put back ID3 and ID4 to test drop behavior
    assert!(reg1
        .register_connection_id(&ext_id_3, None, TEST_TOKEN_3)
        .is_ok());
    assert!(reg2
        .register_connection_id(&ext_id_4, None, TEST_TOKEN_4)
        .is_ok());
    assert_eq!(Some(id1), mapper.lookup_internal_connection_id(&ext_id_3));
    assert_eq!(Some(id2), mapper.lookup_internal_connection_id(&ext_id_4));

    // If a registration is dropped all entries are removed
    drop(reg1);
    assert_eq!(None, mapper.lookup_internal_connection_id(&ext_id_1));
    assert_eq!(None, mapper.lookup_internal_connection_id(&ext_id_2));
    assert_eq!(None, mapper.lookup_internal_connection_id(&ext_id_3));
    assert_eq!(Some(id2), mapper.lookup_internal_connection_id(&ext_id_4));
}

#[test]
fn on_retire_connection_id() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");

    let now = time::now();
    let (mapper, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);

    reg1.set_active_connection_id_limit(2);

    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.16
    //= type=test
    //# Receipt of a RETIRE_CONNECTION_ID frame containing a sequence number
    //# greater than any previously sent to the peer MUST be treated as a
    //# connection error of type PROTOCOL_VIOLATION.
    assert_eq!(
        Some(LocalIdRegistrationError::InvalidSequenceNumber),
        reg1.on_retire_connection_id(1, &ext_id_1, Duration::default(), now)
            .err()
    );

    assert!(reg1
        .register_connection_id(&ext_id_2, None, TEST_TOKEN_2)
        .is_ok());

    let rtt = Duration::from_millis(500);

    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.16
    //= type=test
    //# The sequence number specified in a RETIRE_CONNECTION_ID frame MUST
    //# NOT refer to the Destination Connection ID field of the packet in
    //# which the frame is contained.

    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.16
    //= type=test
    //# The peer MAY treat this as a
    //# connection error of type PROTOCOL_VIOLATION.
    assert_eq!(
        Some(LocalIdRegistrationError::InvalidSequenceNumber),
        reg1.on_retire_connection_id(1, &ext_id_2, Duration::default(), now)
            .err()
    );

    assert!(reg1.on_retire_connection_id(1, &ext_id_1, rtt, now).is_ok());

    assert_eq!(
        PendingRemoval(now + rtt * RTT_MULTIPLIER),
        reg1.get_connection_id_info(&ext_id_2).unwrap().status
    );

    // ID 1 wasn't impacted by the request to retire ID 2
    assert_eq!(
        Active,
        reg1.get_connection_id_info(&ext_id_1).unwrap().status
    );

    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
    //= type=test
    //# An endpoint SHOULD supply a new connection ID when the peer retires a
    //# connection ID.
    assert_eq!(
        connection::id::Interest::New(1),
        reg1.connection_id_interest()
    );

    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
    //= type=test
    //# When an endpoint issues a connection ID, it MUST accept packets that
    //# carry this connection ID for the duration of the connection or until
    //# its peer invalidates the connection ID via a RETIRE_CONNECTION_ID
    //# frame (Section 19.16).

    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
    //= type=test
    //# The endpoint SHOULD continue to
    //# accept the previously issued connection IDs until they are retired by
    //# the peer.
    reg1.unregister_expired_ids(now + rtt * RTT_MULTIPLIER);
    assert!(mapper.lookup_internal_connection_id(&ext_id_2).is_none());
}

#[test]
fn on_retire_connection_id_pending_removal() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");

    let now = time::now() + Duration::from_secs(60);

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);
    reg1.set_active_connection_id_limit(2);

    assert!(reg1
        .register_connection_id(&ext_id_2, Some(now), TEST_TOKEN_2)
        .is_ok());

    reg1.retire_handshake_connection_id();
    reg1.on_timeout(now);

    assert_eq!(
        PendingRetirementConfirmation(None),
        reg1.get_connection_id_info(&ext_id_1).unwrap().status
    );
    assert_eq!(
        PendingRetirementConfirmation(Some(now + EXPIRATION_BUFFER)),
        reg1.get_connection_id_info(&ext_id_2).unwrap().status
    );

    let rtt = Duration::from_millis(500);

    assert!(reg1.on_retire_connection_id(1, &ext_id_1, rtt, now).is_ok());

    assert_eq!(
        PendingRetirementConfirmation(None),
        reg1.get_connection_id_info(&ext_id_1).unwrap().status
    );
    // When the ON_RETIRE_CONNECTION_ID frame is received from the peer, the
    // removal time for the retired connection ID is updated
    assert_eq!(
        PendingRemoval(now + rtt * RTT_MULTIPLIER),
        reg1.get_connection_id_info(&ext_id_2).unwrap().status
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//= type=test
//# An endpoint that initiates migration and requires non-zero-length
//# connection IDs SHOULD ensure that the pool of connection IDs
//# available to its peer allows the peer to use a new connection ID on
//# migration, as the peer will be unable to respond if the pool is
//# exhausted.

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//= type=test
//# An endpoint SHOULD ensure that its peer has a sufficient number of
//# available and unused connection IDs.

//= https://www.rfc-editor.org/rfc/rfc9000#section-9.5
//= type=test
//# To ensure that migration is possible and packets sent on different
//# paths cannot be correlated, endpoints SHOULD provide new connection
//# IDs before peers migrate; see Section 5.1.1.
#[test]
fn connection_id_interest() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);

    // Active connection ID limit starts at 1, so there is no interest initially
    assert_eq!(
        connection::id::Interest::None,
        reg1.connection_id_interest()
    );

    reg1.set_active_connection_id_limit(5);
    assert_eq!(
        MAX_ACTIVE_CONNECTION_ID_LIMIT,
        reg1.active_connection_id_limit as u64
    );

    assert_eq!(
        connection::id::Interest::New(reg1.active_connection_id_limit - 1),
        reg1.connection_id_interest()
    );

    assert!(reg1
        .register_connection_id(&ext_id_2, None, TEST_TOKEN_2)
        .is_ok());

    assert_eq!(
        connection::id::Interest::New(reg1.active_connection_id_limit - 2),
        reg1.connection_id_interest()
    );

    assert!(reg1
        .register_connection_id(&ext_id_3, None, TEST_TOKEN_3)
        .is_ok());

    assert_eq!(
        connection::id::Interest::None,
        reg1.connection_id_interest()
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//= type=test
//# An endpoint MUST NOT
//# provide more connection IDs than the peer's limit.
#[test]
#[should_panic]
fn endpoint_must_not_provide_more_ids_than_peer_limit() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);

    reg1.set_active_connection_id_limit(2);

    assert_eq!(
        connection::id::Interest::New(1),
        reg1.connection_id_interest()
    );

    assert!(reg1
        .register_connection_id(&ext_id_2, None, TEST_TOKEN_2)
        .is_ok());

    // Panics because we are inserting more than the limit
    let _ = reg1.register_connection_id(&ext_id_3, None, TEST_TOKEN_3);
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//= type=test
//# An endpoint MAY
//# send connection IDs that temporarily exceed a peer's limit if the
//# NEW_CONNECTION_ID frame also requires the retirement of any excess,
//# by including a sufficiently large value in the Retire Prior To field.
#[test]
fn endpoint_may_exceed_limit_temporarily() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");

    let now = time::now();

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);
    reg1.set_active_connection_id_limit(2);

    assert_eq!(
        connection::id::Interest::New(1),
        reg1.connection_id_interest()
    );

    assert!(reg1
        .register_connection_id(&ext_id_2, Some(now + EXPIRATION_BUFFER), TEST_TOKEN_2)
        .is_ok());
    reg1.retire_handshake_connection_id();
    reg1.on_timeout(now + EXPIRATION_BUFFER);

    // We can register another ID because the retire_prior_to field retires old IDs
    assert_eq!(
        connection::id::Interest::New(2),
        reg1.connection_id_interest()
    );
    assert!(reg1
        .register_connection_id(&ext_id_3, None, TEST_TOKEN_3)
        .is_ok());
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//= type=test
//# An endpoint MAY also limit the issuance of
//# connection IDs to reduce the amount of per-path state it maintains,
//# such as path validation status, as its peer might interact with it
//# over as many paths as there are issued connection IDs.
#[test]
fn endpoint_may_limit_connection_ids() {
    let ext_id_1 = id(b"id01");
    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);
    reg1.set_active_connection_id_limit(100);

    assert_eq!(
        MAX_ACTIVE_CONNECTION_ID_LIMIT,
        reg1.active_connection_id_limit as u64
    );
}

#[test]
fn on_transmit() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");

    let now = time::now() + Duration::from_secs(60);

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);

    reg1.set_active_connection_id_limit(3);

    assert_eq!(
        transmission::Interest::None,
        reg1.get_transmission_interest()
    );

    assert!(reg1
        .register_connection_id(&ext_id_2, Some(now), TEST_TOKEN_2)
        .is_ok());

    assert_eq!(
        transmission::Interest::NewData,
        reg1.get_transmission_interest()
    );

    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut write_context = MockWriteContext::new(
        time::now(),
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Server,
    );
    reg1.on_transmit(&mut write_context);

    let expected_frame = Frame::NewConnectionId(NewConnectionId {
        sequence_number: VarInt::from_u32(1),
        retire_prior_to: VarInt::from_u32(0),
        connection_id: ext_id_2.as_bytes(),
        stateless_reset_token: TEST_TOKEN_2.as_ref().try_into().unwrap(),
    });

    assert_eq!(
        expected_frame,
        write_context.frame_buffer.pop_front().unwrap().as_frame()
    );

    assert_eq!(
        transmission::Interest::None,
        reg1.get_transmission_interest()
    );

    // Retire everything
    reg1.retire_handshake_connection_id();
    reg1.on_timeout(now);
    assert!(reg1
        .register_connection_id(&ext_id_3, None, TEST_TOKEN_3)
        .is_ok());

    assert_eq!(
        transmission::Interest::NewData,
        reg1.get_transmission_interest()
    );

    // Switch ID 3 to PendingReissue
    reg1.get_connection_id_info_mut(&ext_id_3).unwrap().status = PendingReissue;
    reg1.transmission_interest.clear();

    assert_eq!(
        transmission::Interest::LostData,
        reg1.get_transmission_interest()
    );

    reg1.on_transmit(&mut write_context);

    let expected_frame = Frame::NewConnectionId(NewConnectionId {
        sequence_number: VarInt::from_u32(2),
        retire_prior_to: VarInt::from_u32(2),
        connection_id: ext_id_3.as_bytes(),
        stateless_reset_token: TEST_TOKEN_3.as_ref().try_into().unwrap(),
    });

    assert_eq!(
        expected_frame,
        write_context.frame_buffer.pop_front().unwrap().as_frame()
    );

    assert_eq!(
        transmission::Interest::None,
        reg1.get_transmission_interest()
    );
}

#[test]
fn on_transmit_constrained() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);

    reg1.set_active_connection_id_limit(3);

    assert_eq!(
        transmission::Interest::None,
        reg1.get_transmission_interest()
    );

    assert!(reg1
        .register_connection_id(&ext_id_2, None, TEST_TOKEN_2)
        .is_ok());
    assert!(reg1
        .register_connection_id(&ext_id_3, None, TEST_TOKEN_3)
        .is_ok());

    assert_eq!(
        transmission::Interest::NewData,
        reg1.get_transmission_interest()
    );

    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut write_context = MockWriteContext::new(
        time::now(),
        &mut frame_buffer,
        transmission::Constraint::RetransmissionOnly,
        transmission::Mode::Normal,
        endpoint::Type::Server,
    );
    reg1.on_transmit(&mut write_context);

    // No frame written because only retransmissions are allowed
    assert!(write_context.frame_buffer.is_empty());

    reg1.get_connection_id_info_mut(&ext_id_2).unwrap().status = PendingReissue;
    reg1.transmission_interest.clear();

    assert_eq!(
        transmission::Interest::LostData,
        reg1.get_transmission_interest()
    );

    reg1.on_transmit(&mut write_context);

    // Only the ID pending reissue should be written
    assert_eq!(1, write_context.frame_buffer.len());

    let expected_frame = Frame::NewConnectionId(NewConnectionId {
        sequence_number: VarInt::from_u32(1),
        retire_prior_to: VarInt::from_u32(0),
        connection_id: ext_id_2.as_bytes(),
        stateless_reset_token: TEST_TOKEN_2.as_ref().try_into().unwrap(),
    });

    assert_eq!(
        expected_frame,
        write_context.frame_buffer.pop_front().unwrap().as_frame()
    );

    assert_eq!(
        transmission::Interest::NewData,
        reg1.get_transmission_interest()
    );
}

#[test]
fn on_packet_ack_and_loss() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);

    reg1.set_active_connection_id_limit(3);

    assert!(reg1
        .register_connection_id(&ext_id_2, None, TEST_TOKEN_2)
        .is_ok());

    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut write_context = MockWriteContext::new(
        time::now(),
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Server,
    );

    // Transition ID to PendingAcknowledgement
    let packet_number = write_context.packet_number();
    reg1.on_transmit(&mut write_context);

    // Packet was lost
    reg1.on_packet_loss(&PacketNumberRange::new(packet_number, packet_number));

    assert_eq!(
        PendingReissue,
        reg1.get_connection_id_info(&ext_id_2).unwrap().status
    );

    // Transition ID to PendingAcknowledgement again
    let packet_number = write_context.packet_number();
    reg1.on_transmit(&mut write_context);

    reg1.on_packet_ack(&PacketNumberRange::new(packet_number, packet_number));

    assert_eq!(
        Active,
        reg1.get_connection_id_info(&ext_id_2).unwrap().status
    );
    assert_eq!(
        stateless_reset::Token::ZEROED,
        reg1.get_connection_id_info(&ext_id_2)
            .unwrap()
            .stateless_reset_token
    );
}

#[test]
fn timers() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");

    let (_, mut reg1) = mapper(ext_id_1, None, TEST_TOKEN_1);
    reg1.set_active_connection_id_limit(3);

    // No timer set for the handshake connection ID
    assert_eq!(0, reg1.armed_timer_count());

    let now = time::now();
    let expiration = now + Duration::from_secs(60);

    assert!(reg1
        .register_connection_id(&ext_id_2, Some(expiration), TEST_TOKEN_2)
        .is_ok());

    // Expiration timer is armed based on retire time
    assert_eq!(1, reg1.armed_timer_count());
    assert_eq!(Some(expiration - EXPIRATION_BUFFER), reg1.next_expiration());

    reg1.get_connection_id_info_mut(&ext_id_1)
        .unwrap()
        .retire(Some(now));
    reg1.next_expiration.clear();

    // Expiration timer is armed based on removal time
    assert_eq!(1, reg1.armed_timer_count());
    assert_eq!(Some(now + EXPIRATION_BUFFER), reg1.next_expiration());

    reg1.get_connection_id_info_mut(&ext_id_2)
        .unwrap()
        .retire(Some(now));
    reg1.next_expiration.clear();

    // Expiration timer is armed based on removal time
    assert_eq!(1, reg1.armed_timer_count());
    assert_eq!(Some(now + EXPIRATION_BUFFER), reg1.next_expiration());

    // Unregister CIDs 1 and 2 (sequence numbers 0 and 1)
    reg1.unregister_expired_ids(now + Duration::from_secs(120));

    // No more timers are set
    assert_eq!(0, reg1.armed_timer_count());
}

#[test]
fn on_timeout() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");

    let now = time::now();
    let handshake_expiration = now + Duration::from_secs(60);

    let (_, mut reg1) = mapper(ext_id_1, Some(handshake_expiration), TEST_TOKEN_1);
    reg1.set_active_connection_id_limit(3);

    // Timer set for the handshake connection ID
    assert_eq!(1, reg1.armed_timer_count());

    reg1.retire_handshake_connection_id();

    // Too early, no timer is ready
    reg1.on_timeout(now);

    assert_eq!(Some(handshake_expiration), reg1.next_expiration());
    assert!(reg1.get_connection_id_info(&ext_id_1).is_some());

    // Now the expiration timer is ready
    reg1.on_timeout(handshake_expiration);
    // ID 1 was removed since it expired
    assert!(reg1.get_connection_id_info(&ext_id_1).is_none());
    assert!(reg1.next_expiration().is_none());

    let expiration_2 = now + Duration::from_secs(60);
    let expiration_3 = now + Duration::from_secs(120);

    assert!(reg1
        .register_connection_id(&ext_id_2, Some(expiration_2), TEST_TOKEN_2)
        .is_ok());
    assert!(reg1
        .register_connection_id(&ext_id_3, Some(expiration_3), TEST_TOKEN_3)
        .is_ok());

    // Expiration timer is set based on the retirement time of ID 2
    assert_eq!(
        Some(expiration_2 - EXPIRATION_BUFFER),
        reg1.next_expiration()
    );

    reg1.on_timeout(expiration_2 - EXPIRATION_BUFFER);

    // ID 2 is moved into pending retirement confirmation
    assert_eq!(
        PendingRetirementConfirmation(Some(expiration_2)),
        reg1.get_connection_id_info(&ext_id_2).unwrap().status
    );
    // Expiration timer is set to the expiration time of ID 2
    assert_eq!(Some(expiration_2), reg1.next_expiration());

    reg1.on_timeout(expiration_2);

    assert!(reg1.get_connection_id_info(&ext_id_2).is_none());

    // Expiration timer is set to the retirement time of ID 3
    assert_eq!(
        Some(expiration_3 - EXPIRATION_BUFFER),
        reg1.next_expiration()
    );
}

#[test]
fn retire_handshake_connection_id() {
    let ext_id_1 = id(b"id01");
    let ext_id_2 = id(b"id02");
    let ext_id_3 = id(b"id03");

    let now = time::now();
    let handshake_expiration = now + Duration::from_secs(60);
    let (_, mut reg1) = mapper(ext_id_1, Some(handshake_expiration), TEST_TOKEN_1);

    reg1.set_active_connection_id_limit(3);

    assert!(reg1
        .register_connection_id(&ext_id_2, None, TEST_TOKEN_2)
        .is_ok());
    assert!(reg1
        .register_connection_id(&ext_id_3, None, TEST_TOKEN_3)
        .is_ok());

    reg1.retire_handshake_connection_id();

    assert_eq!(3, reg1.registered_ids.iter().count());

    for id_info in reg1.registered_ids.iter() {
        if id_info.sequence_number == 0 {
            assert!(id_info.is_retired());
            assert_eq!(
                Some(handshake_expiration - EXPIRATION_BUFFER),
                id_info.retirement_time
            );
        } else {
            assert!(!id_info.is_retired())
        }
    }

    // Calling retire_handshake_connection_id again does nothing
    reg1.retire_handshake_connection_id();

    assert_eq!(3, reg1.registered_ids.iter().count());

    for id_info in reg1.registered_ids.iter() {
        if id_info.sequence_number == 0 {
            assert!(id_info.is_retired());
            assert_eq!(
                Some(handshake_expiration - EXPIRATION_BUFFER),
                id_info.retirement_time
            );
        } else {
            assert!(!id_info.is_retired())
        }
    }

    assert!(reg1
        .on_retire_connection_id(0, &ext_id_2, Duration::from_millis(100), now)
        .is_ok());

    // ON_RETIRE_CONNECTION_ID received for the handshake CID should change the status to PendingRemoval
    for id_info in reg1.registered_ids.iter() {
        if id_info.sequence_number == 0 {
            assert!(matches!(id_info.status, PendingRemoval(_)));
        } else {
            assert!(!id_info.is_retired())
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use crate::{
    connection::{
        peer_id_registry::{
            testing::{id, peer_registry},
            PeerIdRegistrationError,
            PeerIdRegistrationError::{
                ExceededActiveConnectionIdLimit, ExceededRetiredConnectionIdLimit,
                InvalidNewConnectionId,
            },
            PeerIdStatus::{
                InUse, InUsePendingNewConnectionId, New, PendingAcknowledgement, PendingRetirement,
                PendingRetirementRetransmission,
            },
            RETIRED_CONNECTION_ID_LIMIT,
        },
        ConnectionIdMapper, InternalConnectionIdGenerator,
    },
    contexts::{
        testing::{MockWriteContext, OutgoingFrameBuffer},
        WriteContext,
    },
    transmission,
    transmission::interest::Provider,
};
use s2n_quic_core::{
    endpoint,
    frame::{new_connection_id::STATELESS_RESET_TOKEN_LEN, Frame, RetireConnectionId},
    packet::number::PacketNumberRange,
    random,
    stateless_reset::token::testing::*,
    time::clock::testing as time,
    transport,
    varint::VarInt,
};

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
//= type=test
//# An endpoint SHOULD limit the number of connection IDs it has retired
//# locally for which RETIRE_CONNECTION_ID frames have not yet been
//# acknowledged.  An endpoint SHOULD allow for sending and tracking a
//# number of RETIRE_CONNECTION_ID frames of at least twice the value of
//# the active_connection_id_limit transport parameter.  An endpoint MUST
//# NOT forget a connection ID without retiring it, though it MAY choose
//# to treat having connection IDs in need of retirement that exceed this
//# limit as a connection error of type CONNECTION_ID_LIMIT_ERROR.
#[test]
fn error_when_exceeding_retired_connection_id_limit() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    // Register 6 more new IDs for a total of 7, with 6 retired
    for i in 2u32..=(RETIRED_CONNECTION_ID_LIMIT + 1).into() {
        assert!(reg
            .on_new_connection_id(
                &id(&i.to_ne_bytes()),
                i,
                i,
                &[i as u8; STATELESS_RESET_TOKEN_LEN].into()
            )
            .is_ok());
    }

    // Retiring one more ID exceeds the limit
    let result = reg.on_new_connection_id(
        &id(b"id08"),
        8,
        8,
        &[8_u8; STATELESS_RESET_TOKEN_LEN].into(),
    );

    assert_eq!(Some(ExceededRetiredConnectionIdLimit), result.err());
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//= type=test
//# After processing a NEW_CONNECTION_ID frame and
//# adding and retiring active connection IDs, if the number of active
//# connection IDs exceeds the value advertised in its
//# active_connection_id_limit transport parameter, an endpoint MUST
//# close the connection with an error of type CONNECTION_ID_LIMIT_ERROR.
#[test]
fn error_when_exceeding_active_connection_id_limit() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    // Register 5 more new IDs with a retire_prior_to of 4, so there will be a total of
    // 6 connection IDs, with 3 retired and 3 active
    for i in 2u32..=6 {
        assert!(reg
            .on_new_connection_id(
                &id(&i.to_ne_bytes()),
                i,
                4,
                &[i as u8; STATELESS_RESET_TOKEN_LEN].into()
            )
            .is_ok());
    }

    // Adding one more ID exceeds the limit
    let result = reg.on_new_connection_id(
        &id(b"id07"),
        8,
        0,
        &[8_u8; STATELESS_RESET_TOKEN_LEN].into(),
    );

    assert_eq!(Some(ExceededActiveConnectionIdLimit), result.err());
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
//= type=test
//# Receipt of the same frame multiple times MUST NOT be treated as a
//# connection error.
#[test]
fn no_error_when_duplicate() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    let id_2 = id(b"id02");
    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_1).is_ok());

    assert_eq!(2, reg.registered_ids.len());
    reg.registered_ids[1].status = PendingRetirement;

    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_1).is_ok());
    assert_eq!(2, reg.registered_ids.len());
    assert_eq!(PendingRetirement, reg.registered_ids[1].status);
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-18.2
//= type=test
//# The value of the active_connection_id_limit parameter MUST be at least 2.
#[test]
fn active_connection_id_limit_must_be_at_least_2() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    let id_2 = id(b"id02");
    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_1).is_ok());

    let id_3 = id(b"id03");
    assert!(reg.on_new_connection_id(&id_3, 2, 0, &TEST_TOKEN_2).is_ok());

    assert_eq!(
        2,
        reg.registered_ids
            .iter()
            .filter(|id_info| id_info.is_active())
            .count()
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
//= type=test
//# If an endpoint receives a NEW_CONNECTION_ID frame that repeats a
//# previously issued connection ID with a different Stateless Reset
//# Token field value or a different Sequence Number field value, or if a
//# sequence number is used for different connection IDs, the endpoint
//# MAY treat that receipt as a connection error of type
//# PROTOCOL_VIOLATION.
#[test]
fn duplicate_new_id_different_token_or_sequence_number() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    let id_2 = id(b"id02");
    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_1).is_ok());

    // Change the sequence number
    let mut result = reg.on_new_connection_id(&id_2, 2, 0, &TEST_TOKEN_1);
    assert_eq!(Some(InvalidNewConnectionId), result.err());

    // Change the stateless reset token
    result = reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_2);
    assert_eq!(Some(InvalidNewConnectionId), result.err());
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
//= type=test
//# If an endpoint receives a NEW_CONNECTION_ID frame that repeats a
//# previously issued connection ID with a different Stateless Reset
//# Token field value or a different Sequence Number field value, or if a
//# sequence number is used for different connection IDs, the endpoint
//# MAY treat that receipt as a connection error of type
//# PROTOCOL_VIOLATION.
#[test]
fn non_duplicate_new_id_same_token_or_sequence_number() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    let id_2 = id(b"id02");
    let id_3 = id(b"id03");
    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_1).is_ok());

    // Same sequence number
    let mut result = reg.on_new_connection_id(&id_3, 1, 0, &TEST_TOKEN_2);
    assert_eq!(Some(InvalidNewConnectionId), result.err());

    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.2
    //= type=test
    //# Endpoints are not required to compare new values
    //# against all previous values, but a duplicate value MAY be treated as
    //# a connection error of type PROTOCOL_VIOLATION.
    // Same stateless reset token
    result = reg.on_new_connection_id(&id_3, 2, 0, &TEST_TOKEN_1);
    assert_eq!(Some(InvalidNewConnectionId), result.err());
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
//= type=test
//# A receiver MUST ignore any Retire Prior To fields that do not
//# increase the largest received Retire Prior To value.
#[test]
fn ignore_retire_prior_to_that_does_not_increase() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    let id_2 = id(b"id02");
    let id_3 = id(b"id03");
    let id_4 = id(b"id04");
    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_2).is_ok());
    assert!(reg.on_new_connection_id(&id_3, 2, 1, &TEST_TOKEN_3).is_ok());
    assert_eq!(1, reg.retire_prior_to);
    assert!(reg.on_new_connection_id(&id_4, 3, 0, &TEST_TOKEN_4).is_ok());
    assert_eq!(1, reg.retire_prior_to);
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
//= type=test
//# Upon receipt of an increased Retire Prior To field, the peer MUST
//# stop using the corresponding connection IDs and retire them with
//# RETIRE_CONNECTION_ID frames before adding the newly provided
//# connection ID to the set of active connection IDs.
#[test]
fn retire_connection_id_when_retire_prior_to_increases() {
    let id_1 = id(b"id01");
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let mut reg =
        mapper.create_client_peer_id_registry(InternalConnectionIdGenerator::new().generate_id());
    reg.register_initial_connection_id(id_1);
    reg.register_initial_stateless_reset_token(TEST_TOKEN_1);

    let id_2 = id(b"id02");
    assert!(reg.on_new_connection_id(&id_2, 1, 1, &TEST_TOKEN_2).is_ok());

    assert_eq!(PendingRetirement, reg.registered_ids[0].status);
    assert_eq!(
        transmission::Interest::NewData,
        reg.get_transmission_interest()
    );

    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut write_context = MockWriteContext::new(
        time::now(),
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Server,
    );
    let packet_number = write_context.packet_number();
    reg.on_transmit(&mut write_context);

    let expected_frame = Frame::RetireConnectionId(RetireConnectionId {
        sequence_number: VarInt::from_u32(0),
    });

    assert_eq!(
        expected_frame,
        write_context.frame_buffer.pop_front().unwrap().as_frame()
    );
    assert_eq!(
        PendingAcknowledgement(packet_number),
        reg.registered_ids[0].status
    );

    assert_eq!(
        transmission::Interest::None,
        reg.get_transmission_interest()
    );

    reg.on_packet_loss(&PacketNumberRange::new(packet_number, packet_number));

    assert_eq!(
        PendingRetirementRetransmission,
        reg.registered_ids[0].status
    );
    assert_eq!(
        transmission::Interest::LostData,
        reg.get_transmission_interest()
    );

    // Transition ID to PendingAcknowledgement again
    let packet_number = write_context.packet_number();
    reg.on_transmit(&mut write_context);

    assert_eq!(
        Some(reg.internal_id),
        mapper.remove_internal_connection_id_by_stateless_reset_token(&TEST_TOKEN_1)
    );

    reg.on_packet_ack(&PacketNumberRange::new(packet_number, packet_number));

    // ID 1 was removed
    assert_eq!(reg.registered_ids.len(), 1);
    assert_eq!(id_2, reg.registered_ids[0].id);
    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
    //= type=test
    //# An endpoint MUST NOT check for any stateless reset tokens associated
    //# with connection IDs it has not used or for connection IDs that have
    //# been retired.
    assert_eq!(
        None,
        mapper.remove_internal_connection_id_by_stateless_reset_token(&TEST_TOKEN_1)
    );

    assert_eq!(
        transmission::Interest::None,
        reg.get_transmission_interest()
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
//= type=test
//# If an endpoint receives a NEW_CONNECTION_ID frame that repeats a
//# previously issued connection ID with a different Stateless Reset
//# Token field value or a different Sequence Number field value, or if a
//# sequence number is used for different connection IDs, the endpoint
//# MAY treat that receipt as a connection error of type
//# PROTOCOL_VIOLATION.
#[test]
fn retire_new_connection_id_if_sequence_number_smaller_than_retire_prior_to() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    let id_2 = id(b"id02");
    assert!(reg
        .on_new_connection_id(&id_2, 10, 10, &TEST_TOKEN_2)
        .is_ok());
    assert_eq!(PendingRetirement, reg.registered_ids[0].status);
    assert_eq!(New, reg.registered_ids[1].status);

    let id_3 = id(b"id03");
    assert!(reg.on_new_connection_id(&id_3, 1, 0, &TEST_TOKEN_3).is_ok());

    assert_eq!(New, reg.registered_ids[1].status);
    assert_eq!(PendingRetirement, reg.registered_ids[2].status);
}

#[test]
fn retire_initial_id_when_new_connection_id_available() {
    let id_1 = id(b"id01");
    let mut reg = peer_registry(id_1, None);

    assert_eq!(InUsePendingNewConnectionId, reg.registered_ids[0].status);

    let id_2 = id(b"id02");
    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_2).is_ok());

    assert_eq!(PendingRetirement, reg.registered_ids[0].status);
}

#[test]
pub fn initial_id_is_active() {
    let id_1 = id(b"id01");
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let reg = mapper
        .create_server_peer_id_registry(InternalConnectionIdGenerator::new().generate_id(), id_1);

    assert!(reg.is_active(&id_1));
}

#[test]
pub fn retired_id_is_not_active() {
    let id_1 = id(b"id01");
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let mut reg = mapper
        .create_server_peer_id_registry(InternalConnectionIdGenerator::new().generate_id(), id_1);

    assert!(reg.is_active(&id_1));
    reg.registered_ids[0].status = PendingRetirement;
    assert!(!reg.is_active(&id_1));
}

#[test]
pub fn unknown_id_is_not_active() {
    let id_1 = id(b"id01");
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let reg = mapper
        .create_server_peer_id_registry(InternalConnectionIdGenerator::new().generate_id(), id_1);

    assert!(reg.is_active(&id_1));
    let id_unknown = id(b"unknown");
    assert!(!reg.is_active(&id_unknown));
}

#[test]
pub fn consume_new_id_should_return_id() {
    let id_1 = id(b"id01");
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let mut reg = mapper
        .create_server_peer_id_registry(InternalConnectionIdGenerator::new().generate_id(), id_1);

    let id_2 = id(b"id02");
    assert!(reg.on_new_connection_id(&id_2, 1, 0, &TEST_TOKEN_2).is_ok());
    reg.registered_ids[1].status = New;

    assert!(reg
        .state
        .lock()
        .unwrap()
        .stateless_reset_map
        .remove(&TEST_TOKEN_2)
        .is_none());
    assert_eq!(Some(id_2), reg.consume_new_id_inner());
    reg.registered_ids[1].status = InUse;
    // this is an indirect way to test that we inserted a reset token when we consumed id_2
    assert!(reg
        .state
        .lock()
        .unwrap()
        .stateless_reset_map
        .remove(&TEST_TOKEN_2)
        .is_some());
}

#[test]
pub fn consume_new_id_should_error_if_no_ids_are_available() {
    let id_1 = id(b"id01");
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let mut reg = mapper
        .create_server_peer_id_registry(InternalConnectionIdGenerator::new().generate_id(), id_1);

    assert_eq!(None, reg.consume_new_id_inner());
}

#[test]
fn error_conversion() {
    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
    //= type=test
    //# If an endpoint receives a NEW_CONNECTION_ID frame that repeats a
    //# previously issued connection ID with a different Stateless Reset
    //# Token field value or a different Sequence Number field value, or if a
    //# sequence number is used for different connection IDs, the endpoint
    //# MAY treat that receipt as a connection error of type
    //# PROTOCOL_VIOLATION.
    let mut transport_error: transport::Error =
        PeerIdRegistrationError::InvalidNewConnectionId.into();
    assert_eq!(
        transport::Error::PROTOCOL_VIOLATION.code,
        transport_error.code
    );

    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
    //= type=test
    //# After processing a NEW_CONNECTION_ID frame and
    //# adding and retiring active connection IDs, if the number of active
    //# connection IDs exceeds the value advertised in its
    //# active_connection_id_limit transport parameter, an endpoint MUST
    //# close the connection with an error of type CONNECTION_ID_LIMIT_ERROR.
    transport_error = PeerIdRegistrationError::ExceededActiveConnectionIdLimit.into();
    assert_eq!(
        transport::Error::CONNECTION_ID_LIMIT_ERROR.code,
        transport_error.code
    );

    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
    //= type=test
    //# An endpoint MUST NOT forget a connection ID without retiring it,
    //# though it MAY choose to treat having connection IDs in need of
    //# retirement that exceed this limit as a connection error of type
    //# CONNECTION_ID_LIMIT_ERROR.
    transport_error = PeerIdRegistrationError::ExceededRetiredConnectionIdLimit.into();
    assert_eq!(
        transport::Error::CONNECTION_ID_LIMIT_ERROR.code,
        transport_error.code
    );
}

#[test]
pub fn client_peer_id_registry_should_not_register_cid() {
    let mut random_generator = random::testing::Generator(123);
    let mut mapper = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server);
    let reg =
        mapper.create_client_peer_id_registry(InternalConnectionIdGenerator::new().generate_id());

    assert!(reg.registered_ids.is_empty());
    assert!(reg.is_empty());
}


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

//! Defines Context traits, which are passed to various lifecycle callbacks
//! within the connection in order to collect data

use crate::{connection::InternalConnectionId, transmission, wakeup_queue::WakeupHandle};
use s2n_codec::encoder::EncoderValue;
use s2n_quic_core::{
    endpoint,
    event::{self, IntoEvent},
    frame::{ack::AckRanges as AckRangesTrait, ack_elicitation::AckElicitation, Ack, FrameTrait},
    packet::number::PacketNumber,
    time::Timestamp,
};

/// Context information that is passed to `on_transmit` calls on Streams
pub trait WriteContext {
    /// Returns the current point of time
    fn current_time(&self) -> Timestamp;

    /// Returns the transmission constraint for the current packet
    fn transmission_constraint(&self) -> transmission::Constraint;

    /// Returns the transmission mode for the current packet
    fn transmission_mode(&self) -> transmission::Mode;

    /// Returns the number of available bytes remaining in the current payload
    fn remaining_capacity(&self) -> usize;

    /// Attempt to write an ack frame.
    ///
    /// If this was successful the number of the packet
    /// that will be used to send the frame will be returned.
    fn write_ack_frame<AckRanges: AckRangesTrait>(
        &mut self,
        ack_frame: &Ack<AckRanges>,
    ) -> Option<PacketNumber> {
        self.write_frame(ack_frame)
    }

    /// Attempt to write a frame.
    ///
    /// If this was successful the number of the packet
    /// that will be used to send the frame will be returned.
    fn write_frame<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>;

    /// Writes a pre-fitted frame.
    ///
    /// Callers should ensure the frame fits within the outgoing buffer when using this function.
    /// The context should panic if otherwise.
    fn write_fitted_frame<Frame>(&mut self, frame: &Frame) -> PacketNumber
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>;

    /// Attempt to write a frame, bypassing congestion controller constraint checks.
    /// If this was successful the number of the packet that will be used to send
    /// the frame will be returned.
    fn write_frame_forced<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>;

    /// Returns the ack elicitation of the current packet
    fn ack_elicitation(&self) -> AckElicitation;

    /// Returns the packet number for the current packet
    fn packet_number(&self) -> PacketNumber;

    /// Returns the local endpoint type (client or server)
    fn local_endpoint_type(&self) -> endpoint::Type;

    /// Returns the length of the packet header in bytes
    fn header_len(&self) -> usize;

    /// Returns the length of the authentication tag in bytes
    fn tag_len(&self) -> usize;
}

/// Enumerates error values for `on_transmit` calls
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OnTransmitError {
    /// It was not possible to write a frame
    CouldNotWriteFrame,
    /// It was not possible to obtain a large enough space for writing a frame
    CouldNotAcquireEnoughSpace,
}

/// Enumerates error values for `on_transmit` calls on connections
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ConnectionOnTransmitError {
    /// It was not possible to obtain a datagram to write into
    NoDatagram,
}

/// The context parameter which is passed from all external API calls
pub struct ConnectionApiCallContext<'a> {
    wakeup_handle: &'a WakeupHandle<InternalConnectionId>,
}

impl<'a> ConnectionApiCallContext<'a> {
    /// Creates an [`ConnectionApiCallContext`] from a [`WakeupHandle`]
    pub fn from_wakeup_handle(wakeup_handle: &'a WakeupHandle<InternalConnectionId>) -> Self {
        Self { wakeup_handle }
    }

    /// Returns a reference to the WakeupHandle
    pub fn wakeup_handle(&mut self) -> &WakeupHandle<InternalConnectionId> {
        self.wakeup_handle
    }
}

#[cfg(test)]
pub mod testing;


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

use crate::contexts::WriteContext;
use alloc::collections::VecDeque;
use s2n_codec::{
    encoder::{EncoderBuffer, EncoderValue},
    DecoderBufferMut,
};
use s2n_quic_core::{
    endpoint,
    event::{self, IntoEvent},
    frame::{
        ack_elicitation::{AckElicitable, AckElicitation},
        FrameMut, FrameTrait,
    },
    packet::number::{PacketNumber, PacketNumberSpace},
    time::Timestamp,
    transmission,
    transmission::{Constraint, Mode},
    varint::VarInt,
};

/// A single frame which had been written.
/// The buffer always stores only a single serialized Frame.
#[derive(Clone, Debug)]
pub struct WrittenFrame {
    pub data: Vec<u8>,
    pub packet_nr: PacketNumber,
}

impl WrittenFrame {
    /// Deserializes a written frame into a quic_frame::Frame type.
    /// panics if deserialization fails.
    pub fn as_frame(&mut self) -> FrameMut {
        let buffer = DecoderBufferMut::new(&mut self.data[..]);
        let (frame, remaining) = buffer
            .decode::<FrameMut>()
            .expect("Buffer contains a valid frame");
        assert_eq!(0, remaining.len());
        frame
    }
}

/// Stores frames which have been written by components
#[derive(Clone, Debug)]
pub struct OutgoingFrameBuffer {
    pub ack_elicitation: AckElicitation,
    /// Frames which had been written so far
    pub frames: VecDeque<WrittenFrame>,
    /// The PacketNumber which will be returned by the next `write_frame()` call
    next_packet_nr: PacketNumber,
    /// If set, this indicates the maximum packet size
    max_buffer_size: Option<usize>,
    /// The remaining space in the current packet
    remaining_packet_space: usize,
    /// If set, the value indicates after how many `write_frame()` will still be
    /// permitted before errors are returned on write. This can be used to simulate
    /// failing write calls.
    error_after_frames: Option<usize>,
}

impl Default for OutgoingFrameBuffer {
    fn default() -> Self {
        OutgoingFrameBuffer {
            ack_elicitation: Default::default(),
            frames: VecDeque::new(),
            next_packet_nr: PacketNumberSpace::ApplicationData
                .new_packet_number(VarInt::from_u8(0)),
            max_buffer_size: None,
            remaining_packet_space: 0,
            error_after_frames: None,
        }
    }
}

impl OutgoingFrameBuffer {
    pub fn new() -> Self {
        Self::default()
    }

    /// Clears all frames which have been fully written
    pub fn clear(&mut self) {
        self.frames.clear();
        self.ack_elicitation = Default::default();
    }

    /// Configures an explicit maximum packet size. If this is configured the
    /// component will write multiple frames into a single packet.
    /// Otherwise an individual packet will get written for each frame.
    pub fn set_max_packet_size(&mut self, max_buffer_size: Option<usize>) {
        // If we go from one mode to another, we should flush any pending
        // packets in between.
        self.flush();
        self.max_buffer_size = max_buffer_size;

        if let Some(max_buffer_size) = max_buffer_size {
            self.remaining_packet_space = max_buffer_size;
        }
    }

    /// Instructs the `FrameBuffer` to only allow `n` frame writes,
    /// and to fail the following write attempt
    pub fn set_error_write_after_n_frames(&mut self, n: usize) {
        self.error_after_frames = Some(n);
    }

    /// Returns the amount of written frames
    pub fn len(&self) -> usize {
        self.frames.len()
    }

    /// Returns false if there are written frames
    pub fn is_empty(&self) -> bool {
        self.frames.is_empty()
    }

    /// Returns the oldest written frame
    pub fn pop_front(&mut self) -> Option<WrittenFrame> {
        self.frames.pop_front()
    }

    /// Flushes any pending packet
    pub fn flush(&mut self) {
        if let Some(max_buffer_size) = self.max_buffer_size {
            if self.remaining_packet_space == max_buffer_size {
                // No frame had been written into this packet
                return;
            }

            // Increment the packet number and reset the amount of available
            // packet space. Thereby frames after that flush will get assigned
            // a fresh packet number.
            self.next_packet_nr = self.next_packet_nr.next().unwrap();
            self.remaining_packet_space = max_buffer_size;
            self.ack_elicitation = Default::default();
        }
    }

    fn remaining_capacity(&self) -> usize {
        if self.max_buffer_size.is_some() {
            self.remaining_packet_space
        } else {
            core::usize::MAX
        }
    }

    fn encode_frame_to_vec<Frame: s2n_codec::EncoderValue>(
        frame: &Frame,
        encoded_size: usize,
    ) -> Vec<u8> {
        // Create a new buffer for the frame
        let mut write_buffer = vec![0u8; encoded_size];
        let mut encoder_buffer = EncoderBuffer::new(&mut write_buffer[..]);
        frame.encode(&mut encoder_buffer);
        write_buffer
    }

    pub fn write_frame<Frame: s2n_codec::EncoderValue + AckElicitable>(
        &mut self,
        frame: &Frame,
    ) -> Option<PacketNumber> {
        if let Some(error_after_frames) = self.error_after_frames {
            if error_after_frames == 0 {
                // No more frames are permitted -> fail the write
                return None;
            }
            // Decrease the amount of permitted writes
            self.error_after_frames = Some(error_after_frames - 1);
        }

        let encoded_size = frame.encoding_size();

        if let Some(max_buffer_size) = self.max_buffer_size {
            if encoded_size > max_buffer_size {
                // This can never be written
                return None;
            }

            if self.remaining_packet_space < encoded_size {
                // Flush the current packet to make space
                self.flush();
            }

            let encoded_frame = Self::encode_frame_to_vec(frame, encoded_size);

            // Store the frame, but don't increase the packet number, since we
            // might be writing more frames after this.
            let packet_nr = self.next_packet_nr;

            self.ack_elicitation |= frame.ack_elicitation();
            self.frames.push_back(WrittenFrame {
                data: encoded_frame,
                packet_nr,
            });

            self.remaining_packet_space -= encoded_size;

            Some(self.next_packet_nr)
        } else {
            // There is no write limit configured. This means we directly store
            // each frame as an individual buffer.
            let encoded_frame = Self::encode_frame_to_vec(frame, encoded_size);

            let packet_nr = self.next_packet_nr;
            self.next_packet_nr = self.next_packet_nr.next().unwrap();

            self.ack_elicitation |= frame.ack_elicitation();
            self.frames.push_back(WrittenFrame {
                data: encoded_frame,
                packet_nr,
            });

            Some(packet_nr)
        }
    }
}

#[derive(Debug)]
pub struct MockWriteContext<'a> {
    pub current_time: Timestamp,
    pub frame_buffer: &'a mut OutgoingFrameBuffer,
    pub transmission_constraint: Constraint,
    pub transmission_mode: Mode,
    pub endpoint: endpoint::Type,
}

impl<'a> MockWriteContext<'a> {
    pub fn new(
        current_time: Timestamp,
        frame_buffer: &'a mut OutgoingFrameBuffer,
        transmission_constraint: Constraint,
        transmission_mode: Mode,
        endpoint: endpoint::Type,
    ) -> MockWriteContext<'a> {
        MockWriteContext {
            current_time,
            frame_buffer,
            transmission_constraint,
            transmission_mode,
            endpoint,
        }
    }
}

impl<'a> WriteContext for MockWriteContext<'a> {
    fn current_time(&self) -> Timestamp {
        self.current_time
    }

    fn transmission_constraint(&self) -> Constraint {
        self.transmission_constraint
    }

    fn transmission_mode(&self) -> Mode {
        self.transmission_mode
    }

    fn remaining_capacity(&self) -> usize {
        self.frame_buffer.remaining_capacity()
    }

    fn write_frame<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        match self.transmission_constraint() {
            transmission::Constraint::AmplificationLimited => {
                unreachable!("frames should not be written when we're amplification limited")
            }
            transmission::Constraint::CongestionLimited => {
                assert!(!frame.is_congestion_controlled());
            }
            transmission::Constraint::RetransmissionOnly => {}
            transmission::Constraint::None => {}
        }
        self.frame_buffer.write_frame(frame)
    }

    fn write_fitted_frame<Frame>(&mut self, frame: &Frame) -> PacketNumber
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        self.write_frame(frame)
            .expect("frame should fit in current buffer")
    }

    fn write_frame_forced<Frame>(&mut self, frame: &Frame) -> Option<PacketNumber>
    where
        Frame: EncoderValue + FrameTrait,
        for<'frame> &'frame Frame: IntoEvent<event::builder::Frame>,
    {
        self.frame_buffer.write_frame(frame)
    }

    fn ack_elicitation(&self) -> AckElicitation {
        self.frame_buffer.ack_elicitation
    }

    fn packet_number(&self) -> PacketNumber {
        self.frame_buffer.next_packet_nr
    }

    fn local_endpoint_type(&self) -> endpoint::Type {
        self.endpoint
    }

    fn header_len(&self) -> usize {
        0
    }

    fn tag_len(&self) -> usize {
        0
    }
}