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

use crate::endpoint;
use alloc::collections::VecDeque;
use core::time::Duration;
use s2n_codec::{Encoder, EncoderBuffer, EncoderValue};
use s2n_quic_core::{
    event,
    inet::ExplicitCongestionNotification,
    io::tx,
    packet,
    packet::ProtectedPacket,
    path::{self, MINIMUM_MTU},
};

#[derive(Debug)]
pub struct Negotiator<Config: endpoint::Config> {
    transmissions: VecDeque<Transmission<Config::PathHandle>>,
    max_peers: usize,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Error;

const SUPPORTED_VERSIONS: &[u32] = &[
    0x1, // Draft 34 / Version 1 (https://github.com/quicwg/base-drafts/wiki/21st-Implementation-Draft)
];

macro_rules! is_supported {
    ($packet:ident, $publisher:ident) => {{
        let supported = SUPPORTED_VERSIONS
            .iter()
            .cloned()
            .any(|v| v == $packet.version);

        if supported {
            //= https://tools.ietf.org/id/draft-marx-qlog-event-definitions-quic-h3-02#5.3.1
            //# Upon receiving a client initial with a supported version, the
            //# server logs this event with server_versions and chosen_version set
            $publisher.on_version_information(event::builder::VersionInformation {
                server_versions: &SUPPORTED_VERSIONS,
                client_versions: &[],
                chosen_version: Some($packet.version),
            });
        } else {
            //= https://tools.ietf.org/id/draft-marx-qlog-event-definitions-quic-h3-02#5.3.1
            //# Upon receiving a client initial with an unsupported version, the
            //# server logs this event with server_versions set and
            //# client_versions to the single-element array containing the
            //# client's attempted version.  The absence of chosen_version implies
            //# no overlap was found.
            $publisher.on_version_information(event::builder::VersionInformation {
                server_versions: &SUPPORTED_VERSIONS,
                client_versions: &[$packet.version],
                chosen_version: None,
            });
        }

        supported
    }};
}

impl<Config: endpoint::Config> Default for Negotiator<Config> {
    fn default() -> Self {
        Self::new(endpoint::DEFAULT_MAX_PEERS)
    }
}

impl<Config: endpoint::Config> Negotiator<Config> {
    pub fn new(max_peers: usize) -> Self {
        Self {
            transmissions: if Config::ENDPOINT_TYPE.is_server() {
                VecDeque::with_capacity(max_peers)
            } else {
                VecDeque::new()
            },
            max_peers,
        }
    }

    pub fn on_packet<Pub: event::EndpointPublisher>(
        &mut self,
        path: &Config::PathHandle,
        payload_len: usize,
        packet: &ProtectedPacket,
        publisher: &mut Pub,
    ) -> Result<(), Error> {
        // always forward packets for clients on to connections
        if Config::ENDPOINT_TYPE.is_client() {
            return Ok(());
        }

        let packet = match packet {
            ProtectedPacket::Initial(packet) => {
                if is_supported!(packet, publisher) {
                    return Ok(());
                }
                packet
            }
            ProtectedPacket::ZeroRtt(packet) => {
                if is_supported!(packet, publisher) {
                    return Ok(());
                }

                //= https://www.rfc-editor.org/rfc/rfc9000#section-6.1
                //# a server that is able to recognize packets as
                //# 0-RTT might choose not to send Version Negotiation packets in
                //# response to 0-RTT packets with the expectation that it will
                //# eventually receive an Initial packet.
                return Err(Error);
            }
            ProtectedPacket::VersionNegotiation(_packet) => {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-6.1
                //# An endpoint MUST NOT send a Version Negotiation packet
                //# in response to receiving a Version Negotiation packet.
                return Ok(());
            }
            _ => return Ok(()),
        };

        //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.2
        //# If a server receives a packet that indicates an unsupported version
        //# and if the packet is large enough to initiate a new connection for
        //# any supported version, the server SHOULD send a Version Negotiation
        //# packet as described in Section 6.1.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-6
        //# A server might not send a Version
        //# Negotiation packet if the datagram it receives is smaller than the
        //# minimum size specified in a different version;
        if payload_len < (MINIMUM_MTU as usize) {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.2
            //# Servers MUST
            //# drop smaller packets that specify unsupported versions.
            return Err(Error);
        }

        {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.2
            //# A server MAY limit the number of packets
            //# to which it responds with a Version Negotiation packet.

            //= https://www.rfc-editor.org/rfc/rfc9000#section-6.1
            //# A server MAY limit the number of Version Negotiation packets it
            //# sends.

            // store the peer's address if we're not at capacity
            if self.transmissions.len() != self.max_peers {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-5.2.2
                //# Servers SHOULD respond with a Version
                //# Negotiation packet, provided that the datagram is sufficiently long.
                self.transmissions
                    .push_back(Transmission::new(*path, packet));
            }
        }

        Err(Error)
    }

    pub fn on_transmit<
        Tx: tx::Queue<Handle = Config::PathHandle>,
        Pub: event::EndpointPublisher,
    >(
        &mut self,
        queue: &mut Tx,
        publisher: &mut Pub,
    ) {
        // clients don't transmit version negotiation packets
        if Config::ENDPOINT_TYPE.is_client() {
            return;
        }

        while let Some(transmission) = self.transmissions.pop_front() {
            match queue.push(&transmission) {
                Ok(tx::Outcome { len, .. }) => {
                    publisher.on_endpoint_packet_sent(event::builder::EndpointPacketSent {
                        packet_header: event::builder::PacketHeader::VersionNegotiation {},
                    });

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

struct Transmission<Path: path::Handle> {
    path: Path,
    // The MINIMUM_MTU size allows for at least 170 supported versions
    packet: [u8; MINIMUM_MTU as usize],
    packet_len: usize,
}

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

impl<Path: path::Handle> Transmission<Path> {
    pub fn new(path: Path, initial_packet: &packet::initial::ProtectedInitial) -> Self {
        let mut packet_buf = [0u8; MINIMUM_MTU as usize];
        let version_packet = packet::version_negotiation::VersionNegotiation::from_initial(
            initial_packet,
            SupportedVersions,
        );

        let mut buffer = EncoderBuffer::new(&mut packet_buf);
        version_packet.encode(&mut buffer);
        let packet_len = buffer.len();

        Self {
            path,
            packet: packet_buf,
            packet_len,
        }
    }
}

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

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

    fn path_handle(&self) -> &Self::Handle {
        &self.path
    }

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

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

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

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

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

#[derive(Clone, Copy, Debug)]
pub struct SupportedVersions;

impl EncoderValue for SupportedVersions {
    fn encode<E: Encoder>(&self, encoder: &mut E) {
        for version in SUPPORTED_VERSIONS {
            encoder.encode(version);
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-6.3
        //# Endpoints MAY add reserved versions to any field where unknown or
        //# unsupported versions are ignored to test that a peer correctly
        //# ignores the value.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-6.3
        //# Endpoints MAY send packets with a reserved version to test that a
        //# peer correctly discards the packet.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-6.3
        //# For a server to use a new version in the future, clients need to
        //# correctly handle unsupported versions.  Some version numbers
        //# (0x?a?a?a?a as defined in Section 15) are reserved for inclusion in
        //# fields that contain version numbers.
        encoder.encode(&0xdadadadau32);
    }
}

#[cfg(any(test, feature = "testing"))]
mod tests {
    use super::*;
    use crate::endpoint::testing;
    use core::mem::size_of;
    use s2n_codec::{DecoderBufferMut, Encoder, EncoderBuffer};
    use s2n_quic_core::{
        connection,
        connection::id::ConnectionInfo,
        event::testing::Publisher,
        inet::{DatagramInfo, SocketAddress},
        packet::{
            handshake::Handshake,
            initial::Initial,
            number::{PacketNumberSpace, TruncatedPacketNumber},
            short::Short,
            version_negotiation::VersionNegotiation,
            zero_rtt::ZeroRtt,
        },
        path::RemoteAddress,
        time::clock::testing as time,
        varint::VarInt,
    };

    type Server = Negotiator<testing::Server>;
    type Client = Negotiator<testing::Client>;

    // TAG to append to packet payloads to ensure they meet the minimum packet size
    // that would be expected had they undergone packet protection.
    const DUMMY_TAG: [u8; 16] = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15];

    fn datagram_info(payload_len: usize) -> (RemoteAddress, DatagramInfo) {
        (
            RemoteAddress::from(SocketAddress::default()),
            DatagramInfo {
                timestamp: time::now(),
                payload_len,
                ecn: Default::default(),
                destination_connection_id: connection::LocalId::TEST_ID,
                source_connection_id: None,
            },
        )
    }

    macro_rules! on_packet {
        (
            $negotiator:ident,
            $publisher:ident,
            $remote_address:expr,
            $payload_len:expr,
            $packet:expr
        ) => {{
            let mut buffer = vec![0u8; 1200];
            let mut encoder = EncoderBuffer::new(&mut buffer);

            encoder.encode(&$packet);

            let len = encoder.len();
            let decoder = DecoderBufferMut::new(&mut buffer[..len]);
            let remote_address = SocketAddress::default();
            let connection_info = ConnectionInfo::new(&remote_address);
            let (packet, _) = ProtectedPacket::decode(decoder, &connection_info, &3).unwrap();
            $negotiator.on_packet(&$remote_address, $payload_len, &packet, $publisher)
        }};
    }

    fn on_initial_packet<C: endpoint::Config>(
        datagram_info: (C::PathHandle, DatagramInfo),
        version: u32,
        negotiator: &mut Negotiator<C>,
        publisher: &mut Publisher,
    ) -> Result<(), Error> {
        on_packet!(
            negotiator,
            publisher,
            datagram_info.0,
            datagram_info.1.payload_len,
            Initial {
                version,
                destination_connection_id: &[1u8, 2, 3][..],
                source_connection_id: &[4u8, 5, 6, 7][..],
                token: &[0u8; 0][..],
                packet_number: pn(PacketNumberSpace::Initial),
                payload: &[1u8, 2, 3, 4, 5][..],
            }
        )
    }

    fn on_future_version_initial_packet<C: endpoint::Config>(
        datagram_info: (C::PathHandle, DatagramInfo),
        version: u32,
        negotiator: &mut Negotiator<C>,
        publisher: &mut Publisher,
    ) -> Result<(), Error> {
        let mut payload = vec![1u8, 2, 3, 4, 5];
        payload.extend_from_slice(&DUMMY_TAG[..]);

        on_packet!(
            negotiator,
            publisher,
            datagram_info.0,
            datagram_info.1.payload_len,
            Initial {
                version,
                // Maximum length connection IDs that may be valid in future versions
                destination_connection_id: &[1u8; size_of::<u8>()][..],
                source_connection_id: &[2u8; size_of::<u8>()][..],
                token: &[0u8; 0][..],
                packet_number: pn(PacketNumberSpace::Initial),
                payload: payload.as_slice(),
            }
        )
    }

    fn on_handshake_packet<C: endpoint::Config>(
        datagram_info: (C::PathHandle, DatagramInfo),
        version: u32,
        negotiator: &mut Negotiator<C>,
        publisher: &mut Publisher,
    ) -> Result<(), Error> {
        let mut payload = vec![1u8, 2, 3, 4, 5];
        payload.extend_from_slice(&DUMMY_TAG[..]);

        on_packet!(
            negotiator,
            publisher,
            datagram_info.0,
            datagram_info.1.payload_len,
            Handshake {
                version,
                destination_connection_id: &[1u8, 2, 3][..],
                source_connection_id: &[4u8, 5, 6][..],
                packet_number: pn(PacketNumberSpace::ApplicationData),
                payload: payload.as_slice(),
            }
        )
    }

    fn on_zerortt_packet<C: endpoint::Config>(
        datagram_info: (C::PathHandle, DatagramInfo),
        version: u32,
        negotiator: &mut Negotiator<C>,
        publisher: &mut Publisher,
    ) -> Result<(), Error> {
        let mut payload = vec![1u8, 2, 3, 4, 5];
        payload.extend_from_slice(&DUMMY_TAG[..]);

        on_packet!(
            negotiator,
            publisher,
            datagram_info.0,
            datagram_info.1.payload_len,
            ZeroRtt {
                version,
                destination_connection_id: &[1u8, 2, 3][..],
                source_connection_id: &[4u8, 5, 6][..],
                packet_number: pn(PacketNumberSpace::ApplicationData),
                payload: payload.as_slice(),
            }
        )
    }

    fn on_version_negotiation_packet<C: endpoint::Config>(
        datagram_info: (C::PathHandle, DatagramInfo),
        negotiator: &mut Negotiator<C>,
        publisher: &mut Publisher,
    ) -> Result<(), Error> {
        on_packet!(
            negotiator,
            publisher,
            datagram_info.0,
            datagram_info.1.payload_len,
            VersionNegotiation {
                tag: 0,
                destination_connection_id: &[1u8, 2, 3][..],
                source_connection_id: &[4u8, 5, 6][..],
                supported_versions: SupportedVersions,
            }
        )
    }

    fn on_short_packet<C: endpoint::Config>(
        datagram_info: (C::PathHandle, DatagramInfo),
        negotiator: &mut Negotiator<C>,
        publisher: &mut Publisher,
    ) -> Result<(), Error> {
        let mut payload = vec![1u8, 2, 3, 4, 5];
        payload.extend_from_slice(&DUMMY_TAG[..]);

        on_packet!(
            negotiator,
            publisher,
            datagram_info.0,
            datagram_info.1.payload_len,
            Short {
                destination_connection_id: &[1u8, 2, 3][..],
                key_phase: Default::default(),
                spin_bit: Default::default(),
                packet_number: pn(PacketNumberSpace::ApplicationData),
                payload: payload.as_slice(),
            }
        )
    }

    fn pn(space: PacketNumberSpace) -> TruncatedPacketNumber {
        let pn = space.new_packet_number(VarInt::default());
        pn.truncate(pn).unwrap()
    }

    const INVALID_VERSION: u32 = 123;

    #[test]
    fn client_test() {
        let mut client = Client::default();
        let mut publisher = Publisher::snapshot();

        assert_eq!(
            on_initial_packet(
                datagram_info(1200),
                INVALID_VERSION,
                &mut client,
                &mut publisher
            ),
            Ok(()),
            "client implementation should always allow version negotiator packets through",
        );

        assert_eq!(
            on_zerortt_packet(
                datagram_info(1200),
                INVALID_VERSION,
                &mut client,
                &mut publisher
            ),
            Ok(()),
            "client implementation should always allow version negotiator packets through",
        );

        assert!(
            client.transmissions.is_empty(),
            "clients should never negotiate at the endpoint level"
        );
    }

    #[test]
    fn server_initial_test() {
        let mut server = Server::default();
        let mut publisher = Publisher::snapshot();

        assert_eq!(
            on_initial_packet(
                datagram_info(1200),
                INVALID_VERSION,
                &mut server,
                &mut publisher
            ),
            Err(Error),
            "server implementations should error on invalid versions"
        );

        assert!(
            !server.transmissions.is_empty(),
            "servers should negotiate with initial packets"
        );
    }

    #[test]
    fn server_future_version_initial_test() {
        let mut server = Server::default();
        let mut publisher = Publisher::snapshot();

        assert_eq!(
            on_future_version_initial_packet(
                datagram_info(1200),
                INVALID_VERSION,
                &mut server,
                &mut publisher
            ),
            Err(Error),
            "server implementations should error on invalid versions"
        );

        assert!(
            !server.transmissions.is_empty(),
            "servers should negotiate with initial packets"
        );
    }

    #[test]
    fn server_zerortt_test() {
        let mut server = Server::default();
        let mut publisher = Publisher::snapshot();

        assert_eq!(
            on_zerortt_packet(
                datagram_info(1200),
                INVALID_VERSION,
                &mut server,
                &mut publisher
            ),
            Err(Error),
            "server implementations should error on invalid versions"
        );

        assert!(
            server.transmissions.is_empty(),
            "servers should not negotiate with zero_rtt packets"
        );
    }

    #[test]
    fn server_undersized_test() {
        let mut server = Server::default();
        let mut publisher = Publisher::snapshot();

        assert_eq!(
            on_initial_packet(
                datagram_info(32),
                INVALID_VERSION,
                &mut server,
                &mut publisher
            ),
            Err(Error),
            "server implementations should error on invalid versions"
        );

        assert!(
            server.transmissions.is_empty(),
            "servers should not negotiate with undersized packets"
        );
    }

    #[test]
    fn server_max_peers_test() {
        let mut server = Server::new(2);
        let mut publisher = Publisher::snapshot();

        for _ in 0..5 {
            assert_eq!(
                on_initial_packet(
                    datagram_info(1200),
                    INVALID_VERSION,
                    &mut server,
                    &mut publisher
                ),
                Err(Error),
                "server implementations should error on invalid versions"
            );
        }

        assert_eq!(
            server.transmissions.len(),
            2,
            "servers should not negotiate with more than the allowed max_peers"
        );
    }

    #[test]
    fn server_other_packets_test() {
        let mut server = Server::default();
        let mut publisher = Publisher::snapshot();

        assert_eq!(
            on_version_negotiation_packet(datagram_info(1200), &mut server, &mut publisher),
            Ok(()),
            "server implementations should allow other packets through to connections"
        );

        assert_eq!(
            on_handshake_packet(
                datagram_info(1200),
                INVALID_VERSION,
                &mut server,
                &mut publisher
            ),
            Ok(()),
            "server implementations should allow other packets through to connections"
        );

        assert_eq!(
            on_short_packet(datagram_info(1200), &mut server, &mut publisher),
            Ok(()),
            "server implementations should allow other packets through to connections"
        );

        assert!(
            server.transmissions.is_empty(),
            "servers should not negotiate with version negotiation packets"
        );
    }
}


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

use crate::{contexts::WriteContext, transmission};
use s2n_quic_core::{
    ct::ConstantTimeEq,
    event, frame,
    time::{timer, Duration, Timer, Timestamp},
};

pub type Data = [u8; frame::path_challenge::DATA_LEN];
const DISABLED_DATA: Data = [0; frame::path_challenge::DATA_LEN];

#[derive(Clone, Debug)]
pub struct Challenge {
    state: State,
    abandon_duration: Duration,
    abandon_timer: Timer,
    data: Data,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum State {
    /// PATH_CHALLENGE is not used for path validation when initiating a new connection
    InitialPathDisabled,

    /// A Challenge frame must be sent. The `u8` represents the remaining number of retries
    RequiresTransmission(u8),

    /// Challenge has been sent and we are awaiting a response until the abandon timer expires
    PendingResponse,

    /// The Challenge has been abandoned due to the abandon_timer
    Abandoned,

    /// When the PATH_CHALLENGE was validated by a PATH_RESPONSE
    Validated,
}

impl transmission::interest::Provider for State {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        match self {
            State::RequiresTransmission(_) => query.on_interest(transmission::Interest::NewData),
            _ => Ok(()),
        }
    }
}

impl Challenge {
    pub fn new(abandon_duration: Duration, data: Data) -> Self {
        Self {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
            //# An endpoint SHOULD NOT probe a new path with packets containing a
            //# PATH_CHALLENGE frame more frequently than it would send an Initial
            //# packet.

            //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
            //# An endpoint MAY send multiple PATH_CHALLENGE frames to guard against
            //# packet loss.

            // Re-transmitting twice guards against packet loss, while remaining
            // below the amplification limit of 3.
            state: State::RequiresTransmission(2),
            abandon_duration,
            abandon_timer: Timer::default(),
            data,
        }
    }

    pub fn disabled() -> Self {
        Self {
            state: State::InitialPathDisabled,
            abandon_duration: Duration::ZERO,
            abandon_timer: Timer::default(),
            data: DISABLED_DATA,
        }
    }

    /// When a PATH_CHALLENGE is transmitted this handles any internal state operations.
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        match self.state {
            State::RequiresTransmission(0) => self.state = State::PendingResponse,
            State::RequiresTransmission(remaining) => {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
                //# However, an endpoint SHOULD NOT send multiple
                //# PATH_CHALLENGE frames in a single packet.
                let frame = frame::PathChallenge { data: &self.data };

                if context.write_frame(&frame).is_some() {
                    let remaining = remaining - 1;
                    self.state = State::RequiresTransmission(remaining);

                    if !self.abandon_timer.is_armed() {
                        self.abandon_timer
                            .set(context.current_time() + self.abandon_duration);
                    }
                }
            }
            _ => {}
        }
    }

    pub fn on_timeout<Pub: event::ConnectionPublisher>(
        &mut self,
        timestamp: Timestamp,
        publisher: &mut Pub,
        path: event::builder::Path,
    ) {
        if self.abandon_timer.poll_expiration(timestamp).is_ready() {
            self.abandon(publisher, path);
        }
    }

    pub fn abandon<Pub: event::ConnectionPublisher>(
        &mut self,
        publisher: &mut Pub,
        path: event::builder::Path,
    ) {
        if self.is_pending() {
            self.state = State::Abandoned;
            self.abandon_timer.cancel();
            publisher.on_path_challenge_updated(event::builder::PathChallengeUpdated {
                path_challenge_status: event::builder::PathChallengeStatus::Abandoned,
                path,
                challenge_data: self.challenge_data(),
            });
        }
    }

    pub fn is_disabled(&self) -> bool {
        matches!(self.state, State::InitialPathDisabled)
    }

    pub fn is_pending(&self) -> bool {
        matches!(
            self.state,
            State::PendingResponse | State::RequiresTransmission(_)
        )
    }

    pub fn on_validated(&mut self, data: &[u8]) -> bool {
        if self.is_pending() && ConstantTimeEq::ct_eq(&self.data[..], data).into() {
            self.state = State::Validated;
            true
        } else {
            false
        }
    }

    pub fn challenge_data(&self) -> &[u8] {
        &self.data
    }
}

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

        Ok(())
    }
}

impl transmission::interest::Provider for Challenge {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.state.transmission_interest(query)
    }
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use super::*;
    use s2n_quic_core::time::{Clock, Duration, NoopClock};

    pub fn helper_challenge() -> Helper {
        let now = NoopClock {}.get_time();
        let abandon_duration = Duration::from_millis(10_000);
        let expected_data: [u8; 8] = [0; 8];

        let challenge = Challenge::new(abandon_duration, expected_data);

        Helper {
            now,
            abandon_duration,
            expected_data,
            challenge,
        }
    }

    #[allow(dead_code)]
    pub struct Helper {
        pub now: Timestamp,
        pub abandon_duration: Duration,
        pub expected_data: Data,
        pub challenge: Challenge,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::contexts::testing::{MockWriteContext, OutgoingFrameBuffer};
    use s2n_quic_core::{
        endpoint,
        time::{Clock, Duration, NoopClock},
    };
    use testing::*;

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
    //= type=test
    //# An endpoint MAY send multiple PATH_CHALLENGE frames to guard against
    //# packet loss.
    #[test]
    fn create_challenge_that_requires_two_transmissions() {
        let helper = helper_challenge();
        assert_eq!(helper.challenge.state, State::RequiresTransmission(2));
    }

    #[test]
    fn create_disabled_challenge() {
        let challenge = Challenge::disabled();
        assert_eq!(challenge.state, State::InitialPathDisabled);
        assert!(!challenge.abandon_timer.is_armed());
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
    //= type=test
    //# An endpoint SHOULD NOT probe a new path with packets containing a
    //# PATH_CHALLENGE frame more frequently than it would send an Initial
    //# packet.
    #[test]
    fn transmit_challenge_only_twice() {
        // Setup:
        let mut helper = helper_challenge();
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        assert_eq!(helper.challenge.state, State::RequiresTransmission(2));

        // Trigger:
        helper.challenge.on_transmit(&mut context);

        // Expectation:
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
        //= type=test
        //# However, an endpoint SHOULD NOT send multiple
        //# PATH_CHALLENGE frames in a single packet.
        assert_eq!(context.frame_buffer.len(), 1);

        assert_eq!(helper.challenge.state, State::RequiresTransmission(1));
        let written_data = match context.frame_buffer.pop_front().unwrap().as_frame() {
            frame::Frame::PathChallenge(frame) => Some(*frame.data),
            _ => None,
        }
        .unwrap();
        assert_eq!(written_data, helper.expected_data);

        // Trigger:
        helper.challenge.on_transmit(&mut context);

        // Expectation:
        assert_eq!(helper.challenge.state, State::RequiresTransmission(0));
        let written_data = match context.frame_buffer.pop_front().unwrap().as_frame() {
            frame::Frame::PathChallenge(frame) => Some(*frame.data),
            _ => None,
        }
        .unwrap();
        assert_eq!(written_data, helper.expected_data);

        // Trigger:
        helper.challenge.on_transmit(&mut context);

        // Expectation:
        assert_eq!(helper.challenge.state, State::PendingResponse);
        assert_eq!(context.frame_buffer.len(), 0);
    }

    #[test]
    fn successful_on_transmit_arms_the_timer() {
        // Setup:
        let mut helper = helper_challenge();
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        assert_eq!(helper.challenge.state, State::RequiresTransmission(2));
        assert!(!helper.challenge.abandon_timer.is_armed());

        // Trigger:
        helper.challenge.on_transmit(&mut context);

        // Expectation:
        assert!(helper.challenge.abandon_timer.is_armed());
    }

    #[test]
    fn maintain_idle_and_dont_transmit_when_pending_response_state() {
        // Setup:
        let mut helper = helper_challenge();
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        helper.challenge.state = State::PendingResponse;
        assert_eq!(helper.challenge.state, State::PendingResponse);

        // Trigger:
        helper.challenge.on_transmit(&mut context);

        // Expectation:
        assert_eq!(helper.challenge.state, State::PendingResponse);
        assert_eq!(context.frame_buffer.len(), 0);
    }

    #[test]
    fn test_on_timeout() {
        let mut helper = helper_challenge();
        let expiration_time = helper.now + helper.abandon_duration;

        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        helper.challenge.on_transmit(&mut context);

        let mut publisher = event::testing::Publisher::snapshot();
        let path = event::builder::Path::test();

        helper.challenge.on_timeout(
            expiration_time - Duration::from_millis(10),
            &mut publisher,
            path,
        );
        assert!(helper.challenge.is_pending());

        let path = event::builder::Path::test();
        helper.challenge.on_timeout(
            expiration_time + Duration::from_millis(10),
            &mut publisher,
            path,
        );
        assert!(!helper.challenge.is_pending());
    }

    #[test]
    fn challenge_must_remains_abandoned_once_abandoned() {
        let mut helper = helper_challenge();
        let expiration_time = helper.now + helper.abandon_duration;

        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        helper.challenge.on_transmit(&mut context);

        let mut publisher = event::testing::Publisher::snapshot();
        let path = event::builder::Path::test();

        // Trigger:
        helper.challenge.on_timeout(
            expiration_time + Duration::from_millis(10),
            &mut publisher,
            path,
        );

        // Expectation:
        assert!(!helper.challenge.is_pending());

        let path = event::builder::Path::test();

        // Trigger:
        helper.challenge.on_timeout(
            expiration_time - Duration::from_millis(10),
            &mut publisher,
            path,
        );

        // Expectation:
        assert!(!helper.challenge.is_pending());
    }

    #[test]
    fn dont_abandon_disabled_state() {
        let mut challenge = Challenge::disabled();
        let now = NoopClock {}.get_time();

        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        challenge.on_transmit(&mut context);

        assert_eq!(challenge.state, State::InitialPathDisabled);

        let mut publisher = event::testing::Publisher::snapshot();
        let path = event::builder::Path::test();

        let large_expiration_time = now + Duration::from_secs(1_000_000);
        challenge.on_timeout(large_expiration_time, &mut publisher, path);
        assert_eq!(challenge.state, State::InitialPathDisabled);
    }

    #[test]
    fn test_is_abandoned() {
        let mut helper = helper_challenge();
        let expiration_time = helper.now + helper.abandon_duration;

        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        helper.challenge.on_transmit(&mut context);

        let mut publisher = event::testing::Publisher::snapshot();
        let path = event::builder::Path::test();

        assert!(helper.challenge.is_pending());

        helper.challenge.on_timeout(
            expiration_time + Duration::from_millis(10),
            &mut publisher,
            path,
        );
        assert!(!helper.challenge.is_pending());
    }

    #[test]
    fn test_on_validate() {
        let mut helper = helper_challenge();

        let wrong_data: [u8; 8] = [5; 8];
        assert!(!helper.challenge.on_validated(&wrong_data));
        assert!(helper.challenge.is_pending());

        assert!(helper.challenge.on_validated(&helper.expected_data));
        assert_eq!(helper.challenge.state, State::Validated);
    }

    #[test]
    fn is_disabled() {
        let challenge = Challenge::disabled();

        assert_eq!(challenge.state, State::InitialPathDisabled);
        assert!(challenge.is_disabled());
    }

    #[test]
    fn dont_validate_disabled_state() {
        let mut helper = helper_challenge();
        helper.challenge.state = State::InitialPathDisabled;

        assert!(!helper.challenge.on_validated(&helper.expected_data));
        assert_eq!(helper.challenge.state, State::InitialPathDisabled);
    }

    #[test]
    fn dont_abandon_a_validated_challenge() {
        let mut helper = helper_challenge();
        helper.challenge.state = State::Validated;
        let mut publisher = event::testing::Publisher::snapshot();
        let path = event::builder::Path::test();

        helper.challenge.abandon(&mut publisher, path);

        assert_eq!(helper.challenge.state, State::Validated);
    }

    #[test]
    fn cancel_abandon_timer_on_abandon() {
        let mut helper = helper_challenge();
        let mut publisher = event::testing::Publisher::snapshot();
        let path = event::builder::Path::test();

        helper.challenge.abandon(&mut publisher, path);

        assert!(!helper.challenge.abandon_timer.is_armed());
    }
}


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

use core::ops::RangeInclusive;
use s2n_quic_core::{
    counter::{Counter, Saturating},
    event,
    event::IntoEvent,
    frame::ack::EcnCounts,
    inet::ExplicitCongestionNotification,
    number::CheckedSub,
    random,
    time::{timer, Duration, Timer, Timestamp},
    transmission,
    varint::VarInt,
};

//= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2
//# If an endpoint has cause to expect that IP packets with an ECT codepoint
//# might be dropped by a faulty network element, the endpoint could set an
//# ECT codepoint for only the first ten outgoing packets on a path, or for
//# a period of three PTOs
const TESTING_PACKET_THRESHOLD: u8 = 10;

// After a failure has been detected, the ecn::Controller will wait this duration
// before testing for ECN support again.
const RETEST_COOL_OFF_DURATION: Duration = Duration::from_secs(60);

// The number of round trip times an ECN capable path will wait before transmitting an ECN-CE marked packet.
const CE_SUPPRESSION_TESTING_RTT_MULTIPLIER: RangeInclusive<u16> = 10..=100;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ValidationOutcome {
    /// The path is ECN capable and congestion was experienced
    ///
    /// Contains the incremental count of packets that experienced congestion
    CongestionExperienced(VarInt),
    /// The path failed validation
    Failed,
    /// The path passed validation
    Passed,
    /// Validation was not performed
    Skipped,
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum State {
    // ECN capability is being tested, tracking the number of ECN marked packets sent
    Testing(u8),
    // ECN capability has been tested, but not validated yet
    Unknown,
    // ECN validation has failed. Validation will be restarted based on the timer
    Failed(Timer),
    // ECN validation has succeeded. CE suppression will be tested based on the timer.
    Capable(Timer),
}

impl IntoEvent<event::builder::EcnState> for &State {
    #[inline]
    fn into_event(self) -> event::builder::EcnState {
        match self {
            State::Testing(_) => event::builder::EcnState::Testing,
            State::Unknown => event::builder::EcnState::Unknown,
            State::Failed(_) => event::builder::EcnState::Failed,
            State::Capable(_) => event::builder::EcnState::Capable,
        }
    }
}

impl Default for State {
    fn default() -> Self {
        State::Testing(0)
    }
}

#[derive(Clone, Debug, Default)]
pub struct Controller {
    state: State,
    // A count of the number of packets with ECN marking lost since
    // the last time a packet with ECN marking was acknowledged.
    black_hole_counter: Counter<u8, Saturating>,
    // The largest acknowledged packet sent with an ECN marking. Used when tracking
    // packets that have been lost for the purpose of detecting a black hole.
    last_acked_ecn_packet_timestamp: Option<Timestamp>,
}

impl Controller {
    /// Restart testing of ECN capability
    pub fn restart<Pub: event::ConnectionPublisher>(
        &mut self,
        path: event::builder::Path,
        publisher: &mut Pub,
    ) {
        if self.state != State::Testing(0) {
            self.change_state(State::Testing(0), path, publisher);
        }
        self.black_hole_counter = Default::default();
    }

    /// Called when the connection timer expires
    pub fn on_timeout<Pub: event::ConnectionPublisher>(
        &mut self,
        now: Timestamp,
        path: event::builder::Path,
        random_generator: &mut dyn random::Generator,
        rtt: Duration,
        publisher: &mut Pub,
    ) {
        match self.state {
            State::Failed(ref mut retest_timer) => {
                if retest_timer.poll_expiration(now).is_ready() {
                    self.restart(path, publisher);
                }
            }
            State::Capable(ref mut ce_suppression_timer) if !ce_suppression_timer.is_armed() => {
                ce_suppression_timer
                    .set(now + Self::next_ce_packet_duration(random_generator, rtt));
            }
            State::Testing(_) | State::Unknown | State::Capable(_) => {}
        }
    }

    /// Gets the ECN marking to use on packets sent to the peer
    pub fn ecn(
        &mut self,
        transmission_mode: transmission::Mode,
        now: Timestamp,
    ) -> ExplicitCongestionNotification {
        if transmission_mode.is_loss_recovery_probing() {
            // Don't mark loss recovery probes as ECN capable in case the ECN
            // marking is causing packet loss
            return ExplicitCongestionNotification::NotEct;
        }

        match self.state {
            //= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.4
            //# On paths with a "testing" or "capable" state, the endpoint
            //# sends packets with an ECT marking -- ECT(0) by default;
            //# otherwise, the endpoint sends unmarked packets.
            State::Testing(_) => ExplicitCongestionNotification::Ect0,
            State::Capable(ref mut ce_suppression_timer) => {
                if ce_suppression_timer.poll_expiration(now).is_ready() {
                    //= https://www.rfc-editor.org/rfc/rfc9002#section-8.3
                    //# A sender can detect suppression of reports by marking occasional
                    //# packets that it sends with an ECN-CE marking.
                    ExplicitCongestionNotification::Ce
                } else {
                    //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.2
                    //# Upon successful validation, an endpoint MAY continue to set an ECT
                    //# codepoint in subsequent packets it sends, with the expectation that
                    //# the path is ECN-capable.
                    ExplicitCongestionNotification::Ect0
                }
            }
            //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.2
            //# If validation fails, then the endpoint MUST disable ECN. It stops setting the ECT
            //# codepoint in IP packets that it sends, assuming that either the network path or
            //# the peer does not support ECN.
            State::Failed(_) | State::Unknown => ExplicitCongestionNotification::NotEct,
        }
    }

    /// Returns a duration based on a randomly generated value in the CE_SUPPRESSION_TESTING_RTT_MULTIPLIER
    /// range multiplied by the given round trip time. This duration represents the amount of time
    /// to wait before an ECN-CE marked packet should be sent, to test if CE reports are being
    /// suppressed by the peer.
    ///
    /// Note: This function performs a modulo operation on the random generated bytes to restrict
    /// the result to the `CE_SUPPRESSION_TESTING_RTT_MULTIPLIER` range. This may introduce a modulo
    /// bias in the resulting count, but does not result in any reduction in security for this
    /// usage. Other usages that require uniform sampling should implement rejection sampling or
    /// other methodologies and not copy this implementation.
    fn next_ce_packet_duration(
        random_generator: &mut dyn random::Generator,
        rtt: Duration,
    ) -> Duration {
        let mut bytes = [0; core::mem::size_of::<u16>()];
        random_generator.public_random_fill(&mut bytes);
        let result = u16::from_le_bytes(bytes);

        let max_variance = (CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.end()
            - CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start())
        .saturating_add(1);
        let result = CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() + result % max_variance;
        result as u32 * rtt
    }

    /// Returns true if the path has been determined to be capable of handling ECN marked packets
    pub fn is_capable(&self) -> bool {
        matches!(self.state, State::Capable(_))
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.2
    //# Network routing and path elements can change mid-connection; an endpoint
    //# MUST disable ECN if validation later fails.
    /// Validate the given `EcnCounts`, updating the current validation state based on the
    /// validation outcome.
    ///
    /// * `newly_acked_ecn_counts` - total ECN counts that were sent on packets newly acknowledged by the peer
    /// * `sent_packet_ecn_counts` - total ECN counts for all outstanding packets, including those newly
    ///                              acknowledged during this validation
    /// * `baseline_ecn_counts` - the ECN counts present in the Ack frame the last time ECN counts were processed
    /// * `ack_frame_ecn_counts` - the ECN counts present in the current Ack frame (if any)
    /// * `now` - the time the Ack frame was received
    #[allow(clippy::too_many_arguments)]
    pub fn validate<Pub: event::ConnectionPublisher>(
        &mut self,
        newly_acked_ecn_counts: EcnCounts,
        sent_packet_ecn_counts: EcnCounts,
        baseline_ecn_counts: EcnCounts,
        ack_frame_ecn_counts: Option<EcnCounts>,
        now: Timestamp,
        rtt: Duration,
        path: event::builder::Path,
        publisher: &mut Pub,
    ) -> ValidationOutcome {
        if matches!(self.state, State::Failed(_)) {
            // Validation had already failed
            return ValidationOutcome::Skipped;
        }

        if ack_frame_ecn_counts.is_none() {
            if newly_acked_ecn_counts.as_option().is_some() {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.1
                //# If an ACK frame newly acknowledges a packet that the endpoint sent with
                //# either the ECT(0) or ECT(1) codepoint set, ECN validation fails if the
                //# corresponding ECN counts are not present in the ACK frame. This check
                //# detects a network element that zeroes the ECN field or a peer that does
                //# not report ECN markings.
                self.fail(now, path, publisher);
                return ValidationOutcome::Failed;
            }

            if baseline_ecn_counts == EcnCounts::default() {
                // Nothing to validate
                return ValidationOutcome::Skipped;
            }
        }

        let congestion_experienced_count = if let Some(incremental_ecn_counts) =
            ack_frame_ecn_counts
                .unwrap_or_default()
                .checked_sub(baseline_ecn_counts)
        {
            if Self::ce_remarking(incremental_ecn_counts, newly_acked_ecn_counts)
                || Self::remarked_to_ect0_or_ect1(incremental_ecn_counts, sent_packet_ecn_counts)
                || Self::ce_suppression(incremental_ecn_counts, newly_acked_ecn_counts)
            {
                self.fail(now, path, publisher);
                return ValidationOutcome::Failed;
            }

            // ce_suppression check above ensures this doesn't underflow
            incremental_ecn_counts.ce_count - newly_acked_ecn_counts.ce_count
        } else {
            // ECN counts decreased from the baseline
            self.fail(now, path, publisher);
            return ValidationOutcome::Failed;
        };

        //= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.4
        //# From the "unknown" state, successful validation of the ECN counts in an ACK frame
        //# (see Section 13.4.2.1) causes the ECN state for the path to become "capable",
        //# unless no marked packet has been acknowledged.
        if matches!(self.state, State::Unknown)
            && newly_acked_ecn_counts.ect_0_count > VarInt::from_u8(0)
        {
            // Arm the ce suppression timer to send a ECN-CE marked packet to test for
            // CE suppression by the peer.
            let mut ce_suppression_timer = Timer::default();
            ce_suppression_timer
                .set(now + *CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() as u32 * rtt);
            self.change_state(State::Capable(ce_suppression_timer), path, publisher);
        }

        if self.is_capable() && congestion_experienced_count > VarInt::ZERO {
            return ValidationOutcome::CongestionExperienced(congestion_experienced_count);
        }

        ValidationOutcome::Passed
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.1
    //# ECN validation also fails if the sum of the increase in ECT(0)
    //# and ECN-CE counts is less than the number of newly acknowledged
    //# packets that were originally sent with an ECT(0) marking.
    #[inline]
    fn ce_remarking(incremental_ecn_counts: EcnCounts, newly_acked_ecn_counts: EcnCounts) -> bool {
        let ect_0_increase = incremental_ecn_counts
            .ect_0_count
            .saturating_add(incremental_ecn_counts.ce_count);
        ect_0_increase < newly_acked_ecn_counts.ect_0_count
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.1
    //# ECN validation can fail if the received total count for either ECT(0) or ECT(1)
    //# exceeds the total number of packets sent with each corresponding ECT codepoint.
    #[inline]
    fn remarked_to_ect0_or_ect1(
        incremental_ecn_counts: EcnCounts,
        sent_packet_ecn_counts: EcnCounts,
    ) -> bool {
        incremental_ecn_counts.ect_0_count > sent_packet_ecn_counts.ect_0_count
            || incremental_ecn_counts.ect_1_count > sent_packet_ecn_counts.ect_1_count
    }

    //= https://www.rfc-editor.org/rfc/rfc9002#section-8.3
    //# A receiver can misreport ECN markings to alter the congestion
    //# response of a sender.  Suppressing reports of ECN-CE markings could
    //# cause a sender to increase their send rate.  This increase could
    //# result in congestion and loss.

    //= https://www.rfc-editor.org/rfc/rfc9002#section-8.3
    //# A sender can detect suppression of reports by marking occasional
    //# packets that it sends with an ECN-CE marking.  If a packet sent with
    //# an ECN-CE marking is not reported as having been CE marked when the
    //# packet is acknowledged, then the sender can disable ECN for that path
    //# by not setting ECN-Capable Transport (ECT) codepoints in subsequent
    //# packets sent on that path [RFC3168].
    #[inline]
    fn ce_suppression(
        incremental_ecn_counts: EcnCounts,
        newly_acked_ecn_counts: EcnCounts,
    ) -> bool {
        incremental_ecn_counts.ce_count < newly_acked_ecn_counts.ce_count
    }

    /// This method gets called when a packet has been sent
    pub fn on_packet_sent<Pub: event::ConnectionPublisher>(
        &mut self,
        ecn: ExplicitCongestionNotification,
        path: event::builder::Path,
        publisher: &mut Pub,
    ) {
        debug_assert!(
            !matches!(ecn, ExplicitCongestionNotification::Ect1),
            "Ect1 is not used"
        );

        if let (true, State::Testing(ref mut packet_count)) = (ecn.using_ecn(), &mut self.state) {
            *packet_count += 1;

            if *packet_count >= TESTING_PACKET_THRESHOLD {
                self.change_state(State::Unknown, path, publisher);
            }
        }
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack(&mut self, time_sent: Timestamp, ecn: ExplicitCongestionNotification) {
        if self.ecn_packet_sent_after_last_acked_ecn_packet(time_sent, ecn) {
            // Reset the black hole counter since a packet with ECN marking
            // has been acknowledged, indicating the path may still be ECN-capable
            self.black_hole_counter = Default::default();
            self.last_acked_ecn_packet_timestamp = Some(time_sent);
        }
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<Pub: event::ConnectionPublisher>(
        &mut self,
        time_sent: Timestamp,
        ecn: ExplicitCongestionNotification,
        now: Timestamp,
        path: event::builder::Path,
        publisher: &mut Pub,
    ) {
        if matches!(self.state, State::Failed(_)) {
            return;
        }

        if self.ecn_packet_sent_after_last_acked_ecn_packet(time_sent, ecn) {
            // An ECN marked packet that was sent after the last
            // acknowledged ECN marked packet has been lost
            self.black_hole_counter += 1;
        }

        if self.black_hole_counter > TESTING_PACKET_THRESHOLD {
            self.fail(now, path, publisher);
        }
    }

    /// Returns true if a packet sent at the given `time_sent` with the given ECN marking
    /// was marked as using ECN and was sent after the last time an ECN marked packet had
    /// been acknowledged.
    fn ecn_packet_sent_after_last_acked_ecn_packet(
        &mut self,
        time_sent: Timestamp,
        ecn: ExplicitCongestionNotification,
    ) -> bool {
        ecn.using_ecn()
            && self
                .last_acked_ecn_packet_timestamp
                .map_or(true, |last_acked| last_acked < time_sent)
    }

    /// Set the state to Failed and arm the retest timer
    fn fail<Pub: event::ConnectionPublisher>(
        &mut self,
        now: Timestamp,
        path: event::builder::Path,
        publisher: &mut Pub,
    ) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.2
        //# Even if validation fails, an endpoint MAY revalidate ECN for the same path at any later
        //# time in the connection. An endpoint could continue to periodically attempt validation.
        let mut retest_timer = Timer::default();
        retest_timer.set(now + RETEST_COOL_OFF_DURATION);
        self.change_state(State::Failed(retest_timer), path, publisher);
        self.black_hole_counter = Default::default();
    }

    fn change_state<Pub: event::ConnectionPublisher>(
        &mut self,
        state: State,
        path: event::builder::Path,
        publisher: &mut Pub,
    ) {
        debug_assert_ne!(self.state, state);

        self.state = state;

        publisher.on_ecn_state_changed(event::builder::EcnStateChanged {
            path,
            state: self.state.into_event(),
        })
    }
}

impl timer::Provider for Controller {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        if let State::Failed(timer) = &self.state {
            timer.timers(query)?
        }
        // The ce suppression timer in State::Capable is not queried here as that
        // timer is passively polled when transmitting and does not require firing
        // precisely.

        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 Manager implementation

use crate::{
    connection::PeerIdRegistry,
    endpoint, path,
    path::{challenge, Path},
    transmission,
};
use s2n_quic_core::{
    ack,
    connection::{self, PeerId},
    event::{self, builder::DatagramDropReason, IntoEvent},
    frame,
    frame::path_validation,
    inet::DatagramInfo,
    packet::number::PacketNumberSpace,
    path::{
        migration::{self, Validator as _},
        Handle as _, Id, MaxMtu,
    },
    random,
    recovery::{
        congestion_controller::{self, Endpoint as _},
        RttEstimator,
    },
    stateless_reset,
    time::{timer, Timestamp},
    transport,
};
use smallvec::SmallVec;

/// The amount of Paths that can be maintained without using the heap.
/// This value is also used to limit the number of connection migrations.
const MAX_ALLOWED_PATHS: usize = 5;

/// The PathManager handles paths for a specific connection.
/// It will handle path validation operations, and track the active path for a connection.
#[derive(Debug)]
pub struct Manager<Config: endpoint::Config> {
    /// Path array
    paths: SmallVec<[Path<Config>; MAX_ALLOWED_PATHS]>,

    /// Registry of `connection::PeerId`s
    pub(crate) peer_id_registry: PeerIdRegistry,

    /// Index to the active path
    active: u8,

    /// Index of last known active and validated path.
    ///
    /// The Path must be validated and also active at some some point to
    /// be set as the last_known_active_validated_path.
    last_known_active_validated_path: Option<u8>,

    /// The current index of a path that is pending packet protection authentication
    ///
    /// This field is used to annotate a new path that is pending packet authentication.
    /// If packet authentication fails then this path index will get reused instead of
    /// appending another to the list. This is used to prevent an off-path attacker from
    /// creating new paths with garbage data and preventing the peer to migrate paths.
    ///
    /// Note that it doesn't prevent an on-path attacker from observing/forwarding
    /// authenticated packets from bogus addresses. Because of the current hard limit
    /// of `MAX_ALLOWED_PATHS`, this will prevent the peer from migrating, if it needs to.
    /// The `paths` data structure will need to be enhanced to include garbage collection
    /// of old paths to overcome this limitation.
    pending_packet_authentication: Option<u8>,
}

impl<Config: endpoint::Config> Manager<Config> {
    pub fn new(initial_path: Path<Config>, peer_id_registry: PeerIdRegistry) -> Self {
        let mut manager = Manager {
            paths: SmallVec::from_elem(initial_path, 1),
            peer_id_registry,
            active: 0,
            last_known_active_validated_path: None,
            pending_packet_authentication: None,
        };
        manager.paths[0].activated = true;
        manager.paths[0].is_active = true;
        manager
    }

    /// Update the active path
    fn update_active_path<Pub: event::ConnectionPublisher>(
        &mut self,
        new_path_id: Id,
        random_generator: &mut dyn random::Generator,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        debug_assert!(new_path_id != path_id(self.active));

        let prev_path_id = self.active_path_id();

        let mut peer_connection_id = self[new_path_id].peer_connection_id;

        // The path's connection id might have retired since we last used it. Check if it is still
        // active, otherwise try and consume a new connection id.
        if !self.peer_id_registry.is_active(&peer_connection_id) {
            // If there are no new connection ids the peer is responsible for
            // providing additional connection ids to continue.
            //
            // Insufficient connection ids should not cause the connection to close.
            // Investigate api after this is used.
            peer_connection_id = self
                .peer_id_registry
                .consume_new_id_for_existing_path(new_path_id, peer_connection_id, publisher)
                .ok_or(
                    // TODO: add an event if active path update fails due to insufficient ids
                    transport::Error::INTERNAL_ERROR,
                )?;
        };
        self[new_path_id].peer_connection_id = peer_connection_id;

        if self.active_path().is_validated() {
            self.last_known_active_validated_path = Some(self.active);
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.3
        //# In response to an apparent migration, endpoints MUST validate the
        //# previously active path using a PATH_CHALLENGE frame.
        //
        // TODO: https://github.com/aws/s2n-quic/issues/711
        // The usage of 'apparent' is vague and its not clear if the previous path should
        // always be validated or only if the new active path is not validated.
        if !self.active_path().is_challenge_pending() {
            self.set_challenge(self.active_path_id(), random_generator);
        }

        self.activate_path(publisher, prev_path_id, new_path_id);

        // Restart ECN validation to check that the path still supports ECN
        let path = self.active_path_mut();
        path.ecn_controller
            .restart(path_event!(path, new_path_id), publisher);
        Ok(())
    }

    /// Return the active path
    #[inline]
    pub fn active_path(&self) -> &Path<Config> {
        &self.paths[self.active as usize]
    }

    /// Return a mutable reference to the active path
    #[inline]
    pub fn active_path_mut(&mut self) -> &mut Path<Config> {
        &mut self.paths[self.active as usize]
    }

    /// Return the Id of the active path
    #[inline]
    pub fn active_path_id(&self) -> Id {
        path_id(self.active)
    }

    pub fn check_active_path_is_synced(&self) {
        if cfg!(debug_assertions) {
            for (idx, path) in self.paths.iter().enumerate() {
                assert_eq!(path.is_active, (self.active == idx as u8));
            }
        }
    }

    pub fn activate_path<Pub: event::ConnectionPublisher>(
        &mut self,
        publisher: &mut Pub,
        prev_path_id: Id,
        new_path_id: Id,
    ) {
        self.check_active_path_is_synced();
        self.active = new_path_id.as_u8();
        self[prev_path_id].is_active = false;
        self[new_path_id].is_active = true;
        self[new_path_id].on_activated();
        self.check_active_path_is_synced();

        let prev_path = &self[prev_path_id];
        let new_path = &self[new_path_id];
        publisher.on_active_path_updated(event::builder::ActivePathUpdated {
            previous: path_event!(prev_path, prev_path_id),
            active: path_event!(new_path, new_path_id),
        });
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3
    //= type=TODO
    //= tracking-issue=714
    //# An endpoint MAY skip validation of a peer address if
    //# that address has been seen recently.
    /// Returns the Path for the provided address if the PathManager knows about it
    #[inline]
    pub fn path(&self, handle: &Config::PathHandle) -> Option<(Id, &Path<Config>)> {
        self.paths
            .iter()
            .enumerate()
            .find(|(_id, path)| Path::eq_by_handle(path, handle))
            .map(|(id, path)| (path_id(id as u8), path))
    }

    /// Returns the Path for the provided address if the PathManager knows about it
    #[inline]
    pub fn path_mut(&mut self, handle: &Config::PathHandle) -> Option<(Id, &mut Path<Config>)> {
        self.paths
            .iter_mut()
            .enumerate()
            .find(|(_id, path)| Path::eq_by_handle(path, handle))
            .map(|(id, path)| (path_id(id as u8), path))
    }

    /// Returns an iterator over all paths pending path_challenge or path_response
    /// transmission.
    pub fn paths_pending_validation(&mut self) -> PathsPendingValidation<Config> {
        PathsPendingValidation::new(self)
    }

    /// Called when a datagram is received on a connection
    /// Upon success, returns a `(Id, bool)` containing the path ID and a boolean that is
    /// true if the path had been amplification limited prior to receiving the datagram
    /// and is now no longer amplification limited.
    ///
    /// This function is called prior to packet authentication. If possible add business
    /// logic to [`Self::on_processed_packet`], which is called after the packet has been
    /// authenticated.
    #[allow(clippy::too_many_arguments)]
    pub fn on_datagram_received<Pub: event::ConnectionPublisher>(
        &mut self,
        path_handle: &Config::PathHandle,
        datagram: &DatagramInfo,
        handshake_confirmed: bool,
        congestion_controller_endpoint: &mut Config::CongestionControllerEndpoint,
        migration_validator: &mut Config::PathMigrationValidator,
        max_mtu: MaxMtu,
        publisher: &mut Pub,
    ) -> Result<(Id, bool), DatagramDropReason> {
        let valid_initial_received = self.valid_initial_received();

        if let Some((id, path)) = self.path_mut(path_handle) {
            let source_cid_changed = datagram.source_connection_id.map_or(false, |scid| {
                scid != path.peer_connection_id && valid_initial_received
            });

            if source_cid_changed {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
                //# Once a client has received a valid Initial packet from the server, it MUST
                //# discard any subsequent packet it receives on that connection with a
                //# different Source Connection ID.

                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
                //# Any further changes to the Destination Connection ID are only
                //# permitted if the values are taken from NEW_CONNECTION_ID frames; if
                //# subsequent Initial packets include a different Source Connection ID,
                //# they MUST be discarded.

                return Err(DatagramDropReason::InvalidSourceConnectionId);
            }

            // update the address if it was resolved
            path.handle.maybe_update(path_handle);

            let unblocked = path.on_bytes_received(datagram.payload_len);
            return Ok((id, unblocked));
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
        //# If a client receives packets from an unknown server address,
        //# the client MUST discard these packets.
        if Config::ENDPOINT_TYPE.is_client() {
            return Err(DatagramDropReason::UnknownServerAddress);
        };

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
        //# The design of QUIC relies on endpoints retaining a stable address
        //# for the duration of the handshake.  An endpoint MUST NOT initiate
        //# connection migration before the handshake is confirmed, as defined
        //# in section 4.1.2 of [QUIC-TLS].
        if !handshake_confirmed {
            return Err(DatagramDropReason::ConnectionMigrationDuringHandshake);
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
        //# If the peer
        //# violates this requirement, the endpoint MUST either drop the incoming
        //# packets on that path without generating a Stateless Reset or proceed
        //# with path validation and allow the peer to migrate.  Generating a
        //# Stateless Reset or closing the connection would allow third parties
        //# in the network to cause connections to close by spoofing or otherwise
        //# manipulating observed traffic.

        self.handle_connection_migration(
            path_handle,
            datagram,
            congestion_controller_endpoint,
            migration_validator,
            max_mtu,
            publisher,
        )
    }

    #[allow(clippy::too_many_arguments)]
    fn handle_connection_migration<Pub: event::ConnectionPublisher>(
        &mut self,
        path_handle: &Config::PathHandle,
        datagram: &DatagramInfo,
        congestion_controller_endpoint: &mut Config::CongestionControllerEndpoint,
        migration_validator: &mut Config::PathMigrationValidator,
        max_mtu: MaxMtu,
        publisher: &mut Pub,
    ) -> Result<(Id, bool), DatagramDropReason> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
        //# Clients are responsible for initiating all migrations.
        debug_assert!(Config::ENDPOINT_TYPE.is_server());

        let remote_address = path_handle.remote_address();
        let local_address = path_handle.local_address();
        let active_local_addr = self.active_path().local_address();
        let active_remote_addr = self.active_path().remote_address();

        // TODO set alpn if available
        let attempt: migration::Attempt = migration::AttemptBuilder {
            active_path: event::builder::Path {
                local_addr: active_local_addr.into_event(),
                local_cid: self.active_path().local_connection_id.into_event(),
                remote_addr: active_remote_addr.into_event(),
                remote_cid: self.active_path().peer_connection_id.into_event(),
                id: self.active_path_id().into_event(),
                is_active: true,
            }
            .into_event(),
            packet: migration::PacketInfoBuilder {
                remote_address: &remote_address,
                local_address: &local_address,
            }
            .into(),
        }
        .into();

        match migration_validator.on_migration_attempt(&attempt) {
            migration::Outcome::Allow => {
                // no-op: allow the migration to continue
            }
            migration::Outcome::Deny(reason) => {
                publisher.on_connection_migration_denied(reason.into_event());
                return Err(DatagramDropReason::RejectedConnectionMigration);
            }
            _ => {
                unimplemented!("unimplemented migration outcome");
            }
        }

        // Determine which index will be used for the newly created path
        //
        // If a previously allocated path failed to contain an authenticated packet, we
        // use that index instead of pushing on to the end.
        let new_path_idx = if let Some(idx) = self.pending_packet_authentication {
            idx as _
        } else {
            let idx = self.paths.len();
            self.pending_packet_authentication = Some(idx as _);
            idx
        };

        // TODO: Support deletion of old paths: https://github.com/aws/s2n-quic/issues/741
        // The current path manager implementation does not delete or reuse indices
        // in the path array. This can result in an unbounded number of paths. To prevent
        // this we limit the max number of paths per connection.
        if new_path_idx >= MAX_ALLOWED_PATHS {
            return Err(DatagramDropReason::PathLimitExceeded);
        }
        let new_path_id = path_id(new_path_idx as u8);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.4
        //= type=TODO
        //# Because port-only changes are commonly the
        //# result of NAT rebinding or other middlebox activity, the endpoint MAY
        //# instead retain its congestion control state and round-trip estimate
        //# in those cases instead of reverting to initial values.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.4
        //# On confirming a peer's ownership of its new address, an endpoint MUST
        //# immediately reset the congestion controller and round-trip time
        //# estimator for the new path to initial values (see Appendices A.3 and
        //# B.3 of [QUIC-RECOVERY]) unless the only change in the peer's address
        //# is its port number.
        // Since we maintain a separate congestion controller and round-trip time
        // estimator for the new path, and they are initialized with initial values,
        // we do not need to reset congestion controller and round-trip time estimator
        // again on confirming the peer's ownership of its new address.
        let rtt = RttEstimator::new(self.active_path().rtt_estimator.max_ack_delay());
        let path_info = congestion_controller::PathInfo::new(&remote_address);
        let cc = congestion_controller_endpoint.new_congestion_controller(path_info);

        let peer_connection_id = {
            if self.active_path().local_connection_id != datagram.destination_connection_id {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-9.5
                //# Similarly, an endpoint MUST NOT reuse a connection ID when sending to
                //# more than one destination address.

                // Peer has intentionally tried to migrate to this new path because they changed
                // their destination_connection_id, so we will change our destination_connection_id as well.
                self.peer_id_registry
                    .consume_new_id_for_new_path()
                    .ok_or(DatagramDropReason::InsufficientConnectionIds)?
            } else {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-9.5
                //# Due to network changes outside
                //# the control of its peer, an endpoint might receive packets from a new
                //# source address with the same Destination Connection ID field value,
                //# in which case it MAY continue to use the current connection ID with
                //# the new remote address while still sending from the same local
                //# address.
                self.active_path().peer_connection_id
            }
        };

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.1
        //# Until a peer's address is deemed valid, an endpoint limits
        //# the amount of data it sends to that address; see Section 8.
        //
        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3
        //# An endpoint MAY send data to an unvalidated peer address, but it MUST
        //# protect against potential attacks as described in Sections 9.3.1 and
        //# 9.3.2.
        //
        // New paths for a Server endpoint start in AmplificationLimited state until they are validated.
        let mut path = Path::new(
            *path_handle,
            peer_connection_id,
            datagram.destination_connection_id,
            rtt,
            cc,
            true,
            max_mtu,
        );

        let unblocked = path.on_bytes_received(datagram.payload_len);

        let active_path = self.active_path();
        let active_path_id = self.active_path_id();
        publisher.on_path_created(event::builder::PathCreated {
            active: path_event!(active_path, active_path_id),
            new: path_event!(path, new_path_id),
        });

        publisher.on_mtu_updated(event::builder::MtuUpdated {
            path_id: new_path_id.into_event(),
            mtu: path.mtu_controller.mtu() as u16,
            cause: MtuUpdatedCause::NewPath,
        });

        // create a new path
        if new_path_idx < self.paths.len() {
            self.paths[new_path_idx] = path;
        } else {
            self.paths.push(path);
        }

        Ok((new_path_id, unblocked))
    }

    fn set_challenge(&mut self, path_id: Id, random_generator: &mut dyn random::Generator) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
        //# The endpoint MUST use unpredictable data in every PATH_CHALLENGE
        //# frame so that it can associate the peer's response with the
        //# corresponding PATH_CHALLENGE.
        let mut data: challenge::Data = [0; 8];
        random_generator.public_random_fill(&mut data);

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.4
        //# Endpoints SHOULD abandon path validation based on a timer.
        //
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.4
        //# When
        //# setting this timer, implementations are cautioned that the new path
        //# could have a longer round-trip time than the original.  A value of
        //# three times the larger of the current PTO or the PTO for the new path
        //# (using kInitialRtt, as defined in [QUIC-RECOVERY]) is RECOMMENDED.
        let abandon_duration = self[path_id].pto_period(PacketNumberSpace::ApplicationData);
        let abandon_duration = 3 * abandon_duration.max(
            self.active_path()
                .pto_period(PacketNumberSpace::ApplicationData),
        );

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9
        //# An endpoint MUST
        //# perform path validation (Section 8.2) if it detects any change to a
        //# peer's address, unless it has previously validated that address.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.6.3
        //# Servers SHOULD initiate path validation to the client's new address
        //# upon receiving a probe packet from a different address.
        let challenge = challenge::Challenge::new(abandon_duration, data);
        self[path_id].set_challenge(challenge);
    }

    /// Returns true if a valid initial packet has been received
    pub fn valid_initial_received(&self) -> bool {
        if Config::ENDPOINT_TYPE.is_server() {
            // Since the path manager is owned by a connection, and a connection can only exist
            // on the server if a valid initial has been received, we immediately return true
            true
        } else {
            // A QUIC client uses a randomly generated value as the Initial Connection Id
            // until it receives a packet from the Server. Upon receiving a Server packet,
            // the Client switches to using the new Destination Connection Id. The
            // PeerIdRegistry is expected to be empty until the first Server initial packet.
            !self.peer_id_registry.is_empty()
        }
    }

    /// Writes any frames the path manager wishes to transmit to the given context
    #[inline]
    pub fn on_transmit<W: transmission::WriteContext>(&mut self, context: &mut W) {
        self.peer_id_registry.on_transmit(context)

        // TODO Add in per-path constraints based on whether a Challenge needs to be
        // transmitted.
    }

    /// Called when packets are acknowledged
    #[inline]
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        self.peer_id_registry.on_packet_ack(ack_set);
    }

    /// Called when packets are lost
    #[inline]
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        self.peer_id_registry.on_packet_loss(ack_set);
    }

    #[inline]
    pub fn on_path_challenge(
        &mut self,
        path_id: Id,
        challenge: &frame::path_challenge::PathChallenge,
    ) {
        self[path_id].on_path_challenge(challenge.data);
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.3
    //# Path validation succeeds when a PATH_RESPONSE frame is received that
    //# contains the data that was sent in a previous PATH_CHALLENGE frame.
    //# A PATH_RESPONSE frame received on any network path validates the path
    //# on which the PATH_CHALLENGE was sent.
    #[inline]
    pub fn on_path_response<Pub: event::ConnectionPublisher>(
        &mut self,
        response: &frame::PathResponse,
        publisher: &mut Pub,
    ) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
        //# A PATH_RESPONSE frame MUST be sent on the network path where the
        //# PATH_CHALLENGE frame was received.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
        //# This requirement MUST NOT be enforced by the endpoint that initiates
        //# path validation, as that would enable an attack on migration; see
        //# Section 9.3.3.
        //
        // The 'attack on migration' refers to the following scenario:
        // If the packet forwarded by the off-attacker is received before the
        // genuine packet, the genuine packet will be discarded as a duplicate
        // and path validation will fail.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.3
        //# A PATH_RESPONSE frame received on any network path validates the path
        //# on which the PATH_CHALLENGE was sent.

        for (id, path) in self.paths.iter_mut().enumerate() {
            if path.on_path_response(response.data) {
                let id = id as u64;
                publisher.on_path_challenge_updated(event::builder::PathChallengeUpdated {
                    path_challenge_status: event::builder::PathChallengeStatus::Validated,
                    path: path_event!(path, id),
                    challenge_data: path.challenge.challenge_data().into_event(),
                });
                // A path was validated so check if it becomes the new
                // last_known_active_validated_path
                if path.is_activated() {
                    self.last_known_active_validated_path = Some(id as u8);
                }
                break;
            }
        }
    }

    /// Process a packet and update internal state.
    ///
    /// Check if the packet is a non-probing (path validation) packet and attempt to
    /// update the active path for the connection.
    pub fn on_processed_packet<Pub: event::ConnectionPublisher>(
        &mut self,
        path_id: Id,
        source_connection_id: Option<PeerId>,
        path_validation_probing: path_validation::Probe,
        random_generator: &mut dyn random::Generator,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
        //# A client MUST change the Destination Connection ID it uses for
        //# sending packets in response to only the first received Initial or
        //# Retry packet.
        if !self.valid_initial_received() {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
            //# Until a packet is received from the server, the client MUST
            //# use the same Destination Connection ID value on all packets in this
            //# connection.
            //
            // This is the first Server packet so start using the newly provided
            // connection id form the Server.
            assert!(Config::ENDPOINT_TYPE.is_client());
            if let Some(source_connection_id) = source_connection_id {
                self[path_id].peer_connection_id = source_connection_id;
                self.peer_id_registry
                    .register_initial_connection_id(source_connection_id);
            }
        }

        // Remove the temporary status after successfully processing a packet
        if self.pending_packet_authentication == Some(path_id.as_u8()) {
            self.pending_packet_authentication = None;

            // We can finally arm the challenge after authenticating the packet
            self.set_challenge(path_id, random_generator);
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-9.2
        //# An endpoint can migrate a connection to a new local address by
        //# sending packets containing non-probing frames from that address.
        if !path_validation_probing.is_probing() && self.active_path_id() != path_id {
            self.update_active_path(path_id, random_generator, publisher)?;

            //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3
            //# After changing the address to which it sends non-probing packets, an
            //# endpoint can abandon any path validation for other addresses.
            //
            // Abandon other path validations only if the active path is validated since an
            // attacker could block all path validation attempts simply by forwarding packets.
            if self.active_path().is_validated() {
                self.abandon_all_path_challenges(publisher);
            } else if !self.active_path().is_challenge_pending() {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3
                //# If the recipient permits the migration, it MUST send subsequent
                //# packets to the new peer address and MUST initiate path validation
                //# (Section 8.2) to verify the peer's ownership of the address if
                //# validation is not already underway.
                self.set_challenge(self.active_path_id(), random_generator);
            }
        }
        Ok(())
    }

    #[inline]
    fn abandon_all_path_challenges<Pub: event::ConnectionPublisher>(
        &mut self,
        publisher: &mut Pub,
    ) {
        for (idx, path) in self.paths.iter_mut().enumerate() {
            let path_id = idx as u64;
            path.abandon_challenge(publisher, path_id);
        }
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
    //# Tokens are
    //# invalidated when their associated connection ID is retired via a
    //# RETIRE_CONNECTION_ID frame (Section 19.16).
    pub fn on_connection_id_retire(&self, _connection_id: &connection::LocalId) {
        // TODO invalidate any tokens issued under this connection id
    }

    /// Called when a NEW_CONNECTION_ID frame is received from the peer
    pub fn on_new_connection_id<Pub: event::ConnectionPublisher>(
        &mut self,
        connection_id: &connection::PeerId,
        sequence_number: u32,
        retire_prior_to: u32,
        stateless_reset_token: &stateless_reset::Token,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        // Retire and register connection ID
        self.peer_id_registry.on_new_connection_id(
            connection_id,
            sequence_number,
            retire_prior_to,
            stateless_reset_token,
        )?;

        //= 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.
        let active_path_connection_id = self.active_path().peer_connection_id;

        if !self.peer_id_registry.is_active(&active_path_connection_id) {
            self.active_path_mut().peer_connection_id = self
                .peer_id_registry
                .consume_new_id_for_existing_path(
                    self.active_path_id(),
                    active_path_connection_id,
                    publisher,
                )
                .expect(
                    "since we are only checking the active path and new id was delivered \
                    via the new_connection_id frames, there will always be a new id available \
                    to consume if necessary",
                );
        }

        Ok(())
    }

    pub fn on_timeout<Pub: event::ConnectionPublisher>(
        &mut self,
        timestamp: Timestamp,
        random_generator: &mut dyn random::Generator,
        publisher: &mut Pub,
    ) -> Result<(), connection::Error> {
        for (id, path) in self.paths.iter_mut().enumerate() {
            path.on_timeout(timestamp, path_id(id as u8), random_generator, publisher);
        }

        if self.active_path().failed_validation() {
            match self.last_known_active_validated_path {
                Some(last_known_active_validated_path) => {
                    //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.2
                    //# To protect the connection from failing due to such a spurious
                    //# migration, an endpoint MUST revert to using the last validated peer
                    //# address when validation of a new peer address fails.
                    let prev_path_id = path_id(self.active);
                    let new_path_id = path_id(last_known_active_validated_path);
                    self.activate_path(publisher, prev_path_id, new_path_id);
                    self.last_known_active_validated_path = None;
                }
                None => {
                    //= https://www.rfc-editor.org/rfc/rfc9000#section-9
                    //# When an endpoint has no validated path on which to send packets, it
                    //# MAY discard connection state.

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-9
                    //= type=TODO
                    //= tracking-issue=713
                    //# An endpoint capable of connection
                    //# migration MAY wait for a new path to become available before
                    //# discarding connection state.

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.2
                    //# If an endpoint has no state about the last validated peer address, it
                    //# MUST close the connection silently by discarding all connection
                    //# state.

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-10
                    //# An endpoint MAY discard connection state if it does not have a
                    //# validated path on which it can send packets; see Section 8.2
                    return Err(connection::Error::no_valid_path());
                }
            }
        }

        Ok(())
    }

    /// Notifies the path manager of the connection closing event
    pub fn on_closing(&mut self) {
        self.active_path_mut().on_closing();
        // TODO clean up other paths
    }

    /// true if ALL paths are amplification_limited
    #[inline]
    pub fn is_amplification_limited(&self) -> bool {
        self.paths
            .iter()
            .all(|path| path.transmission_constraint().is_amplification_limited())
    }

    /// true if ANY of the paths can transmit
    #[inline]
    pub fn can_transmit(&self, interest: transmission::Interest) -> bool {
        self.paths.iter().any(|path| {
            let constraint = path.transmission_constraint();
            interest.can_transmit(constraint)
        })
    }

    #[inline]
    pub fn transmission_constraint(&self) -> transmission::Constraint {
        // Return the lowest constraint which will ensure we don't get blocked on transmission by a single path
        self.paths
            .iter()
            .map(|path| path.transmission_constraint())
            .min()
            .unwrap_or(transmission::Constraint::None)
    }

    /// Returns the maximum size the UDP payload can reach for any probe packet.
    #[inline]
    pub fn max_mtu(&self) -> MaxMtu {
        let value = self.active_path().max_mtu();

        // This value is the same for each path so just return the active value
        if cfg!(debug_assertions) {
            for path in self.paths.iter() {
                assert_eq!(value, path.max_mtu());
            }
        }

        value
    }
}

#[inline]
fn path_id(id: u8) -> path::Id {
    // Safety: The path::Manager is responsible for managing path ID and is thus
    //         responsible for using them safely
    unsafe { path::Id::new(id) }
}

impl<Config: endpoint::Config> timer::Provider for Manager<Config> {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        for path in self.paths.iter() {
            path.timers(query)?;
        }

        Ok(())
    }
}

/// Iterate over all paths that have an interest in sending PATH_CHALLENGE
/// or PATH_RESPONSE frames.
///
/// This abstraction allows for iterating over pending paths while also
/// having mutable access to the Manager.
pub struct PathsPendingValidation<'a, Config: endpoint::Config> {
    index: u8,
    path_manager: &'a mut Manager<Config>,
}

impl<'a, Config: endpoint::Config> PathsPendingValidation<'a, Config> {
    pub fn new(path_manager: &'a mut Manager<Config>) -> Self {
        Self {
            index: 0,
            path_manager,
        }
    }

    #[inline]
    pub fn next_path(&mut self) -> Option<(Id, &mut Manager<Config>)> {
        loop {
            let path = self.path_manager.paths.get(self.index as usize)?;

            // Advance the index otherwise this will continue to process the
            // same path.
            self.index += 1;

            if path.is_challenge_pending() || path.is_response_pending() {
                return Some((path_id(self.index - 1), self.path_manager));
            }
        }
    }
}

impl<Config: endpoint::Config> transmission::interest::Provider for Manager<Config> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.peer_id_registry.transmission_interest(query)?;

        for path in self.paths.iter() {
            // query PATH_CHALLENGE and PATH_RESPONSE interest for each path
            path.transmission_interest(query)?;
        }

        Ok(())
    }
}

impl<Config: endpoint::Config> core::ops::Index<Id> for Manager<Config> {
    type Output = Path<Config>;

    #[inline]
    fn index(&self, id: Id) -> &Self::Output {
        &self.paths[id.as_u8() as usize]
    }
}

impl<Config: endpoint::Config> core::ops::IndexMut<Id> for Manager<Config> {
    #[inline]
    fn index_mut(&mut self, id: Id) -> &mut Self::Output {
        &mut self.paths[id.as_u8() as usize]
    }
}

macro_rules! path_event {
    ($path:ident, $path_id:ident) => {{
        event::builder::Path {
            local_addr: $path.local_address().into_event(),
            local_cid: $path.local_connection_id.into_event(),
            remote_addr: $path.remote_address().into_event(),
            remote_cid: $path.peer_connection_id.into_event(),
            id: $path_id.into_event(),
            is_active: $path.is_active(),
        }
    }};
}

pub(crate) use path_event;
use s2n_quic_core::event::builder::MtuUpdatedCause;

#[cfg(test)]
mod tests;

#[cfg(test)]
mod fuzz_target;