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

use crate::{
    ack::AckManager,
    connection::{self, ConnectionTransmissionContext, ProcessingError},
    endpoint, path,
    path::{path_event, Path},
    processed_packet::ProcessedPacket,
    recovery,
    space::{CryptoStream, HandshakeStatus, PacketSpace, TxPacketNumbers},
    transmission,
};
use core::{fmt, marker::PhantomData};
use s2n_codec::EncoderBuffer;
use s2n_quic_core::{
    connection::PeerId,
    crypto::{tls, CryptoSuite, InitialKey},
    event::{self, ConnectionPublisher as _, IntoEvent},
    frame::{ack::AckRanges, crypto::CryptoRef, Ack, ConnectionClose},
    inet::DatagramInfo,
    packet::{
        encoding::{PacketEncoder, PacketEncodingError},
        initial::{CleartextInitial, Initial, ProtectedInitial},
        number::{PacketNumber, PacketNumberRange, PacketNumberSpace, SlidingWindow},
    },
    time::{timer, Timestamp},
    transport,
};
use smallvec::SmallVec;

pub struct InitialSpace<Config: endpoint::Config> {
    pub ack_manager: AckManager,
    //= https://www.rfc-editor.org/rfc/rfc9001#section-4
    //# If QUIC needs to retransmit that data, it MUST use
    //# the same keys even if TLS has already updated to newer keys.
    pub key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialKey,
    pub header_key:
        <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialHeaderKey,
    //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9
    //# If packets from a lower encryption level contain
    //# CRYPTO frames, frames that retransmit that data MUST be sent at the
    //# same encryption level.
    pub crypto_stream: CryptoStream,
    pub tx_packet_numbers: TxPacketNumbers,
    pub received_hello_message: bool,
    //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.3
    //# Subsequent Initial packets from the client include the connection ID
    //# and token values from the Retry packet.
    retry_token: Vec<u8>,
    processed_packet_numbers: SlidingWindow,
    recovery_manager: recovery::Manager<Config>,
}

impl<Config: endpoint::Config> fmt::Debug for InitialSpace<Config> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("InitialSpace")
            .field("ack_manager", &self.ack_manager)
            .field("tx_packet_numbers", &self.tx_packet_numbers)
            .field("processed_packet_numbers", &self.processed_packet_numbers)
            .field("recovery_manager", &self.recovery_manager)
            .finish()
    }
}

impl<Config: endpoint::Config> InitialSpace<Config> {
    pub fn new(
        key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialKey,
        header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialHeaderKey,
        now: Timestamp,
        ack_manager: AckManager,
    ) -> Self {
        Self {
            ack_manager,
            key,
            header_key,
            crypto_stream: CryptoStream::new(),
            tx_packet_numbers: TxPacketNumbers::new(PacketNumberSpace::Initial, now),
            received_hello_message: false,
            retry_token: Vec::new(),
            processed_packet_numbers: SlidingWindow::default(),
            recovery_manager: recovery::Manager::new(PacketNumberSpace::Initial),
        }
    }

    /// This method gets called when a Retry packet is processed.
    ///
    /// Reset the TLS stack and recover state when the first Retry packet is processed.
    /// Also regenerate the Initial keys based on the new retry_source_connection_id.
    pub fn on_retry_packet<Pub: event::ConnectionPublisher>(
        &mut self,
        path: &mut path::Path<Config>,
        path_id: path::Id,
        retry_source_connection_id: &PeerId,
        retry_token: &[u8],
        publisher: &mut Pub,
    ) {
        debug_assert!(Config::ENDPOINT_TYPE.is_client());
        self.retry_token = retry_token.to_vec();

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

        self.key = initial_key;
        self.header_key = initial_header_key;

        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.5.3
        //# Other than updating the Destination Connection ID and Token fields,
        //# the Initial packet sent by the client is subject to the same
        //# restrictions as the first Initial packet.  A client MUST use the same
        //# cryptographic handshake message it included in this packet.
        self.crypto_stream.on_retry_packet();

        // Reset the recovery state; discarding any previous Initial packets that
        // might have been sent/lost.
        self.recovery_manager
            .on_retry_packet(path, path_id, publisher);
    }

    /// Returns true if the packet number has already been processed
    pub fn is_duplicate<Pub: event::ConnectionPublisher>(
        &self,
        packet_number: PacketNumber,
        path_id: path::Id,
        path: &path::Path<Config>,
        publisher: &mut Pub,
    ) -> bool {
        let packet_check = self.processed_packet_numbers.check(packet_number);
        if let Err(error) = packet_check {
            publisher.on_duplicate_packet(event::builder::DuplicatePacket {
                packet_header: event::builder::PacketHeader::new(
                    packet_number,
                    publisher.quic_version(),
                ),
                path: path_event!(path, path_id),
                error: error.into_event(),
            });
        }
        match packet_check {
            Ok(()) => false,
            Err(_) => true,
        }
    }

    pub fn on_transmit<'a>(
        &mut self,
        context: &mut ConnectionTransmissionContext<Config>,
        transmission_constraint: transmission::Constraint,
        handshake_status: &HandshakeStatus,
        buffer: EncoderBuffer<'a>,
    ) -> Result<(transmission::Outcome, EncoderBuffer<'a>), PacketEncodingError<'a>> {
        let mut packet_number = self.tx_packet_numbers.next();

        if self.recovery_manager.requires_probe() {
            //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.4
            //# If the sender wants to elicit a faster acknowledgement on PTO, it can
            //# skip a packet number to eliminate the acknowledgment delay.

            // TODO Does this interact negatively with persistent congestion detection, which
            //      relies on consecutive packet numbers?
            packet_number = packet_number.next().unwrap();
        }

        let packet_number_encoder = self.packet_number_encoder();
        let mut outcome = transmission::Outcome::default();

        let destination_connection_id = context.path().peer_connection_id;
        let payload = transmission::Transmission {
            config: <PhantomData<Config>>::default(),
            outcome: &mut outcome,
            packet_number,
            payload: transmission::early::Payload {
                ack_manager: &mut self.ack_manager,
                crypto_stream: &mut self.crypto_stream,
                packet_number_space: PacketNumberSpace::Initial,
                recovery_manager: &mut self.recovery_manager,
            },
            timestamp: context.timestamp,
            transmission_constraint,
            transmission_mode: context.transmission_mode,
            tx_packet_numbers: &mut self.tx_packet_numbers,
            path_id: context.path_id,
            publisher: context.publisher,
            packet_interceptor: context.packet_interceptor,
        };

        let packet = Initial {
            version: context.quic_version,
            destination_connection_id,
            source_connection_id: context.path_manager[context.path_id].local_connection_id,
            token: self.retry_token.as_slice(),
            packet_number,
            payload,
        };

        let (_protected_packet, buffer) = packet.encode_packet(
            &self.key,
            &self.header_key,
            packet_number_encoder,
            context.min_packet_len,
            buffer,
        )?;

        let time_sent = context.timestamp;
        let path_id = context.path_id;
        let (recovery_manager, mut recovery_context) =
            self.recovery(handshake_status, path_id, context.path_manager);
        recovery_manager.on_packet_sent(
            packet_number,
            outcome,
            time_sent,
            context.ecn,
            context.transmission_mode,
            None,
            &mut recovery_context,
            context.publisher,
        );

        context
            .publisher
            .on_packet_sent(event::builder::PacketSent {
                packet_header: event::builder::PacketHeader::new(
                    packet_number,
                    context.publisher.quic_version(),
                ),
                packet_len: outcome.bytes_sent,
            });

        Ok((outcome, buffer))
    }

    pub(super) fn on_transmit_close<'a>(
        &mut self,
        context: &mut ConnectionTransmissionContext<Config>,
        connection_close: &ConnectionClose,
        buffer: EncoderBuffer<'a>,
    ) -> Result<(transmission::Outcome, EncoderBuffer<'a>), PacketEncodingError<'a>> {
        let packet_number = self.tx_packet_numbers.next();

        let packet_number_encoder = self.packet_number_encoder();
        let mut outcome = transmission::Outcome::default();

        let destination_connection_id = context.path().peer_connection_id;
        let payload = transmission::Transmission {
            config: <PhantomData<Config>>::default(),
            outcome: &mut outcome,
            packet_number,
            payload: transmission::connection_close::Payload {
                connection_close,
                packet_number_space: PacketNumberSpace::Initial,
            },
            timestamp: context.timestamp,
            transmission_constraint: transmission::Constraint::None,
            transmission_mode: transmission::Mode::Normal,
            tx_packet_numbers: &mut self.tx_packet_numbers,
            path_id: context.path_id,
            publisher: context.publisher,
            packet_interceptor: context.packet_interceptor,
        };

        let packet = Initial {
            version: context.quic_version,
            destination_connection_id,
            source_connection_id: context.path_manager[context.path_id].local_connection_id,
            token: &[0u8; 0][..],
            packet_number,
            payload,
        };

        let (_protected_packet, buffer) = packet.encode_packet(
            &self.key,
            &self.header_key,
            packet_number_encoder,
            context.min_packet_len,
            buffer,
        )?;

        context
            .publisher
            .on_packet_sent(event::builder::PacketSent {
                packet_header: event::builder::PacketHeader::new(
                    packet_number,
                    context.publisher.quic_version(),
                ),
                packet_len: outcome.bytes_sent,
            });

        Ok((outcome, buffer))
    }

    /// Signals the connection was previously blocked by anti-amplification limits
    /// but is now no longer limited.
    pub fn on_amplification_unblocked(
        &mut self,
        path: &Path<Config>,
        timestamp: Timestamp,
        is_handshake_confirmed: bool,
    ) {
        debug_assert!(
            Config::ENDPOINT_TYPE.is_server(),
            "Clients are never in an anti-amplification state"
        );

        //= https://www.rfc-editor.org/rfc/rfc9002#appendix-A.6
        //# When a server is blocked by anti-amplification limits, receiving a
        //# datagram unblocks it, even if none of the packets in the datagram are
        //# successfully processed.  In such a case, the PTO timer will need to
        //# be re-armed.
        self.recovery_manager
            .update_pto_timer(path, timestamp, is_handshake_confirmed);
    }

    /// Called when the connection timer expired
    pub fn on_timeout<Pub: event::ConnectionPublisher>(
        &mut self,
        handshake_status: &HandshakeStatus,
        path_id: path::Id,
        path_manager: &mut path::Manager<Config>,
        random_generator: &mut Config::RandomGenerator,
        timestamp: Timestamp,
        publisher: &mut Pub,
    ) {
        self.ack_manager.on_timeout(timestamp);

        let (recovery_manager, mut context) =
            self.recovery(handshake_status, path_id, path_manager);
        recovery_manager.on_timeout(timestamp, random_generator, &mut context, publisher);
    }

    /// Called before the Initial packet space is discarded
    pub fn on_discard<Pub: event::ConnectionPublisher>(
        &mut self,
        path: &mut Path<Config>,
        path_id: path::Id,
        publisher: &mut Pub,
    ) {
        publisher.on_key_space_discarded(event::builder::KeySpaceDiscarded {
            space: event::builder::KeySpace::Initial,
        });
        self.recovery_manager
            .on_packet_number_space_discarded(path, path_id, publisher);
    }

    pub fn requires_probe(&self) -> bool {
        self.recovery_manager.requires_probe()
    }

    /// Returns the Packet Number to be used when decoding incoming packets
    pub fn packet_number_decoder(&self) -> PacketNumber {
        self.ack_manager.largest_received_packet_number_acked()
    }

    /// Returns the Packet Number to be used when encoding outgoing packets
    fn packet_number_encoder(&self) -> PacketNumber {
        self.tx_packet_numbers.largest_sent_packet_number_acked()
    }

    fn recovery<'a>(
        &'a mut self,
        handshake_status: &'a HandshakeStatus,
        path_id: path::Id,
        path_manager: &'a mut path::Manager<Config>,
    ) -> (
        &'a mut recovery::Manager<Config>,
        RecoveryContext<'a, Config>,
    ) {
        (
            &mut self.recovery_manager,
            RecoveryContext {
                ack_manager: &mut self.ack_manager,
                crypto_stream: &mut self.crypto_stream,
                tx_packet_numbers: &mut self.tx_packet_numbers,
                handshake_status,
                config: PhantomData,
                path_id,
                path_manager,
            },
        )
    }

    /// Validate packets in the Initial packet space
    pub fn validate_and_decrypt_packet<'a, Pub: event::ConnectionPublisher>(
        &self,
        protected: ProtectedInitial<'a>,
        path_id: path::Id,
        path: &path::Path<Config>,
        publisher: &mut Pub,
    ) -> Result<CleartextInitial<'a>, ProcessingError> {
        let packet_number_decoder = self.packet_number_decoder();
        let packet = protected
            .unprotect(&self.header_key, packet_number_decoder)
            .map_err(|err| {
                publisher.on_packet_dropped(event::builder::PacketDropped {
                    reason: event::builder::PacketDropReason::UnprotectFailed {
                        space: event::builder::KeySpace::Initial,
                        path: path_event!(path, path_id),
                    },
                });
                err
            })?;

        if self.is_duplicate(packet.packet_number, path_id, path, publisher) {
            return Err(ProcessingError::DuplicatePacket);
        }

        let packet_header =
            event::builder::PacketHeader::new(packet.packet_number, publisher.quic_version());
        let decrypted = packet.decrypt(&self.key).map_err(|err| {
            publisher.on_packet_dropped(event::builder::PacketDropped {
                reason: event::builder::PacketDropReason::DecryptionFailed {
                    packet_header,
                    path: path_event!(path, path_id),
                },
            });
            err
        })?;

        if Config::ENDPOINT_TYPE.is_client() && !decrypted.token.is_empty() {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
            //# Initial packets sent by the server MUST set the Token Length field
            //# to 0; clients that receive an Initial packet with a non-zero Token
            //# Length field MUST either discard the packet or generate a
            //# connection error of type PROTOCOL_VIOLATION.
            publisher.on_packet_dropped(event::builder::PacketDropped {
                reason: event::builder::PacketDropReason::NonEmptyRetryToken {
                    path: path_event!(path, path_id),
                },
            });
            return Err(ProcessingError::NonEmptyRetryToken);
        }

        Ok(decrypted)
    }

    #[inline]
    fn parse_client_hello<Pub: event::ConnectionPublisher>(
        &mut self,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        debug_assert!(Config::ENDPOINT_TYPE.is_server());

        if let Some(payload) = self.parse_hello(tls::HandshakeType::ClientHello)? {
            publisher.on_tls_client_hello(event::builder::TlsClientHello { payload: &payload });
        }

        Ok(())
    }

    #[inline]
    fn parse_server_hello<Pub: event::ConnectionPublisher>(
        &mut self,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        debug_assert!(Config::ENDPOINT_TYPE.is_client());

        if let Some(payload) = self.parse_hello(tls::HandshakeType::ServerHello)? {
            publisher.on_tls_server_hello(event::builder::TlsServerHello { payload: &payload });
        }

        Ok(())
    }

    #[inline]
    fn parse_hello(
        &mut self,
        msg_type: tls::HandshakeType,
    ) -> Result<Option<SmallVec<[&[u8]; 16]>>, transport::Error> {
        debug_assert!(
            !self.received_hello_message,
            "should only be called before the hello is parsed"
        );
        debug_assert_eq!(
            self.crypto_stream.rx.consumed_len(),
            0,
            "should not consume any crypto data"
        );

        let chunks = self.crypto_stream.rx.iter();
        let total_received_len = self.crypto_stream.rx.total_received_len();

        let mut chunks = chunks.peekable();
        let empty_chunk = &[][..];
        let header_chunk = chunks.peek().unwrap_or(&empty_chunk);

        // TODO make this configurable:
        //      https://github.com/aws/s2n-quic/issues/1001
        const MAX_HELLO_SIZE: u64 = 2 << 16;

        let outcome =
            <<Config::TLSEndpoint as tls::Endpoint>::Session as tls::Session>::parse_hello(
                msg_type,
                header_chunk,
                total_received_len,
                MAX_HELLO_SIZE,
            )?;

        if let Some(offsets) = outcome {
            // record that we've received the hello message
            self.received_hello_message = true;

            let payload = offsets.trim_chunks(chunks).collect();
            Ok(Some(payload))
        } else {
            // we are still waiting on more data
            Ok(None)
        }
    }
}

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

        Ok(())
    }
}

impl<Config: endpoint::Config> transmission::interest::Provider for InitialSpace<Config> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.ack_manager.transmission_interest(query)?;
        self.crypto_stream.transmission_interest(query)?;
        self.recovery_manager.transmission_interest(query)?;
        Ok(())
    }
}

impl<Config: endpoint::Config> connection::finalization::Provider for InitialSpace<Config> {
    fn finalization_status(&self) -> connection::finalization::Status {
        // there's nothing in here that hold up finalizing a connection
        connection::finalization::Status::Idle
    }
}

struct RecoveryContext<'a, Config: endpoint::Config> {
    ack_manager: &'a mut AckManager,
    crypto_stream: &'a mut CryptoStream,
    tx_packet_numbers: &'a mut TxPacketNumbers,
    handshake_status: &'a HandshakeStatus,
    config: PhantomData<Config>,
    path_id: path::Id,
    path_manager: &'a mut path::Manager<Config>,
}

impl<'a, Config: endpoint::Config> recovery::Context<Config> for RecoveryContext<'a, Config> {
    const ENDPOINT_TYPE: endpoint::Type = Config::ENDPOINT_TYPE;

    fn is_handshake_confirmed(&self) -> bool {
        self.handshake_status.is_confirmed()
    }

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

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

    fn path_by_id(&self, path_id: path::Id) -> &path::Path<Config> {
        &self.path_manager[path_id]
    }

    fn path_mut_by_id(&mut self, path_id: path::Id) -> &mut path::Path<Config> {
        &mut self.path_manager[path_id]
    }

    fn path_id(&self) -> path::Id {
        self.path_id
    }

    fn validate_packet_ack(
        &mut self,
        timestamp: Timestamp,
        packet_number_range: &PacketNumberRange,
    ) -> Result<(), transport::Error> {
        self.tx_packet_numbers
            .on_packet_ack(timestamp, packet_number_range)
    }

    fn on_new_packet_ack<Pub: event::ConnectionPublisher>(
        &mut self,
        packet_number_range: &PacketNumberRange,
        _publisher: &mut Pub,
    ) {
        self.crypto_stream.on_packet_ack(packet_number_range);
    }

    fn on_packet_ack(&mut self, timestamp: Timestamp, packet_number_range: &PacketNumberRange) {
        self.ack_manager
            .on_packet_ack(timestamp, packet_number_range);
    }

    fn on_packet_loss<Pub: event::ConnectionPublisher>(
        &mut self,
        packet_number_range: &PacketNumberRange,
        _publisher: &mut Pub,
    ) {
        self.crypto_stream.on_packet_loss(packet_number_range);
        self.ack_manager.on_packet_loss(packet_number_range);
    }

    fn on_rtt_update(&mut self) {}
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
//# The payload of an Initial packet includes a CRYPTO frame (or frames)
//# containing a cryptographic handshake message, ACK frames, or both.
//# PING, PADDING, and CONNECTION_CLOSE frames of type 0x1c are also
//# permitted.  An endpoint that receives an Initial packet containing
//# other frames can either discard the packet as spurious or treat it as
//# a connection error.
impl<Config: endpoint::Config> PacketSpace<Config> for InitialSpace<Config> {
    const INVALID_FRAME_ERROR: &'static str = "invalid frame in initial space";

    fn handle_crypto_frame<Pub: event::ConnectionPublisher>(
        &mut self,
        frame: CryptoRef,
        _datagram: &DatagramInfo,
        _path: &mut Path<Config>,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        self.crypto_stream.on_crypto_frame(frame)?;

        // try to parse out the hello message if we haven't yet
        if !self.received_hello_message {
            match Config::ENDPOINT_TYPE {
                endpoint::Type::Server => self.parse_client_hello(publisher)?,
                endpoint::Type::Client => self.parse_server_hello(publisher)?,
            }
        }

        Ok(())
    }

    fn handle_ack_frame<A: AckRanges, Pub: event::ConnectionPublisher>(
        &mut self,
        frame: Ack<A>,
        timestamp: Timestamp,
        path_id: path::Id,
        path_manager: &mut path::Manager<Config>,
        packet_number: PacketNumber,
        handshake_status: &mut HandshakeStatus,
        _local_id_registry: &mut connection::LocalIdRegistry,
        random_generator: &mut Config::RandomGenerator,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        let (recovery_manager, mut context) =
            self.recovery(handshake_status, path_id, path_manager);
        recovery_manager.on_ack_frame(
            timestamp,
            frame,
            packet_number,
            random_generator,
            &mut context,
            publisher,
        )
    }

    fn handle_connection_close_frame(
        &mut self,
        frame: ConnectionClose,
        _timestamp: Timestamp,
        _path: &mut Path<Config>,
    ) -> Result<(), transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-17.2.2
        //# CONNECTION_CLOSE frames of type 0x1c are also
        //# permitted.

        if frame.tag() != 0x1c {
            return Err(transport::Error::PROTOCOL_VIOLATION);
        }

        Ok(())
    }

    fn on_processed_packet<Pub: event::ConnectionPublisher>(
        &mut self,
        processed_packet: ProcessedPacket,
        path_id: path::Id,
        path: &Path<Config>,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        self.ack_manager.on_processed_packet(
            &processed_packet,
            path_event!(path, path_id),
            publisher,
        );
        self.processed_packet_numbers
            .insert(processed_packet.packet_number)
            .expect("packet number was already checked");
        Ok(())
    }
}


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

use core::{task::Poll, time::Duration};
use s2n_quic_core::time::{timer, Timer, Timestamp};

#[derive(Debug)]
pub struct KeepAlive {
    enabled: bool,
    period: Duration,
    timer: Timer,
}

impl KeepAlive {
    pub fn new(max_idle_timeout: Option<Duration>, max_period: Duration) -> Self {
        let period = if let Some(max_idle_timeout) = max_idle_timeout {
            // send a ping frame at 3/4 max idle timeout to ensure it is delivered in time
            (max_idle_timeout * 3 / 4).min(max_period)
        } else {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-10.1.2
            //# A connection will time out if no packets are sent or received for a
            //# period longer than the time negotiated using the max_idle_timeout
            //# transport parameter; see Section 10.  However, state in middleboxes
            //# might time out earlier than that.  Though REQ-5 in [RFC4787]
            //# recommends a 2-minute timeout interval, experience shows that sending
            //# packets every 30 seconds is necessary to prevent the majority of
            //# middleboxes from losing state for UDP flows [GATEWAY].

            // Even though we don't have an idle timeout, we should still have a default
            // keep-alive period to ensure middleboxes don't drop their UDP flow
            max_period
        };

        Self {
            enabled: false,
            period,
            timer: Timer::default(),
        }
    }

    #[inline]
    pub fn update(&mut self, enabled: bool) {
        self.enabled = enabled;
    }

    #[inline]
    pub fn reset(&mut self, now: Timestamp) {
        self.timer.set(now + self.period)
    }

    #[inline]
    pub fn on_timeout(&mut self, now: Timestamp) -> Poll<()> {
        if !self.enabled {
            return Poll::Pending;
        }

        let res = self.timer.poll_expiration(now);

        if res.is_ready() {
            self.reset(now);
        }

        res
    }

    #[inline]
    pub fn period(&self) -> Duration {
        self.period
    }
}

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


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

use crate::{
    ack,
    ack::AckManager,
    connection, endpoint, path,
    path::{path_event, Path},
    processed_packet::ProcessedPacket,
    stream::Manager as _,
    transmission,
};
use bytes::Bytes;
use core::{
    fmt,
    task::{Poll, Waker},
};
use s2n_codec::DecoderBufferMut;
use s2n_quic_core::{
    application::ServerName,
    connection::{limits::Limits, InitialId, PeerId},
    crypto::{tls, tls::Session, CryptoSuite, Key},
    event::{self, IntoEvent},
    frame::{
        ack::AckRanges, crypto::CryptoRef, datagram::DatagramRef, stream::StreamRef, Ack,
        ConnectionClose, DataBlocked, HandshakeDone, MaxData, MaxStreamData, MaxStreams,
        NewConnectionId, NewToken, PathChallenge, PathResponse, ResetStream, RetireConnectionId,
        StopSending, StreamDataBlocked, StreamsBlocked,
    },
    inet::DatagramInfo,
    packet::number::{PacketNumber, PacketNumberSpace},
    time::{timer, Timestamp},
    transport,
};

mod application;
mod crypto_stream;
pub(crate) mod datagram;
mod handshake;
mod handshake_status;
mod initial;
mod keep_alive;
mod session_context;
mod tx_packet_numbers;

pub(crate) use application::ApplicationSpace;
pub(crate) use crypto_stream::CryptoStream;
pub(crate) use handshake::HandshakeSpace;
pub(crate) use handshake_status::HandshakeStatus;
pub(crate) use initial::InitialSpace;
pub(crate) use session_context::SessionContext;
pub(crate) use tx_packet_numbers::TxPacketNumbers;

struct SessionInfo<Config: endpoint::Config> {
    session: <Config::TLSEndpoint as tls::Endpoint>::Session,
    initial_cid: InitialId,
}

pub struct PacketSpaceManager<Config: endpoint::Config> {
    session_info: Option<SessionInfo<Config>>,
    retry_cid: Option<Box<PeerId>>,
    initial: Option<Box<InitialSpace<Config>>>,
    handshake: Option<Box<HandshakeSpace<Config>>>,
    application: Option<Box<ApplicationSpace<Config>>>,
    zero_rtt_crypto:
        Option<Box<<<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::ZeroRttKey>>,
    handshake_status: HandshakeStatus,
    /// Server Name Indication
    pub server_name: Option<ServerName>,
    //= https://www.rfc-editor.org/rfc/rfc9000#section-7
    //# Endpoints MUST explicitly negotiate an application protocol.

    //= https://www.rfc-editor.org/rfc/rfc9001#section-8.1
    //# Unless
    //# another mechanism is used for agreeing on an application protocol,
    //# endpoints MUST use ALPN for this purpose.
    pub application_protocol: Bytes,
}

impl<Config: endpoint::Config> fmt::Debug for PacketSpaceManager<Config> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("PacketSpaceManager")
            .field("initial", &self.initial)
            .field("handshake", &self.handshake)
            .field("application", &self.application)
            .field("handshake_status", &self.handshake_status)
            .finish()
    }
}

macro_rules! packet_space_api {
    ($ty:ty, $field:ident, $get_mut:ident $(, $discard:ident $(, $assert_discard:ident)? )?) => {
        #[allow(dead_code)]
        pub fn $field(&self) -> Option<&$ty> {
            self.$field
                .as_ref()
                .map(Box::as_ref)
        }

        pub fn $get_mut(&mut self) -> Option<(&mut $ty, &mut HandshakeStatus)> {
            let space = self.$field
                .as_mut()
                .map(Box::as_mut)?;
            Some((space, &mut self.handshake_status))
        }

        $(
            pub fn $discard<Pub: event::ConnectionPublisher>(
                &mut self,
                path: &mut Path<Config>,
                path_id: path::Id,
                publisher: &mut Pub,
            ) {
                //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.2
                //# When Initial or Handshake keys are discarded, the PTO and loss
                //# detection timers MUST be reset, because discarding keys indicates
                //# forward progress and the loss detection timer might have been set for
                //# a now discarded packet number space.
                if let Some(mut space) = self.$field.take() {
                    // reset the PTO backoff value as part of resetting the PTO timer.
                    path.reset_pto_backoff();

                    space.on_discard(path, path_id, publisher);
                }

                //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.1
                //# Endpoints MUST NOT send
                //# Initial packets after this point.
                // By discarding a space, we are no longer capable of sending packets with those
                // keys.

                debug_assert!(self.$field.is_none(), "space should have been discarded");
                $(
                    debug_assert!(self.$assert_discard.is_none(), "space should have been discarded");
                )?
            }
        )?
    };
}

impl<Config: endpoint::Config> PacketSpaceManager<Config> {
    pub fn new<Pub: event::ConnectionPublisher>(
        initial_cid: InitialId,
        session: <Config::TLSEndpoint as tls::Endpoint>::Session,
        initial_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialKey,
        header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::InitialHeaderKey,
        now: Timestamp,
        publisher: &mut Pub,
    ) -> Self {
        let ack_manager = AckManager::new(PacketNumberSpace::Initial, ack::Settings::EARLY);

        publisher.on_key_update(event::builder::KeyUpdate {
            key_type: event::builder::KeyType::Initial,
            cipher_suite: initial_key.cipher_suite().into_event(),
        });
        Self {
            session_info: Some(SessionInfo {
                session,
                initial_cid,
            }),
            retry_cid: None,
            initial: Some(Box::new(InitialSpace::new(
                initial_key,
                header_key,
                now,
                ack_manager,
            ))),
            handshake: None,
            application: None,
            zero_rtt_crypto: None,
            handshake_status: HandshakeStatus::default(),
            server_name: None,
            application_protocol: Bytes::new(),
        }
    }

    packet_space_api!(InitialSpace<Config>, initial, initial_mut, discard_initial);

    packet_space_api!(
        HandshakeSpace<Config>,
        handshake,
        handshake_mut,
        discard_handshake,
        initial
    );

    packet_space_api!(ApplicationSpace<Config>, application, application_mut);

    #[allow(dead_code)] // 0RTT hasn't been started yet
    pub fn zero_rtt_crypto(
        &self,
    ) -> Option<&<<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::ZeroRttKey> {
        self.zero_rtt_crypto.as_ref().map(Box::as_ref)
    }

    pub fn discard_zero_rtt_crypto(&mut self) {
        self.zero_rtt_crypto = None;
    }

    #[allow(clippy::too_many_arguments)]
    pub fn poll_crypto<Pub: event::ConnectionPublisher>(
        &mut self,
        path_manager: &mut path::Manager<Config>,
        local_id_registry: &mut connection::LocalIdRegistry,
        limits: &mut Limits,
        now: Timestamp,
        waker: &Waker,
        publisher: &mut Pub,
        datagram: &mut Config::DatagramEndpoint,
    ) -> Poll<Result<(), transport::Error>> {
        if let Some(session_info) = self.session_info.as_mut() {
            let mut context: SessionContext<Config, Pub> = SessionContext {
                now,
                initial_cid: &session_info.initial_cid,
                retry_cid: self.retry_cid.as_deref(),
                initial: &mut self.initial,
                handshake: &mut self.handshake,
                application: &mut self.application,
                zero_rtt_crypto: &mut self.zero_rtt_crypto,
                path_manager,
                handshake_status: &mut self.handshake_status,
                local_id_registry,
                limits,
                server_name: &mut self.server_name,
                application_protocol: &mut self.application_protocol,
                waker,
                publisher,
                datagram,
            };

            match session_info.session.poll(&mut context)? {
                Poll::Ready(_success) => {
                    // The TLS session and retry_cid is no longer needed
                    self.session_info = None;
                    self.retry_cid = None;
                }
                Poll::Pending => return Poll::Pending,
            };
        }

        Poll::Ready(Ok(()))
    }

    /// Called when the connection timer expired
    pub fn on_timeout<Pub: event::ConnectionPublisher>(
        &mut self,
        local_id_registry: &mut connection::LocalIdRegistry,
        path_manager: &mut path::Manager<Config>,
        random_generator: &mut Config::RandomGenerator,
        timestamp: Timestamp,
        publisher: &mut Pub,
    ) {
        let path_id = path_manager.active_path_id();
        let path = path_manager.active_path_mut();

        // ensure the backoff doesn't grow too quickly
        let max_backoff = path.pto_backoff * 2;

        if let Some((space, handshake_status)) = self.initial_mut() {
            space.on_timeout(
                handshake_status,
                path_id,
                path_manager,
                random_generator,
                timestamp,
                publisher,
            )
        }
        if let Some((space, handshake_status)) = self.handshake_mut() {
            space.on_timeout(
                handshake_status,
                path_id,
                path_manager,
                random_generator,
                timestamp,
                publisher,
            )
        }
        if let Some((space, handshake_status)) = self.application_mut() {
            space.on_timeout(
                path_manager,
                handshake_status,
                local_id_registry,
                random_generator,
                timestamp,
                publisher,
            )
        }

        let path = path_manager.active_path_mut();
        path.pto_backoff = path.pto_backoff.min(max_backoff);
    }

    /// Signals the connection was previously blocked by anti-amplification limits
    /// but is now no longer limited.
    pub fn on_amplification_unblocked(&mut self, path: &Path<Config>, timestamp: Timestamp) {
        if let Some((space, handshake_status)) = self.initial_mut() {
            space.on_amplification_unblocked(path, timestamp, handshake_status.is_confirmed());
        }

        if let Some((space, handshake_status)) = self.handshake_mut() {
            space.on_amplification_unblocked(path, timestamp, handshake_status.is_confirmed());
        }

        if let Some((space, handshake_status)) = self.application_mut() {
            space.on_amplification_unblocked(path, timestamp, handshake_status.is_confirmed());
        }
    }

    pub fn requires_probe(&self) -> bool {
        core::iter::empty()
            .chain(self.initial.iter().map(|space| space.requires_probe()))
            .chain(self.handshake.iter().map(|space| space.requires_probe()))
            .chain(self.application.iter().map(|space| space.requires_probe()))
            .any(|requires_probe| requires_probe)
    }

    pub fn is_handshake_confirmed(&self) -> bool {
        self.handshake_status.is_confirmed()
    }

    pub fn is_handshake_complete(&self) -> bool {
        self.handshake_status.is_complete()
    }

    pub(crate) fn on_transmit_close(
        &mut self,
        early_connection_close: &ConnectionClose,
        connection_close: &ConnectionClose,
        context: &mut connection::ConnectionTransmissionContext<Config>,
        packet_buffer: &mut endpoint::PacketBuffer,
    ) -> Option<Bytes> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.2.3
        //# When sending a CONNECTION_CLOSE frame, the goal is to ensure that the
        //# peer will process the frame.  Generally, this means sending the frame
        //# in a packet with the highest level of packet protection to avoid the
        //# packet being discarded.
        let mut can_send_initial = self.initial.is_some();
        let mut can_send_handshake = self.handshake.is_some();
        let can_send_application = self.application.is_some();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.2.3
        //# After the handshake is confirmed (see
        //# Section 4.1.2 of [QUIC-TLS]), an endpoint MUST send any
        //# CONNECTION_CLOSE frames in a 1-RTT packet.
        if self.is_handshake_confirmed() {
            can_send_initial = false;
            can_send_handshake = false;
            debug_assert!(
                can_send_application,
                "if the handshake is confirmed, 1rtt keys should be available"
            );
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.2.3
        //# A client will always know whether the server has Handshake keys
        //# (see Section 17.2.2.1), but it is possible that a server does not
        //# know whether the client has Handshake keys.  Under these
        //# circumstances, a server SHOULD send a CONNECTION_CLOSE frame in
        //# both Handshake and Initial packets to ensure that at least one of
        //# them is processable by the client.
        if can_send_handshake {
            match Config::ENDPOINT_TYPE {
                endpoint::Type::Client => {
                    // if we are a client and have handshake keys, we know the server
                    // has handshake keys as well, so no need to transmit in initial.
                    can_send_initial = false;
                }
                endpoint::Type::Server => {
                    // try to send an initial packet if the space is still available
                    //
                    // Note: this assignment isn't actually needed; it's mostly to make
                    //       the code easier to follow
                    can_send_initial &= true;
                }
            }
        }

        packet_buffer.write(|buffer| {
            macro_rules! write_packet {
                ($buffer:expr, $space:ident, $check:expr, $frame:expr) => {
                    if let Some((space, _handshake_status)) = self.$space().filter(|_| $check) {
                        let result = space.on_transmit_close(context, &$frame, $buffer);

                        match result {
                            Ok((outcome, buffer)) => {
                                *context.outcome += outcome;
                                buffer
                            }
                            Err(err) => err.take_buffer(),
                        }
                    } else {
                        $buffer
                    }
                };
            }

            let buffer = write_packet!(
                buffer,
                initial_mut,
                can_send_initial,
                early_connection_close
            );
            let buffer = write_packet!(
                buffer,
                handshake_mut,
                can_send_handshake,
                early_connection_close
            );
            write_packet!(
                buffer,
                application_mut,
                can_send_application,
                connection_close
            )
        })
    }

    pub fn close<Pub: event::ConnectionPublisher>(
        &mut self,
        error: connection::Error,
        path: &mut path::Path<Config>,
        path_id: path::Id,
        publisher: &mut Pub,
    ) {
        self.session_info = None;
        self.retry_cid = None;
        self.discard_initial(path, path_id, publisher);
        self.discard_handshake(path, path_id, publisher);
        self.discard_zero_rtt_crypto();

        // Don't discard the application space until the application has read the error
        if let Some((application, _handshake_status)) = self.application_mut() {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-10.2
            //# A CONNECTION_CLOSE frame
            //# causes all streams to immediately become closed; open streams can be
            //# assumed to be implicitly reset.

            // Close all streams with the derived error
            application.stream_manager.close(error);
        }
    }

    pub fn on_retry_packet(&mut self, retry_source_connection_id: PeerId) {
        debug_assert!(Config::ENDPOINT_TYPE.is_client());
        self.retry_cid = Some(Box::new(retry_source_connection_id));
    }

    pub fn retry_cid(&self) -> Option<&PeerId> {
        self.retry_cid.as_deref()
    }
}

impl<Config: endpoint::Config> timer::Provider for PacketSpaceManager<Config> {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# When ack-eliciting packets in multiple packet number spaces are in
        //# flight, the timer MUST be set to the earlier value of the Initial and
        //# Handshake packet number spaces.

        if let Some(space) = self.application.as_ref() {
            space.timers(query)?;
        }
        if let Some(space) = self.handshake.as_ref() {
            space.timers(query)?;
        }
        if let Some(space) = self.initial.as_ref() {
            space.timers(query)?;
        }
        Ok(())
    }
}

impl<Config: endpoint::Config> transmission::interest::Provider for PacketSpaceManager<Config> {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        if let Some(space) = self.application.as_ref() {
            space.transmission_interest(query)?;
        }

        self.handshake_status.transmission_interest(query)?;

        if let Some(space) = self.initial.as_ref() {
            space.transmission_interest(query)?;
        }

        if let Some(space) = self.handshake.as_ref() {
            space.transmission_interest(query)?;
        }

        Ok(())
    }
}

impl<Config: endpoint::Config> connection::finalization::Provider for PacketSpaceManager<Config> {
    fn finalization_status(&self) -> connection::finalization::Status {
        core::iter::empty()
            .chain(self.initial.iter().map(|space| space.finalization_status()))
            .chain(
                self.handshake
                    .iter()
                    .map(|space| space.finalization_status()),
            )
            .chain(
                self.application
                    .iter()
                    .map(|space| space.finalization_status()),
            )
            .sum()
    }
}

macro_rules! default_frame_handler {
    ($name:ident, $frame:ty) => {
        fn $name(&mut self, frame: $frame) -> Result<(), transport::Error> {
            Err(transport::Error::PROTOCOL_VIOLATION
                .with_reason(Self::INVALID_FRAME_ERROR)
                .with_frame_type(frame.tag().into()))
        }
    };
}

pub trait PacketSpace<Config: endpoint::Config> {
    const INVALID_FRAME_ERROR: &'static str;

    fn handle_crypto_frame<Pub: event::ConnectionPublisher>(
        &mut self,
        frame: CryptoRef,
        datagram: &DatagramInfo,
        path: &mut Path<Config>,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error>;

    #[allow(clippy::too_many_arguments)]
    fn handle_ack_frame<A: AckRanges, Pub: event::ConnectionPublisher>(
        &mut self,
        frame: Ack<A>,
        timestamp: Timestamp,
        path_id: path::Id,
        path_manager: &mut path::Manager<Config>,
        packet_number: PacketNumber,
        handshake_status: &mut HandshakeStatus,
        local_id_registry: &mut connection::LocalIdRegistry,
        random_generator: &mut Config::RandomGenerator,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error>;

    fn handle_connection_close_frame(
        &mut self,
        frame: ConnectionClose,
        timestamp: Timestamp,
        path: &mut Path<Config>,
    ) -> Result<(), transport::Error>;

    fn handle_handshake_done_frame<Pub: event::ConnectionPublisher>(
        &mut self,
        frame: HandshakeDone,
        _timestamp: Timestamp,
        _path: &mut Path<Config>,
        _local_id_registry: &mut connection::LocalIdRegistry,
        _handshake_status: &mut HandshakeStatus,
        _publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        Err(transport::Error::PROTOCOL_VIOLATION
            .with_reason(Self::INVALID_FRAME_ERROR)
            .with_frame_type(frame.tag().into()))
    }

    fn handle_retire_connection_id_frame(
        &mut self,
        frame: RetireConnectionId,
        _datagram: &DatagramInfo,
        _path: &mut Path<Config>,
        _local_id_registry: &mut connection::LocalIdRegistry,
    ) -> Result<(), transport::Error> {
        Err(transport::Error::PROTOCOL_VIOLATION
            .with_reason(Self::INVALID_FRAME_ERROR)
            .with_frame_type(frame.tag().into()))
    }

    fn handle_new_connection_id_frame<Pub: event::ConnectionPublisher>(
        &mut self,
        frame: NewConnectionId,
        _datagram: &DatagramInfo,
        _path_manager: &mut path::Manager<Config>,
        _publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        Err(transport::Error::PROTOCOL_VIOLATION
            .with_reason(Self::INVALID_FRAME_ERROR)
            .with_frame_type(frame.tag().into()))
    }

    fn handle_path_response_frame<Pub: event::ConnectionPublisher>(
        &mut self,
        frame: PathResponse,
        _path_manager: &mut path::Manager<Config>,
        _publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        Err(transport::Error::PROTOCOL_VIOLATION
            .with_reason(Self::INVALID_FRAME_ERROR)
            .with_frame_type(frame.tag().into()))
    }

    fn handle_path_challenge_frame(
        &mut self,
        frame: PathChallenge,
        _path_id: path::Id,
        _path_manager: &mut path::Manager<Config>,
    ) -> Result<(), transport::Error> {
        Err(transport::Error::PROTOCOL_VIOLATION
            .with_reason(Self::INVALID_FRAME_ERROR)
            .with_frame_type(frame.tag().into()))
    }

    fn handle_stream_frame(
        &mut self,
        frame: StreamRef,
        _packet: &mut ProcessedPacket,
    ) -> Result<(), transport::Error> {
        Err(transport::Error::PROTOCOL_VIOLATION
            .with_reason(Self::INVALID_FRAME_ERROR)
            .with_frame_type(frame.tag().into()))
    }

    fn handle_datagram_frame(
        &mut self,
        _path: s2n_quic_core::event::api::Path<'_>,
        frame: DatagramRef,
    ) -> Result<(), transport::Error> {
        Err(transport::Error::PROTOCOL_VIOLATION
            .with_reason(Self::INVALID_FRAME_ERROR)
            .with_frame_type(frame.tag().into()))
    }

    default_frame_handler!(handle_data_blocked_frame, DataBlocked);
    default_frame_handler!(handle_max_data_frame, MaxData);
    default_frame_handler!(handle_max_stream_data_frame, MaxStreamData);
    default_frame_handler!(handle_max_streams_frame, MaxStreams);
    default_frame_handler!(handle_reset_stream_frame, ResetStream);
    default_frame_handler!(handle_stop_sending_frame, StopSending);
    default_frame_handler!(handle_stream_data_blocked_frame, StreamDataBlocked);
    default_frame_handler!(handle_streams_blocked_frame, StreamsBlocked);
    default_frame_handler!(handle_new_token_frame, NewToken);

    fn on_processed_packet<Pub: event::ConnectionPublisher>(
        &mut self,
        processed_packet: ProcessedPacket,
        path_id: path::Id,
        path: &Path<Config>,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error>;

    // TODO: Reduce arguments, https://github.com/aws/s2n-quic/issues/312
    #[allow(clippy::too_many_arguments)]
    fn handle_cleartext_payload<'a, Pub: event::ConnectionPublisher>(
        &mut self,
        packet_number: PacketNumber,
        payload: DecoderBufferMut<'a>,
        datagram: &'a DatagramInfo,
        path_id: path::Id,
        path_manager: &mut path::Manager<Config>,
        handshake_status: &mut HandshakeStatus,
        local_id_registry: &mut connection::LocalIdRegistry,
        random_generator: &mut Config::RandomGenerator,
        publisher: &mut Pub,
        packet_interceptor: &mut Config::PacketInterceptor,
    ) -> Result<ProcessedPacket<'a>, connection::Error> {
        use s2n_quic_core::{
            frame::{Frame, FrameMut},
            varint::VarInt,
        };

        let mut payload = {
            use s2n_quic_core::packet::interceptor::{Interceptor, Packet};

            // intercept the payload after it is decrypted, but before we process the frames
            packet_interceptor.intercept_rx_payload(
                &publisher.subject(),
                &Packet {
                    number: packet_number,
                    timestamp: datagram.timestamp,
                },
                payload,
            )
        };

        let mut processed_packet = ProcessedPacket::new(packet_number, datagram);

        macro_rules! on_frame_processed {
            ($frame:ident) => {{
                let frame_type = $frame.tag();
                processed_packet.on_processed_frame(&$frame);
                move |err: transport::Error| err.with_frame_type(VarInt::from_u8(frame_type))
            }};
        }

        while !payload.is_empty() {
            let (frame, remaining) = payload
                .decode::<FrameMut>()
                .map_err(transport::Error::from)?;

            let path = &path_manager[path_id];
            publisher.on_frame_received(event::builder::FrameReceived {
                packet_header: event::builder::PacketHeader::new(
                    packet_number,
                    publisher.quic_version(),
                ),
                path: path_event!(path, path_id),
                frame: frame.into_event(),
            });

            match frame {
                Frame::Padding(frame) => {
                    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.1
                    //# A PADDING frame (type=0x00) has no semantic value.  PADDING frames
                    //# can be used to increase the size of a packet.  Padding can be used to
                    //# increase an Initial packet to the minimum required size or to provide
                    //# protection against traffic analysis for protected packets.
                    let _ = on_frame_processed!(frame);
                }
                Frame::Ping(frame) => {
                    //= https://www.rfc-editor.org/rfc/rfc9000#section-19.2
                    //# Endpoints can use PING frames (type=0x01) to verify that their peers
                    //# are still alive or to check reachability to the peer.
                    let _ = on_frame_processed!(frame);
                }
                Frame::Crypto(frame) => {
                    let on_error = on_frame_processed!(frame);

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-7.5
                    //# Packets containing
                    //# discarded CRYPTO frames MUST be acknowledged because the packet has
                    //# been received and processed by the transport even though the CRYPTO
                    //# frame was discarded.

                    self.handle_crypto_frame(
                        frame.into(),
                        datagram,
                        &mut path_manager[path_id],
                        publisher,
                    )
                    .map_err(on_error)?;
                }
                Frame::Ack(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_ack_frame(
                        frame,
                        datagram.timestamp,
                        path_id,
                        path_manager,
                        packet_number,
                        handshake_status,
                        local_id_registry,
                        random_generator,
                        publisher,
                    )
                    .map_err(on_error)?;
                }
                Frame::ConnectionClose(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_connection_close_frame(
                        frame,
                        datagram.timestamp,
                        &mut path_manager[path_id],
                    )
                    .map_err(on_error)?;

                    // skip processing any other frames and return an error
                    return Err(frame.into());
                }
                Frame::Stream(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_stream_frame(frame.into(), &mut processed_packet)
                        .map_err(on_error)?;
                }
                Frame::Datagram(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_datagram_frame(
                        path_event!(path, path_id).into_event(),
                        frame.into(),
                    )
                    .map_err(on_error)?;
                }
                Frame::DataBlocked(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_data_blocked_frame(frame).map_err(on_error)?;
                }
                Frame::MaxData(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_max_data_frame(frame).map_err(on_error)?;
                }
                Frame::MaxStreamData(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_max_stream_data_frame(frame).map_err(on_error)?;
                }
                Frame::MaxStreams(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_max_streams_frame(frame).map_err(on_error)?;
                }
                Frame::ResetStream(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_reset_stream_frame(frame).map_err(on_error)?;
                }
                Frame::StopSending(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_stop_sending_frame(frame).map_err(on_error)?;
                }
                Frame::StreamDataBlocked(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_stream_data_blocked_frame(frame)
                        .map_err(on_error)?;
                }
                Frame::StreamsBlocked(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_streams_blocked_frame(frame).map_err(on_error)?;
                }
                Frame::NewToken(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_new_token_frame(frame).map_err(on_error)?;
                }
                Frame::NewConnectionId(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_new_connection_id_frame(frame, datagram, path_manager, publisher)
                        .map_err(on_error)?;
                }
                Frame::RetireConnectionId(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_retire_connection_id_frame(
                        frame,
                        datagram,
                        &mut path_manager[path_id],
                        local_id_registry,
                    )
                    .map_err(on_error)?;
                }
                Frame::PathChallenge(frame) => {
                    let on_error = on_frame_processed!(frame);

                    //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.3
                    //# An endpoint that receives a PATH_CHALLENGE on an active path SHOULD
                    //# send a non-probing packet in response.
                    if path_manager.active_path_id() == path_id {
                        processed_packet.path_challenge_on_active_path = true;
                    }
                    self.handle_path_challenge_frame(frame, path_id, path_manager)
                        .map_err(on_error)?;
                }
                Frame::PathResponse(frame) => {
                    let on_error = on_frame_processed!(frame);

                    self.handle_path_response_frame(frame, path_manager, publisher)
                        .map_err(on_error)?;
                }
                Frame::HandshakeDone(frame) => {
                    let on_error = on_frame_processed!(frame);
                    self.handle_handshake_done_frame(
                        frame,
                        datagram.timestamp,
                        &mut path_manager[path_id],
                        local_id_registry,
                        handshake_status,
                        publisher,
                    )
                    .map_err(on_error)?;
                }
            }

            payload = remaining;
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-12.4
        //# The payload of a packet that contains frames MUST contain at least
        //# one frame, and MAY contain multiple frames and multiple frame types.
        //# An endpoint MUST treat receipt of a packet containing no frames as a
        //# connection error of type PROTOCOL_VIOLATION.
        if processed_packet.frames == 0 {
            return Err(transport::Error::PROTOCOL_VIOLATION
                .with_reason("packet contained no frames")
                .into());
        }

        path_manager.on_processed_packet(
            path_id,
            datagram.source_connection_id,
            processed_packet.path_validation_probing,
            random_generator,
            publisher,
        )?;

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.1
        //# A packet MUST NOT be acknowledged until packet protection has been
        //# successfully removed and all frames contained in the packet have been
        //# processed.  For STREAM frames, this means the data has been enqueued
        //# in preparation to be received by the application protocol, but it
        //# does not require that data be delivered and consumed.
        //#
        //# Once the packet has been fully processed, a receiver acknowledges
        //# receipt by sending one or more ACK frames containing the packet
        //# number of the received packet.

        self.on_processed_packet(processed_packet, path_id, &path_manager[path_id], publisher)?;

        Ok(processed_packet)
    }
}


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

use crate::{
    ack::AckManager,
    connection::{self, limits::Limits},
    endpoint, path,
    space::{
        datagram, keep_alive::KeepAlive, ApplicationSpace, HandshakeSpace, HandshakeStatus,
        InitialSpace,
    },
    stream,
};
use bytes::Bytes;
use core::{ops::Not, task::Waker};
use s2n_codec::{DecoderBuffer, DecoderValue};
use s2n_quic_core::{
    ack,
    application::ServerName,
    connection::{InitialId, PeerId},
    crypto,
    crypto::{tls, CryptoSuite, Key},
    ct::ConstantTimeEq,
    datagram::{ConnectionInfo, Endpoint},
    event,
    event::IntoEvent,
    packet::number::PacketNumberSpace,
    time::Timestamp,
    transport::{
        self,
        parameters::{
            ActiveConnectionIdLimit, ClientTransportParameters, DatagramLimits,
            InitialFlowControlLimits, InitialSourceConnectionId, MaxAckDelay,
            ServerTransportParameters,
        },
    },
};

pub struct SessionContext<'a, Config: endpoint::Config, Pub: event::ConnectionPublisher> {
    pub now: Timestamp,
    pub initial_cid: &'a InitialId,
    pub retry_cid: Option<&'a PeerId>,
    pub path_manager: &'a mut path::Manager<Config>,
    pub initial: &'a mut Option<Box<InitialSpace<Config>>>,
    pub handshake: &'a mut Option<Box<HandshakeSpace<Config>>>,
    pub application: &'a mut Option<Box<ApplicationSpace<Config>>>,
    pub zero_rtt_crypto: &'a mut Option<
        Box<<<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::ZeroRttKey>,
    >,
    pub handshake_status: &'a mut HandshakeStatus,
    pub local_id_registry: &'a mut connection::LocalIdRegistry,
    pub limits: &'a mut Limits,
    pub server_name: &'a mut Option<ServerName>,
    pub application_protocol: &'a mut Bytes,
    pub waker: &'a Waker,
    pub publisher: &'a mut Pub,
    pub datagram: &'a mut Config::DatagramEndpoint,
}

impl<'a, Config: endpoint::Config, Pub: event::ConnectionPublisher>
    SessionContext<'a, Config, Pub>
{
    // This is called by the client
    fn on_server_params(
        &mut self,
        decoder: DecoderBuffer,
    ) -> Result<
        (
            InitialFlowControlLimits,
            ActiveConnectionIdLimit,
            DatagramLimits,
            MaxAckDelay,
        ),
        transport::Error,
    > {
        debug_assert!(Config::ENDPOINT_TYPE.is_client());

        let (peer_parameters, remaining) =
            ServerTransportParameters::decode(decoder).map_err(|_| {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.4
                //# An endpoint SHOULD treat receipt of
                //# duplicate transport parameters as a connection error of type
                //# TRANSPORT_PARAMETER_ERROR.
                transport::Error::TRANSPORT_PARAMETER_ERROR
                    .with_reason("Invalid transport parameters")
            })?;

        debug_assert_eq!(remaining.len(), 0);
        self.publisher.on_transport_parameters_received(
            event::builder::TransportParametersReceived {
                transport_parameters: peer_parameters.into_event(),
            },
        );

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
        //# An endpoint MUST treat the following as a connection error of type
        //# TRANSPORT_PARAMETER_ERROR or PROTOCOL_VIOLATION:
        self.validate_initial_source_connection_id(
            &peer_parameters.initial_source_connection_id,
            self.path_manager
                .active_path()
                .peer_connection_id
                .as_bytes(),
        )?;

        match (self.retry_cid, peer_parameters.retry_source_connection_id) {
            (Some(retry_packet_value), Some(transport_params_value)) => {
                if retry_packet_value
                    .as_bytes()
                    .ct_eq(transport_params_value.as_bytes())
                    .not()
                    .into()
                {
                    return Err(transport::Error::TRANSPORT_PARAMETER_ERROR
                        .with_reason("retry_source_connection_id mismatch"));
                }
            }
            (Some(_), None) => {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
                //# *  absence of the retry_source_connection_id transport parameter from
                //# the server after receiving a Retry packet,
                return Err(transport::Error::TRANSPORT_PARAMETER_ERROR.with_reason(
                    "retry_source_connection_id transport parameter absent \
                    after receiving a Retry packet from the server",
                ));
            }
            (None, Some(_)) => {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
                //# *  presence of the retry_source_connection_id transport parameter
                //# when no Retry packet was received, or
                return Err(transport::Error::TRANSPORT_PARAMETER_ERROR.with_reason(
                    "retry_source_connection_id transport parameter present \
                    when no Retry packet was received",
                ));
            }
            (None, None) => {}
        }

        if let Some(peer_value) = peer_parameters.original_destination_connection_id {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
            //# The values provided by a peer for these transport parameters MUST
            //# match the values that an endpoint used in the Destination and Source
            //# Connection ID fields of Initial packets that it sent (and received,
            //# for servers).  Endpoints MUST validate that received transport
            //# parameters match received connection ID values.
            if peer_value
                .as_bytes()
                .ct_eq(self.initial_cid.as_bytes())
                .not()
                .into()
            {
                return Err(transport::Error::TRANSPORT_PARAMETER_ERROR
                    .with_reason("original_destination_connection_id mismatch"));
            }
        } else {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
            //# An endpoint MUST treat the absence of the
            //# initial_source_connection_id transport parameter from either endpoint
            //# or the absence of the original_destination_connection_id transport
            //# parameter from the server as a connection error of type
            //# TRANSPORT_PARAMETER_ERROR.
            return Err(transport::Error::TRANSPORT_PARAMETER_ERROR
                .with_reason("missing original_destination_connection_id"));
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-10.3
        //# Servers can also issue a stateless_reset_token transport parameter during the
        //# handshake that applies to the connection ID that it selected during
        //# the handshake.  These exchanges are protected by encryption, so only
        //# client and server know their value.  Note that clients cannot use the
        //# stateless_reset_token transport parameter because their transport
        //# parameters do not have confidentiality protection.
        if let Some(stateless_reset_token) = peer_parameters.stateless_reset_token {
            self.path_manager
                .peer_id_registry
                .register_initial_stateless_reset_token(stateless_reset_token);
        }

        // Load the peer's transport parameters into the connection's limits
        self.limits.load_peer(&peer_parameters);

        let initial_flow_control_limits = peer_parameters.flow_control_limits();
        let active_connection_id_limit = peer_parameters.active_connection_id_limit;
        let datagram_limits = peer_parameters.datagram_limits();

        Ok((
            initial_flow_control_limits,
            active_connection_id_limit,
            datagram_limits,
            peer_parameters.max_ack_delay,
        ))
    }

    // This is called by the server
    fn on_client_params(
        &mut self,
        decoder: DecoderBuffer,
    ) -> Result<
        (
            InitialFlowControlLimits,
            ActiveConnectionIdLimit,
            DatagramLimits,
            MaxAckDelay,
        ),
        transport::Error,
    > {
        debug_assert!(Config::ENDPOINT_TYPE.is_server());

        let (peer_parameters, remaining) =
            ClientTransportParameters::decode(decoder).map_err(|_| {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-7.4
                //# An endpoint SHOULD treat receipt of
                //# duplicate transport parameters as a connection error of type
                //# TRANSPORT_PARAMETER_ERROR.
                transport::Error::TRANSPORT_PARAMETER_ERROR
                    .with_reason("Invalid transport parameters")
            })?;

        debug_assert_eq!(remaining.len(), 0);
        self.publisher.on_transport_parameters_received(
            event::builder::TransportParametersReceived {
                transport_parameters: peer_parameters.into_event(),
            },
        );

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
        //# An endpoint MUST treat the following as a connection error of type
        //# TRANSPORT_PARAMETER_ERROR or PROTOCOL_VIOLATION:
        self.validate_initial_source_connection_id(
            &peer_parameters.initial_source_connection_id,
            self.path_manager
                .active_path()
                .peer_connection_id
                .as_bytes(),
        )?;

        // Load the peer's transport parameters into the connection's limits
        self.limits.load_peer(&peer_parameters);

        let initial_flow_control_limits = peer_parameters.flow_control_limits();
        let active_connection_id_limit = peer_parameters.active_connection_id_limit;
        let datagram_limits = peer_parameters.datagram_limits();

        Ok((
            initial_flow_control_limits,
            active_connection_id_limit,
            datagram_limits,
            peer_parameters.max_ack_delay,
        ))
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
    //# Each endpoint includes the value of the Source Connection ID field
    //# from the first Initial packet it sent in the
    //# initial_source_connection_id transport parameter
    //
    // When the endpoint is a Server this is the peer's connection id.
    //
    // When the endpoint is a Client, this is the randomly generated
    // initial_connection_id which is locally generated for the first Initial packet.
    fn validate_initial_source_connection_id(
        &self,
        peer_value: &Option<InitialSourceConnectionId>,
        expected_value: &[u8],
    ) -> Result<(), transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
        //# * a mismatch between values received from a peer in these transport
        //# parameters and the value sent in the corresponding Destination or
        //# Source Connection ID fields of Initial packets.
        if let Some(peer_value) = peer_value {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
            //# The values provided by a peer for these transport parameters MUST
            //# match the values that an endpoint used in the Destination and Source
            //# Connection ID fields of Initial packets that it sent (and received,
            //# for servers).  Endpoints MUST validate that received transport
            //# parameters match received connection ID values.
            if peer_value.as_bytes().ct_eq(expected_value).not().into() {
                return Err(transport::Error::TRANSPORT_PARAMETER_ERROR
                    .with_reason("initial_source_connection_id mismatch"));
            }
        } else {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.3
            //# An endpoint MUST treat the absence of the
            //# initial_source_connection_id transport parameter from either endpoint
            //# or the absence of the original_destination_connection_id transport
            //# parameter from the server as a connection error of type
            //# TRANSPORT_PARAMETER_ERROR.
            return Err(transport::Error::TRANSPORT_PARAMETER_ERROR
                .with_reason("missing initial_source_connection_id"));
        }

        Ok(())
    }
}

impl<'a, Config: endpoint::Config, Pub: event::ConnectionPublisher>
    tls::Context<<Config::TLSEndpoint as tls::Endpoint>::Session>
    for SessionContext<'a, Config, Pub>
{
    fn on_handshake_keys(
        &mut self,
        key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::HandshakeKey,
        header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::HandshakeHeaderKey,
    ) -> Result<(), transport::Error> {
        if self.handshake.is_some() {
            return Err(transport::Error::INTERNAL_ERROR
                .with_reason("handshake keys initialized more than once"));
        }

        // After receiving handshake keys, the initial crypto stream should be completely
        // finished
        if let Some(space) = self.initial.as_mut() {
            space.crypto_stream.finish()?;
        }

        let ack_manager = AckManager::new(PacketNumberSpace::Handshake, ack::Settings::EARLY);

        let cipher_suite = key.cipher_suite().into_event();
        *self.handshake = Some(Box::new(HandshakeSpace::new(
            key,
            header_key,
            self.now,
            ack_manager,
        )));
        self.publisher.on_key_update(event::builder::KeyUpdate {
            key_type: event::builder::KeyType::Handshake,
            cipher_suite,
        });
        Ok(())
    }

    fn on_zero_rtt_keys(
        &mut self,
        key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::ZeroRttKey,
        _header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::ZeroRttHeaderKey,
        _application_parameters: tls::ApplicationParameters,
    ) -> Result<(), transport::Error> {
        if self.zero_rtt_crypto.is_some() {
            return Err(transport::Error::INTERNAL_ERROR
                .with_reason("zero rtt keys initialized more than once"));
        }

        let cipher_suite = key.cipher_suite().into_event();

        // TODO: also store the header_key https://github.com/aws/s2n-quic/issues/319
        *self.zero_rtt_crypto = Some(Box::new(key));

        self.publisher.on_key_update(event::builder::KeyUpdate {
            key_type: event::builder::KeyType::ZeroRtt,
            cipher_suite,
        });
        Ok(())
    }

    fn on_one_rtt_keys(
        &mut self,
        key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::OneRttKey,
        header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::OneRttHeaderKey,
        application_parameters: tls::ApplicationParameters,
    ) -> Result<(), transport::Error> {
        if self.application.is_some() {
            return Err(transport::Error::INTERNAL_ERROR
                .with_reason("application keys initialized more than once"));
        }

        if Config::ENDPOINT_TYPE.is_client() {
            //= https://www.rfc-editor.org/rfc/rfc9001#section-4.9.3
            //# Therefore, a client SHOULD discard 0-RTT keys as soon as it installs
            //# 1-RTT keys as they have no use after that moment.

            *self.zero_rtt_crypto = None;
        }

        // Parse transport parameters
        let param_decoder = DecoderBuffer::new(application_parameters.transport_parameters);
        let (peer_flow_control_limits, active_connection_id_limit, datagram_limits, max_ack_delay) =
            match Config::ENDPOINT_TYPE {
                endpoint::Type::Client => self.on_server_params(param_decoder)?,
                endpoint::Type::Server => self.on_client_params(param_decoder)?,
            };

        self.local_id_registry
            .set_active_connection_id_limit(active_connection_id_limit.as_u64());

        let stream_manager = <Config::StreamManager as stream::Manager>::new(
            self.limits,
            Config::ENDPOINT_TYPE,
            self.limits.initial_flow_control_limits(),
            peer_flow_control_limits,
        );

        let ack_manager = AckManager::new(
            PacketNumberSpace::ApplicationData,
            self.limits.ack_settings(),
        );

        let keep_alive = KeepAlive::new(
            self.limits.max_idle_timeout(),
            self.limits.max_keep_alive_period(),
        );

        let conn_info =
            ConnectionInfo::new(datagram_limits.max_datagram_payload, self.waker.clone());
        let (datagram_sender, datagram_receiver) = self.datagram.create_connection(&conn_info);
        let datagram_manager = datagram::Manager::new(
            datagram_sender,
            datagram_receiver,
            datagram_limits.max_datagram_payload,
        );

        self.path_manager
            .active_path_mut()
            .rtt_estimator
            .on_max_ack_delay(max_ack_delay);

        let cipher_suite = key.cipher_suite().into_event();
        let max_mtu = self.path_manager.max_mtu();
        *self.application = Some(Box::new(ApplicationSpace::new(
            key,
            header_key,
            self.now,
            stream_manager,
            ack_manager,
            keep_alive,
            max_mtu,
            datagram_manager,
        )));
        self.publisher.on_key_update(event::builder::KeyUpdate {
            key_type: event::builder::KeyType::OneRtt { generation: 0 },
            cipher_suite,
        });

        Ok(())
    }

    fn on_server_name(&mut self, server_name: ServerName) -> Result<(), transport::Error> {
        self.publisher
            .on_server_name_information(event::builder::ServerNameInformation {
                chosen_server_name: &server_name,
            });
        *self.server_name = Some(server_name);

        Ok(())
    }

    fn on_application_protocol(
        &mut self,
        application_protocol: Bytes,
    ) -> Result<(), transport::Error> {
        self.publisher.on_application_protocol_information(
            event::builder::ApplicationProtocolInformation {
                chosen_application_protocol: &application_protocol,
            },
        );
        *self.application_protocol = application_protocol;

        Ok(())
    }

    fn on_handshake_complete(&mut self) -> Result<(), transport::Error> {
        // After the handshake is complete, the handshake crypto stream should be completely
        // finished
        if let Some(space) = self.handshake.as_mut() {
            space.crypto_stream.finish()?;
        }

        if self.application_protocol.is_empty() {
            //= https://www.rfc-editor.org/rfc/rfc9001#section-8.1
            //# When using ALPN, endpoints MUST immediately close a connection (see
            //# Section 10.2 of [QUIC-TRANSPORT]) with a no_application_protocol TLS
            //# alert (QUIC error code 0x178; see Section 4.8) if an application
            //# protocol is not negotiated.

            //= https://www.rfc-editor.org/rfc/rfc9001#section-8.1
            //# While [ALPN] only specifies that servers
            //# use this alert, QUIC clients MUST use error 0x178 to terminate a
            //# connection when ALPN negotiation fails.
            let err = crypto::CryptoError::NO_APPLICATION_PROTOCOL
                .with_reason("Missing ALPN protocol")
                .into();
            return Err(err);
        }

        self.handshake_status
            .on_handshake_complete(Config::ENDPOINT_TYPE, self.publisher);

        if let Some(application) = self.application.as_mut() {
            if Config::ENDPOINT_TYPE.is_server() {
                // All of the other spaces are discarded by the time the handshake is complete so
                // we only need to notify the application space
                //
                //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.2
                //# the TLS handshake is considered confirmed at the
                //# server when the handshake completes.
                application.on_handshake_confirmed(
                    self.path_manager.active_path(),
                    self.local_id_registry,
                    self.now,
                );
            }
            Ok(())
        } else {
            Err(transport::Error::INTERNAL_ERROR
                .with_reason("handshake cannot be completed without application keys"))
        }
    }

    fn receive_initial(&mut self, max_len: Option<usize>) -> Option<Bytes> {
        let space = self.initial.as_deref_mut()?;

        // don't pass the buffer until we have a full hello message
        if !space.received_hello_message {
            return None;
        }

        space
            .crypto_stream
            .rx
            .pop_watermarked(max_len.unwrap_or(usize::MAX))
            .map(|bytes| bytes.freeze())
    }

    fn receive_handshake(&mut self, max_len: Option<usize>) -> Option<Bytes> {
        self.handshake
            .as_deref_mut()?
            .crypto_stream
            .rx
            .pop_watermarked(max_len.unwrap_or(usize::MAX))
            .map(|bytes| bytes.freeze())
    }

    fn receive_application(&mut self, _max_len: Option<usize>) -> Option<Bytes> {
        // Application doesn't currently have a buffer
        None
    }

    fn can_send_initial(&self) -> bool {
        self.initial
            .as_ref()
            .map(|space| space.crypto_stream.can_send())
            .unwrap_or_default()
    }

    fn send_initial(&mut self, transmission: Bytes) {
        self.initial
            .as_mut()
            .expect("can_send_initial should be called before sending")
            .crypto_stream
            .tx
            .push(transmission);
    }

    fn can_send_handshake(&self) -> bool {
        self.handshake
            .as_ref()
            .map(|space| space.crypto_stream.can_send())
            .unwrap_or_default()
    }

    fn send_handshake(&mut self, transmission: Bytes) {
        self.handshake
            .as_mut()
            .expect("can_send_handshake should be called before sending")
            .crypto_stream
            .tx
            .push(transmission);
    }

    fn can_send_application(&self) -> bool {
        false
    }

    fn send_application(&mut self, _transmission: Bytes) {
        unimplemented!("application level crypto frames cannot currently be sent")
    }

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


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

use s2n_quic_core::{
    ack,
    packet::number::{PacketNumber, PacketNumberSpace},
    time::Timestamp,
    transport,
    varint::VarInt,
};

/// Context for tracking transmission of packet numbers
#[derive(Debug)]
pub struct TxPacketNumbers {
    largest_sent_acked: (PacketNumber, Timestamp),
    next: PacketNumber,
}

impl TxPacketNumbers {
    pub fn new(packet_space: PacketNumberSpace, now: Timestamp) -> Self {
        let initial_packet_number = packet_space.new_packet_number(VarInt::from_u8(0));
        Self {
            largest_sent_acked: (initial_packet_number, now),
            next: initial_packet_number,
        }
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(
        &mut self,
        timestamp: Timestamp,
        ack_set: &A,
    ) -> Result<(), transport::Error> {
        let largest = ack_set.largest();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.1
        //# An endpoint SHOULD treat receipt of an acknowledgment for a packet it
        //# did not send as a connection error of type PROTOCOL_VIOLATION, if it
        //# is able to detect the condition.

        if largest >= self.next {
            return Err(transport::Error::PROTOCOL_VIOLATION
                .with_reason("received an ACK for a packet that was not sent"));
        }

        // record the largest packet acked
        if largest > self.largest_sent_acked.0 {
            self.largest_sent_acked = (largest, timestamp);
        }

        Ok(())
    }

    /// Called after a packet is transmitted with a given packet number
    pub fn on_transmit(&mut self, packet_number: PacketNumber) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-12.3
        //# A QUIC endpoint MUST NOT reuse a packet number within the same packet
        //# number space in one connection.  If the packet number for sending
        //# reaches 2^62-1, the sender MUST close the connection without sending
        //# a CONNECTION_CLOSE frame or any further packets; an endpoint MAY send
        //# a Stateless Reset (Section 10.3) in response to further packets that
        //# it receives.

        // Assuming a constant TX rate of 10,000 packets/second, it would take
        // (2^62 - 1)packets / 10,000packets-per-second  ~= 4.6e14seconds = ~14,000,000years to overflow.
        // It's probably OK to just panic as there will be other things to worry about at that point.

        self.next = packet_number.next().expect("packet number overflowed");
    }

    /// Returns the next packet number in the sequence
    pub fn next(&self) -> PacketNumber {
        self.next
    }

    /// Returns the largest packet number that the endpoint sent that
    /// was ACKed by the peer
    pub fn largest_sent_packet_number_acked(&self) -> PacketNumber {
        self.largest_sent_acked.0
    }
}


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

use crate::{
    contexts::testing::*,
    endpoint,
    space::HandshakeStatus,
    transmission::{self, interest::Provider},
};
use bolero::{check, generator::*};
use s2n_quic_core::{ack, event::testing::Publisher, packet, time::clock::testing as time};

#[derive(Debug)]
struct Oracle {
    pending: bool,
    complete: bool,
    confirmed: bool,
    endpoint_type: endpoint::Type,

    // SERVER
    handshake_done_sent: bool,
    handshake_done_delivered: bool,

    // Client
    handshake_done_received: bool,
}

impl Oracle {
    fn new(endpoint_type: endpoint::Type) -> Self {
        Oracle {
            pending: true,
            complete: false,
            confirmed: false,
            endpoint_type,
            handshake_done_sent: false,
            handshake_done_delivered: false,
            handshake_done_received: false,
        }
    }

    fn can_transmit(&self) -> bool {
        // if the server endpoint is complete but has not yet
        // delivered the HANDSHAKE_DONE
        if self.endpoint_type.is_server() && self.complete && !self.handshake_done_delivered {
            return true;
        }

        false
    }

    fn on_handshake_complete(&mut self) {
        self.pending = false;
        self.complete = true;
        if self.endpoint_type.is_server() {
            self.confirmed = true;
        }
    }

    fn on_transmit(&mut self) {
        if self.endpoint_type.is_client() {
            return;
        }

        self.handshake_done_sent = true;
    }

    fn on_packet_ack(&mut self, ack_handshake_done: bool) {
        if self.endpoint_type.is_client() {
            return;
        }

        // check that the HANDSHAKE_DONE frame was sent
        if self.complete && self.handshake_done_sent && ack_handshake_done {
            self.handshake_done_delivered = true;
        }
    }

    fn on_packet_loss(&mut self, lost_handshake_done: bool) {
        if self.endpoint_type.is_client() {
            return;
        }

        // check that the HANDSHAKE_DONE frame was sent but not yet delivered
        if self.handshake_done_sent && !self.handshake_done_delivered && lost_handshake_done {
            self.handshake_done_sent = false;
        }
    }

    fn on_handshake_done_received(&mut self) {
        if self.endpoint_type.is_server() {
            return;
        }

        if self.complete {
            self.confirmed = true;
            self.handshake_done_received = true;
        }
    }
}

#[derive(Debug, TypeGenerator)]
enum Operation {
    Complete,

    // SERVER: below are server specific operations
    TransmitHandshakeDone {},
    AckHandshakeDonePacket { ack_handshake_done: bool },
    LostHandshakeDonePacket { lost_handshake_done: bool },

    // Client: below are client specific operations
    ReceivedHandshakeDone {},
}

struct Model {
    subject: HandshakeStatus,
    oracle: Oracle,
}

impl Model {
    fn new(endpoint_type: endpoint::Type) -> Self {
        Model {
            subject: HandshakeStatus::InProgress,
            oracle: Oracle::new(endpoint_type),
        }
    }

    fn apply(&mut self, operation: &Operation) {
        match operation {
            Operation::Complete => self.on_complete(),
            Operation::TransmitHandshakeDone {} => self.packet_transmit(),
            Operation::AckHandshakeDonePacket { ack_handshake_done } => {
                self.packet_acked(*ack_handshake_done)
            }
            Operation::LostHandshakeDonePacket {
                lost_handshake_done,
            } => self.packet_loss(*lost_handshake_done),
            Operation::ReceivedHandshakeDone {} => self.on_handshake_done_received(),
        }
    }

    fn on_complete(&mut self) {
        if !self.oracle.complete {
            self.subject
                .on_handshake_complete(self.oracle.endpoint_type, &mut Publisher::no_snapshot());
            self.oracle.on_handshake_complete();
        }
    }

    fn packet_transmit(&mut self) {
        if self.oracle.can_transmit() {
            assert!(matches!(
                self.subject,
                HandshakeStatus::ServerCompleteConfirmed(_)
            ));

            let mut frame_buffer = OutgoingFrameBuffer::new();
            let mut context = MockWriteContext::new(
                time::now(),
                &mut frame_buffer,
                // always allow for transmission for the purpose of this test
                transmission::Constraint::None,
                transmission::Mode::Normal,
                self.oracle.endpoint_type,
            );
            self.subject.on_transmit(&mut context);
            // verify that the subject wrote a frame
            assert!(!frame_buffer.is_empty());

            self.oracle.on_transmit();
        }
    }

    fn packet_acked(&mut self, ack_handshake_done: bool) {
        self.subject.on_packet_ack(
            &AckSetMock(ack_handshake_done),
            &mut Publisher::no_snapshot(),
        );
        self.oracle.on_packet_ack(ack_handshake_done);
    }

    fn packet_loss(&mut self, lost_handshake_done: bool) {
        self.subject.on_packet_loss(
            &AckSetMock(lost_handshake_done),
            &mut Publisher::no_snapshot(),
        );

        // perform some checks before calling `oracle.on_packet_loss`
        if self.oracle.endpoint_type.is_server() {
            // check that the HANDSHAKE_DONE frame was sent but not yet delivered
            if self.oracle.handshake_done_sent && !self.oracle.handshake_done_delivered {
                assert!(matches!(
                    self.subject,
                    HandshakeStatus::ServerCompleteConfirmed(_)
                ));

                if lost_handshake_done {
                    if let HandshakeStatus::ServerCompleteConfirmed(flag) = &self.subject {
                        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.3
                        //= type=test
                        //# The HANDSHAKE_DONE frame MUST be retransmitted until it is
                        //# acknowledged.
                        assert!(flag.has_transmission_interest());
                    }
                }
            }
        }

        self.oracle.on_packet_loss(lost_handshake_done);
    }

    fn on_handshake_done_received(&mut self) {
        self.subject
            .on_handshake_done_received(&mut Publisher::no_snapshot());
        self.oracle.on_handshake_done_received();
    }

    fn invariants(&self) {
        assert_eq!(self.subject.is_complete(), self.oracle.complete);
        assert_eq!(self.subject.is_confirmed(), self.oracle.confirmed);

        assert_eq!(
            self.oracle.pending,
            matches!(self.subject, HandshakeStatus::InProgress)
        );
        if matches!(self.subject, HandshakeStatus::Confirmed) {
            assert!(self.oracle.complete);
            assert!(self.oracle.confirmed);
        }

        // SERVER
        if self.oracle.endpoint_type.is_server() {
            if self.oracle.complete {
                assert!(self.subject.is_confirmed());
            }

            if self.oracle.handshake_done_sent {
                assert!(!matches!(self.subject, HandshakeStatus::InProgress));
                if let HandshakeStatus::ServerCompleteConfirmed(flag) = &self.subject {
                    // HANDSHAKE_DONE should either needs transmission or is in-flight
                    assert!(!flag.is_idle());
                }
            }
            if self.oracle.handshake_done_delivered {
                assert!(matches!(self.subject, HandshakeStatus::Confirmed));
            }
        }

        // CLIENT
        if self.oracle.endpoint_type.is_client() {
            // If handshake is complete but awaiting confirmation
            if self.subject.is_complete() && !self.oracle.handshake_done_received {
                assert!(!self.subject.is_confirmed());
                assert!(matches!(self.subject, HandshakeStatus::ClientComplete));
            }

            if self.oracle.handshake_done_received {
                assert!(self.subject.is_confirmed());
            }
        }
    }
}

struct AckSetMock(bool);
impl ack::Set for AckSetMock {
    fn contains(&self, _packet_number: packet::number::PacketNumber) -> bool {
        self.0
    }

    fn smallest(&self) -> packet::number::PacketNumber {
        todo!("unused for this test")
    }

    fn largest(&self) -> packet::number::PacketNumber {
        todo!("unused for this test")
    }
}

#[test]
fn handshake_status_fuzz() {
    check!()
        .with_type::<Vec<Operation>>()
        .for_each(|operations| {
            let mut server_model = Model::new(endpoint::Type::Server);
            let mut client_model = Model::new(endpoint::Type::Client);

            for operation in operations.iter() {
                server_model.apply(operation);
                client_model.apply(operation);
            }

            server_model.invariants();
            client_model.invariants();
        });
}