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

//! This module contains the Path implementation

use crate::{
    connection,
    contexts::WriteContext,
    endpoint,
    endpoint::Type,
    recovery::{congestion_controller, CongestionController, RttEstimator},
    transmission::{self, Mode},
};
use s2n_quic_core::{
    counter::{Counter, Saturating},
    event::{self, IntoEvent},
    frame, packet, random,
    time::{timer, Timestamp},
};

mod challenge;
pub(crate) mod ecn;
mod manager;
pub(crate) mod mtu;

pub use challenge::*;
pub use manager::*;

/// re-export core
pub use s2n_quic_core::path::*;

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    /// Path has no transmission limitations
    Validated,

    /// Path has not been validated and is subject to amplification limits
    AmplificationLimited {
        tx_allowance: Counter<u32, Saturating>,
    },
}

#[derive(Debug)]
pub struct Path<Config: endpoint::Config> {
    /// The peer's socket address
    pub handle: Config::PathHandle,
    /// The connection id of the peer
    pub peer_connection_id: connection::PeerId,
    /// The local connection id which the peer sends to
    pub local_connection_id: connection::LocalId,
    /// The path owns the roundtrip between peers
    pub rtt_estimator: RttEstimator,
    /// The congestion controller for the path
    pub congestion_controller: <Config::CongestionControllerEndpoint as congestion_controller::Endpoint>::CongestionController,
    /// Probe timeout backoff multiplier
    pub pto_backoff: u32,
    /// Tracks whether this path has passed Address or Path validation
    state: State,
    /// Controller for determining the maximum transmission unit of the path
    pub mtu_controller: mtu::Controller,
    /// Controller for determining the ECN capability of the path
    pub ecn_controller: ecn::Controller,

    /// True if the path has been validated by the peer
    peer_validated: bool,

    /// Challenge sent to the peer in a PATH_CHALLENGE
    challenge: Challenge,

    /// Received a Challenge and should echo back data in PATH_RESPONSE
    response_data: Option<challenge::Data>,

    /// True if the path is currently or at some point been an active path.
    ///
    /// A path becomes an active path if it receives a non-path-validation-probing
    /// packet. `activated` is a one way state used to mark paths that have been the
    /// active path at some point in the connection. This parameter is used to
    /// determine if the path should become the last_known_active_validated_path.
    activated: bool,

    /// True if the path is currently active
    is_active: bool,
}

impl<Config: endpoint::Config> Clone for Path<Config> {
    fn clone(&self) -> Self {
        Self {
            handle: self.handle,
            peer_connection_id: self.peer_connection_id,
            local_connection_id: self.local_connection_id,
            rtt_estimator: self.rtt_estimator,
            congestion_controller: self.congestion_controller.clone(),
            pto_backoff: self.pto_backoff,
            state: self.state,
            mtu_controller: self.mtu_controller.clone(),
            ecn_controller: self.ecn_controller.clone(),
            peer_validated: self.peer_validated,
            challenge: self.challenge.clone(),
            response_data: self.response_data,
            activated: self.activated,
            is_active: self.is_active,
        }
    }
}

/// A Path holds the local and peer socket addresses, connection ids, and state. It can be
/// validated or pending validation.
impl<Config: endpoint::Config> Path<Config> {
    #[allow(clippy::too_many_arguments)]
    pub fn new(
        handle: Config::PathHandle,
        peer_connection_id: connection::PeerId,
        local_connection_id: connection::LocalId,
        rtt_estimator: RttEstimator,
        congestion_controller: <Config::CongestionControllerEndpoint as congestion_controller::Endpoint>::CongestionController,
        peer_validated: bool,
        max_mtu: MaxMtu,
    ) -> Path<Config> {
        let state = match Config::ENDPOINT_TYPE {
            Type::Server => {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
                //# If the client IP address has changed, the
                //# server MUST adhere to the anti-amplification limit; see Section 8.
                // Start each path in State::AmplificationLimited until it has been validated.
                State::AmplificationLimited {
                    tx_allowance: Default::default(),
                }
            }
            //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
            //# Clients are only constrained by the congestion controller.
            Type::Client => State::Validated,
        };
        let peer_socket_address = handle.remote_address();
        Path {
            handle,
            peer_connection_id,
            local_connection_id,
            rtt_estimator,
            congestion_controller,
            pto_backoff: INITIAL_PTO_BACKOFF,
            state,
            mtu_controller: mtu::Controller::new(max_mtu, &peer_socket_address),
            ecn_controller: ecn::Controller::default(),
            peer_validated,
            challenge: Challenge::disabled(),
            response_data: None,
            activated: false,
            is_active: false,
        }
    }

    #[inline]
    pub fn remote_address(&self) -> RemoteAddress {
        self.handle.remote_address()
    }

    #[inline]
    pub fn local_address(&self) -> LocalAddress {
        self.handle.local_address()
    }

    #[inline]
    pub fn set_challenge(&mut self, challenge: Challenge) {
        self.challenge = challenge;
    }

    #[inline]
    pub fn abandon_challenge<Pub: event::ConnectionPublisher>(
        &mut self,
        publisher: &mut Pub,
        path_id: u64,
    ) {
        self.challenge
            .abandon(publisher, path_event!(self, path_id));
    }

    #[inline]
    pub fn is_active(&self) -> bool {
        self.is_active
    }

    /// Called when bytes have been transmitted on this path
    #[inline]
    pub fn on_bytes_transmitted(&mut self, bytes: usize) {
        if bytes == 0 {
            return;
        }

        debug_assert_ne!(
            self.clamp_mtu(bytes, transmission::Mode::Normal),
            0,
            "path should not transmit when amplification limited; tried to transmit {bytes}"
        );

        if let State::AmplificationLimited { tx_allowance, .. } = &mut self.state {
            *tx_allowance -= bytes as u32
        }
    }

    /// Called when bytes have been received on this path
    /// Returns true if receiving these bytes unblocked the
    /// path from being amplification limited
    #[inline]
    pub fn on_bytes_received(&mut self, bytes: usize) -> bool {
        let was_at_amplification_limit = self.at_amplification_limit();

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
        //# For the purposes of
        //# avoiding amplification prior to address validation, servers MUST
        //# count all of the payload bytes received in datagrams that are
        //# uniquely attributed to a single connection.
        //
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
        //# Prior to validating the client address, servers MUST NOT send more
        //# than three times as many bytes as the number of bytes they have
        //# received.
        //
        if let State::AmplificationLimited { tx_allowance } = &mut self.state {
            *tx_allowance += bytes.saturating_mul(3) as u32;
        }

        was_at_amplification_limit && !self.at_amplification_limit()
    }

    #[inline]
    pub fn on_timeout<Pub: event::ConnectionPublisher>(
        &mut self,
        timestamp: Timestamp,
        path_id: Id,
        random_generator: &mut dyn random::Generator,
        publisher: &mut Pub,
    ) {
        self.challenge
            .on_timeout(timestamp, publisher, path_event!(self, path_id));
        self.mtu_controller.on_timeout(timestamp);
        self.ecn_controller.on_timeout(
            timestamp,
            path_event!(self, path_id),
            random_generator,
            self.rtt_estimator.smoothed_rtt(),
            publisher,
        );
    }

    /// Returns true if this path is able to transmit packets at the given timestamp
    #[inline]
    pub fn can_transmit(&self, timestamp: Timestamp) -> bool {
        !self.at_amplification_limit()
            && self
                .congestion_controller
                .earliest_departure_time()
                .map_or(true, |edt| edt.has_elapsed(timestamp))
    }

    /// Only PATH_CHALLENGE and PATH_RESPONSE frames should be transmitted here.
    #[inline]
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
        //# A PATH_RESPONSE frame MUST be sent on the network path where the
        //# PATH_CHALLENGE frame was received.
        if let Some(response_data) = &mut self.response_data {
            let frame = frame::PathResponse {
                data: response_data,
            };
            if context.write_frame(&frame).is_some() {
                //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
                //# An endpoint MUST NOT send more than one PATH_RESPONSE frame in
                //# response to one PATH_CHALLENGE frame; see Section 13.3.
                self.response_data = None;
            }
        }

        self.challenge.on_transmit(context)
    }

    /// Check if path validation was attempted and failed.
    #[inline]
    pub fn failed_validation(&self) -> bool {
        // PATH_CHALLENGE is not used for validating the initial path and is disabled. Check if
        // the challenge is disabled before executing the following block since there won't be
        // a last_known_validated_path.
        !self.challenge.is_disabled() && !self.is_validated() && !self.is_challenge_pending()
    }

    #[inline]
    pub fn is_challenge_pending(&self) -> bool {
        self.challenge.is_pending()
    }

    #[inline]
    pub fn is_response_pending(&self) -> bool {
        self.response_data.is_some()
    }

    #[inline]
    pub fn on_path_challenge(&mut self, response: &challenge::Data) {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
        //# On receiving a PATH_CHALLENGE frame, an endpoint MUST respond by
        //# echoing the data contained in the PATH_CHALLENGE frame in a
        //# PATH_RESPONSE frame.
        self.response_data = Some(*response);
    }

    /// Validates the path if the PATH_RESPONSE data matches the PATH_CHALLENGE data
    /// and returns if the path was validated.
    #[inline]
    pub fn on_path_response(&mut self, response: &[u8]) -> bool {
        if self.challenge.on_validated(response) {
            self.on_validated();

            return true;

            //= https://www.rfc-editor.org/rfc/rfc9000#section-9.3
            //= type=TODO
            //# After verifying a new client address, the server SHOULD send new
            //# address validation tokens (Section 8) to the client.
        }

        false
    }

    /// Called when a handshake packet is received.
    ///
    /// Receiving a handshake packet acts as path validation for the initial path
    #[inline]
    pub fn on_handshake_packet(&mut self) {
        self.on_validated();
    }

    /// Checks if the peer has started using a different destination Connection Id.
    ///
    /// The CleartextShort packet guarantees the packet has been validated
    /// (authenticated and de-duped).
    pub fn on_process_local_connection_id<Pub: event::ConnectionPublisher>(
        &mut self,
        path_id: Id,
        packet: &packet::short::CleartextShort<'_>,
        local_connection_id: &connection::LocalId,
        publisher: &mut Pub,
    ) {
        debug_assert_eq!(
            packet.destination_connection_id(),
            local_connection_id.as_ref()
        );

        if &self.local_connection_id != local_connection_id {
            publisher.on_connection_id_updated(event::builder::ConnectionIdUpdated {
                path_id: path_id.into_event(),
                cid_consumer: endpoint::Location::Remote,
                previous: self.local_connection_id.into_event(),
                current: local_connection_id.into_event(),
            });
            self.local_connection_id = *local_connection_id;
        }
    }

    /// Called when the path is validated
    #[inline]
    fn on_validated(&mut self) {
        self.state = State::Validated;

        if self.is_peer_validated() {
            self.on_fully_validated();
        }
    }

    /// Returns whether this path has passed address validation
    #[inline]
    pub fn is_validated(&self) -> bool {
        self.state == State::Validated
    }

    /// The path received a non-path-validation-probing packet so mark it as activated.
    #[inline]
    pub fn on_activated(&mut self) {
        self.activated = true;
    }

    /// Returns if the path is currently or at some point been an active path.
    #[inline]
    pub fn is_activated(&self) -> bool {
        self.activated
    }

    /// Marks the path as peer validated
    #[inline]
    pub fn on_peer_validated(&mut self) {
        self.peer_validated = true;

        if self.is_validated() {
            self.on_fully_validated();
        }
    }

    /// Returns whether this path has been validated by the peer
    #[inline]
    pub fn is_peer_validated(&self) -> bool {
        self.peer_validated
    }

    /// Called when the path has been validated locally, and also by the peer
    fn on_fully_validated(&mut self) {
        // Enable the mtu controller to allow for PMTU discovery
        self.mtu_controller.enable()
    }

    #[inline]
    pub fn mtu(&self, transmission_mode: transmission::Mode) -> usize {
        match transmission_mode {
            // Use the minimum MTU for loss recovery probes to allow detection of packets
            // lost when the previously confirmed path MTU is no longer supported.
            //
            // The priority during PathValidationOnly is to validate the path, so the
            // minimum MTU is used to avoid packet loss due to MTU limits.
            Mode::LossRecoveryProbing | Mode::PathValidationOnly => MINIMUM_MTU as usize,
            // When MTU Probing, clamp to the size of the MTU we are attempting to validate
            Mode::MtuProbing => self.mtu_controller.probed_sized(),
            // Otherwise use the confirmed MTU
            Mode::Normal => self.mtu_controller.mtu(),
        }
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-14.1
    //# The server MUST also limit the number of bytes it sends before
    //# validating the address of the client; see Section 8.

    //= https://www.rfc-editor.org/rfc/rfc9000#section-14.2
    //# All QUIC
    //# packets that are not sent in a PMTU probe SHOULD be sized to fit
    //# within the maximum datagram size to avoid the datagram being
    //# fragmented or dropped [RFC8085].

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //# A PL MUST NOT send a datagram (other than a probe
    //# packet) with a size at the PL that is larger than the current
    //# PLPMTU.

    /// Clamps payload sizes to the current MTU for the path
    ///
    /// # Panics
    ///
    /// Panics if this is called when the path is amplification limited
    #[inline]
    pub fn clamp_mtu(&self, requested_size: usize, transmission_mode: transmission::Mode) -> usize {
        debug_assert!(
            !self.at_amplification_limit(),
            "amplification limits should be checked before clamping MTU values"
        );

        requested_size.min(self.mtu(transmission_mode))
    }

    #[inline]
    pub fn transmission_constraint(&self) -> transmission::Constraint {
        if self.at_amplification_limit() {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
            //# Prior to validating the client address, servers MUST NOT send more
            //# than three times as many bytes as the number of bytes they have
            //# received.
            transmission::Constraint::AmplificationLimited
        } else if self.congestion_controller.is_congestion_limited() {
            if self.congestion_controller.requires_fast_retransmission() {
                //= https://www.rfc-editor.org/rfc/rfc9002#section-7.3.2
                //# If the congestion window is reduced immediately, a
                //# single packet can be sent prior to reduction.  This speeds up loss
                //# recovery if the data in the lost packet is retransmitted and is
                //# similar to TCP as described in Section 5 of [RFC6675].
                transmission::Constraint::RetransmissionOnly
            } else {
                //= https://www.rfc-editor.org/rfc/rfc9002#section-7
                //# An endpoint MUST NOT send a packet if it would cause bytes_in_flight
                //# (see Appendix B.2) to be larger than the congestion window, unless
                //# the packet is sent on a PTO timer expiration (see Section 6.2) or
                //# when entering recovery (see Section 7.3.2).
                transmission::Constraint::CongestionLimited
            }
        } else {
            transmission::Constraint::None
        }
    }

    /// Returns whether this path should be limited according to connection establishment amplification limits
    ///
    /// Note: As long as the path has _any_ TX credits we don't consider it to be amplification-limited.
    ///       This may result in sending slightly more than 3x bytes but networking infrastructure mostly
    ///       cares about the number of packets rather than bytes.
    #[inline]
    pub fn at_amplification_limit(&self) -> bool {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
        //# Prior to validating the client address, servers MUST NOT send more
        //# than three times as many bytes as the number of bytes they have
        //# received.
        match self.state {
            State::Validated => false,
            State::AmplificationLimited { tx_allowance } => tx_allowance == 0,
        }
    }

    /// Returns the current PTO period
    #[inline]
    pub fn pto_period(
        &self,
        space: s2n_quic_core::packet::number::PacketNumberSpace,
    ) -> core::time::Duration {
        self.rtt_estimator.pto_period(self.pto_backoff, space)
    }

    /// Resets the PTO backoff to the initial value
    #[inline]
    pub fn reset_pto_backoff(&mut self) {
        self.pto_backoff = INITIAL_PTO_BACKOFF;
    }

    /// Marks the path as closing
    pub fn on_closing(&mut self) {
        // Revert the path state to AmplificationLimited so we can control the number
        // of packets sent back with anti-amplification limits
        match &self.state {
            // keep the current amplification limits
            State::AmplificationLimited { .. } => {}
            State::Validated => {
                self.state = State::AmplificationLimited {
                    tx_allowance: Counter::new(MINIMUM_MTU as u32 * 3),
                };
            }
        }
    }

    #[inline]
    pub fn max_mtu(&self) -> MaxMtu {
        self.mtu_controller.max_mtu()
    }

    /// Returns `true` if the congestion window does not have sufficient space for a packet of
    /// size `mtu` considering the current bytes in flight and the additional `bytes_sent` provided
    #[inline]
    pub fn is_congestion_limited(&self, bytes_sent: usize) -> bool {
        let cwnd = self.congestion_controller.congestion_window();
        let bytes_in_flight = self
            .congestion_controller
            .bytes_in_flight()
            .saturating_add(bytes_sent as u32);
        let mtu = self.mtu(transmission::Mode::Normal) as u32;

        cwnd.saturating_sub(bytes_in_flight) < mtu
    }

    // Compare a Path based on its PathHandle.
    //
    // Currently the local_address on the Client connection is unknown and set to
    // a default un-specified value; therefore only the remote_address is used
    // to compare Paths.
    fn eq_by_handle(&self, handle: &Config::PathHandle) -> bool {
        if self.handle.local_address().port() == 0 {
            // Only compare the remote address if we haven't updated the local yet
            s2n_quic_core::path::Handle::eq(&self.handle.remote_address(), &handle.remote_address())
        } else {
            self.handle.eq(handle)
        }
    }
}

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

        Ok(())
    }
}

impl<Config: endpoint::Config> transmission::interest::Provider for Path<Config> {
    /// Indicate if the path is interested in transmitting PATH_CHALLENGE or
    /// PATH_RESPONSE frames.
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
        //# An endpoint MUST NOT delay transmission of a
        //# packet containing a PATH_RESPONSE frame unless constrained by
        //# congestion control.
        if self.response_data.is_some() {
            query.on_new_data()?;
        }

        self.challenge.transmission_interest(query)?;

        Ok(())
    }
}

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use crate::{
        endpoint,
        path::{Path, DEFAULT_MAX_MTU},
    };
    use core::time::Duration;
    use s2n_quic_core::{connection, recovery::RttEstimator};

    pub fn helper_path_server() -> Path<endpoint::testing::Server> {
        Path::new(
            Default::default(),
            connection::PeerId::try_from_bytes(&[]).unwrap(),
            connection::LocalId::TEST_ID,
            RttEstimator::new(Duration::from_millis(30)),
            Default::default(),
            true,
            DEFAULT_MAX_MTU,
        )
    }

    pub fn helper_path_client() -> Path<endpoint::testing::Client> {
        Path::new(
            Default::default(),
            connection::PeerId::try_from_bytes(&[]).unwrap(),
            connection::LocalId::TEST_ID,
            RttEstimator::new(Duration::from_millis(30)),
            Default::default(),
            false,
            DEFAULT_MAX_MTU,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        contexts::testing::{MockWriteContext, OutgoingFrameBuffer},
        endpoint::testing::Server as Config,
        path,
        path::{challenge::testing::helper_challenge, testing, testing::helper_path_client},
    };
    use core::time::Duration;
    use s2n_quic_core::{
        connection, endpoint,
        event::testing::Publisher,
        recovery::{CongestionController, RttEstimator},
        time::{Clock, NoopClock},
        transmission,
    };

    type Path = super::Path<Config>;

    #[test]
    fn response_data_should_only_be_sent_once() {
        // Setup:
        let mut path = testing::helper_path_server();
        let now = NoopClock {}.get_time();

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

        // set response data
        let expected_data: [u8; 8] = [0; 8];
        path.on_path_challenge(&expected_data);
        assert_eq!(path.response_data.unwrap(), expected_data);

        // Trigger:
        path.on_transmit(&mut context); // send response data

        // Expectation:
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
        //= type=test
        //# An endpoint MUST NOT send more than one PATH_RESPONSE frame in
        //# response to one PATH_CHALLENGE frame; see Section 13.3.
        assert!(path.response_data.is_none());

        assert_eq!(context.frame_buffer.len(), 1);
        let written_data = match context.frame_buffer.pop_front().unwrap().as_frame() {
            frame::Frame::PathResponse(frame) => Some(*frame.data),
            _ => None,
        };
        assert_eq!(written_data.unwrap(), expected_data);
    }

    #[test]
    fn on_timeout_should_set_challenge_to_none_on_challenge_abandonment() {
        // Setup:
        let mut publisher = Publisher::snapshot();
        let mut path = testing::helper_path_server();
        let helper_challenge = helper_challenge();
        let expiration_time = helper_challenge.now + helper_challenge.abandon_duration;
        path.set_challenge(helper_challenge.challenge);

        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper_challenge.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        path.on_transmit(&mut context); // send challenge and arm timer

        // Expectation:
        assert!(path.is_challenge_pending());
        assert!(path.challenge.is_pending());

        // Trigger:
        path.on_timeout(
            expiration_time + Duration::from_millis(10),
            path::Id::test_id(),
            &mut random::testing::Generator(123),
            &mut publisher,
        );

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

    #[test]
    fn is_challenge_pending_should_return_false_if_challenge_is_not_set() {
        // Setup:
        let mut path = testing::helper_path_server();
        let helper_challenge = helper_challenge();

        // Expectation:
        assert!(!path.is_challenge_pending());
        assert!(!path.challenge.is_pending());

        // Trigger:
        path.set_challenge(helper_challenge.challenge);

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

    #[test]
    fn first_path_in_disabled_state_cant_fail_validation() {
        // Setup:
        let path = testing::helper_path_server();

        // Expectation:
        assert!(path.challenge.is_disabled());
        assert!(!path.is_challenge_pending());
        assert!(!path.is_validated());

        assert!(!path.failed_validation());
    }

    #[test]
    fn failed_validation() {
        // Setup:
        let mut publisher = Publisher::snapshot();
        let mut path = testing::helper_path_server();
        let helper_challenge = helper_challenge();

        path.set_challenge(helper_challenge.challenge);
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut context = MockWriteContext::new(
            helper_challenge.now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Client,
        );
        path.on_transmit(&mut context); // send challenge and arm timer

        let expiration_time = helper_challenge.now + helper_challenge.abandon_duration;

        // Trigger:
        path.on_timeout(
            expiration_time + Duration::from_millis(10),
            path::Id::test_id(),
            &mut random::testing::Generator(123),
            &mut publisher,
        );

        // Expectation:
        assert!(!path.challenge.is_disabled());
        assert!(!path.is_challenge_pending());
        assert!(!path.is_validated());

        assert!(path.failed_validation());
    }

    #[test]
    fn abandon_challenge() {
        // Setup:
        let mut path = testing::helper_path_server();
        let helper_challenge = helper_challenge();
        path.set_challenge(helper_challenge.challenge);
        let mut publisher = event::testing::Publisher::snapshot();

        // Trigger:
        path.abandon_challenge(&mut publisher, 0);

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

    #[test]
    fn on_path_challenge_should_set_response_data() {
        // Setup:
        let mut path = testing::helper_path_server();

        // Expectation:
        assert!(path.response_data.is_none());

        // Trigger:
        let expected_data: [u8; 8] = [0; 8];
        path.on_path_challenge(&expected_data);

        // Expectation:
        assert!(path.response_data.is_some());
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
    //= type=test
    //# On receiving a PATH_CHALLENGE frame, an endpoint MUST respond by
    //# echoing the data contained in the PATH_CHALLENGE frame in a
    //# PATH_RESPONSE frame.
    #[test]
    fn on_path_challenge_should_replace_response_data() {
        // Setup:
        let mut path = testing::helper_path_server();
        let expected_data: [u8; 8] = [0; 8];

        // Trigger 1:
        path.on_path_challenge(&expected_data);

        // Expectation 1:
        assert_eq!(path.response_data.unwrap(), expected_data);

        // Trigger 2:
        let new_expected_data: [u8; 8] = [1; 8];
        path.on_path_challenge(&new_expected_data);

        // Expectation 2:
        assert_ne!(expected_data, new_expected_data);
        assert_eq!(path.response_data.unwrap(), new_expected_data);
    }

    #[test]
    fn validate_path_response_should_only_validate_if_challenge_is_set() {
        // Setup:
        let mut path = testing::helper_path_server();
        let helper_challenge = helper_challenge();

        // Expectation:
        assert!(!path.is_validated());

        // Trigger:
        path.set_challenge(helper_challenge.challenge);
        assert!(path.on_path_response(&helper_challenge.expected_data));

        // Expectation:
        assert!(path.is_validated());
    }

    #[test]
    fn on_validated_should_change_state_to_validated_and_clear_challenge() {
        // Setup:
        let mut path = testing::helper_path_server();
        let helper_challenge = helper_challenge();
        path.set_challenge(helper_challenge.challenge);

        assert!(!path.is_validated());
        assert!(path.challenge.is_pending());

        // Trigger:
        path.on_validated();

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

    #[test]
    fn on_validated_when_already_validated_does_nothing() {
        // Setup:
        let mut path = testing::helper_path_server();
        path.set_challenge(helper_challenge().challenge);
        path.on_validated();

        // Trigger:
        path.on_validated();

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

    #[test]
    fn amplification_limit_test() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# If the client IP address has changed, the server MUST
        //# adhere to the anti-amplification limit; see Section 8.
        // This is tested here by verifying a new Path starts in State::AmplificationLimited

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
        //= type=test
        //# For the purposes of
        //# avoiding amplification prior to address validation, servers MUST
        //# count all of the payload bytes received in datagrams that are
        //# uniquely attributed to a single connection.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1.4
        //= type=test
        //# If the client IP address has changed, the server MUST
        //# adhere to the anti-amplification limit; see Section 8.
        // This tests the IP change because a new path is created when a new peer_address is
        // detected. This new path should always start in State::Pending.
        let mut path = testing::helper_path_server();

        // Verify we enforce the amplification limit if we can't send
        // at least 1 minimum sized packet
        let mut unblocked = path.on_bytes_received(1200);
        assert!(unblocked);
        path.on_bytes_transmitted((1200 * 2) + 1);

        // we round up to the nearest mtu
        assert!(!path.at_amplification_limit());
        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::None
        );

        unblocked = path.on_bytes_received(1200);
        assert!(!path.at_amplification_limit());
        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::None
        );
        // If we were not amplification limited previously, receiving
        // more bytes doesn't unblock
        assert!(!unblocked);

        path.on_bytes_transmitted((1200 * 6) + 1);
        assert!(path.at_amplification_limit());
        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::AmplificationLimited
        );
        unblocked = path.on_bytes_received(1200);
        assert!(unblocked);

        path.on_validated();
        path.on_bytes_transmitted(24);
        // Validated paths should always be able to transmit
        assert!(!path.at_amplification_limit());
        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::None
        );

        // If we were already not amplification limited, receiving
        // more bytes doesn't unblock
        unblocked = path.on_bytes_received(1200);
        assert!(!unblocked);

        // Clients are not amplification limited
        let path = helper_path_client();
        assert!(path.is_validated());
    }

    #[test]
    fn amplification_limited_mtu_test() {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-8.1
        //= type=test
        //# Prior to validating the client address, servers MUST NOT send more
        //# than three times as many bytes as the number of bytes they have
        //# received.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-14.1
        //= type=test
        //# The server MUST also limit the number of bytes it sends before
        //# validating the address of the client; see Section 8.

        //= https://www.rfc-editor.org/rfc/rfc9000#section-14.2
        //= type=test
        //# All QUIC
        //# packets that are not sent in a PMTU probe SHOULD be sized to fit
        //# within the maximum datagram size to avoid the datagram being
        //# fragmented or dropped [RFC8085].

        //= https://www.rfc-editor.org/rfc/rfc8899#section-3
        //= type=test
        //# A PL MUST NOT send a datagram (other than a probe
        //# packet) with a size at the PL that is larger than the current
        //# PLPMTU.

        // TODO this would work better as a fuzz test

        for &transmission_mode in &[
            Mode::Normal,
            Mode::PathValidationOnly,
            Mode::MtuProbing,
            Mode::LossRecoveryProbing,
        ] {
            let mut path = testing::helper_path_server();
            // Verify we can transmit up to the mtu
            let mtu = path.mtu(transmission_mode);

            path.on_bytes_received(3);
            path.on_bytes_transmitted(8);

            assert_eq!(path.clamp_mtu(1, transmission_mode), 1);
            assert_eq!(path.clamp_mtu(10, transmission_mode), 10);
            assert_eq!(path.clamp_mtu(1800, transmission_mode), mtu);

            path.on_bytes_transmitted(1);
            // Verify we can't transmit any more bytes
            assert!(path.at_amplification_limit());

            path.on_bytes_received(1);
            // Verify we can transmit up to 3 more bytes
            assert_eq!(path.clamp_mtu(1, transmission_mode), 1);
            assert_eq!(path.clamp_mtu(10, transmission_mode), 10);
            assert_eq!(path.clamp_mtu(1800, transmission_mode), mtu);

            path.on_validated();
            // Validated paths should always be able to transmit
            assert_eq!(path.clamp_mtu(4, transmission_mode), 4);
        }
    }

    #[test]
    fn clamp_mtu_for_validated_path() {
        let mut path = testing::helper_path_server();
        path.on_validated();
        let mtu = 1472;
        let probed_size = 1500;
        path.mtu_controller = mtu::testing::test_controller(mtu, probed_size);

        assert_eq!(
            path.mtu_controller.mtu(),
            path.clamp_mtu(10000, transmission::Mode::Normal)
        );
        assert_eq!(
            MINIMUM_MTU as usize,
            path.clamp_mtu(10000, transmission::Mode::PathValidationOnly)
        );
        assert_eq!(
            MINIMUM_MTU as usize,
            path.clamp_mtu(10000, transmission::Mode::LossRecoveryProbing)
        );
        assert_eq!(
            path.mtu_controller.probed_sized(),
            path.clamp_mtu(10000, transmission::Mode::MtuProbing)
        );
    }

    #[test]
    fn path_mtu() {
        let mut path = testing::helper_path_server();
        path.on_bytes_received(1);
        let mtu = 1472;
        let probed_size = 1500;
        path.mtu_controller = mtu::testing::test_controller(mtu, probed_size);

        assert_eq!(
            path.mtu_controller.mtu(),
            path.mtu(transmission::Mode::Normal)
        );
        assert_eq!(
            MINIMUM_MTU as usize,
            path.mtu(transmission::Mode::PathValidationOnly)
        );
        assert_eq!(
            MINIMUM_MTU as usize,
            path.mtu(transmission::Mode::LossRecoveryProbing)
        );
        assert_eq!(
            path.mtu_controller.probed_sized(),
            path.mtu(transmission::Mode::MtuProbing)
        );
    }

    #[test]
    #[should_panic]
    fn clamp_mtu_when_tx_more_than_rx() {
        let mut path = testing::helper_path_server();
        let mtu = 1472;
        let probed_size = 1500;
        path.mtu_controller = mtu::testing::test_controller(mtu, probed_size);

        assert_eq!(0, path.clamp_mtu(10000, transmission::Mode::Normal));
    }

    #[test]
    fn peer_validated_test() {
        let mut path = testing::helper_path_client();

        assert!(!path.is_peer_validated());

        path.on_peer_validated();

        assert!(path.is_peer_validated());
    }

    #[test]
    fn transmission_constraint_test() {
        let mut path = Path::new(
            Default::default(),
            connection::PeerId::try_from_bytes(&[]).unwrap(),
            connection::LocalId::TEST_ID,
            RttEstimator::new(Duration::from_millis(30)),
            Default::default(),
            false,
            DEFAULT_MAX_MTU,
        );
        let now = NoopClock.get_time();
        let random = &mut random::testing::Generator::default();
        let mut publisher = event::testing::Publisher::snapshot();
        let mut publisher =
            congestion_controller::PathPublisher::new(&mut publisher, path::Id::test_id());
        path.on_validated();

        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::None
        );

        // Fill up the congestion controller
        let packet_info = path.congestion_controller.on_packet_sent(
            now,
            path.congestion_controller.congestion_window() as usize,
            None,
            &path.rtt_estimator,
            &mut publisher,
        );

        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::CongestionLimited
        );

        // Lose a byte to enter recovery
        path.congestion_controller.on_packet_lost(
            1,
            packet_info,
            false,
            false,
            random,
            now,
            &mut publisher,
        );
        path.congestion_controller.requires_fast_retransmission = true;

        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::RetransmissionOnly
        );

        // Lose remaining bytes
        path.congestion_controller.on_packet_lost(
            path.congestion_controller.congestion_window(),
            packet_info,
            false,
            false,
            random,
            now,
            &mut publisher,
        );
        path.congestion_controller.requires_fast_retransmission = false;

        // Since we are no longer congestion limited, there is no transmission constraint
        assert_eq!(
            path.transmission_constraint(),
            transmission::Constraint::None
        );
    }

    #[test]
    fn on_closing_validated_path() {
        let mut path = testing::helper_path_server();
        path.on_validated();
        assert!(path.is_validated());

        // Trigger:
        path.on_closing();

        // Expectation:
        match path.state {
            path::State::Validated => panic!("transition to AmplificationLimited when closing"),
            path::State::AmplificationLimited { tx_allowance } => {
                assert_eq!(*tx_allowance, (MINIMUM_MTU * 3) as u32)
            }
        }
    }

    // Maintain amplification limits if already in AmplificationLimited state
    #[test]
    fn on_closing_not_validated_path() {
        let mut path = testing::helper_path_server();
        assert!(!path.is_validated());

        // Trigger:
        path.on_closing();

        // Expectation:
        match path.state {
            path::State::Validated => panic!("transition to AmplificationLimited when closing"),
            path::State::AmplificationLimited { tx_allowance } => {
                assert_eq!(*tx_allowance, 0)
            }
        }
    }

    #[test]
    fn is_congestion_limited() {
        let mut path = testing::helper_path_client();
        let mtu = path.mtu_controller.mtu() as u32;

        path.congestion_controller.congestion_window = 12000;
        path.congestion_controller.bytes_in_flight = 12000 - 500 - mtu;

        // There is room for an MTU sized packet after including the 500 bytes, so the path is not congestion limited
        assert!(!path.is_congestion_limited(500));

        // There isn't room for an MTU sized packet after including the 501 bytes, so the path is congestion limited
        assert!(path.is_congestion_limited(501));
    }
}


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

use crate::{
    contexts::WriteContext,
    path,
    path::{MaxMtu, MINIMUM_MTU},
    recovery::congestion_controller,
    transmission,
};
use core::time::Duration;
use s2n_codec::EncoderValue;
use s2n_quic_core::{
    counter::{Counter, Saturating},
    event,
    event::{builder::MtuUpdatedCause, IntoEvent},
    frame,
    inet::SocketAddress,
    packet::number::PacketNumber,
    path::{IPV4_MIN_HEADER_LEN, IPV6_MIN_HEADER_LEN, UDP_HEADER_LEN},
    recovery::CongestionController,
    time::{timer, Timer, Timestamp},
};

#[derive(Clone, Debug, PartialEq, Eq)]
enum State {
    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.2
    //# The DISABLED state is the initial state before probing has started.
    Disabled,
    /// SEARCH_REQUESTED is used to indicate a probe packet has been requested
    /// to be transmitted, but has not been transmitted yet.
    SearchRequested,
    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.2
    //# The SEARCHING state is the main probing state.
    Searching(PacketNumber, Timestamp),
    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.2
    //# The SEARCH_COMPLETE state indicates that a search has completed.
    SearchComplete,
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-14.3
//# Endpoints SHOULD set the initial value of BASE_PLPMTU (Section 5.1 of
//# [DPLPMTUD]) to be consistent with QUIC's smallest allowed maximum
//# datagram size.

//= https://www.rfc-editor.org/rfc/rfc8899#section-5.1.2
//# When using IPv4, there is no currently equivalent size specified,
//# and a default BASE_PLPMTU of 1200 bytes is RECOMMENDED.
const BASE_PLPMTU: u16 = MINIMUM_MTU;

//= https://www.rfc-editor.org/rfc/rfc8899#section-5.1.2
//# The MAX_PROBES is the maximum value of the PROBE_COUNT
//# counter (see Section 5.1.3).  MAX_PROBES represents the limit for
//# the number of consecutive probe attempts of any size.  Search
//# algorithms benefit from a MAX_PROBES value greater than 1 because
//# this can provide robustness to isolated packet loss.  The default
//# value of MAX_PROBES is 3.
const MAX_PROBES: u8 = 3;

/// The minimum length of the data field of a packet sent over an
/// Ethernet is 1500 octets, thus the maximum length of an IP datagram
/// sent over an Ethernet is 1500 octets.
/// See https://www.rfc-editor.org/rfc/rfc894.txt
const ETHERNET_MTU: u16 = 1500;

/// If the next value to probe is within the PROBE_THRESHOLD bytes of
/// the current Path MTU, probing will be considered complete.
const PROBE_THRESHOLD: u16 = 20;

/// When the black_hole_counter exceeds this threshold, on_black_hole_detected will be
/// called to reduce the MTU to the BASE_PLPMTU. The black_hole_counter is incremented when
/// a burst of consecutive packets is lost that starts with a packet that is:
///      1) not an MTU probe
///      2) larger than the BASE_PLPMTU
///      3) sent after the largest MTU-sized acknowledged packet number
/// This is a possible indication that the path cannot support the MTU that was previously confirmed.
const BLACK_HOLE_THRESHOLD: u8 = 3;

/// After a black hole has been detected, the mtu::Controller will wait this duration
/// before probing for a larger MTU again.
const BLACK_HOLE_COOL_OFF_DURATION: Duration = Duration::from_secs(60);

//= https://www.rfc-editor.org/rfc/rfc8899#section-5.1.1
//# The PMTU_RAISE_TIMER is configured to the period a
//# sender will continue to use the current PLPMTU, after which it
//# reenters the Search Phase.  This timer has a period of 600
//# seconds, as recommended by PLPMTUD [RFC4821].
const PMTU_RAISE_TIMER_DURATION: Duration = Duration::from_secs(600);

#[derive(Clone, Debug)]
pub struct Controller {
    state: State,
    //= https://www.rfc-editor.org/rfc/rfc8899#section-2
    //# The Packetization Layer PMTU is an estimate of the largest size
    //# of PL datagram that can be sent by a path, controlled by PLPMTUD
    plpmtu: u16,
    /// The maximum size any packet can reach
    max_mtu: MaxMtu,
    /// The maximum size the UDP payload can reach for any probe packet.
    max_udp_payload: u16,
    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.1.3
    //# The PROBED_SIZE is the size of the current probe packet
    //# as determined at the PL.  This is a tentative value for the
    //# PLPMTU, which is awaiting confirmation by an acknowledgment.
    probed_size: u16,
    /// The maximum size datagram to probe for. In contrast to the max_udp_payload,
    /// this value will decrease if probes are not acknowledged.
    max_probe_size: u16,
    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.1.3
    //# The PROBE_COUNT is a count of the number of successive
    //# unsuccessful probe packets that have been sent.
    probe_count: u8,
    /// A count of the number of packets with a size > MINIMUM_MTU lost since
    /// the last time a packet with size equal to the current MTU was acknowledged.
    black_hole_counter: Counter<u8, Saturating>,
    /// The largest acknowledged packet with size >= the plpmtu. Used when tracking
    /// packets that have been lost for the purpose of detecting a black hole.
    largest_acked_mtu_sized_packet: Option<PacketNumber>,
    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.1.1
    //# The PMTU_RAISE_TIMER is configured to the period a
    //# sender will continue to use the current PLPMTU, after which it
    //# reenters the Search Phase.
    pmtu_raise_timer: Timer,
}

impl Controller {
    /// Construct a new mtu::Controller with the given `max_mtu` and `peer_socket_address`
    ///
    /// The UDP header length and IP header length will be subtracted from `max_mtu` to
    /// determine the max_udp_payload used for limiting the payload length of probe packets.
    /// max_mtu is the maximum allowed mtu, e.g. for jumbo frames this value is expected to
    /// be over 9000.
    pub fn new(max_mtu: MaxMtu, peer_socket_address: &SocketAddress) -> Self {
        let min_ip_header_len = match peer_socket_address {
            SocketAddress::IpV4(_) => IPV4_MIN_HEADER_LEN,
            SocketAddress::IpV6(_) => IPV6_MIN_HEADER_LEN,
        };
        let max_udp_payload = u16::from(max_mtu) - UDP_HEADER_LEN - min_ip_header_len;
        debug_assert!(
            max_udp_payload >= BASE_PLPMTU,
            "max_udp_payload must be at least {BASE_PLPMTU}"
        );

        // The UDP payload size for the most likely MTU is based on standard Ethernet MTU minus
        // the minimum length IP headers (without IPv4 options or IPv6 extensions) and UPD header
        let initial_probed_size =
            (ETHERNET_MTU - UDP_HEADER_LEN - min_ip_header_len).min(max_udp_payload);

        Self {
            state: State::Disabled,
            plpmtu: BASE_PLPMTU,
            probed_size: initial_probed_size,
            max_mtu,
            max_udp_payload,
            max_probe_size: max_udp_payload,
            probe_count: 0,
            black_hole_counter: Default::default(),
            largest_acked_mtu_sized_packet: None,
            pmtu_raise_timer: Timer::default(),
        }
    }

    /// Enable path MTU probing
    pub fn enable(&mut self) {
        if self.state != State::Disabled {
            return;
        }

        // TODO: Look up current MTU in a cache. If there is a cache hit
        //       move directly to SearchComplete and arm the PMTU raise timer.
        //       Otherwise, start searching for a larger PMTU immediately
        self.request_new_search(None);
    }

    /// Called when the connection timer expires
    pub fn on_timeout(&mut self, now: Timestamp) {
        if self.pmtu_raise_timer.poll_expiration(now).is_ready() {
            self.request_new_search(None);
        }
    }

    //= https://www.rfc-editor.org/rfc/rfc8899#section-4.2
    //# When
    //# supported, this mechanism MAY also be used by DPLPMTUD to acknowledge
    //# reception of a probe packet.
    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<CC: CongestionController, Pub: event::ConnectionPublisher>(
        &mut self,
        packet_number: PacketNumber,
        sent_bytes: u16,
        congestion_controller: &mut CC,
        path_id: path::Id,
        publisher: &mut Pub,
    ) {
        if self.state == State::Disabled || !packet_number.space().is_application_data() {
            return;
        }

        if sent_bytes >= self.plpmtu
            && self
                .largest_acked_mtu_sized_packet
                .map_or(true, |pn| packet_number > pn)
        {
            // Reset the black hole counter since a packet the size of the current MTU or larger
            // has been acknowledged, indicating the path can still support the current MTU
            self.black_hole_counter = Default::default();
            self.largest_acked_mtu_sized_packet = Some(packet_number);
        }

        if let State::Searching(probe_packet_number, transmit_time) = self.state {
            if packet_number == probe_packet_number {
                self.plpmtu = self.probed_size;
                // A new MTU has been confirmed, notify the congestion controller
                congestion_controller.on_mtu_update(
                    self.plpmtu,
                    &mut congestion_controller::PathPublisher::new(publisher, path_id),
                );

                publisher.on_mtu_updated(event::builder::MtuUpdated {
                    path_id: path_id.into_event(),
                    mtu: self.plpmtu,
                    cause: MtuUpdatedCause::ProbeAcknowledged,
                });

                self.update_probed_size();

                //= https://www.rfc-editor.org/rfc/rfc8899#section-8
                //# To avoid excessive load, the interval between individual probe
                //# packets MUST be at least one RTT, and the interval between rounds of
                //# probing is determined by the PMTU_RAISE_TIMER.

                // Subsequent probe packets are sent based on the round trip transmission and
                // acknowledgement/loss of a packet, so the interval will be at least 1 RTT.
                self.request_new_search(Some(transmit_time));
            }
        }
    }

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //# The PL is REQUIRED to be
    //# robust in the case where probe packets are lost due to other
    //# reasons (including link transmission error, congestion).
    /// This method gets called when a packet loss is reported
    #[allow(clippy::too_many_arguments)]
    pub fn on_packet_loss<CC: CongestionController, Pub: event::ConnectionPublisher>(
        &mut self,
        packet_number: PacketNumber,
        lost_bytes: u16,
        new_loss_burst: bool,
        now: Timestamp,
        congestion_controller: &mut CC,
        path_id: path::Id,
        publisher: &mut Pub,
    ) {
        // MTU probes are only sent in application data space
        if !packet_number.space().is_application_data() {
            return;
        }

        match &self.state {
            State::Disabled => {}
            State::Searching(probe_pn, _) if *probe_pn == packet_number => {
                // The MTU probe was lost
                if self.probe_count == MAX_PROBES {
                    // We've sent MAX_PROBES without acknowledgement, so
                    // attempt a smaller probe size
                    self.max_probe_size = self.probed_size;
                    self.update_probed_size();
                    self.request_new_search(None);
                } else {
                    // Try the same probe size again
                    self.state = State::SearchRequested
                }
            }
            State::Searching(_, _) | State::SearchComplete | State::SearchRequested => {
                if (BASE_PLPMTU + 1..=self.plpmtu).contains(&lost_bytes)
                    && self
                        .largest_acked_mtu_sized_packet
                        .map_or(true, |pn| packet_number > pn)
                    && new_loss_burst
                {
                    // A non-probe packet larger than the BASE_PLPMTU that was sent after the last
                    // acknowledged MTU-sized packet has been lost
                    self.black_hole_counter += 1;
                }

                if self.black_hole_counter > BLACK_HOLE_THRESHOLD {
                    self.on_black_hole_detected(now, congestion_controller, path_id, publisher);
                }
            }
        }
    }

    /// Queries the component for any outgoing frames that need to get sent
    ///
    /// This method assumes that no other data (other than the packet header) has been written
    /// to the supplied `WriteContext`. This necessitates the caller ensuring the probe packet
    /// written by this method to be in its own connection transmission.
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) {
        if self.state != State::SearchRequested || !context.transmission_mode().is_mtu_probing() {
            //= https://www.rfc-editor.org/rfc/rfc8899#section-5.2
            //# When used with an acknowledged PL (e.g., SCTP), DPLPMTUD SHOULD NOT continue to
            //# generate PLPMTU probes in this state.
            return;
        }

        // Each packet contains overhead in the form of a packet header and an authentication tag.
        // This overhead contributes to the overall size of the packet, so the payload we write
        // to the packet will account for this overhead to reach the target probed size.
        let probe_payload_size =
            self.probed_size as usize - context.header_len() - context.tag_len();

        if context.remaining_capacity() < probe_payload_size {
            // There isn't enough capacity in the buffer to write the datagram we
            // want to probe, so we've reached the maximum pmtu and the search is complete.
            self.state = State::SearchComplete;
            return;
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-14.4
        //# Endpoints could limit the content of PMTU probes to PING and PADDING
        //# frames, since packets that are larger than the current maximum
        //# datagram size are more likely to be dropped by the network.

        //= https://www.rfc-editor.org/rfc/rfc8899#section-3
        //# Probe loss recovery: It is RECOMMENDED to use probe packets that
        //# do not carry any user data that would require retransmission if
        //# lost.

        //= https://www.rfc-editor.org/rfc/rfc8899#section-4.1
        //# DPLPMTUD MAY choose to use only one of these methods to simplify the
        //# implementation.

        context.write_frame(&frame::Ping);
        let padding_size = probe_payload_size - frame::Ping.encoding_size();
        if let Some(packet_number) = context.write_frame(&frame::Padding {
            length: padding_size,
        }) {
            self.probe_count += 1;
            self.state = State::Searching(packet_number, context.current_time());
        }
    }

    /// Gets the currently validated maximum transmission unit, not including IP or UDP header len
    pub fn mtu(&self) -> usize {
        self.plpmtu as usize
    }

    /// Returns the maximum size any packet can reach
    pub fn max_mtu(&self) -> MaxMtu {
        self.max_mtu
    }

    /// Gets the MTU currently being probed for
    pub fn probed_sized(&self) -> usize {
        self.probed_size as usize
    }

    /// Sets `probed_size` to the next MTU size to probe for based on a binary search
    fn update_probed_size(&mut self) {
        //= https://www.rfc-editor.org/rfc/rfc8899#section-5.3.2
        //# Implementations SHOULD select the set of probe packet sizes to
        //# maximize the gain in PLPMTU from each search step.
        self.probed_size = self.plpmtu + ((self.max_probe_size - self.plpmtu) / 2);
    }

    /// Requests a new search to be initiated
    ///
    /// If `last_probe_time` is supplied, the PMTU Raise Timer will be armed as
    /// necessary if the probed_size is already within the PROBE_THRESHOLD
    /// of the current PLPMTU
    fn request_new_search(&mut self, last_probe_time: Option<Timestamp>) {
        if self.probed_size - self.plpmtu >= PROBE_THRESHOLD {
            self.probe_count = 0;
            self.state = State::SearchRequested;
        } else {
            // The next probe size is within the threshold of the current MTU
            // so its not worth additional probing.
            self.state = State::SearchComplete;

            if let Some(last_probe_time) = last_probe_time {
                self.arm_pmtu_raise_timer(last_probe_time + PMTU_RAISE_TIMER_DURATION);
            }
        }
    }

    /// Called when an excessive number of packets larger than the BASE_PLPMTU have been lost
    fn on_black_hole_detected<CC: CongestionController, Pub: event::ConnectionPublisher>(
        &mut self,
        now: Timestamp,
        congestion_controller: &mut CC,
        path_id: path::Id,
        publisher: &mut Pub,
    ) {
        self.black_hole_counter = Default::default();
        self.largest_acked_mtu_sized_packet = None;
        // Reset the plpmtu back to the BASE_PLPMTU and notify the congestion controller
        self.plpmtu = BASE_PLPMTU;
        congestion_controller.on_mtu_update(
            BASE_PLPMTU,
            &mut congestion_controller::PathPublisher::new(publisher, path_id),
        );
        // Cancel any current probes
        self.state = State::SearchComplete;
        // Arm the PMTU raise timer to try a larger MTU again after a cooling off period
        self.arm_pmtu_raise_timer(now + BLACK_HOLE_COOL_OFF_DURATION);

        publisher.on_mtu_updated(event::builder::MtuUpdated {
            path_id: path_id.into_event(),
            mtu: self.plpmtu,
            cause: MtuUpdatedCause::Blackhole,
        })
    }

    /// Arm the PMTU Raise Timer if there is still room to increase the
    /// MTU before hitting the max plpmtu
    fn arm_pmtu_raise_timer(&mut self, timestamp: Timestamp) {
        // Reset the max_probe_size to the max_udp_payload to allow for larger probe sizes
        self.max_probe_size = self.max_udp_payload;
        self.update_probed_size();

        if self.probed_size - self.plpmtu >= PROBE_THRESHOLD {
            // There is still some room to try a larger MTU again,
            // so arm the pmtu raise timer
            self.pmtu_raise_timer.set(timestamp);
        }
    }
}

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

        Ok(())
    }
}

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

#[cfg(any(test, feature = "testing"))]
pub mod testing {
    use crate::path::mtu::{test, Controller};

    /// Creates a new mtu::Controller with the given mtu and probed size
    pub fn test_controller(mtu: u16, probed_size: u16) -> Controller {
        let mut controller = test::new_controller(u16::max_value());
        controller.plpmtu = mtu;
        controller.probed_size = probed_size;
        controller
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::contexts::testing::{MockWriteContext, OutgoingFrameBuffer};
    use s2n_quic_core::{
        endpoint,
        event::testing::Publisher,
        frame::Frame,
        packet::number::PacketNumberSpace,
        recovery::congestion_controller::testing::mock::CongestionController,
        time::{clock::testing::now, timer::Provider as _},
        varint::VarInt,
    };
    use std::{convert::TryInto, net::SocketAddr};

    /// Creates a new mtu::Controller with an IPv4 address and the given `max_mtu`
    pub fn new_controller(max_mtu: u16) -> Controller {
        let addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
        Controller::new(max_mtu.try_into().unwrap(), &addr.into())
    }

    /// Creates an application space packet number with the given value
    fn pn(nr: usize) -> PacketNumber {
        PacketNumberSpace::ApplicationData.new_packet_number(VarInt::new(nr as u64).unwrap())
    }

    #[test]
    fn base_plpmtu_is_1200() {
        //= https://www.rfc-editor.org/rfc/rfc8899#section-5.1.2
        //= type=test
        //# When using
        //# IPv4, there is no currently equivalent size specified, and a
        //# default BASE_PLPMTU of 1200 bytes is RECOMMENDED.
        assert_eq!(BASE_PLPMTU, 1200);
    }

    #[test]
    #[should_panic]
    fn new_max_mtu_too_small() {
        new_controller(BASE_PLPMTU + UDP_HEADER_LEN + IPV4_MIN_HEADER_LEN - 1);
    }

    #[test]
    fn new_max_mtu_smaller_than_common_mtu() {
        let mut controller = new_controller(BASE_PLPMTU + UDP_HEADER_LEN + IPV4_MIN_HEADER_LEN + 1);
        assert_eq!(BASE_PLPMTU + 1, controller.probed_size);

        controller.enable();
        assert_eq!(State::SearchComplete, controller.state);
    }

    #[test]
    fn new_ipv4() {
        let addr: SocketAddr = "127.0.0.1:443".parse().unwrap();
        let controller = Controller::new(1600.try_into().unwrap(), &addr.into());
        assert_eq!(
            1600 - UDP_HEADER_LEN - IPV4_MIN_HEADER_LEN,
            controller.max_udp_payload
        );
        assert_eq!(
            1600 - UDP_HEADER_LEN - IPV4_MIN_HEADER_LEN,
            controller.max_probe_size
        );
        assert_eq!(BASE_PLPMTU as usize, controller.mtu());
        assert_eq!(0, controller.probe_count);
        assert_eq!(State::Disabled, controller.state);
        assert!(!controller.pmtu_raise_timer.is_armed());
        assert_eq!(
            ETHERNET_MTU - UDP_HEADER_LEN - IPV4_MIN_HEADER_LEN,
            controller.probed_size
        );
    }

    #[test]
    fn new_ipv6() {
        let addr: SocketAddr = "[2001:0db8:85a3:0001:0002:8a2e:0370:7334]:9000"
            .parse()
            .unwrap();
        let controller = Controller::new(2000.try_into().unwrap(), &addr.into());
        assert_eq!(
            2000 - UDP_HEADER_LEN - IPV6_MIN_HEADER_LEN,
            controller.max_udp_payload
        );
        assert_eq!(
            2000 - UDP_HEADER_LEN - IPV6_MIN_HEADER_LEN,
            controller.max_probe_size
        );
        assert_eq!(BASE_PLPMTU as usize, controller.mtu());
        assert_eq!(0, controller.probe_count);
        assert_eq!(State::Disabled, controller.state);
        assert!(!controller.pmtu_raise_timer.is_armed());
        assert_eq!(
            ETHERNET_MTU - UDP_HEADER_LEN - IPV6_MIN_HEADER_LEN,
            controller.probed_size
        );
    }

    #[test]
    fn enable_already_enabled() {
        let mut controller = new_controller(1500);
        assert_eq!(State::Disabled, controller.state);
        controller.enable();
        assert_eq!(State::SearchRequested, controller.state);
        controller.state = State::SearchComplete;
        controller.enable();
        assert_eq!(State::SearchComplete, controller.state);
    }

    #[test]
    fn enable() {
        let mut controller = new_controller(1500);
        assert_eq!(State::Disabled, controller.state);
        controller.enable();
        assert_eq!(State::SearchRequested, controller.state);
    }

    //= https://www.rfc-editor.org/rfc/rfc8899#section-4.2
    //= type=test
    //# When
    //# supported, this mechanism MAY also be used by DPLPMTUD to acknowledge
    //# reception of a probe packet.
    #[test]
    fn on_packet_ack_within_threshold() {
        let mut controller = new_controller(1472 + PROBE_THRESHOLD * 2);
        let max_udp_payload = controller.max_udp_payload;
        let pn = pn(1);
        let mut cc = CongestionController::default();
        let now = now();
        let mut publisher = Publisher::snapshot();
        controller.state = State::Searching(pn, now);
        controller.probed_size = BASE_PLPMTU;
        controller.max_probe_size = BASE_PLPMTU + PROBE_THRESHOLD * 2 - 1;

        controller.on_packet_ack(
            pn,
            BASE_PLPMTU,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );

        assert_eq!(
            BASE_PLPMTU + (max_udp_payload - BASE_PLPMTU) / 2,
            controller.probed_size
        );
        assert_eq!(1, cc.on_mtu_update);
        assert_eq!(State::SearchComplete, controller.state);
        assert!(controller.pmtu_raise_timer.is_armed());
        assert_eq!(
            Some(now + PMTU_RAISE_TIMER_DURATION),
            controller.next_expiration()
        );

        // Enough time passes that its time to try raising the PMTU again
        let now = now + PMTU_RAISE_TIMER_DURATION;
        controller.on_timeout(now);

        assert_eq!(State::SearchRequested, controller.state);
        assert_eq!(
            BASE_PLPMTU + (max_udp_payload - BASE_PLPMTU) / 2,
            controller.probed_size
        );
    }

    #[test]
    fn on_packet_ack_within_threshold_of_max_plpmtu() {
        let mut controller = new_controller(1472 + (PROBE_THRESHOLD * 2 - 1));
        let max_udp_payload = controller.max_udp_payload;
        let pn = pn(1);
        let mut cc = CongestionController::default();
        let now = now();
        let mut publisher = Publisher::snapshot();
        controller.state = State::Searching(pn, now);

        controller.on_packet_ack(
            pn,
            controller.probed_size,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );

        assert_eq!(1472 + (max_udp_payload - 1472) / 2, controller.probed_size);
        assert_eq!(1, cc.on_mtu_update);
        assert_eq!(State::SearchComplete, controller.state);
        assert!(!controller.pmtu_raise_timer.is_armed());
    }

    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.3.2
    //= type=test
    //# Implementations SHOULD select the set of probe packet sizes to
    //# maximize the gain in PLPMTU from each search step.

    //= https://www.rfc-editor.org/rfc/rfc8899#section-8
    //= type=test
    //# To avoid excessive load, the interval between individual probe
    //# packets MUST be at least one RTT, and the interval between rounds of
    //# probing is determined by the PMTU_RAISE_TIMER.
    #[test]
    fn on_packet_ack_search_requested() {
        let mut controller = new_controller(1500 + (PROBE_THRESHOLD * 2));
        let max_udp_payload = controller.max_udp_payload;
        let pn = pn(1);
        let mut cc = CongestionController::default();
        let now = now();
        let mut publisher = Publisher::snapshot();
        controller.state = State::Searching(pn, now);

        controller.on_packet_ack(
            pn,
            controller.probed_size,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );

        assert_eq!(1472 + (max_udp_payload - 1472) / 2, controller.probed_size);
        assert_eq!(1, cc.on_mtu_update);
        assert_eq!(State::SearchRequested, controller.state);
        assert!(!controller.pmtu_raise_timer.is_armed());
    }

    #[test]
    fn on_packet_ack_resets_black_hole_counter() {
        let mut controller = new_controller(1500 + (PROBE_THRESHOLD * 2));
        let pnum = pn(3);
        let mut cc = CongestionController::default();
        let mut publisher = Publisher::snapshot();
        controller.enable();

        controller.black_hole_counter += 1;
        // ack a packet smaller than the plpmtu
        controller.on_packet_ack(
            pnum,
            controller.plpmtu - 1,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );
        assert_eq!(controller.black_hole_counter, 1);
        assert_eq!(None, controller.largest_acked_mtu_sized_packet);

        // ack a packet the size of the plpmtu
        controller.on_packet_ack(
            pnum,
            controller.plpmtu,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );
        assert_eq!(controller.black_hole_counter, 0);
        assert_eq!(Some(pnum), controller.largest_acked_mtu_sized_packet);

        controller.black_hole_counter += 1;

        // ack an older packet
        let pnum_2 = pn(2);
        controller.on_packet_ack(
            pnum_2,
            controller.plpmtu,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );
        assert_eq!(controller.black_hole_counter, 1);
        assert_eq!(Some(pnum), controller.largest_acked_mtu_sized_packet);
    }

    #[test]
    fn on_packet_ack_disabled_controller() {
        let mut controller = new_controller(1500 + (PROBE_THRESHOLD * 2));
        let pnum = pn(3);
        let mut cc = CongestionController::default();
        let mut publisher = Publisher::snapshot();

        controller.black_hole_counter += 1;
        controller.largest_acked_mtu_sized_packet = Some(pnum);

        let pn = pn(10);
        controller.on_packet_ack(
            pn,
            controller.plpmtu,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );

        assert_eq!(State::Disabled, controller.state);
        assert_eq!(controller.black_hole_counter, 1);
        assert_eq!(Some(pnum), controller.largest_acked_mtu_sized_packet);
    }

    #[test]
    fn on_packet_ack_not_application_space() {
        let mut controller = new_controller(1500 + (PROBE_THRESHOLD * 2));
        let pnum = pn(3);
        let mut cc = CongestionController::default();
        let mut publisher = Publisher::snapshot();
        controller.enable();

        controller.black_hole_counter += 1;
        controller.largest_acked_mtu_sized_packet = Some(pnum);

        // on_packet_ack will be called with packet numbers from Initial and Handshake space,
        // so it should not fail in this scenario.
        let pn = PacketNumberSpace::Handshake.new_packet_number(VarInt::from_u8(10));
        controller.on_packet_ack(
            pn,
            controller.plpmtu,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );
        assert_eq!(controller.black_hole_counter, 1);
        assert_eq!(Some(pnum), controller.largest_acked_mtu_sized_packet);
    }

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //= type=test
    //# The PL is REQUIRED to be
    //# robust in the case where probe packets are lost due to other
    //# reasons (including link transmission error, congestion).
    #[test]
    fn on_packet_loss() {
        let mut controller = new_controller(1500);
        let max_udp_payload = controller.max_udp_payload;
        let pn = pn(1);
        let mut cc = CongestionController::default();
        let now = now();
        let mut publisher = Publisher::snapshot();
        controller.state = State::Searching(pn, now);
        let probed_size = controller.probed_size;

        controller.on_packet_loss(
            pn,
            controller.probed_size,
            false,
            now,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );

        assert_eq!(0, cc.on_mtu_update);
        assert_eq!(max_udp_payload, controller.max_probe_size);
        assert_eq!(probed_size, controller.probed_size);
        assert_eq!(State::SearchRequested, controller.state);
    }

    #[test]
    fn on_packet_loss_max_probes() {
        let mut controller = new_controller(1500);
        let max_udp_payload = controller.max_udp_payload;
        let pn = pn(1);
        let mut cc = CongestionController::default();
        let now = now();
        let mut publisher = Publisher::snapshot();
        controller.state = State::Searching(pn, now);
        controller.probe_count = MAX_PROBES;
        assert_eq!(max_udp_payload, controller.max_probe_size);

        controller.on_packet_loss(
            pn,
            controller.probed_size,
            false,
            now,
            &mut cc,
            path::Id::test_id(),
            &mut publisher,
        );

        assert_eq!(0, cc.on_mtu_update);
        assert_eq!(1472, controller.max_probe_size);
        assert_eq!(
            BASE_PLPMTU + (1472 - BASE_PLPMTU) / 2,
            controller.probed_size
        );
        assert_eq!(State::SearchRequested, controller.state);
    }

    #[test]
    fn on_packet_loss_black_hole() {
        let mut controller = new_controller(1500);
        let mut cc = CongestionController::default();
        let now = now();
        let mut publisher = Publisher::snapshot();
        controller.plpmtu = 1472;
        controller.enable();

        for i in 0..BLACK_HOLE_THRESHOLD + 1 {
            let pn = pn(i as usize);

            // Losing a packet the size of the BASE_PLPMTU should not increase the black_hole_counter
            controller.on_packet_loss(
                pn,
                BASE_PLPMTU,
                true,
                now,
                &mut cc,
                path::Id::test_id(),
                &mut publisher,
            );
            assert_eq!(controller.black_hole_counter, i);

            // Losing a packet larger than the PLPMTU should not increase the black_hole_counter
            controller.on_packet_loss(
                pn,
                controller.plpmtu + 1,
                true,
                now,
                &mut cc,
                path::Id::test_id(),
                &mut publisher,
            );
            assert_eq!(controller.black_hole_counter, i);

            // Losing a packet that does not start a new loss burst should not increase the black_hole_counter
            controller.on_packet_loss(
                pn,
                BASE_PLPMTU + 1,
                false,
                now,
                &mut cc,
                path::Id::test_id(),
                &mut publisher,
            );
            assert_eq!(controller.black_hole_counter, i);

            controller.on_packet_loss(
                pn,
                BASE_PLPMTU + 1,
                true,
                now,
                &mut cc,
                path::Id::test_id(),
                &mut publisher,
            );
            if i < BLACK_HOLE_THRESHOLD {
                assert_eq!(controller.black_hole_counter, i + 1);
            }
        }

        assert_eq!(controller.black_hole_counter, 0);
        assert_eq!(None, controller.largest_acked_mtu_sized_packet);
        assert_eq!(1, cc.on_mtu_update);
        assert_eq!(BASE_PLPMTU, controller.plpmtu);
        assert_eq!(State::SearchComplete, controller.state);
        assert_eq!(
            Some(now + BLACK_HOLE_COOL_OFF_DURATION),
            controller.pmtu_raise_timer.next_expiration()
        );
    }

    #[test]
    fn on_packet_loss_disabled_controller() {
        let mut controller = new_controller(1500);
        let mut cc = CongestionController::default();
        let now = now();
        let mut publisher = Publisher::snapshot();

        for i in 0..BLACK_HOLE_THRESHOLD + 1 {
            let pn = pn(i as usize);
            assert_eq!(controller.black_hole_counter, 0);
            controller.on_packet_loss(
                pn,
                BASE_PLPMTU + 1,
                false,
                now,
                &mut cc,
                path::Id::test_id(),
                &mut publisher,
            );
        }

        assert_eq!(State::Disabled, controller.state);
        assert_eq!(controller.black_hole_counter, 0);
        assert_eq!(0, cc.on_mtu_update);
    }

    #[test]
    fn on_packet_loss_not_application_space() {
        let mut controller = new_controller(1500);
        let mut cc = CongestionController::default();
        let mut publisher = Publisher::snapshot();

        // test the loss in each state
        for state in vec![
            State::Disabled,
            State::SearchRequested,
            State::Searching(pn(1), now()),
            State::SearchComplete,
        ] {
            controller.state = state;
            for i in 0..BLACK_HOLE_THRESHOLD + 1 {
                // on_packet_loss may be called with packet numbers from Initial and Handshake space
                // so it should not fail in this scenario.
                let pn = PacketNumberSpace::Initial.new_packet_number(VarInt::from_u8(i));
                controller.on_packet_loss(
                    pn,
                    BASE_PLPMTU + 1,
                    false,
                    now(),
                    &mut cc,
                    path::Id::test_id(),
                    &mut publisher,
                );
                assert_eq!(controller.black_hole_counter, 0);
                assert_eq!(0, cc.on_mtu_update);
            }
        }
    }

    //= https://www.rfc-editor.org/rfc/rfc8899#section-5.2
    //= type=test
    //# When used with an
    //# acknowledged PL (e.g., SCTP), DPLPMTUD SHOULD NOT continue to
    //# generate PLPMTU probes in this state.
    #[test]
    fn on_transmit_search_not_requested() {
        let mut controller = new_controller(1500);
        controller.state = State::SearchComplete;
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut write_context = MockWriteContext::new(
            now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::MtuProbing,
            endpoint::Type::Server,
        );

        controller.on_transmit(&mut write_context);
        assert!(frame_buffer.is_empty());
        assert_eq!(State::SearchComplete, controller.state);
    }

    #[test]
    fn on_transmit_not_mtu_probing() {
        let mut controller = new_controller(1500);
        controller.state = State::SearchRequested;
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut write_context = MockWriteContext::new(
            now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Server,
        );

        controller.on_transmit(&mut write_context);
        assert!(frame_buffer.is_empty());
        assert_eq!(State::SearchRequested, controller.state);

        controller.state = State::SearchComplete;
        let mut frame_buffer = OutgoingFrameBuffer::new();
        let mut write_context = MockWriteContext::new(
            now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::Normal,
            endpoint::Type::Server,
        );

        controller.on_transmit(&mut write_context);
        assert!(frame_buffer.is_empty());
        assert_eq!(State::SearchComplete, controller.state);
    }

    #[test]
    fn on_transmit_no_capacity() {
        let mut controller = new_controller(1500);
        controller.state = State::SearchRequested;
        let mut frame_buffer = OutgoingFrameBuffer::new();
        frame_buffer.set_max_packet_size(Some(controller.probed_size as usize - 1));
        let mut write_context = MockWriteContext::new(
            now(),
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::MtuProbing,
            endpoint::Type::Server,
        );

        controller.on_transmit(&mut write_context);
        assert!(frame_buffer.is_empty());
        assert_eq!(State::SearchComplete, controller.state);
    }

    //= https://www.rfc-editor.org/rfc/rfc9000#section-14.4
    //= type=test
    //# Endpoints could limit the content of PMTU probes to PING and PADDING
    //# frames, since packets that are larger than the current maximum
    //# datagram size are more likely to be dropped by the network.

    //= https://www.rfc-editor.org/rfc/rfc8899#section-3
    //= type=test
    //# Probe loss recovery: It is RECOMMENDED to use probe packets that
    //# do not carry any user data that would require retransmission if
    //# lost.

    //= https://www.rfc-editor.org/rfc/rfc8899#section-4.1
    //= type=test
    //# DPLPMTUD MAY choose to use only one of these methods to simplify the
    //# implementation.
    #[test]
    fn on_transmit() {
        let mut controller = new_controller(1500);
        controller.state = State::SearchRequested;
        let now = now();
        let mut frame_buffer = OutgoingFrameBuffer::new();
        frame_buffer.set_max_packet_size(Some(controller.probed_size as usize));
        let mut write_context = MockWriteContext::new(
            now,
            &mut frame_buffer,
            transmission::Constraint::None,
            transmission::Mode::MtuProbing,
            endpoint::Type::Server,
        );
        let packet_number = write_context.packet_number();

        controller.on_transmit(&mut write_context);
        assert_eq!(0, write_context.remaining_capacity());
        assert_eq!(
            Frame::Ping(frame::Ping),
            write_context.frame_buffer.pop_front().unwrap().as_frame()
        );
        assert_eq!(
            Frame::Padding(frame::Padding {
                length: controller.probed_size as usize - 1
            }),
            write_context.frame_buffer.pop_front().unwrap().as_frame()
        );
        assert_eq!(State::Searching(packet_number, now), controller.state);
    }
}