// 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::{datagram, keep_alive::KeepAlive, HandshakeStatus, PacketSpace, TxPacketNumbers},
    stream::Manager as _,
    sync::flag,
    transmission,
    transmission::interest::Provider,
};
use core::{convert::TryInto, fmt, marker::PhantomData};
use once_cell::sync::OnceCell;
use s2n_codec::EncoderBuffer;
use s2n_quic_core::{
    crypto::{application::KeySet, limited, tls, CryptoSuite},
    event::{self, ConnectionPublisher as _, 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::{
        encoding::{PacketEncoder, PacketEncodingError},
        number::{PacketNumber, PacketNumberRange, PacketNumberSpace, SlidingWindow},
        short::{CleartextShort, ProtectedShort, Short, SpinBit},
    },
    path::MaxMtu,
    time::{timer, Timestamp},
    transport,
};

pub struct ApplicationSpace<Config: endpoint::Config> {
    /// Transmission Packet numbers
    pub tx_packet_numbers: TxPacketNumbers,
    /// Ack manager
    pub ack_manager: AckManager,
    /// All streams that are managed through this connection
    pub stream_manager: Config::StreamManager,
    /// The current state of the Spin bit
    /// TODO: Spin me
    pub spin_bit: SpinBit,
    /// The crypto suite for application data
    /// TODO: What about ZeroRtt?
    //= https://www.rfc-editor.org/rfc/rfc9001#section-6.3
    //# For this reason, endpoints MUST be able to retain two sets of packet
    //# protection keys for receiving packets: the current and the next.

    //= https://www.rfc-editor.org/rfc/rfc9001#section-6.1
    //# An endpoint MUST NOT initiate a key update prior to having confirmed
    //# the handshake (Section 4.1.2).
    key_set: KeySet<<<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::OneRttKey>,
    header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::OneRttHeaderKey,

    ping: flag::Ping,
    keep_alive: KeepAlive,
    processed_packet_numbers: SlidingWindow,
    recovery_manager: recovery::Manager<Config>,
    pub datagram_manager: datagram::Manager<Config>,
}

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

impl<Config: endpoint::Config> ApplicationSpace<Config> {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::OneRttKey,
        header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::OneRttHeaderKey,
        now: Timestamp,
        stream_manager: Config::StreamManager,
        ack_manager: AckManager,
        keep_alive: KeepAlive,
        max_mtu: MaxMtu,
        datagram_manager: datagram::Manager<Config>,
    ) -> Self {
        let key_set = KeySet::new(key, Self::key_limits(max_mtu));

        Self {
            tx_packet_numbers: TxPacketNumbers::new(PacketNumberSpace::ApplicationData, now),
            ack_manager,
            spin_bit: SpinBit::Zero,
            stream_manager,
            key_set,
            header_key,
            ping: flag::Ping::default(),
            keep_alive,
            processed_packet_numbers: SlidingWindow::default(),
            recovery_manager: recovery::Manager::new(PacketNumberSpace::ApplicationData),
            datagram_manager,
        }
    }

    /// 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: &mut 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 timestamp = context.timestamp;
        let transmission_mode = context.transmission_mode;
        let min_packet_len = context.min_packet_len;
        let bytes_progressed = self.stream_manager.outgoing_bytes_progressed();

        let payload = transmission::Transmission {
            config: <PhantomData<Config>>::default(),
            outcome: &mut outcome,
            packet_number,
            payload: transmission::application::Payload::<Config>::new(
                context.path_id,
                context.path_manager,
                context.local_id_registry,
                context.transmission_mode,
                &mut self.ack_manager,
                handshake_status,
                &mut self.ping,
                &mut self.stream_manager,
                &mut self.recovery_manager,
                &mut self.datagram_manager,
            ),
            timestamp,
            transmission_constraint,
            transmission_mode,
            tx_packet_numbers: &mut self.tx_packet_numbers,
            path_id: context.path_id,
            publisher: context.publisher,
            packet_interceptor: context.packet_interceptor,
        };

        let spin_bit = self.spin_bit;
        let header_key = &self.header_key;
        let (_protected_packet, buffer) =
            self.key_set
                .encrypt_packet(buffer, |buffer, key, key_phase| {
                    let packet = Short {
                        spin_bit,
                        key_phase,
                        destination_connection_id,
                        packet_number,
                        payload,
                    };
                    packet.encode_packet(
                        key,
                        header_key,
                        packet_number_encoder,
                        min_packet_len,
                        buffer,
                    )
                })?;

        outcome.bytes_progressed +=
            (self.stream_manager.outgoing_bytes_progressed() - bytes_progressed).as_u64() as usize;

        let app_limited = self.is_app_limited(context.path(), outcome.bytes_sent);

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

        // reset the keep alive timer after sending an ack-eliciting packet
        if outcome.ack_elicitation.is_ack_eliciting() {
            self.keep_alive.reset(timestamp);
        }

        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::ApplicationData,
            },
            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 spin_bit = self.spin_bit;
        let min_packet_len = context.min_packet_len;
        let header_key = &self.header_key;
        let (_protected_packet, buffer) =
            self.key_set
                .encrypt_packet(buffer, |buffer, key, key_phase| {
                    let packet = Short {
                        spin_bit,
                        key_phase,
                        destination_connection_id,
                        packet_number,
                        payload,
                    };
                    packet.encode_packet(
                        key,
                        header_key,
                        packet_number_encoder,
                        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);
    }

    /// Signals the handshake is confirmed
    pub fn on_handshake_confirmed(
        &mut self,
        path: &Path<Config>,
        local_id_registry: &mut connection::LocalIdRegistry,
        timestamp: Timestamp,
    ) {
        // Retire the local connection ID used during the handshake to reduce linkability
        local_id_registry.retire_handshake_connection_id();

        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# A sender SHOULD restart its PTO timer every time an ack-eliciting
        //# packet is sent or acknowledged, or when Initial or Handshake keys are
        //# discarded (Section 4.9 of [QUIC-TLS]).

        //= https://www.rfc-editor.org/rfc/rfc9002#section-6.2.1
        //# An endpoint MUST NOT set its PTO timer for the Application Data
        //# packet number space until the handshake is confirmed.

        // Since we maintain a separate PTO timer for each packet space, we don't have to update
        // it when the Initial or Handshake keys are discarded. However, we do need to update the
        // PTO timer when the handshake is confirmed, as the Application space PTO timer is not
        // started until the handshake is confirmed.
        self.recovery_manager
            .update_pto_timer(path, timestamp, true)
    }

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

        let (recovery_manager, mut context) = self.recovery(
            handshake_status,
            local_id_registry,
            path_manager.active_path_id(),
            path_manager,
        );

        recovery_manager.on_timeout(timestamp, random_generator, &mut context, publisher);

        self.stream_manager.on_timeout(timestamp);

        if self.keep_alive.on_timeout(timestamp).is_ready() {
            publisher.on_keep_alive_timer_expired(event::builder::KeepAliveTimerExpired {
                timeout: self.keep_alive.period(),
            });

            // send a ping after timing out
            self.ping();
        }
    }

    /// Returns `true` if the recovery manager for this packet space requires a probe
    /// packet to be sent.
    pub fn requires_probe(&self) -> bool {
        self.recovery_manager.requires_probe()
    }

    pub fn ping(&mut self) {
        self.ping.send()
    }

    pub fn keep_alive(&mut self, enabled: bool) {
        self.keep_alive.update(enabled);
    }

    /// 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 mut HandshakeStatus,
        local_id_registry: &'a mut connection::LocalIdRegistry,
        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,
                handshake_status,
                ping: &mut self.ping,
                stream_manager: &mut self.stream_manager,
                local_id_registry,
                path_id,
                path_manager,
                tx_packet_numbers: &mut self.tx_packet_numbers,
            },
        )
    }

    /// Returns `true` if sending is limited by the application and not the congestion controller
    ///
    /// Sending is app limited if the application is not fully utilizing the available
    /// congestion window currently and there is no more application data remaining to send.
    fn is_app_limited(&self, path: &Path<Config>, bytes_sent: usize) -> bool {
        !path.is_congestion_limited(bytes_sent) && !self.has_transmission_interest()
    }

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

        let packet_number = packet.packet_number;
        let packet_header =
            event::builder::PacketHeader::new(packet.packet_number, publisher.quic_version());
        let decrypted = self.key_set.decrypt_packet(
            packet,
            largest_acked,
            //= https://www.rfc-editor.org/rfc/rfc9001#section-6.3
            //# For a short period after a key
            //# update completes, up to the PTO, endpoints MAY defer generation of
            //# the next set of receive packet protection keys.  This allows
            //# endpoints to retain only two sets of receive keys; see Section 6.5.

            //= https://www.rfc-editor.org/rfc/rfc9001#section-6.5
            //# An endpoint MAY allow a period of approximately the Probe Timeout
            //# (PTO; see [QUIC-RECOVERY]) after promoting the next set of receive
            //# keys to be current before it creates the subsequent set of packet
            //# protection keys.
            datagram.timestamp
                + path
                    .rtt_estimator
                    .pto_period(1, PacketNumberSpace::ApplicationData),
        );
        match decrypted {
            Ok((_, Some(generation))) => {
                publisher.on_key_update(event::builder::KeyUpdate {
                    key_type: event::builder::KeyType::OneRtt { generation },
                    cipher_suite: self.key_set.cipher_suite().into_event(),
                });
            }
            Ok(_) => {}
            Err(_) => {
                publisher.on_packet_dropped(event::builder::PacketDropped {
                    reason: event::builder::PacketDropReason::DecryptionFailed {
                        packet_header,
                        path: path_event!(path, path_id),
                    },
                });
            }
        }
        //= https://www.rfc-editor.org/rfc/rfc9001#section-9.5
        //# For authentication to be
        //# free from side channels, the entire process of header protection
        //# removal, packet number recovery, and packet protection removal MUST
        //# be applied together without timing and other side channels.

        // We perform decryption prior to checking for duplicate to avoid short-circuiting
        // and maintain constant-time operation.
        if self.is_duplicate(packet_number, path_id, path, publisher) {
            return Err(ProcessingError::DuplicatePacket);
        }

        if decrypted.is_ok() {
            // reset the keep alive timer after receiving a packet
            self.keep_alive.reset(datagram.timestamp);
        }

        decrypted.map(|x| x.0)
    }

    fn key_limits(max_mtu: MaxMtu) -> limited::Limits {
        let mut limits = limited::Limits::default();

        limits.max_mtu = max_mtu;

        // AEAD optimizations are currently in the testing phase so make them opt-in at runtime
        limits.sealer_optimization_threshold = {
            static THRESHOLD: OnceCell<u64> = OnceCell::new();

            *THRESHOLD.get_or_init(|| {
                std::env::var("S2N_UNSTABLE_CRYPTO_OPT_TX")
                    .ok()
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(u64::MAX)
            })
        };

        limits.opener_optimization_threshold = {
            static THRESHOLD: OnceCell<u64> = OnceCell::new();

            *THRESHOLD.get_or_init(|| {
                std::env::var("S2N_UNSTABLE_CRYPTO_OPT_RX")
                    .ok()
                    .and_then(|v| v.parse().ok())
                    .unwrap_or(u64::MAX)
            })
        };

        limits
    }
}

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

        Ok(())
    }
}

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

impl<Config: endpoint::Config> connection::finalization::Provider for ApplicationSpace<Config> {
    fn finalization_status(&self) -> connection::finalization::Status {
        self.stream_manager.finalization_status()
    }
}

struct RecoveryContext<'a, Config: endpoint::Config> {
    ack_manager: &'a mut AckManager,
    handshake_status: &'a mut HandshakeStatus,
    ping: &'a mut flag::Ping,
    stream_manager: &'a mut Config::StreamManager,
    local_id_registry: &'a mut connection::LocalIdRegistry,
    path_id: path::Id,
    path_manager: &'a mut path::Manager<Config>,
    tx_packet_numbers: &'a mut TxPacketNumbers,
}

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.handshake_status
            .on_packet_ack(packet_number_range, publisher);
        self.ping.on_packet_ack(packet_number_range);
        self.stream_manager.on_packet_ack(packet_number_range);
        self.local_id_registry.on_packet_ack(packet_number_range);
        self.path_manager.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.ack_manager.on_packet_loss(packet_number_range);
        self.handshake_status
            .on_packet_loss(packet_number_range, publisher);
        self.ping.on_packet_loss(packet_number_range);
        self.stream_manager.on_packet_loss(packet_number_range);
        self.local_id_registry.on_packet_loss(packet_number_range);
        self.path_manager.on_packet_loss(packet_number_range);
    }

    fn on_rtt_update(&mut self) {
        // Update the stream manager if this RTT update was for the active path
        if self.path_manager.active_path_id() == self.path_id {
            self.stream_manager
                .on_rtt_update(&self.path_manager.active_path().rtt_estimator)
        }
    }
}

impl<Config: endpoint::Config> PacketSpace<Config> for ApplicationSpace<Config> {
    const INVALID_FRAME_ERROR: &'static str = "invalid frame in application space";

    fn handle_crypto_frame<Pub: event::ConnectionPublisher>(
        &mut self,
        _frame: CryptoRef,
        _datagram: &DatagramInfo,
        _path: &mut Path<Config>,
        _publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.5
        //# Once the handshake completes, if an endpoint is unable to buffer all
        //# data in a CRYPTO frame, it MAY discard that CRYPTO frame and all
        //# CRYPTO frames received in the future, or it MAY close the connection
        //# with a CRYPTO_BUFFER_EXCEEDED error code.

        // we currently just discard CRYPTO frames post-handshake
        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 path = &mut path_manager[path_id];
        path.on_peer_validated();
        let (recovery_manager, mut context) =
            self.recovery(handshake_status, local_id_registry, 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> {
        Ok(())
    }

    fn handle_stream_frame(
        &mut self,
        frame: StreamRef,
        packet: &mut ProcessedPacket,
    ) -> Result<(), transport::Error> {
        let bytes_progressed = self.stream_manager.incoming_bytes_progressed();

        self.stream_manager.on_data(&frame)?;

        packet.bytes_progressed +=
            (self.stream_manager.incoming_bytes_progressed() - bytes_progressed).as_u64() as usize;

        Ok(())
    }

    fn handle_datagram_frame(
        &mut self,
        path: s2n_quic_core::event::api::Path<'_>,
        frame: DatagramRef,
    ) -> Result<(), transport::Error> {
        self.datagram_manager.on_datagram_frame(path, frame);
        Ok(())
    }

    fn handle_data_blocked_frame(&mut self, frame: DataBlocked) -> Result<(), transport::Error> {
        self.stream_manager.on_data_blocked(frame)
    }

    fn handle_max_data_frame(&mut self, frame: MaxData) -> Result<(), transport::Error> {
        self.stream_manager.on_max_data(frame)
    }

    fn handle_max_stream_data_frame(
        &mut self,
        frame: MaxStreamData,
    ) -> Result<(), transport::Error> {
        self.stream_manager.on_max_stream_data(&frame)
    }

    fn handle_max_streams_frame(&mut self, frame: MaxStreams) -> Result<(), transport::Error> {
        self.stream_manager.on_max_streams(&frame)
    }

    fn handle_reset_stream_frame(&mut self, frame: ResetStream) -> Result<(), transport::Error> {
        self.stream_manager.on_reset_stream(&frame)
    }

    fn handle_stop_sending_frame(&mut self, frame: StopSending) -> Result<(), transport::Error> {
        self.stream_manager.on_stop_sending(&frame)
    }

    fn handle_stream_data_blocked_frame(
        &mut self,
        frame: StreamDataBlocked,
    ) -> Result<(), transport::Error> {
        self.stream_manager.on_stream_data_blocked(&frame)
    }

    fn handle_streams_blocked_frame(
        &mut self,
        frame: StreamsBlocked,
    ) -> Result<(), transport::Error> {
        self.stream_manager.on_streams_blocked(&frame)
    }

    fn handle_new_token_frame(&mut self, frame: NewToken) -> Result<(), transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.7
        //# A server MUST treat receipt
        //# of a NEW_TOKEN frame as a connection error of type
        //# PROTOCOL_VIOLATION.
        if Config::ENDPOINT_TYPE.is_server() {
            return Err(transport::Error::PROTOCOL_VIOLATION
                .with_reason(Self::INVALID_FRAME_ERROR)
                .with_frame_type(frame.tag().into()));
        }
        // TODO add support for NEW_TOKEN_FRAMEs on the client
        Ok(())
    }

    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> {
        if path_manager.active_path().peer_connection_id.is_empty() {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-19.15
            //# An endpoint that is sending packets with a zero-length Destination
            //# Connection ID MUST treat receipt of a NEW_CONNECTION_ID frame as a
            //# connection error of type PROTOCOL_VIOLATION.
            return Err(transport::Error::PROTOCOL_VIOLATION);
        }

        let peer_id = connection::PeerId::try_from_bytes(frame.connection_id)
            .expect("Length is validated when decoding the frame");
        let sequence_number = frame
            .sequence_number
            .as_u64()
            .try_into()
            .map_err(|_err| transport::Error::PROTOCOL_VIOLATION)?;
        let retire_prior_to = frame
            .retire_prior_to
            .as_u64()
            .try_into()
            .map_err(|_err| transport::Error::PROTOCOL_VIOLATION)?;
        let stateless_reset_token = (*frame.stateless_reset_token).into();

        path_manager.on_new_connection_id(
            &peer_id,
            sequence_number,
            retire_prior_to,
            &stateless_reset_token,
            publisher,
        )
    }

    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> {
        let sequence_number = frame
            .sequence_number
            .as_u64()
            .try_into()
            .map_err(|_err| transport::Error::PROTOCOL_VIOLATION)?;

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

        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.16
        //# The peer MAY treat this as a
        //# connection error of type PROTOCOL_VIOLATION.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.16
        //# Receipt of a RETIRE_CONNECTION_ID frame containing a sequence number
        //# greater than any previously sent to the peer MUST be treated as a
        //# connection error of type PROTOCOL_VIOLATION.
        local_id_registry
            .on_retire_connection_id(
                sequence_number,
                &datagram.destination_connection_id,
                path.rtt_estimator.smoothed_rtt(),
                datagram.timestamp,
            )
            .map_err(|err| transport::Error::PROTOCOL_VIOLATION.with_reason(err.message()))
    }

    fn handle_path_challenge_frame(
        &mut self,
        frame: PathChallenge,
        path_id: path::Id,
        path_manager: &mut path::Manager<Config>,
    ) -> Result<(), transport::Error> {
        path_manager.on_path_challenge(path_id, &frame);
        Ok(())
    }

    fn handle_path_response_frame<Pub: event::ConnectionPublisher>(
        &mut self,
        frame: PathResponse,
        path_manager: &mut path::Manager<Config>,
        publisher: &mut Pub,
    ) -> Result<(), transport::Error> {
        path_manager.on_path_response(&frame, publisher);
        Ok(())
    }

    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> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-19.20
        //# A server MUST
        //# treat receipt of a HANDSHAKE_DONE frame as a connection error of type
        //# PROTOCOL_VIOLATION.

        if Config::ENDPOINT_TYPE.is_server() {
            return Err(transport::Error::PROTOCOL_VIOLATION
                .with_reason("Clients MUST NOT send HANDSHAKE_DONE frames")
                .with_frame_type(frame.tag().into()));
        }

        handshake_status.on_handshake_done_received(publisher);

        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.2
        //# At the
        //# client, the handshake is considered confirmed when a HANDSHAKE_DONE
        //# frame is received.
        self.on_handshake_confirmed(path, local_id_registry, timestamp);

        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 crate::{
    sync::data_sender::{self, DataSender, OutgoingDataFlowController},
    transmission,
};
use s2n_quic_core::{
    ack, buffer::ReceiveBuffer, frame::crypto::CryptoRef, transport, varint::VarInt,
};

pub type TxCryptoStream = DataSender<CryptoFlowController, data_sender::writer::Crypto>;

/// Serializes and writes `Crypto` frames
#[derive(Debug, Default)]
pub struct CryptoFlowController {}

/// There is no control flow for crypto data
impl OutgoingDataFlowController for CryptoFlowController {
    fn acquire_flow_control_window(&mut self, _end_offset: VarInt) -> VarInt {
        VarInt::MAX
    }

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

    fn clear_blocked(&mut self) {}

    fn finish(&mut self) {}
}

#[derive(Debug)]
pub struct CryptoStream {
    pub tx: TxCryptoStream,
    pub rx: ReceiveBuffer,
    is_finished: bool,
}

const TX_MAX_BUFFER_CAPACITY: u32 = 4096;

impl Default for CryptoStream {
    fn default() -> Self {
        Self::new()
    }
}

impl CryptoStream {
    pub fn new() -> Self {
        Self {
            tx: TxCryptoStream::new(Default::default(), TX_MAX_BUFFER_CAPACITY),
            rx: ReceiveBuffer::default(),
            is_finished: false,
        }
    }

    pub fn can_send(&self) -> bool {
        !self.is_finished && self.tx.available_buffer_space() > 0
    }

    pub fn finish(&mut self) -> Result<(), transport::Error> {
        self.is_finished = true;
        self.tx.finish();

        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.3
        //# When TLS
        //# provides keys for a higher encryption level, if there is data from
        //# a previous encryption level that TLS has not consumed, this MUST
        //# be treated as a connection error of type PROTOCOL_VIOLATION.
        if self.rx.is_empty() {
            Ok(())
        } else {
            Err(transport::Error::PROTOCOL_VIOLATION)
        }
    }

    pub fn on_crypto_frame(&mut self, frame: CryptoRef) -> Result<(), transport::Error> {
        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.3
        //# *  If the packet is from a previously installed encryption level, it
        //# MUST NOT contain data that extends past the end of previously
        //# received data in that flow.  Implementations MUST treat any
        //# violations of this requirement as a connection error of type
        //# PROTOCOL_VIOLATION.

        if self.is_finished && frame.offset + frame.data.len() > self.rx.total_received_len() {
            return Err(transport::Error::PROTOCOL_VIOLATION);
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.5
        //= type=TODO
        //= tracking-issue=356
        //= feature=Crypto buffer limits
        //# Implementations MUST support buffering at least 4096 bytes of data
        //# received in out-of-order CRYPTO frames.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-7.5
        //= type=TODO
        //= tracking-issue=356
        //= feature=Crypto buffer limits
        //# Endpoints MAY choose to
        //# allow more data to be buffered during the handshake.

        //TODO we need to limit the buffer size here

        self.rx.write_at(frame.offset, frame.data).map_err(|_| {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-7.5
            //# If an endpoint does not expand its buffer, it MUST close
            //# the connection with a CRYPTO_BUFFER_EXCEEDED error code.

            transport::Error::CRYPTO_BUFFER_EXCEEDED
        })?;

        Ok(())
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        self.tx.on_packet_ack(ack_set);
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        self.tx.on_packet_loss(ack_set);
    }

    /// This method gets called when a Retry packet is processed.
    pub fn on_retry_packet(&mut self) {
        self.tx.on_all_lost();
    }
}

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


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

use crate::{
    endpoint,
    stream::Manager as _,
    transmission::{
        interest::{self, Provider},
        WriteContext,
    },
};
use core::task::Poll;
use s2n_codec::EncoderValue;
use s2n_quic_core::{
    datagram::{Endpoint, ReceiveContext, Receiver, Sender, WriteError},
    frame::{self, datagram::DatagramRef},
    query,
    varint::VarInt,
};

// Contains the datagram sender and receiver implementations.
//
// Used to call datagram callbacks during packet transmission and
// packet processing.
pub struct Manager<Config: endpoint::Config> {
    pub sender: <<Config as endpoint::Config>::DatagramEndpoint as Endpoint>::Sender,
    pub receiver: <<Config as endpoint::Config>::DatagramEndpoint as Endpoint>::Receiver,
    max_datagram_payload: u64,
}

impl<Config: endpoint::Config> Manager<Config> {
    pub fn new(
        sender: <<Config as endpoint::Config>::DatagramEndpoint as Endpoint>::Sender,
        receiver: <<Config as endpoint::Config>::DatagramEndpoint as Endpoint>::Receiver,
        max_datagram_payload: u64,
    ) -> Self {
        Self {
            sender,
            receiver,
            max_datagram_payload,
        }
    }

    /// A callback that allows users to write datagrams directly to the packet.
    pub fn on_transmit<W: WriteContext>(
        &mut self,
        context: &mut W,
        stream_manager: &mut Config::StreamManager,
        datagrams_prioritized: bool,
    ) {
        let mut packet = Packet {
            context,
            has_pending_streams: stream_manager.has_pending_streams(),
            datagrams_prioritized,
            max_datagram_payload: self.max_datagram_payload,
        };
        self.sender.on_transmit(&mut packet);
    }

    // A callback that allows users to access datagrams directly after they are
    // received.
    pub fn on_datagram_frame(
        &mut self,
        path: s2n_quic_core::event::api::Path<'_>,
        datagram: DatagramRef,
    ) {
        let context = ReceiveContext::new(path);
        self.receiver.on_datagram(&context, datagram.data);
    }

    pub fn datagram_mut(&mut self, query: &mut dyn query::QueryMut) -> Poll<()> {
        // Try to execute the query on the sender side. If that fails, try the receiver side.
        match query.execute_mut(&mut self.sender) {
            query::ControlFlow::Continue => {
                query.execute_mut(&mut self.receiver);
            }
            query::ControlFlow::Break => (),
        }

        if self.has_transmission_interest() {
            Poll::Ready(())
        } else {
            Poll::Pending
        }
    }
}

impl<Config: endpoint::Config> interest::Provider for Manager<Config> {
    #[inline]
    fn transmission_interest<Q: interest::Query>(&self, query: &mut Q) -> interest::Result {
        if self.sender.has_transmission_interest() {
            query.on_new_data()?;
        }
        Ok(())
    }
}

struct Packet<'a, C: WriteContext> {
    context: &'a mut C,
    has_pending_streams: bool,
    datagrams_prioritized: bool,
    max_datagram_payload: u64,
}

impl<'a, C: WriteContext> s2n_quic_core::datagram::Packet for Packet<'a, C> {
    /// Returns the remaining space in the packet
    fn remaining_capacity(&self) -> usize {
        let space = self.context.remaining_capacity();
        // Remove the frame type length and the maximum length value
        space
            .saturating_sub(frame::datagram::DATAGRAM_TAG.encoding_size())
            .saturating_sub(
                VarInt::new(space as u64)
                    .unwrap_or(VarInt::MAX)
                    .encoding_size(),
            )
    }

    /// Writes a single datagram to a packet
    fn write_datagram(&mut self, data: &[u8]) -> Result<(), WriteError> {
        self.write_datagram_vectored(&[data])
    }

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

        let remaining_capacity = self.context.remaining_capacity();
        let is_last_frame =
            remaining_capacity == frame::datagram::DATAGRAM_TAG.encoding_size() + data_len;
        let frame = frame::Datagram {
            is_last_frame,
            data,
        };
        self.context
            .write_frame(&frame)
            .ok_or(WriteError::ExceedsPacketCapacity)?;

        Ok(())
    }

    /// Returns whether or not there is reliable data ready to send
    fn has_pending_streams(&self) -> bool {
        self.has_pending_streams
    }

    /// Returns whether or not datagrams are prioritized in this packet or not
    fn datagrams_prioritized(&self) -> bool {
        self.datagrams_prioritized
    }
}


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

use crate::{
    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::{
    crypto::{tls, CryptoSuite},
    event::{self, ConnectionPublisher as _, IntoEvent},
    frame::{ack::AckRanges, crypto::CryptoRef, Ack, ConnectionClose},
    inet::DatagramInfo,
    packet::{
        encoding::{PacketEncoder, PacketEncodingError},
        handshake::{CleartextHandshake, Handshake, ProtectedHandshake},
        number::{PacketNumber, PacketNumberRange, PacketNumberSpace, SlidingWindow},
    },
    time::{timer, Timestamp},
    transport,
};

pub struct HandshakeSpace<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>::HandshakeKey,
    pub header_key:
        <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::HandshakeHeaderKey,
    //= 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,
    processed_packet_numbers: SlidingWindow,
    recovery_manager: recovery::Manager<Config>,
}

impl<Config: endpoint::Config> fmt::Debug for HandshakeSpace<Config> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("HandshakeSpace")
            .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> HandshakeSpace<Config> {
    pub fn new(
        key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::HandshakeKey,
        header_key: <<Config::TLSEndpoint as tls::Endpoint>::Session as CryptoSuite>::HandshakeHeaderKey,
        now: Timestamp,
        ack_manager: AckManager,
    ) -> Self {
        Self {
            ack_manager,
            key,
            header_key,
            crypto_stream: CryptoStream::new(),
            tx_packet_numbers: TxPacketNumbers::new(PacketNumberSpace::Handshake, now),
            processed_packet_numbers: SlidingWindow::default(),
            recovery_manager: recovery::Manager::new(PacketNumberSpace::Handshake),
        }
    }

    /// 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::Handshake,
                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 = Handshake {
            version: context.quic_version,
            destination_connection_id,
            source_connection_id: context.path_manager[context.path_id].local_connection_id,
            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::Handshake,
            },
            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 = Handshake {
            version: context.quic_version,
            destination_connection_id,
            source_connection_id: context.path_manager[context.path_id].local_connection_id,
            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 Handshake 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::Handshake,
        });
        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 Handshake packet space
    pub fn validate_and_decrypt_packet<'a, Pub: event::ConnectionPublisher>(
        &self,
        protected: ProtectedHandshake<'a>,
        path_id: path::Id,
        path: &path::Path<Config>,
        publisher: &mut Pub,
    ) -> Result<CleartextHandshake<'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::Handshake,
                        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
        })?;
        Ok(decrypted)
    }
}

impl<Config: endpoint::Config> timer::Provider for HandshakeSpace<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 HandshakeSpace<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 HandshakeSpace<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.4
//# The payload of this packet contains CRYPTO frames and could contain
//# PING, PADDING, or ACK frames.  Handshake packets MAY contain
//# CONNECTION_CLOSE frames of type 0x1c.  Endpoints MUST treat receipt
//# of Handshake packets with other frames as a connection error.
impl<Config: endpoint::Config> PacketSpace<Config> for HandshakeSpace<Config> {
    const INVALID_FRAME_ERROR: &'static str = "invalid frame in handshake 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)?;

        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 path = &mut path_manager[path_id];
        path.on_peer_validated();
        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.4
        //# Handshake packets MAY contain
        //# CONNECTION_CLOSE frames of type 0x1c.

        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 crate::{contexts::WriteContext, endpoint, sync::flag, transmission};
use s2n_quic_core::{
    ack,
    event::{self, ConnectionPublisher},
    frame::HandshakeDone,
    packet::number::PacketNumber,
};

pub type Flag = flag::Flag<HandshakeDoneWriter>;

/// The handshake status is used to track the handshake progress.
///
/// As the handshake proceeds, it unlocks certain progress on the connection i.e.
/// discarding keys, sending and receiving 1-rtt packets.
///
/// The handshake status transitions from Pending -> Complete -> Confirmed. However
/// the progress differs between the Client and Server. The chart below captures
/// these different requirements.
///
/// |        | Complete         | Confirmed                  |
/// |--------|------------------|----------------------------|
/// | server | TLS-completes    | TLS-completes              |
/// |        |                  |                            |
/// | client | TLS-completes    | HANDSHAKE_DONE received    |
/// |        |                  | or 1-rtt acked             |
///
/// TLS-completes: TLS stack reports the handshake as complete. This happens when
///                the TLS stack has sent the Finished message and verified peer's
///                Finished message.
///
/// The major differences between the Client and Server include:
/// - the handshake is complete and confirmed once the TLS-completes on the Server.
/// - the Server is required to send a HANDSHAKE_DONE frame once the handshake completes.
/// - the Client must wait for a HANDSHAKE_DONE (or an acked 1-rtt packet) to 'Confirm'
/// the handshake.
///
/// Note: s2n-quic does not implement the optional 1-rtt acked requirement.
#[derive(Debug, Default)]
pub enum HandshakeStatus {
    /// Awaiting handshake completion
    #[default]
    InProgress,

    /// Client handshake Complete
    ///
    /// Transient state while client awaits HANDSHAKE_DONE
    ClientComplete,

    /// Server handshake Complete
    ///
    /// Transient state where server awaits sending a HANDSHAKE_DONE
    ServerCompleteConfirmed(Flag),

    /// Terminal state requiring no further action
    Confirmed,
}

impl HandshakeStatus {
    /// Returns `true` if the handshake has been completed
    #[inline]
    pub fn is_complete(&self) -> bool {
        // The handshake is complete once its not Pending
        !matches!(self, HandshakeStatus::InProgress)
    }

    /// Returns `true` if the handshake has been confirmed
    pub fn is_confirmed(&self) -> bool {
        match self {
            HandshakeStatus::InProgress | HandshakeStatus::ClientComplete => false,
            HandshakeStatus::ServerCompleteConfirmed(_) => {
                //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.2
                //# the TLS handshake is considered confirmed at the
                //# server when the handshake completes
                true
            }
            HandshakeStatus::Confirmed => true,
        }
    }

    /// This method is called on the client when the HANDSHAKE_DONE
    /// frame has been received
    pub fn on_handshake_done_received<Pub: ConnectionPublisher>(&mut self, publisher: &mut Pub) {
        if let HandshakeStatus::ClientComplete = self {
            publisher.on_handshake_status_updated(event::builder::HandshakeStatusUpdated {
                status: event::builder::HandshakeStatus::HandshakeDoneAcked,
            });
            publisher.on_handshake_status_updated(event::builder::HandshakeStatusUpdated {
                status: event::builder::HandshakeStatus::Confirmed,
            });
            //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.2
            //# At the client, the handshake is
            //# considered confirmed when a HANDSHAKE_DONE frame is received.
            *self = HandshakeStatus::Confirmed;
        }
    }

    /// This method is called after the TLS handshake has been completed
    pub fn on_handshake_complete<Pub: ConnectionPublisher>(
        &mut self,
        endpoint_type: endpoint::Type,
        publisher: &mut Pub,
    ) {
        debug_assert!(
            matches!(self, Self::InProgress),
            "on_handshake_complete should only be called once."
        );
        publisher.on_handshake_status_updated(event::builder::HandshakeStatusUpdated {
            status: event::builder::HandshakeStatus::Complete,
        });

        if endpoint_type.is_server() {
            publisher.on_handshake_status_updated(event::builder::HandshakeStatusUpdated {
                status: event::builder::HandshakeStatus::Confirmed,
            });
            //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.2
            //# The server MUST send a HANDSHAKE_DONE
            //# frame as soon as the handshake is complete.
            let mut flag = Flag::default();
            flag.send();
            *self = HandshakeStatus::ServerCompleteConfirmed(flag);
        } else {
            *self = HandshakeStatus::ClientComplete;
        }
    }

    /// Used for tracking when the HANDSHAKE_DONE frame has been delivered
    /// to the peer.
    pub fn on_packet_ack<A: ack::Set, Pub: event::ConnectionPublisher>(
        &mut self,
        ack_set: &A,
        publisher: &mut Pub,
    ) {
        if let HandshakeStatus::ServerCompleteConfirmed(flag) = self {
            // The server is required to re-transmit the frame until it is
            // acknowledged by the peer. Once it is delivered, the state
            // can transition to Confirmed.
            if flag.on_packet_ack(ack_set) {
                publisher.on_handshake_status_updated(event::builder::HandshakeStatusUpdated {
                    status: event::builder::HandshakeStatus::HandshakeDoneAcked,
                });
                *self = HandshakeStatus::Confirmed;
            }
        }
    }

    /// Used for tracking when the HANDSHAKE_DONE frame needs to be
    /// re-transmitted.
    pub fn on_packet_loss<A: ack::Set, Pub: event::ConnectionPublisher>(
        &mut self,
        ack_set: &A,
        publisher: &mut Pub,
    ) {
        if let HandshakeStatus::ServerCompleteConfirmed(flag) = self {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-13.3
            //# The HANDSHAKE_DONE frame MUST be retransmitted until it is
            //# acknowledged.
            if flag.on_packet_loss(ack_set) {
                publisher.on_handshake_status_updated(event::builder::HandshakeStatusUpdated {
                    status: event::builder::HandshakeStatus::HandshakeDoneLost,
                });
            }
        }
    }

    /// Queries if any HANDSHAKE_DONE frames need to get sent
    pub fn on_transmit<C: WriteContext>(&mut self, context: &mut C) {
        if let HandshakeStatus::ServerCompleteConfirmed(flag) = self {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-19.20
            //# A HANDSHAKE_DONE frame can only be sent by the server.
            let _ = flag.on_transmit(context);
        }
    }
}

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

impl flag::Writer for HandshakeDoneWriter {
    fn write_frame<W: WriteContext>(&mut self, context: &mut W) -> Option<PacketNumber> {
        debug_assert!(
            context.local_endpoint_type().is_server(),
            "Only servers should transmit HANDSHAKE_DONE frames"
        );
        context.write_frame(&HandshakeDone)
    }
}

impl transmission::interest::Provider for HandshakeStatus {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        if let HandshakeStatus::ServerCompleteConfirmed(flag) = self {
            flag.transmission_interest(query)
        } else {
            Ok(())
        }
    }
}

#[cfg(test)]
mod fuzz_target;

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{contexts::testing::*, transmission::interest::Provider};
    use s2n_quic_core::{endpoint, event::testing::Publisher, time::clock::testing as time};

    #[test]
    fn server_test() {
        let mut publisher = Publisher::snapshot();
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            time::now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Server,
        );

        let mut status = HandshakeStatus::default();

        assert!(!status.is_confirmed());
        assert!(!status.is_complete());
        assert_eq!(
            status.get_transmission_interest(),
            transmission::Interest::default(),
            "status should not express interest in default state"
        );

        status.on_transmit(&mut context);
        assert!(
            context.frame_buffer.is_empty(),
            "status should not transmit in default state"
        );

        //= https://www.rfc-editor.org/rfc/rfc9001#section-4.1.2
        //= type=test
        //# the TLS handshake is considered confirmed at the
        //# server when the handshake completes.
        status.on_handshake_complete(endpoint::Type::Server, &mut publisher);
        assert!(status.is_confirmed());
        assert!(status.is_complete());

        assert_eq!(
            status.get_transmission_interest(),
            transmission::Interest::NewData,
            "status should express interest in deliver after handshake complete"
        );

        status.on_transmit(&mut context);

        let packet_number = context
            .frame_buffer
            .pop_front()
            .expect("status should write HANDSHAKE_DONE frames")
            .packet_nr;

        status.on_packet_ack(&packet_number, &mut publisher);
        assert!(status.is_confirmed());

        assert!(
            !status.has_transmission_interest(),
            "status should not express interest after complete",
        );

        status.on_transmit(&mut context);
        assert!(
            context.frame_buffer.is_empty(),
            "status should not transmit after complete"
        );
    }

    #[test]
    fn client_test() {
        let mut publisher = Publisher::snapshot();
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            time::now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );

        let mut status = HandshakeStatus::default();

        assert!(!status.is_complete());
        assert!(!status.is_confirmed());

        assert!(
            !status.has_transmission_interest(),
            "client does not transmit a HANDSHAKE_DONE"
        );

        status.on_transmit(&mut context);
        assert!(
            context.frame_buffer.is_empty(),
            "client does not transmit a HANDSHAKE_DONE"
        );

        // the handshake must be complete prior to being confirmed
        status.on_handshake_done_received(&mut publisher);
        assert!(!status.is_complete());
        assert!(!status.is_confirmed());

        status.on_handshake_complete(endpoint::Type::Client, &mut publisher);
        assert!(status.is_complete());

        assert!(
            !status.has_transmission_interest(),
            "client does not transmit a HANDSHAKE_DONE"
        );

        status.on_transmit(&mut context);
        assert!(
            context.frame_buffer.is_empty(),
            "client does not transmit a HANDSHAKE_DONE"
        );

        // confirm the client handshake
        status.on_handshake_done_received(&mut publisher);
        assert!(status.is_confirmed());

        // try calling it multiple times
        status.on_handshake_done_received(&mut publisher);
        assert!(status.is_confirmed());
    }
}