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

//! This module contains the implementation of QUIC `Connections` and their management

use crate::{
    connection::{
        self, connection_interests::ConnectionInterests, id::ConnectionInfo,
        internal_connection_id::InternalConnectionId, local_id_registry::LocalIdRegistrationError,
        ConnectionIdMapper, Parameters as ConnectionParameters, ProcessingError,
    },
    contexts::ConnectionOnTransmitError,
    endpoint,
    path::{self, path_event},
    stream,
};
use bytes::Bytes;
use core::task::{Context, Poll};
use s2n_codec::DecoderBufferMut;
use s2n_quic_core::{
    application,
    application::ServerName,
    event::{self, builder::DatagramDropReason, supervisor, ConnectionPublisher, IntoEvent},
    inet::{DatagramInfo, SocketAddress},
    io::tx,
    packet::{
        handshake::ProtectedHandshake,
        initial::{CleartextInitial, ProtectedInitial},
        retry::ProtectedRetry,
        short::ProtectedShort,
        version_negotiation::ProtectedVersionNegotiation,
        zero_rtt::ProtectedZeroRtt,
        ProtectedPacket,
    },
    path::{Handle as _, MaxMtu},
    query,
    time::Timestamp,
};

/// A trait which represents an internally used `Connection`
pub trait ConnectionTrait: 'static + Send + Sized {
    /// Static configuration of a connection
    type Config: endpoint::Config;

    /// Creates a new `Connection` instance with the given configuration
    fn new(parameters: ConnectionParameters<Self::Config>) -> Result<Self, connection::Error>;

    /// Returns the Connections internal ID
    fn internal_connection_id(&self) -> InternalConnectionId;

    /// Returns whether the connection is in the handshake state
    fn is_handshaking(&self) -> bool;

    /// Initiates closing the connection as described in
    /// https://www.rfc-editor.org/rfc/rfc9000#section-10
    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,
    );

    /// Marks a connection which advertised itself as having completed the handshake
    /// (via [`ConnectionInterests`]) as accepted. After this call the `accept` interest should
    /// no longer be signalled.
    fn mark_as_accepted(&mut self);

    /// Generates and registers new connection IDs using the given `ConnectionIdFormat` and
    /// `StatelessResetTokenGenerator`
    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<(), LocalIdRegistrationError>;

    /// Queries the connection for outgoing packets
    fn on_transmit<Tx>(
        &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<(), ConnectionOnTransmitError>
    where
        Tx: tx::Queue<Handle = <Self::Config as endpoint::Config>::PathHandle>;

    /// Handles all timeouts on the `Connection`.
    ///
    /// `timestamp` passes the current time.
    fn on_timeout(
        &mut self,
        connection_id_mapper: &mut 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>;

    /// Handles all external wakeups on the [`Connection`].
    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>;

    // Packet handling

    /// Is called when an initial packet had been received
    #[allow(clippy::too_many_arguments)]
    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>;

    /// Is called when an unprotected initial packet had been received
    #[allow(clippy::too_many_arguments)]
    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>;

    /// Is called when a handshake packet had been received
    #[allow(clippy::too_many_arguments)]
    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>;

    /// 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>;

    /// 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>;

    /// 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>;

    /// 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>;

    /// Notifies a connection it has received a datagram from a peer
    #[allow(clippy::too_many_arguments)]
    fn on_datagram_received(
        &mut self,
        path_handle: &<Self::Config as endpoint::Config>::PathHandle,
        datagram: &DatagramInfo,
        congestion_controller_endpoint: &mut <Self::Config as endpoint::Config>::CongestionControllerEndpoint,
        migration_validator: &mut <Self::Config as endpoint::Config>::PathMigrationValidator,
        max_mtu: MaxMtu,
        subscriber: &mut <Self::Config as endpoint::Config>::EventSubscriber,
    ) -> Result<path::Id, DatagramDropReason>;

    /// Returns the Connections interests
    fn interests(&self) -> ConnectionInterests;

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

    /// Handles reception of a single QUIC packet
    #[allow(clippy::too_many_arguments)]
    fn handle_packet(
        &mut self,
        datagram: &DatagramInfo,
        path_id: path::Id,
        packet: ProtectedPacket,
        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> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.1
        //# If a client receives a packet that uses a different version than it
        //# initially selected, it MUST discard that packet.
        if let Some(version) = packet.version() {
            if version != self.quic_version() {
                self.with_event_publisher(
                    datagram.timestamp,
                    Some(path_id),
                    subscriber,
                    |publisher, path| {
                        publisher.on_packet_dropped(event::builder::PacketDropped {
                            reason: event::builder::PacketDropReason::VersionMismatch {
                                version,
                                path: path_event!(path, path_id),
                            },
                        })
                    },
                );
                return Ok(());
            }
        }

        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.4
        //# An endpoint SHOULD continue
        //# to respond to packets that can be processed during this time.
        // We make a best effort to process all of the packet spaces we have available. There isn't
        // any special logic required to meet this requirement as each packet is handled
        // independently.

        match packet {
            ProtectedPacket::Short(packet) => self.handle_short_packet(
                datagram,
                path_id,
                packet,
                random_generator,
                subscriber,
                packet_interceptor,
            ),
            ProtectedPacket::VersionNegotiation(packet) => self.handle_version_negotiation_packet(
                datagram,
                path_id,
                packet,
                subscriber,
                packet_interceptor,
            ),
            ProtectedPacket::Initial(packet) => self.handle_initial_packet(
                datagram,
                path_id,
                packet,
                random_generator,
                subscriber,
                packet_interceptor,
                datagram_endpoint,
            ),
            ProtectedPacket::ZeroRtt(packet) => self.handle_zero_rtt_packet(
                datagram,
                path_id,
                packet,
                subscriber,
                packet_interceptor,
            ),
            ProtectedPacket::Handshake(packet) => self.handle_handshake_packet(
                datagram,
                path_id,
                packet,
                random_generator,
                subscriber,
                packet_interceptor,
                datagram_endpoint,
            ),
            ProtectedPacket::Retry(packet) => {
                self.handle_retry_packet(datagram, path_id, packet, subscriber, packet_interceptor)
            }
        }
    }

    /// This is called to handle the remaining and yet undecoded packets inside
    /// a datagram.
    #[allow(clippy::too_many_arguments)]
    fn handle_remaining_packets(
        &mut self,
        path_handle: &<Self::Config as endpoint::Config>::PathHandle,
        datagram: &DatagramInfo,
        path_id: path::Id,
        connection_id_validator: &<Self::Config as endpoint::Config>::ConnectionIdFormat,
        mut payload: DecoderBufferMut,
        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<(), connection::Error> {
        let remote_address = path_handle.remote_address();
        let connection_info = ConnectionInfo::new(&remote_address);

        while !payload.is_empty() {
            if let Ok((packet, remaining)) =
                ProtectedPacket::decode(payload, &connection_info, connection_id_validator)
            {
                payload = remaining;

                //= https://www.rfc-editor.org/rfc/rfc9000#section-12.2
                //# Senders MUST NOT coalesce QUIC packets
                //# with different connection IDs into a single UDP datagram.  Receivers
                //# SHOULD ignore any subsequent packets with a different Destination
                //# Connection ID than the first packet in the datagram.
                if datagram.destination_connection_id.as_bytes()
                    != packet.destination_connection_id()
                {
                    self.with_event_publisher(
                        datagram.timestamp,
                        Some(path_id),
                        subscriber,
                        |publisher, path| {
                            publisher.on_packet_dropped(event::builder::PacketDropped {
                                reason: event::builder::PacketDropReason::ConnectionIdMismatch {
                                    packet_cid: packet.destination_connection_id(),
                                    path: path_event!(path, path_id),
                                },
                            })
                        },
                    );
                    break;
                }

                let result = self.handle_packet(
                    datagram,
                    path_id,
                    packet,
                    random_generator,
                    subscriber,
                    packet_interceptor,
                    datagram_endpoint,
                );

                if let Err(ProcessingError::ConnectionError(err)) = result {
                    // CryptoErrors returned as a result of a packet failing decryption will be
                    // silently discarded, but this method could return an error on protocol
                    // violations which would result in shutting down the connection anyway. In this
                    // case this will return early without processing the remaining packets.
                    if !payload.is_empty() {
                        self.with_event_publisher(
                            datagram.timestamp,
                            Some(path_id),
                            subscriber,
                            |publisher, path| {
                                publisher.on_packet_dropped(event::builder::PacketDropped {
                                    reason: event::builder::PacketDropReason::ConnectionError {
                                        path: path_event!(path, path_id),
                                    },
                                })
                            },
                        );
                    }
                    return Err(err);
                }
            } else {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-12.2
                //# Every QUIC packet that is coalesced into a single UDP datagram is
                //# separate and complete.  The receiver of coalesced QUIC packets MUST
                //# individually process each QUIC packet and separately acknowledge
                //# them, as if they were received as the payload of different UDP
                //# datagrams.  For example, if decryption fails (because the keys are
                //# not available or for any other reason), the receiver MAY either
                //# discard or buffer the packet for later processing and MUST attempt to
                //# process the remaining packets.
                //
                // We choose to discard the rest of the datagram on parsing errors since it
                // would be difficult to recover from an invalid packet.
                self.with_event_publisher(
                    datagram.timestamp,
                    Some(path_id),
                    subscriber,
                    |publisher, path| {
                        publisher.on_packet_dropped(event::builder::PacketDropped {
                            reason: event::builder::PacketDropReason::DecodingFailed {
                                path: path_event!(path, path_id),
                            },
                        })
                    },
                );

                break;
            }
        }

        Ok(())
    }

    fn poll_stream_request(
        &mut self,
        stream_id: stream::StreamId,
        request: &mut stream::ops::Request,
        context: Option<&Context>,
    ) -> Result<stream::ops::Response, stream::StreamError>;

    fn poll_accept_stream(
        &mut self,
        stream_type: Option<stream::StreamType>,
        context: &Context,
    ) -> Poll<Result<Option<stream::StreamId>, connection::Error>>;

    fn poll_open_stream(
        &mut self,
        stream_type: stream::StreamType,
        open_token: &mut connection::OpenToken,
        context: &Context,
    ) -> Poll<Result<stream::StreamId, connection::Error>>;

    fn application_close(&mut self, error: Option<application::Error>);

    fn server_name(&self) -> Option<ServerName>;

    fn application_protocol(&self) -> Bytes;

    fn ping(&mut self) -> Result<(), connection::Error>;

    fn keep_alive(&mut self, enabled: bool) -> Result<(), connection::Error>;

    fn local_address(&self) -> Result<SocketAddress, connection::Error>;

    fn remote_address(&self) -> Result<SocketAddress, connection::Error>;

    fn error(&self) -> Option<connection::Error>;

    fn query_event_context(&self, query: &mut dyn query::Query);

    fn query_event_context_mut(&mut self, query: &mut dyn query::QueryMut);

    fn datagram_mut(&mut self, query: &mut dyn query::QueryMut);

    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>,
        );
}

/// A lock that synchronizes connection state between the QUIC endpoint thread and application
pub trait Lock<T>: 'static + Send + Sync {
    type Error;

    /// Creates a connection lock
    fn new(value: T) -> Self;

    /// Obtains a read-only reference to the inner connection
    fn read<F: FnOnce(&T) -> R, R>(&self, f: F) -> Result<R, Self::Error>;

    /// Obtains a mutable reference to the inner connection
    fn write<F: FnOnce(&mut T) -> R, R>(&self, f: F) -> Result<R, Self::Error>;
}

#[cfg(feature = "std")]
impl<T: 'static + Send> Lock<T> for std::sync::Mutex<T> {
    type Error = ();

    fn new(value: T) -> Self {
        std::sync::Mutex::new(value)
    }

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

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


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

//! Defines error types for packet handling inside connections

use s2n_quic_core::crypto::CryptoError;

/// Errors that can occur during packet reception
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum PacketHandlingError {
    /// A crypto error occurred
    CryptoError(CryptoError),
}

impl From<CryptoError> for PacketHandlingError {
    fn from(error: CryptoError) -> Self {
        Self::CryptoError(error)
    }
}


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

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Status {
    Idle,
    Draining,
    Final,
}

impl Status {
    #[allow(dead_code)]
    pub fn is_idle(self) -> bool {
        matches!(self, Self::Idle)
    }

    #[allow(dead_code)]
    pub fn is_draining(self) -> bool {
        matches!(self, Self::Draining)
    }

    #[allow(dead_code)]
    pub fn is_final(self) -> bool {
        matches!(self, Self::Final)
    }
}

impl core::ops::Add for Status {
    type Output = Self;

    fn add(self, rhs: Self) -> Self {
        // all components need to be in their final state to finalize
        match (self, rhs) {
            // in idle the component doesn't report any finalization status
            (Self::Idle, rhs) => rhs,
            (lhs, Self::Idle) => lhs,

            // a draining component holds up finaliztion
            (Self::Draining, _) => Self::Draining,
            (_, Self::Draining) => Self::Draining,

            // only return final if both components are final
            (Self::Final, Self::Final) => Self::Final,
        }
    }
}

impl core::ops::AddAssign for Status {
    fn add_assign(&mut self, rhs: Self) {
        *self = *self + rhs;
    }
}

impl core::iter::Sum for Status {
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        let mut interest = Self::Idle;

        for item in iter {
            interest += item;
        }

        interest
    }
}

pub trait Provider {
    fn finalization_status(&self) -> Status;
}


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

//! Defines the internal connection ID

/// A connection ID which is stable and internally identifies a connection over
/// the whole lifetime of an endpoint.
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct InternalConnectionId(u64);

/// Generates internal connection IDs.
///
/// Connection IDs are only internal if they are generated by the same generator.
pub struct InternalConnectionIdGenerator {
    next: u64,
}

impl InternalConnectionIdGenerator {
    /// Returns a new InternalConnectionIdGenerator
    pub fn new() -> Self {
        Self { next: 0 }
    }

    /// Returns a new InternalConnectionId
    pub fn generate_id(&mut self) -> InternalConnectionId {
        let result = self.next;
        // This should never happen in production.
        // If the generator generates 1 million IDs per second,
        // the number will overflow after 213503982 days of operation.
        // Even with 1 billion IDs per second - which requires the CPU to spend
        // a significant amount of time generating IDs - it will take years to
        // overflow.
        self.next = self.next.checked_add(1).expect("Connection ID overflow");

        InternalConnectionId(result)
    }
}

impl From<InternalConnectionId> for u64 {
    fn from(internal_conn_id: InternalConnectionId) -> u64 {
        internal_conn_id.0
    }
}

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

    #[test]
    fn connection_id_generator_test() {
        let mut id_generator = InternalConnectionIdGenerator::new();

        for i in 0..128u64 {
            let id = id_generator.generate_id();
            assert_eq!(InternalConnectionId(i), id);
        }
    }
}


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

use crate::{
    connection::{
        connection_id_mapper::ConnectionIdMapperState, local_id_registry::LocalIdStatus::*,
        InternalConnectionId,
    },
    contexts::WriteContext,
    transmission,
};
use core::convert::TryInto;
use s2n_quic_core::{
    ack, connection, frame,
    memo::Memo,
    packet::number::PacketNumber,
    stateless_reset,
    time::{timer, Duration, Timer, Timestamp},
};
use smallvec::SmallVec;
use std::sync::{Arc, Mutex};

/// The amount of ConnectionIds we can register without dynamic memory allocation
const NR_STATIC_REGISTRABLE_IDS: usize = 5;

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
//# 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.
/// Limit on the number of connection IDs issued to the peer to reduce the amount
/// of per-path state maintained. Increasing this value allows peers to probe
/// more paths simultaneously at the expense of additional state to maintain.
const MAX_ACTIVE_CONNECTION_ID_LIMIT: u64 = 3;

/// Buffer to allow time for a peer to process and retire an expiring connection ID
/// before the connection ID actually expires.
///
/// When a local connection ID is retired a NEW_CONNECTION_ID frame is sent to the peer
/// containing a new connection ID to use as well as an increased "retire prior to" value.
/// The peer is required to send back a RETIRE_CONNECTION_ID frame retiring the connection ID(s)
/// indicated by the "retire prior to" value. When the RETIRE_CONNECTION_ID frame is
/// received, the connection ID is removed from use. The `EXPIRATION_BUFFER` is meant to
/// ensure the peer has enough time to cease using the connection ID before it is permanently
/// removed. 30 seconds should be enough time, even in extremely slow networks, for the
/// peer to flush all packets created with the old connection ID and to receive and start using
/// the new connection ID, at the minimal cost of maintaining state of an extra connection ID for
/// a brief period. Setting this value too low increases the risk of a packet being received with
/// a removed connection ID, resulting in a stateless reset that terminates the connection.
///
/// The value is not dynamically based on RTT as this would introduce a moving retirement time
/// target that would complicate arming the expiration timer and determining which connection IDs
/// should be retired.
const EXPIRATION_BUFFER: Duration = Duration::from_secs(30);

/// When a RETIRE_CONNECTION_ID frame is received, we don't want to remove the connection ID
/// immediately, since some packets with the just retired ID may still be received due to packet
/// reordering. The `RTT_MULTIPLIER` is multiplied by the smoothed RTT to calculate a removal
/// time that gives sufficient time for reordered packets to be processed.
const RTT_MULTIPLIER: u32 = 3;

/// A registration at the [`ConnectionIdMapper`].
///
/// It allows to add and remove external QUIC Connection IDs which are mapped to
/// internal IDs.
#[derive(Debug)]
pub struct LocalIdRegistry {
    /// The internal connection ID for this registration
    internal_id: InternalConnectionId,
    /// The shared state between mapper and registration
    state: Arc<Mutex<ConnectionIdMapperState>>,
    /// The connection IDs which are currently registered at the ConnectionIdMapper
    registered_ids: RegisteredIds,
    /// The sequence number to use the next time a new connection ID is registered
    next_sequence_number: u32,
    /// The current sequence number below which all connection IDs are considered retired
    retire_prior_to: u32,
    /// The maximum number of connection IDs to give to the peer
    active_connection_id_limit: u8,
    /// Memoized query to track retiring and expired connection IDs
    next_expiration: Memo<Option<Timestamp>, RegisteredIds>,
    /// Memoized query to track if there is any ACK interest
    ack_interest: Memo<bool, RegisteredIds>,
    /// Memoized query to track if there is any transmission interest
    transmission_interest: Memo<transmission::Interest, RegisteredIds>,
    /// Memoized query to track the number of active CIDs
    active_id_count: Memo<u8, RegisteredIds>,
}

type RegisteredIds = SmallVec<[LocalIdInfo; NR_STATIC_REGISTRABLE_IDS]>;

#[derive(Debug)]
struct LocalIdInfo {
    id: connection::LocalId,
    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
    //# Each Connection ID has an associated sequence number to assist in
    //# detecting when NEW_CONNECTION_ID or RETIRE_CONNECTION_ID frames refer
    //# to the same value.
    sequence_number: u32,
    retirement_time: Option<Timestamp>,
    stateless_reset_token: stateless_reset::Token,
    status: LocalIdStatus,
}

impl LocalIdInfo {
    // The time the connection ID next needs to change status (either to
    // PendingRetirementConfirmation, or removal altogether)
    fn next_status_change_time(&self) -> Option<Timestamp> {
        self.removal_time().or(self.retirement_time)
    }

    // Returns true if the connection ID should be moved to PendingRemoval
    fn is_retire_ready(&self, timestamp: Timestamp) -> bool {
        !self.is_retired()
            && self.retirement_time.map_or(false, |retirement_time| {
                retirement_time.has_elapsed(timestamp)
            })
    }

    // Returns true if the connection ID has been retired and is pending removal
    fn is_retired(&self) -> bool {
        matches!(
            self.status,
            PendingRetirementConfirmation(_) | PendingRemoval(_)
        )
    }

    // Returns true if the connection ID should no longer be used
    fn is_expired(&self, timestamp: Timestamp) -> bool {
        self.removal_time()
            .map_or(false, |removal_time| removal_time.has_elapsed(timestamp))
    }

    // The time this connection ID should be removed
    #[inline]
    fn removal_time(&self) -> Option<Timestamp> {
        match self.status {
            PendingRetirementConfirmation(removal_time) => removal_time,
            PendingRemoval(removal_time) => Some(removal_time),
            _ => None,
        }
    }

    // Changes the status of the connection ID to PendingRetirementConfirmation with a removal
    // time incorporating the EXPIRATION_BUFFER
    fn retire(&mut self, timestamp: Option<Timestamp>) {
        debug_assert!(!self.is_retired());
        self.status = PendingRetirementConfirmation(timestamp.map(|time| time + EXPIRATION_BUFFER))
    }

    #[inline]
    fn transmission_interest(&self) -> transmission::Interest {
        match self.status {
            PendingIssuance => transmission::Interest::NewData,
            PendingReissue => transmission::Interest::LostData,
            _ => transmission::Interest::None,
        }
    }

    /// Returns true if this status counts towards the active_connection_id_limit
    #[inline]
    fn counts_towards_limit(&self) -> bool {
        !matches!(
            self.status,
            PendingRetirementConfirmation(_) | PendingRemoval(_)
        )
    }
}

/// The current status of the connection ID.
#[derive(Debug, PartialEq, Eq)]
pub enum LocalIdStatus {
    /// New Connection IDs are put in the `PendingIssuance` status
    /// upon creation until a NEW_CONNECTION_ID frame has been sent
    /// to the peer to communicate the new connection ID.
    PendingIssuance,
    /// If the packet containing the NEW_CONNECTION_ID is lost, the
    /// connection ID is put into PendingReissue status.
    PendingReissue,
    /// Once a NEW_CONNECTION_ID frame is transmitted, the connection
    /// ID waits for acknowledgement of the packet.
    PendingAcknowledgement(PacketNumber),
    /// Once a connection ID has been communicated to the peer  and the
    /// peer has acknowledged the ID, it enters the `Active` status.
    /// The initial connection ID starts in this status.
    Active,
    /// Connection IDs are put in the `PendingRetirementConfirmation` status
    /// upon retirement, until confirmation of the retirement
    /// is received from the peer. If the optional removal_time indicated in this status
    /// is exceeded, the connection ID will be removed without confirmation from
    /// the peer.
    PendingRetirementConfirmation(Option<Timestamp>),
    /// Connection IDs are put in the `PendingRemoval` status
    /// when the peer has confirmed the retirement by sending a
    /// RETIRE_CONNECTION_ID_FRAME. This status exists to allow for a brief
    /// period before the Id is removed to account for packet reordering.
    PendingRemoval(Timestamp),
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum LocalIdRegistrationError {
    /// The Connection ID had already been registered
    ConnectionIdInUse,
    /// An invalid sequence number was specified
    InvalidSequenceNumber,
}

impl LocalIdRegistrationError {
    pub fn message(&self) -> &'static str {
        match self {
            LocalIdRegistrationError::ConnectionIdInUse => "Connection ID already in use",
            LocalIdRegistrationError::InvalidSequenceNumber => "Invalid sequence number",
        }
    }
}

impl Drop for LocalIdRegistry {
    fn drop(&mut self) {
        if let Ok(mut guard) = self.state.lock() {
            // Unregister all previously registered IDs
            for id_info in &self.registered_ids {
                guard.local_id_map.remove(&id_info.id);
            }

            // Also clean up the initial ID if it had not already been removed
            guard.initial_id_map.remove(&self.internal_id);
        }
    }
}

impl LocalIdRegistry {
    /// Constructs a new `LocalIdRegistry` and registers the provided `handshake_connection_id`
    pub(crate) fn new(
        internal_id: InternalConnectionId,
        state: Arc<Mutex<ConnectionIdMapperState>>,
        handshake_connection_id: &connection::LocalId,
        handshake_connection_id_expiration_time: Option<Timestamp>,
        stateless_reset_token: stateless_reset::Token,
    ) -> Self {
        let mut registry = Self {
            internal_id,
            state,
            registered_ids: SmallVec::new(),
            next_sequence_number: 0,
            retire_prior_to: 0,
            // Initialize to 1 until we know the actual limit
            // from the peer transport parameters
            active_connection_id_limit: 1,
            next_expiration: Memo::new(|ids| {
                ids.iter()
                    .filter_map(|id_info| id_info.next_status_change_time())
                    .min()
            }),
            ack_interest: Memo::new(|ids| {
                for id in ids.iter() {
                    if matches!(id.status, PendingAcknowledgement(_)) {
                        return true;
                    }
                }

                false
            }),
            transmission_interest: Memo::new(|ids| {
                let mut interest = transmission::Interest::None;

                for id_info in ids {
                    interest = interest.max(id_info.transmission_interest());
                }

                interest
            }),
            active_id_count: Memo::new(|ids| {
                let mut count = 0;
                for id in ids {
                    if id.counts_towards_limit() {
                        count += 1;
                    }
                }
                count
            }),
        };

        let _ = registry.register_connection_id(
            handshake_connection_id,
            handshake_connection_id_expiration_time,
            stateless_reset_token,
        );

        let handshake_connection_id_info = registry
            .registered_ids
            .iter_mut()
            .next()
            .expect("initial id added above");

        // The handshake connection ID is sent in the Initial packet,
        // so it starts in the `Active` status.
        handshake_connection_id_info.status = Active;
        registry.transmission_interest.clear();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
        //# The sequence number of the initial connection ID is 0.
        debug_assert_eq!(handshake_connection_id_info.sequence_number, 0);

        registry.check_consistency();

        registry
    }

    /// Returns the associated internal connection ID
    pub fn internal_connection_id(&self) -> InternalConnectionId {
        self.internal_id
    }

    /// Sets the active connection id limit
    pub fn set_active_connection_id_limit(&mut self, active_connection_id_limit: u64) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
        //# 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.
        self.active_connection_id_limit =
            MAX_ACTIVE_CONNECTION_ID_LIMIT.min(active_connection_id_limit) as u8;
    }

    /// Registers a connection ID mapping at the mapper with an optional expiration
    /// timestamp. Returns the sequence number of the connection ID.
    ///
    /// This will return an error if the provided ConnectionId has already been
    /// registered or is already used by a different internal connection.
    pub fn register_connection_id(
        &mut self,
        id: &connection::LocalId,
        expiration: Option<Timestamp>,
        stateless_reset_token: stateless_reset::Token,
    ) -> Result<(), LocalIdRegistrationError> {
        if self.registered_ids.iter().any(|id_info| id_info.id == *id) {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1
            //# As a trivial example, this means the same connection ID
            //# MUST NOT be issued more than once on the same connection.
            return Err(LocalIdRegistrationError::ConnectionIdInUse);
        }

        self.validate_new_connection_id(stateless_reset_token);

        // Try to insert into the global map
        self.state
            .lock()
            .expect("should succeed unless the lock is poisoned")
            .local_id_map
            .try_insert(id, self.internal_id)
            .map_err(|_| LocalIdRegistrationError::ConnectionIdInUse)?;

        let sequence_number = self.next_sequence_number;
        let retirement_time = expiration.map(|expiration| expiration - EXPIRATION_BUFFER);

        // Track the inserted connection ID info
        self.registered_ids.push(LocalIdInfo {
            id: *id,
            sequence_number,
            retirement_time,
            stateless_reset_token,
            status: PendingIssuance,
        });
        self.active_id_count.clear();
        self.transmission_interest.clear();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
        //# The sequence number on
        //# each newly issued connection ID MUST increase by 1.
        self.next_sequence_number += 1;

        // If we are provided an expiration, update the timers
        if expiration.is_some() {
            self.next_expiration.clear();
        }

        self.check_consistency();

        Ok(())
    }

    /// Unregisters connection IDs that have expired
    fn unregister_expired_ids(&mut self, timestamp: Timestamp) {
        {
            let mut mapper_state = self
                .state
                .lock()
                .expect("should succeed unless the lock is poisoned");

            self.registered_ids.retain(|id_info| {
                if id_info.is_expired(timestamp) {
                    let remove_result = mapper_state.local_id_map.remove(&id_info.id);
                    debug_assert!(
                        remove_result.is_some(),
                        "Connection ID should have been stored in mapper"
                    );

                    // clear all of the memoized values
                    self.ack_interest.clear();
                    self.transmission_interest.clear();
                    self.active_id_count.clear();
                    self.next_expiration.clear();

                    false // Don't retain
                } else {
                    true // Retain
                }
            });
        }

        self.check_consistency();
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
    //# 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
    //# The endpoint SHOULD continue to
    //# accept the previously issued connection IDs until they are retired by
    //# the peer.
    /// Handles the retirement of a sequence_number received from a RETIRE_CONNECTION_ID frame
    pub fn on_retire_connection_id(
        &mut self,
        sequence_number: u32,
        destination_connection_id: &connection::LocalId,
        rtt: Duration,
        timestamp: Timestamp,
    ) -> Result<(), LocalIdRegistrationError> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.16
        //# 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.
        if sequence_number >= self.next_sequence_number {
            return Err(LocalIdRegistrationError::InvalidSequenceNumber);
        }

        let id_info = self
            .registered_ids
            .iter_mut()
            // Filter out IDs that are already PendingRemoval, indicating this was a duplicate
            // RETIRE_CONNECTION_ID frame
            .filter(|id_info| !matches!(id_info.status, PendingRemoval(_)))
            .find(|id_info| id_info.sequence_number == sequence_number);

        if let Some(id_info) = id_info {
            if id_info.id == *destination_connection_id {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-19.16
                //# 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
                //# The peer MAY treat this as a
                //# connection error of type PROTOCOL_VIOLATION.
                return Err(LocalIdRegistrationError::InvalidSequenceNumber);
            }

            // Calculate a removal time based on RTT to give sufficient time for out of
            // order packets using the retired connection ID to be received
            let removal_time = timestamp + rtt * RTT_MULTIPLIER;

            id_info.status = PendingRemoval(removal_time);

            // clear all of the memoized values
            self.ack_interest.clear();
            self.transmission_interest.clear();
            self.active_id_count.clear();
            self.next_expiration.clear();
        }

        self.check_consistency();

        Ok(())
    }

    /// Returns the mappers interest in new connection IDs
    #[inline]
    pub fn connection_id_interest(&self) -> connection::id::Interest {
        let active_connection_id_count = self.active_id_count.get(&self.registered_ids);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
        //# 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-5.1.1
        //# An endpoint MUST NOT
        //# provide more connection IDs than the peer's limit.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
        //# An endpoint SHOULD supply a new connection ID when the peer retires a
        //# connection ID.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.5
        //# 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.
        let new_connection_id_count = self.active_connection_id_limit - active_connection_id_count;

        if new_connection_id_count > 0 {
            self.check_active_connection_id_limit(
                active_connection_id_count,
                new_connection_id_count,
            );
            connection::id::Interest::New(new_connection_id_count)
        } else {
            connection::id::Interest::None
        }
    }

    /// Handles timeouts on the registration
    ///
    /// `timestamp` passes the current time.
    pub fn on_timeout(&mut self, timestamp: Timestamp) {
        if self.timer().poll_expiration(timestamp).is_ready() {
            for id_info in self
                .registered_ids
                .iter_mut()
                .filter(|id_info| id_info.is_retire_ready(timestamp))
            {
                id_info.retire(Some(timestamp));
                self.retire_prior_to = self.retire_prior_to.max(id_info.sequence_number + 1);

                // clear all of the memoized values
                self.ack_interest.clear();
                self.transmission_interest.clear();
                self.active_id_count.clear();
                self.next_expiration.clear();
            }

            self.unregister_expired_ids(timestamp);
        }

        self.check_consistency();
    }

    /// Writes any NEW_CONNECTION_ID frames necessary to the given context
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        let constraint = context.transmission_constraint();

        if !self
            .transmission_interest
            .get(&self.registered_ids)
            .can_transmit(constraint)
        {
            return;
        }

        for mut id_info in self
            .registered_ids
            .iter_mut()
            .filter(|id_info| id_info.transmission_interest().can_transmit(constraint))
        {
            if let Some(packet_number) = context.write_frame(&frame::NewConnectionId {
                sequence_number: id_info.sequence_number.into(),
                retire_prior_to: self.retire_prior_to.into(),
                connection_id: id_info.id.as_bytes(),
                stateless_reset_token: id_info
                    .stateless_reset_token
                    .as_ref()
                    .try_into()
                    .expect("Length is already checked"),
            }) {
                id_info.status = PendingAcknowledgement(packet_number);
                self.transmission_interest.clear();
                self.ack_interest.clear();
            }
        }

        self.check_consistency();
    }

    /// Activates connection IDs that were pending acknowledgement
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        if !self.ack_interest.get(&self.registered_ids) {
            return;
        }

        for mut id_info in self.registered_ids.iter_mut() {
            if let PendingAcknowledgement(packet_number) = id_info.status {
                if ack_set.contains(packet_number) {
                    id_info.status = Active;
                    // Once the NEW_CONNECTION_ID is acknowledged, we don't need the
                    // stateless reset token anymore.
                    id_info.stateless_reset_token = stateless_reset::Token::ZEROED;

                    self.ack_interest.clear();
                }
            }
        }

        self.check_consistency();
    }

    /// Moves connection IDs pending acknowledgement into pending reissue
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        if !self.ack_interest.get(&self.registered_ids) {
            return;
        }

        for mut id_info in self.registered_ids.iter_mut() {
            if let PendingAcknowledgement(packet_number) = id_info.status {
                if ack_set.contains(packet_number) {
                    id_info.status = PendingReissue;

                    self.ack_interest.clear();
                    self.transmission_interest.clear();
                }
            }
        }

        self.check_consistency();
    }

    /// Requests the peer to retire the connection id used during the handshake
    pub fn retire_handshake_connection_id(&mut self) {
        if let Some(handshake_id_info) = self
            .registered_ids
            .iter_mut()
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
            //# The sequence number of the initial connection ID is 0.
            .find(|id_info| id_info.sequence_number == 0 && !id_info.is_retired())
        {
            // Request the peer to retire the handshake CID immediately by incrementing retire_prior_to,
            // but schedule the removal of the handshake CID for its regularly scheduled retirement
            // time, since some peers may not be capable of retiring the handshake CID. This allows the
            // handshake CID to remain in use until the peer is ready to retire it or the handshake CID
            // expires, whichever comes first. If the handshake CID does not have a retirement time,
            // it will not be removed unless the peer sends a RETIRE_CONNECTION_ID frame.
            handshake_id_info.retire(handshake_id_info.retirement_time);

            // The retire_prior_to number is incremented to trigger the peer to send a
            // RETIRE_CONNECTION_ID frame if possible.
            self.retire_prior_to = self
                .retire_prior_to
                .max(handshake_id_info.sequence_number + 1);

            self.active_id_count.clear();
            self.next_expiration.clear();
            self.transmission_interest.clear();
        }

        self.check_consistency();
    }

    /// Validate that the current expiration timer is based on the next status change time
    fn check_consistency(&self) {
        if cfg!(debug_assertions) {
            self.next_expiration.check_consistency(&self.registered_ids);
            self.ack_interest.check_consistency(&self.registered_ids);
            self.transmission_interest
                .check_consistency(&self.registered_ids);
            self.active_id_count.check_consistency(&self.registered_ids);
        }
    }

    fn check_active_connection_id_limit(&self, active_count: u8, new_count: u8) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
        //# 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.
        if cfg!(debug_assertions) {
            let retired_count = self
                .registered_ids
                .iter()
                .filter(|id_info| id_info.sequence_number < self.retire_prior_to)
                .count() as u8;
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
            //# 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.
            assert!(
                (active_count + new_count).saturating_sub(retired_count)
                    <= self.active_connection_id_limit
            );
        }
    }

    #[inline]
    fn timer(&self) -> Timer {
        Timer::from(self.next_expiration.get(&self.registered_ids))
    }

    fn validate_new_connection_id(&self, new_token: stateless_reset::Token) {
        if cfg!(debug_assertions) {
            let active_count = self.active_id_count.get(&self.registered_ids);
            assert!(
                active_count < self.active_connection_id_limit,
                "Attempted to register more connection IDs than the active connection id limit: {}",
                self.active_connection_id_limit
            );

            assert!(
                !self
                    .registered_ids
                    .iter()
                    .map(|id_info| id_info.stateless_reset_token)
                    .any(|token| token == new_token),
                "Registered a duplicate stateless reset token"
            );
        }
    }
}

impl timer::Provider for LocalIdRegistry {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        query.on_timer(&self.timer())?;
        Ok(())
    }
}

impl transmission::interest::Provider for LocalIdRegistry {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        let interest = self.transmission_interest.get(&self.registered_ids);
        query.on_interest(interest)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests;


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

//! This module contains the implementation of QUIC `Connections` and their management

use crate::{
    endpoint, path::MaxMtu, recovery::congestion_controller, space::PacketSpaceManager,
    wakeup_queue::WakeupHandle,
};
use s2n_quic_core::{connection, event, event::supervisor, time::Timestamp};

mod api;
mod api_provider;
mod close_sender;
mod connection_container;
mod connection_id_mapper;
mod connection_impl;
mod connection_interests;
mod connection_timers;
mod connection_trait;
mod errors;
pub(crate) mod finalization;
mod internal_connection_id;
pub(crate) mod local_id_registry;
pub(crate) mod open_token;
pub(crate) mod peer_id_registry;
pub(crate) mod transmission;

pub(crate) use api_provider::{ConnectionApi, ConnectionApiProvider};
pub(crate) use connection_container::{ConnectionContainer, ConnectionContainerIterationResult};
pub(crate) use connection_id_mapper::ConnectionIdMapper;
pub(crate) use connection_interests::ConnectionInterests;
pub(crate) use connection_timers::ConnectionTimers;
pub(crate) use connection_trait::ConnectionTrait as Trait;
pub(crate) use internal_connection_id::{InternalConnectionId, InternalConnectionIdGenerator};
pub(crate) use local_id_registry::LocalIdRegistry;
pub(crate) use peer_id_registry::PeerIdRegistry;
pub(crate) use transmission::{ConnectionTransmission, ConnectionTransmissionContext};

pub use api::Connection;
pub use connection_impl::ConnectionImpl as Implementation;
pub use connection_trait::Lock;
pub use open_token::Pair as OpenToken;
/// re-export core
pub use s2n_quic_core::connection::*;

/// Parameters which are passed to a Connection.
/// These are unique per created connection.
pub struct Parameters<'a, Cfg: endpoint::Config> {
    /// The [`Connection`]s internal identifier
    pub internal_connection_id: InternalConnectionId,
    /// The local ID registry which should be utilized by the connection
    pub local_id_registry: LocalIdRegistry,
    /// The peer ID registry which should be utilized by the connection
    pub peer_id_registry: PeerIdRegistry,
    /// The last utilized remote Connection ID
    pub peer_connection_id: PeerId,
    /// The last utilized local Connection ID
    pub local_connection_id: LocalId,
    /// The path handle on which the connection was created
    pub path_handle: Cfg::PathHandle,
    /// The space manager created for the connection
    pub space_manager: PacketSpaceManager<Cfg>,
    /// A struct which triggers a wakeup for the given connection
    ///
    /// This should be called from the application task
    pub wakeup_handle: WakeupHandle<InternalConnectionId>,
    /// The initial congestion controller for the connection
    pub congestion_controller: <Cfg::CongestionControllerEndpoint as congestion_controller::Endpoint>::CongestionController,
    /// The time the connection is being created
    pub timestamp: Timestamp,
    /// The QUIC protocol version which is used for this particular connection
    pub quic_version: u32,
    /// The limits that were advertised to the peer
    pub limits: connection::Limits,
    /// The largest maximum transmission unit (MTU) that can be sent on a path
    pub max_mtu: MaxMtu,
    /// The context that should be passed to all related connection events
    pub event_context: <Cfg::EventSubscriber as event::Subscriber>::ConnectionContext,
    /// The context passed to the connection supervisor
    pub supervisor_context: &'a supervisor::Context<'a>,
    // The datagram provider for the endpoint
    pub datagram_endpoint: &'a mut Cfg::DatagramEndpoint,
    /// The event subscriber for the endpoint
    pub event_subscriber: &'a mut Cfg::EventSubscriber,
}


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

use core::num::NonZeroU64;

/// An opaque token issued to each connection handle which allows the stream
/// controller to track any pending open requests.
///
/// Each connection handle must have a unique instance of a token so their wakers are correctly
/// tracked.
#[derive(Debug)]
pub struct Pair {
    /// Stores the token for the unidirectional stream type
    pub(crate) unidirectional: Token,
    /// Stores the token for the bididirectional stream type
    pub(crate) bidirectional: Token,
}

impl Pair {
    /// Creates a new open token
    ///
    /// This should be held on the connection handle and should be presented each time
    /// `poll_open_stream` is called.
    #[inline]
    pub const fn new() -> Self {
        Self {
            unidirectional: Token::new(),
            bidirectional: Token::new(),
        }
    }
}

#[derive(Debug)]
pub(crate) struct Token(Option<NonZeroU64>);

impl Token {
    #[inline]
    pub const fn new() -> Self {
        Self(None)
    }

    /// Returns the index that the caller's waker should be stored, given the expired token
    #[inline]
    pub fn index(&self, expired_token: &Self) -> Option<usize> {
        let base = expired_token.0.map_or(0, |v| v.get());
        // since we're dealing with NonZeroU64, add one to get the correct index
        let base = base + 1;
        let v = self.0?;
        let v = v.get().checked_sub(base)?;
        Some(v as usize)
    }

    /// Expires a `count` number of tokens
    ///
    /// This should be called each time at least one waker is woken and removed from the waker list
    #[inline]
    pub fn expire(&mut self, count: usize) {
        if let Some(v) = self.0.as_mut() {
            *v = unsafe {
                // Safety: non-zero N + count is always non-zero
                NonZeroU64::new_unchecked(v.get() + count as u64)
            };
        } else if let Some(v) = NonZeroU64::new(count as _) {
            *self = Self(Some(v));
        }
    }

    /// Resets the token state
    #[inline]
    pub fn clear(&mut self) {
        *self = Self(None);
    }
}

#[derive(Debug)]
pub(crate) struct Counter(NonZeroU64);

impl Counter {
    pub const fn new() -> Self {
        Self(unsafe {
            // Safety: 1 is always non-zero
            NonZeroU64::new_unchecked(1)
        })
    }

    /// Returns the next open token for a connection
    #[inline]
    pub fn next(&mut self) -> Token {
        let v = self.0;
        let next = Token(Some(v));
        self.0 = unsafe {
            // Safety: N + 1 is always non-zero
            NonZeroU64::new_unchecked(v.get() + 1)
        };
        next
    }
}

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

    #[test]
    fn open_token_test() {
        let mut caller_tokens = [Token::new(), Token::new(), Token::new()];
        let mut expired_token = Token::new();
        let mut token_counter = Counter::new();

        assert_eq!(
            caller_tokens[0].index(&expired_token),
            None,
            "empty caller token with no expired tokens"
        );

        // give the caller a token
        caller_tokens[0] = token_counter.next();
        assert_eq!(
            caller_tokens[0].index(&expired_token),
            Some(0),
            "issued caller with no expired tokens should be index 0"
        );

        // expire the issued token
        expired_token.expire(1);
        assert_eq!(
            caller_tokens[0].index(&expired_token),
            None,
            "issued caller token should now be expired"
        );

        for (idx, token) in caller_tokens.iter_mut().enumerate() {
            *token = token_counter.next();
            assert_eq!(
                token.index(&expired_token),
                Some(idx),
                "issued caller should be indexed in order",
            );
        }

        // expire just the first token
        expired_token.expire(1);

        assert_eq!(
            caller_tokens[0].index(&expired_token),
            None,
            "first caller token should now be expired"
        );

        for (idx, token) in caller_tokens[1..].iter().enumerate() {
            assert_eq!(
                token.index(&expired_token),
                Some(idx),
                "caller {} should not be expired",
                idx + 1,
            );
        }
    }
}


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

use crate::{
    connection::{
        connection_id_mapper::ConnectionIdMapperState,
        peer_id_registry::{
            PeerIdRegistrationError::{
                ExceededActiveConnectionIdLimit, ExceededRetiredConnectionIdLimit,
                InvalidNewConnectionId,
            },
            PeerIdStatus::{
                InUse, InUsePendingNewConnectionId, New, PendingAcknowledgement, PendingRetirement,
                PendingRetirementRetransmission,
            },
        },
        InternalConnectionId,
    },
    path,
    transmission::{self, WriteContext},
};
use s2n_quic_core::{
    ack, connection, endpoint,
    event::{self, IntoEvent},
    frame,
    memo::Memo,
    packet::number::PacketNumber,
    stateless_reset, transport,
};
use smallvec::SmallVec;
use std::sync::{Arc, Mutex};

/// The amount of ConnectionIds we can register without dynamic memory allocation
const NR_STATIC_REGISTRABLE_IDS: usize = 5;

//= https://www.rfc-editor.org/rfc/rfc9000#section-18.2
//# This is an integer value
//# specifying the maximum number of connection IDs from the peer that
//# an endpoint is willing to store.  This value includes the
//# connection ID received during the handshake, that received in the
//# preferred_address transport parameter, and those received in
//# NEW_CONNECTION_ID frames.  The value of the
//# active_connection_id_limit parameter MUST be at least 2.
// A value of 3 is sufficient for a client to probe multiple paths and for a server to respond
// to connection migrations from a client.
pub const ACTIVE_CONNECTION_ID_LIMIT: u8 = 3;

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
//# 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.
const RETIRED_CONNECTION_ID_LIMIT: u8 = ACTIVE_CONNECTION_ID_LIMIT * 2;

#[derive(Debug)]
pub struct PeerIdRegistry {
    /// The internal connection ID for this registration
    internal_id: InternalConnectionId,
    /// The shared state between mapper and registration
    state: Arc<Mutex<ConnectionIdMapperState>>,
    /// The connection IDs which are currently registered
    registered_ids: RegisteredIds,
    /// The largest retire prior to value that has been received from the peer
    retire_prior_to: u32,
    /// Memoized query to track if there is any ACK interest
    ack_interest: Memo<bool, RegisteredIds>,
    /// Memoized query to track if there is any transmission interest
    transmission_interest: Memo<transmission::Interest, RegisteredIds>,
}

type RegisteredIds = SmallVec<[PeerIdInfo; NR_STATIC_REGISTRABLE_IDS]>;

#[derive(Debug, Clone)]
struct PeerIdInfo {
    id: connection::PeerId,
    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
    //# Each Connection ID has an associated sequence number to assist in
    //# detecting when NEW_CONNECTION_ID or RETIRE_CONNECTION_ID frames refer
    //# to the same value.
    sequence_number: u32,
    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
    //# A 128-bit value that will be used for a stateless reset when the
    //# associated connection ID is used.
    stateless_reset_token: Option<stateless_reset::Token>,
    // The current status of the connection ID
    status: PeerIdStatus,
}

impl PeerIdInfo {
    /// Returns true if this ID is ready to be retired
    fn is_retire_ready(&self, retire_prior_to: u32) -> bool {
        self.is_active() && self.sequence_number < retire_prior_to
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
    //# 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.
    fn validate_new_connection_id(
        &self,
        new_id: &connection::PeerId,
        stateless_reset_token: &stateless_reset::Token,
        sequence_number: u32,
    ) -> Result<bool, PeerIdRegistrationError> {
        let reset_token_is_equal = self
            .stateless_reset_token
            .map_or(false, |token| token == *stateless_reset_token);
        let sequence_number_is_equal = self.sequence_number == sequence_number;

        if self.id == *new_id {
            if !reset_token_is_equal || !sequence_number_is_equal {
                return Err(InvalidNewConnectionId);
            }

            // This was a valid duplicate new connection ID
            return Ok(true);
        } else if sequence_number_is_equal || reset_token_is_equal {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.2
            //# 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.
            return Err(InvalidNewConnectionId);
        }

        // This was a valid non-duplicate new connection ID
        Ok(false)
    }

    /// Returns true if this PeerId may be used to send packets to the peer
    fn is_active(&self) -> bool {
        matches!(self.status, New | InUse | InUsePendingNewConnectionId)
    }

    /// Returns true if the status of this ID allows for transmission
    /// based on the transmission constraint
    fn transmission_interest(&self) -> transmission::Interest {
        match self.status {
            PendingRetirementRetransmission => transmission::Interest::LostData,
            PendingRetirement => transmission::Interest::NewData,
            _ => transmission::Interest::None,
        }
    }
}

/// The current status of the connection ID.
#[derive(Clone, Debug, PartialEq)]
enum PeerIdStatus {
    /// Connection IDs received in NEW_CONNECTION_ID frames start in the `New` status.
    New,
    /// Once a connection ID is used on a path it moves to the `InUse` status.
    InUse,
    /// The initial connection ID used during the handshake is in use, but will be retired
    /// as soon as a NEW_CONNECTION_ID frame is received from the peer.
    InUsePendingNewConnectionId,
    /// Once a connection ID will no longer be used, it enters the `PendingRetirement` status,
    /// triggering a RETIRE_CONNECTION_ID frame to be sent.
    PendingRetirement,
    /// If the packet that sent the RETIRE_CONNECTION_ID frame was declared lost, the connection
    /// moves to the `PendingRetirementRetransmission` status to allow for faster retransmission
    /// of the lost frame.
    PendingRetirementRetransmission,
    /// Once the RETIRE_CONNECTION_ID frame has been sent, the connection ID enters
    /// `PendingAcknowledgement` status, tracking the packet number of the packet that transmitted
    /// the retire frame. When acknowledgement of that packet is received, the connection ID is
    /// removed.
    PendingAcknowledgement(PacketNumber),
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum PeerIdRegistrationError {
    /// The NEW_CONNECTION_ID frame was invalid
    InvalidNewConnectionId,
    /// The active_connection_id_limit was exceeded
    ExceededActiveConnectionIdLimit,
    /// Too many connection IDs are pending retirement
    ExceededRetiredConnectionIdLimit,
}

impl PeerIdRegistrationError {
    fn message(&self) -> &'static str {
        match self {
            PeerIdRegistrationError::InvalidNewConnectionId => {
                "The new connection ID had an invalid sequence_number or stateless_reset_token"
            }
            PeerIdRegistrationError::ExceededActiveConnectionIdLimit => {
                "The active_connection_id_limit has been exceeded"
            }
            PeerIdRegistrationError::ExceededRetiredConnectionIdLimit => {
                "Too many connection IDs have been retired without acknowledgement from the peer"
            }
        }
    }
}

impl From<PeerIdRegistrationError> for transport::Error {
    fn from(err: PeerIdRegistrationError) -> Self {
        let transport_error = match err {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
            //# 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.
            PeerIdRegistrationError::InvalidNewConnectionId => transport::Error::PROTOCOL_VIOLATION,
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
            //# 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.
            PeerIdRegistrationError::ExceededActiveConnectionIdLimit => {
                transport::Error::CONNECTION_ID_LIMIT_ERROR
            }
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
            //# 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.
            PeerIdRegistrationError::ExceededRetiredConnectionIdLimit => {
                transport::Error::CONNECTION_ID_LIMIT_ERROR
            }
        };
        transport_error.with_reason(err.message())
    }
}

impl Drop for PeerIdRegistry {
    fn drop(&mut self) {
        if let Ok(mut guard) = self.state.lock() {
            // Stop tracking all associated stateless reset tokens
            for token in self
                .registered_ids
                .iter()
                .flat_map(|id_info| id_info.stateless_reset_token)
            {
                guard.stateless_reset_map.remove(&token);
            }
        }
    }
}

impl PeerIdRegistry {
    /// Constructs a new `PeerIdRegistry`.
    pub(crate) fn new(
        internal_id: InternalConnectionId,
        state: Arc<Mutex<ConnectionIdMapperState>>,
    ) -> Self {
        Self {
            internal_id,
            state,
            registered_ids: SmallVec::new(),
            retire_prior_to: 0,
            ack_interest: Memo::new(|ids| {
                for id in ids.iter() {
                    if matches!(id.status, PendingAcknowledgement(_)) {
                        return true;
                    }
                }

                false
            }),
            transmission_interest: Memo::new(|ids| {
                let mut interest = transmission::Interest::None;

                for id_info in ids {
                    interest = interest.max(id_info.transmission_interest());
                }

                interest
            }),
        }
    }

    /// Used to register the initial peer DestinationConnectionId.
    ///
    /// For a Server endpoint this happens immediately after creation of the
    /// PeerIdRegistry, since the ClientHello includes a SourceConnectionId.
    /// A Client endpoint must however wait for the initial Server response
    /// to populate this value.
    pub(crate) fn register_initial_connection_id(&mut self, peer_id: connection::PeerId) {
        debug_assert!(self.is_empty());

        self.registered_ids.push(PeerIdInfo {
            id: peer_id,
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
            //# The sequence number of the initial connection ID is 0.
            sequence_number: 0,
            stateless_reset_token: None,
            // Start the initial PeerId in ActivePendingNewConnectionId so the ID used
            // during the handshake is rotated as soon as the peer sends a new connection ID
            status: PeerIdStatus::InUsePendingNewConnectionId,
        });

        self.check_consistency();
    }

    /// Used to register the initial peer stateless reset token that applies to the connection ID
    /// the server selected during the handshake.
    ///
    /// This method is only used on the client to register a stateless token from a peer server,
    /// as clients cannot transmit a stateless reset token in their transport parameters due to
    /// lack of confidentiality protection.
    pub(crate) fn register_initial_stateless_reset_token(
        &mut self,
        stateless_reset_token: stateless_reset::Token,
    ) {
        debug_assert!(!self.is_empty());

        if let Some(peer_id_info) = self.registered_ids.get_mut(0) {
            debug_assert_eq!(None, peer_id_info.stateless_reset_token);
            peer_id_info.stateless_reset_token = Some(stateless_reset_token);
        }

        self.state
            .lock()
            .expect("should succeed unless the lock is poisoned")
            .stateless_reset_map
            .insert(stateless_reset_token, self.internal_id);
    }

    /// Check if registered_ids is empty.
    ///
    /// This is only expected to be true when an endpoint creates a new
    /// peer_id_registry.
    pub(crate) fn is_empty(&self) -> bool {
        self.registered_ids.is_empty()
    }

    /// Handles a new connection ID received from a NEW_CONNECTION_ID frame
    pub fn on_new_connection_id(
        &mut self,
        new_id: &connection::PeerId,
        sequence_number: u32,
        retire_prior_to: u32,
        stateless_reset_token: &stateless_reset::Token,
    ) -> Result<(), PeerIdRegistrationError> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
        //# A receiver MUST ignore any Retire Prior To fields that do not
        //# increase the largest received Retire Prior To value.
        self.retire_prior_to = self.retire_prior_to.max(retire_prior_to);

        let mut active_id_count = 0;
        let mut is_duplicate = false;
        let mut id_pending_new_connection_id = None;

        // Iterate over all registered IDs, retiring any as necessary
        for id_info in self.registered_ids.iter_mut() {
            is_duplicate |= id_info.validate_new_connection_id(
                new_id,
                stateless_reset_token,
                sequence_number,
            )?;

            if id_info.is_retire_ready(self.retire_prior_to) {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
                //# 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.
                id_info.status = PendingRetirement;
                self.transmission_interest.clear();
            }

            if id_info.is_active() {
                active_id_count += 1;
            }

            if id_pending_new_connection_id.is_none()
                && id_info.status == InUsePendingNewConnectionId
            {
                id_pending_new_connection_id = Some(id_info);
            }
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
        //# Receipt of the same frame multiple times MUST NOT be treated as a
        //# connection error.
        if !is_duplicate {
            let mut new_id_info = PeerIdInfo {
                id: *new_id,
                sequence_number,
                stateless_reset_token: Some(*stateless_reset_token),
                status: New,
            };

            //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
            //# An endpoint that receives a NEW_CONNECTION_ID frame with a sequence
            //# number smaller than the Retire Prior To field of a previously
            //# received NEW_CONNECTION_ID frame MUST send a corresponding
            //# RETIRE_CONNECTION_ID frame that retires the newly received connection
            //# ID, unless it has already done so for that sequence number.
            if new_id_info.is_retire_ready(self.retire_prior_to) {
                new_id_info.status = PendingRetirement;
                self.transmission_interest.clear();
            }

            if new_id_info.is_active() {
                active_id_count += 1;

                if let Some(mut id_pending_new_connection_id) = id_pending_new_connection_id {
                    // If there was an ID pending new connection ID, it can be moved to PendingRetirement
                    // now that we know we aren't processing a duplicate NEW_CONNECTION_ID and the
                    // new connection ID wasn't immediately retired.
                    id_pending_new_connection_id.status = PendingRetirement;
                    self.transmission_interest.clear();
                    // We retired one active connection ID
                    active_id_count -= 1;
                }
            }

            self.registered_ids.push(new_id_info);

            self.check_active_connection_id_limit(active_id_count)?;
        }

        // Duplicate NEW_CONNECTION_ID frames may not change the sequence number or
        // stateless reset token, but the RFC does not specify the behavior if the
        // retire prior to value changes. This means the number of retired connection
        // IDs may have changed even if the NEW_CONNECTION_ID frame was a duplicate,
        // so we will validate the retired id count regardless of the duplicate status.
        let retired_id_count = self.registered_ids.len() - active_id_count;

        self.check_retired_connection_id_limit(retired_id_count)?;

        self.check_consistency();

        Ok(())
    }

    /// Writes any RETIRE_CONNECTION_ID frames necessary to the given context
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        let constraint = context.transmission_constraint();

        if !self
            .transmission_interest
            .get(&self.registered_ids)
            .can_transmit(constraint)
        {
            return;
        }

        for id_info in self
            .registered_ids
            .iter_mut()
            .filter(|id_info| id_info.transmission_interest().can_transmit(constraint))
        {
            if let Some(packet_number) = context.write_frame(&frame::RetireConnectionId {
                sequence_number: id_info.sequence_number.into(),
            }) {
                id_info.status = PendingAcknowledgement(packet_number);
                self.transmission_interest.clear();
                self.ack_interest.clear();
            }
        }

        self.check_consistency();
    }

    /// Removes connection IDs that were pending acknowledgement
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        if !self.ack_interest.get(&self.registered_ids) {
            return;
        }

        let mut mapper_state = self
            .state
            .lock()
            .expect("should succeed unless the lock is poisoned");

        self.registered_ids.retain(|id_info| {
            if let PendingAcknowledgement(packet_number) = id_info.status {
                if ack_set.contains(packet_number) {
                    if let Some(token) = id_info.stateless_reset_token {
                        // Stop tracking the stateless reset token
                        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
                        //# 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.
                        mapper_state.stateless_reset_map.remove(&token);
                    }

                    self.ack_interest.clear();

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
                    //# An endpoint MUST NOT forget a connection ID without retiring it
                    // Don't retain the ID since the retirement was acknowledged
                    return false;
                }
            }

            // Retain IDs that weren't PendingAcknowledgement or weren't acknowledged
            true
        });

        self.check_consistency();
    }

    /// Sets the retransmit flag to true for connection IDs pending acknowledgement with a lost
    /// packet number
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        if !self.ack_interest.get(&self.registered_ids) {
            return;
        }

        for id_info in self.registered_ids.iter_mut() {
            if let PendingAcknowledgement(packet_number) = id_info.status {
                if ack_set.contains(packet_number) {
                    id_info.status = PendingRetirementRetransmission;
                    self.ack_interest.clear();
                    self.transmission_interest.clear();
                }
            }
        }

        self.check_consistency();
    }

    /// Checks if the peer_id exists and if it is active.
    pub fn is_active(&self, peer_id: &connection::PeerId) -> bool {
        self.registered_ids
            .iter()
            .any(|id_info| peer_id == &id_info.id && id_info.is_active())
    }

    /// Tries to consume a new peer_id if one is available.
    ///
    /// Register the stateless reset token once a connection ID is in use.
    fn consume_new_id_inner(&mut self) -> Option<connection::PeerId> {
        for id_info in self.registered_ids.iter_mut() {
            if id_info.status == New {
                // Start tracking the stateless reset token
                //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3.1
                //# 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.
                if let Some(token) = id_info.stateless_reset_token {
                    self.state
                        .lock()
                        .expect("should succeed unless the lock is poisoned")
                        .stateless_reset_map
                        .insert(token, self.internal_id);
                }

                // Consume the new id
                id_info.status = InUse;
                return Some(id_info.id);
            }
        }

        None
    }

    /// Tries to consume a new peer_id if one is available for an existing path.
    pub fn consume_new_id_for_existing_path<Pub: event::ConnectionPublisher>(
        &mut self,
        path_id: path::Id,
        current_peer_connection_id: connection::PeerId,
        publisher: &mut Pub,
    ) -> Option<connection::PeerId> {
        let new_id = self.consume_new_id_inner();
        if let Some(new_id) = new_id {
            debug_assert_ne!(current_peer_connection_id, new_id);

            publisher.on_connection_id_updated(event::builder::ConnectionIdUpdated {
                path_id: path_id.into_event(),
                cid_consumer: endpoint::Location::Local,
                previous: current_peer_connection_id.into_event(),
                current: new_id.into_event(),
            });
        }
        new_id
    }

    /// Tries to consume a new peer_id if one is available for a new path.
    pub fn consume_new_id_for_new_path(&mut self) -> Option<connection::PeerId> {
        self.consume_new_id_inner()
    }

    // Validate that the ACTIVE_CONNECTION_ID_LIMIT has not been exceeded
    fn check_active_connection_id_limit(
        &self,
        active_id_count: usize,
    ) -> Result<(), PeerIdRegistrationError> {
        debug_assert_eq!(
            active_id_count,
            self.registered_ids
                .iter()
                .filter(|id_info| id_info.is_active())
                .count()
        );

        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.1
        //# 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.
        if active_id_count > ACTIVE_CONNECTION_ID_LIMIT as usize {
            return Err(ExceededActiveConnectionIdLimit);
        }

        Ok(())
    }

    // Validate that the RETIRED_CONNECTION_ID_LIMIT has not been exceeded
    fn check_retired_connection_id_limit(
        &self,
        retired_id_count: usize,
    ) -> Result<(), PeerIdRegistrationError> {
        debug_assert_eq!(
            retired_id_count,
            self.registered_ids
                .iter()
                .filter(|id_info| !id_info.is_active())
                .count()
        );

        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
        //# 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.
        if retired_id_count > RETIRED_CONNECTION_ID_LIMIT as usize {
            return Err(ExceededRetiredConnectionIdLimit);
        }

        Ok(())
    }

    fn check_consistency(&self) {
        if cfg!(debug_assertions) {
            self.ack_interest.check_consistency(&self.registered_ids);
            self.transmission_interest
                .check_consistency(&self.registered_ids);

            let before_count = self.registered_ids.len();
            let mut registered_id_copy = self.registered_ids.to_vec();
            registered_id_copy.sort_by_key(|id_info| id_info.sequence_number);
            registered_id_copy.dedup_by_key(|id_info| id_info.sequence_number);
            assert_eq!(before_count, registered_id_copy.len());
            registered_id_copy.sort_by_key(|id_info| id_info.id);
            registered_id_copy.dedup_by_key(|id_info| id_info.id);
            assert_eq!(before_count, registered_id_copy.len());
            registered_id_copy.sort_by_key(|id_info| {
                id_info
                    .stateless_reset_token
                    .map(|token| token.into_inner())
            });
            registered_id_copy.dedup_by_key(|id_info| id_info.stateless_reset_token);
            assert_eq!(before_count, registered_id_copy.len());
        }
    }
}

impl transmission::interest::Provider for PeerIdRegistry {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        let interest = self.transmission_interest.get(&self.registered_ids);
        query.on_interest(interest)?;
        Ok(())
    }
}

#[cfg(test)]
mod tests;

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use crate::connection::{ConnectionIdMapper, InternalConnectionIdGenerator, PeerIdRegistry};
    use s2n_quic_core::{connection, endpoint, random, stateless_reset};

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

    // Helper function to easily create a PeerIdRegistry
    pub(crate) fn peer_registry(
        initial_id: connection::PeerId,
        stateless_reset_token: Option<stateless_reset::Token>,
    ) -> PeerIdRegistry {
        let mut random_generator = random::testing::Generator(123);

        let mut registry = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server)
            .create_client_peer_id_registry(InternalConnectionIdGenerator::new().generate_id());
        registry.register_initial_connection_id(initial_id);
        if let Some(token) = stateless_reset_token {
            registry.register_initial_stateless_reset_token(token);
        }
        registry
    }
}