// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
use super::*;
use s2n_quic_core::{
    event::{builder::Path, testing::Publisher},
    time::{clock::testing as time, timer::Provider},
    varint::VarInt,
};
use std::ops::Deref;

/// Return ECN counts with the given counts
fn helper_ecn_counts(ect0: u8, ect1: u8, ce: u8) -> EcnCounts {
    EcnCounts {
        ect_0_count: VarInt::from_u8(ect0),
        ect_1_count: VarInt::from_u8(ect1),
        ce_count: VarInt::from_u8(ce),
    }
}

#[test]
fn default() {
    let controller = Controller::default();
    assert_eq!(0, *controller.black_hole_counter.deref());
    assert_eq!(State::Testing(0), controller.state);
    assert_eq!(None, controller.last_acked_ecn_packet_timestamp);
}

#[test]
fn restart() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Failed(Timer::default()),
        ..Default::default()
    };
    controller.black_hole_counter += 1;

    controller.restart(Path::test(), &mut publisher);

    assert_eq!(State::Testing(0), controller.state);
    assert_eq!(0, *controller.black_hole_counter.deref());
}

#[test]
fn restart_already_in_testing_0() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Testing(0),
        ..Default::default()
    };
    controller.black_hole_counter += 1;

    controller.restart(Path::test(), &mut publisher);

    assert_eq!(State::Testing(0), controller.state);
    assert_eq!(0, *controller.black_hole_counter.deref());
}

#[test]
fn on_timeout_failed() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    controller.fail(now, Path::test(), &mut publisher);

    if let State::Failed(timer) = &controller.state {
        assert!(timer.is_armed());
    } else {
        panic!("State should be Failed");
    }

    assert_eq!(0, *controller.black_hole_counter.deref());

    let now = now + RETEST_COOL_OFF_DURATION - Duration::from_secs(1);

    // Too soon
    controller.on_timeout(
        now,
        Path::test(),
        &mut random::testing::Generator(123),
        Duration::default(),
        &mut publisher,
    );

    if let State::Failed(timer) = &controller.state {
        assert!(timer.is_armed());
    } else {
        panic!("State should be Failed");
    }

    assert_eq!(0, *controller.black_hole_counter.deref());

    let now = now + Duration::from_secs(1);
    controller.on_timeout(
        now,
        Path::test(),
        &mut random::testing::Generator(123),
        Duration::default(),
        &mut publisher,
    );

    assert_eq!(State::Testing(0), controller.state);
    assert_eq!(0, *controller.black_hole_counter.deref());
}

#[test]
fn on_timeout_capable() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    let rtt = Duration::from_millis(50);
    let ce_suppression_time = now + *CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() as u32 * rtt;
    let mut ce_suppression_timer = Timer::default();
    ce_suppression_timer.set(ce_suppression_time);
    controller.state = State::Capable(ce_suppression_timer);

    let now = now + rtt;

    // Too soon
    controller.on_timeout(
        now,
        Path::test(),
        &mut random::testing::Generator(123),
        rtt,
        &mut publisher,
    );

    if let State::Capable(timer) = &controller.state {
        assert!(timer.is_armed());
        assert_eq!(Some(ce_suppression_time), timer.next_expiration());
    } else {
        panic!("State should be Capable");
    }

    // Timer is no longer armed
    controller.state = State::Capable(Timer::default());

    controller.on_timeout(
        now,
        Path::test(),
        &mut random::testing::Generator(123),
        rtt,
        &mut publisher,
    );

    if let State::Capable(timer) = &controller.state {
        assert!(timer.is_armed());
        assert!(timer.next_expiration().unwrap() > now);
    } else {
        panic!("State should be Capable");
    }
}

#[test]
fn ecn() {
    let mut publisher = Publisher::snapshot();
    let now = time::now();

    for &transmission_mode in &[
        transmission::Mode::Normal,
        transmission::Mode::MtuProbing,
        transmission::Mode::PathValidationOnly,
    ] {
        let mut controller = Controller::default();
        assert!(controller.ecn(transmission_mode, now).using_ecn());

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.2
        //= type=test
        //# Upon successful validation, an endpoint MAY continue to set an ECT
        //# codepoint in subsequent packets it sends, with the expectation that
        //# the path is ECN-capable.
        let mut ce_suppression_timer = Timer::default();
        ce_suppression_timer.set(now + Duration::from_secs(10));
        controller.state = State::Capable(ce_suppression_timer);
        assert!(controller.ecn(transmission_mode, now).using_ecn());
        if let State::Capable(ref timer) = controller.state {
            assert!(timer.is_armed());
        } else {
            panic!("State should be Capable");
        }

        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.2
        //= type=test
        //# If validation fails, then the endpoint MUST disable ECN. It stops setting the ECT
        //# codepoint in IP packets that it sends, assuming that either the network path or
        //# the peer does not support ECN.
        controller.fail(time::now(), Path::test(), &mut publisher);
        assert!(!controller.ecn(transmission_mode, now).using_ecn());

        controller.state = State::Unknown;
        assert!(!controller.ecn(transmission_mode, now).using_ecn());
    }
}

#[test]
fn ecn_ce_suppression() {
    let now = time::now();

    for &transmission_mode in &[
        transmission::Mode::Normal,
        transmission::Mode::MtuProbing,
        transmission::Mode::PathValidationOnly,
    ] {
        let mut controller = Controller::default();
        assert!(controller.ecn(transmission_mode, now).using_ecn());

        let mut ce_suppression_timer = Timer::default();
        ce_suppression_timer.set(now);
        controller.state = State::Capable(ce_suppression_timer);
        assert!(controller
            .ecn(transmission_mode, now)
            .congestion_experienced());
        if let State::Capable(timer) = controller.state {
            assert!(!timer.is_armed());
        } else {
            panic!("State should be Capable");
        }
    }
}

#[test]
fn ecn_loss_recovery_probing() {
    let now = time::now();

    for state in vec![
        State::Capable(Timer::default()),
        State::Testing(0),
        State::Unknown,
        State::Failed(Timer::default()),
    ] {
        let mut controller = Controller {
            state,
            ..Default::default()
        };
        assert!(!controller
            .ecn(transmission::Mode::LossRecoveryProbing, now)
            .using_ecn());
    }
}

#[test]
fn is_capable() {
    for state in vec![
        State::Testing(0),
        State::Unknown,
        State::Failed(Timer::default()),
    ] {
        let controller = Controller {
            state,
            ..Default::default()
        };
        assert!(!controller.is_capable());
    }

    let controller = Controller {
        state: State::Capable(Timer::default()),
        ..Default::default()
    };
    assert!(controller.is_capable());
}

#[test]
fn validate_already_failed() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    controller.fail(now, Path::test(), &mut publisher);
    let outcome = controller.validate(
        EcnCounts::default(),
        EcnCounts::default(),
        EcnCounts::default(),
        None,
        now + Duration::from_secs(5),
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    if let State::Failed(timer) = &controller.state {
        assert!(timer.is_armed());
        assert_eq!(
            controller.next_expiration(),
            Some(now + RETEST_COOL_OFF_DURATION)
        );
    } else {
        panic!("State should be Failed");
    }
    assert_eq!(ValidationOutcome::Skipped, outcome);
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.1
//= type=test
//# If an ACK frame newly acknowledges a packet that the endpoint sent with
//# either the ECT(0) or ECT(1) codepoint set, ECN validation fails if the
//# corresponding ECN counts are not present in the ACK frame. This check
//# detects a network element that zeroes the ECN field or a peer that does
//# not report ECN markings.
#[test]
fn validate_ecn_counts_not_in_ack() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    let expected_ecn_counts = helper_ecn_counts(1, 0, 0);
    let outcome = controller.validate(
        expected_ecn_counts,
        EcnCounts::default(),
        EcnCounts::default(),
        None,
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Failed, outcome);
    assert!(matches!(controller.state, State::Failed(_)));
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.1
//= type=test
//# ECN validation also fails if the sum of the increase in ECT(0)
//# and ECN-CE counts is less than the number of newly acknowledged
//# packets that were originally sent with an ECT(0) marking.
#[test]
fn validate_ecn_ce_remarking() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    let expected_ecn_counts = helper_ecn_counts(1, 0, 0);
    let sent_packet_ecn_counts = helper_ecn_counts(1, 0, 0);
    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(EcnCounts::default()),
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Failed, outcome);
    assert!(matches!(controller.state, State::Failed(_)));
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-13.4.2.1
//= type=test
//# ECN validation can fail if the received total count for either ECT(0) or ECT(1)
//# exceeds the total number of packets sent with each corresponding ECT codepoint.
#[test]
fn validate_ect_0_remarking() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    let expected_ecn_counts = helper_ecn_counts(1, 0, 0);
    let sent_packet_ecn_counts = helper_ecn_counts(1, 0, 0);
    let ack_frame_ecn_counts = helper_ecn_counts(1, 1, 0);
    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(ack_frame_ecn_counts),
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Failed, outcome);
    assert!(matches!(controller.state, State::Failed(_)));
}

#[test]
fn validate_ect_0_remarking_after_restart() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    let expected_ecn_counts = helper_ecn_counts(1, 0, 0);
    let ack_frame_ecn_counts = helper_ecn_counts(0, 3, 0);
    let baseline_ecn_counts = helper_ecn_counts(0, 2, 0);
    let sent_packet_ecn_counts = helper_ecn_counts(1, 0, 0);
    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        baseline_ecn_counts,
        Some(ack_frame_ecn_counts),
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Failed, outcome);
    assert!(matches!(controller.state, State::Failed(_)));
}

#[test]
fn validate_no_ecn_counts() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Unknown,
        ..Default::default()
    };
    let now = time::now();
    let outcome = controller.validate(
        EcnCounts::default(),
        EcnCounts::default(),
        EcnCounts::default(),
        None,
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Skipped, outcome);
    assert_eq!(State::Unknown, controller.state);
}

#[test]
fn validate_ecn_decrease() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    let baseline_ecn_counts = helper_ecn_counts(1, 0, 0);
    let outcome = controller.validate(
        EcnCounts::default(),
        EcnCounts::default(),
        baseline_ecn_counts,
        None,
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Failed, outcome);
    assert!(matches!(controller.state, State::Failed(_)));
}

//= https://www.rfc-editor.org/rfc/rfc9000#appendix-A.4
//= type=test
//# From the "unknown" state, successful validation of the ECN counts in an ACK frame
//# (see Section 13.4.2.1) causes the ECN state for the path to become "capable",
//# unless no marked packet has been acknowledged.
#[test]
fn validate_no_marked_packets_acked() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Unknown,
        ..Default::default()
    };
    let now = time::now();
    let outcome = controller.validate(
        EcnCounts::default(),
        EcnCounts::default(),
        EcnCounts::default(),
        Some(EcnCounts::default()),
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Passed, outcome);
    assert_eq!(State::Unknown, controller.state);
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-8.3
//= type=test
//# A sender can detect suppression of reports by marking occasional packets that it
//# sends with an ECN-CE marking. If a packet sent with an ECN-CE marking is not
//# reported as having been CE marked when the packet is acknowledged, then the
//# sender can disable ECN for that path by not setting ECN-Capable Transport (ECT)
//# codepoints in subsequent packets sent on that path [RFC3168].
#[test]
fn validate_ce_suppression_remarked_to_not_ect() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    // We sent one ECT0 and one CE
    let newly_acked_ecn_counts = helper_ecn_counts(1, 0, 1);
    let sent_packet_ecn_counts = helper_ecn_counts(1, 0, 1);
    // The peer suppressed the CE count
    let ack_frame_ecn_counts = helper_ecn_counts(1, 0, 0);
    let outcome = controller.validate(
        newly_acked_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(ack_frame_ecn_counts),
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Failed, outcome);
    assert!(matches!(controller.state, State::Failed(_)));
}

//= https://www.rfc-editor.org/rfc/rfc9002#section-8.3
//= type=test
//# A sender can detect suppression of reports by marking occasional packets that it
//# sends with an ECN-CE marking. If a packet sent with an ECN-CE marking is not
//# reported as having been CE marked when the packet is acknowledged, then the
//# sender can disable ECN for that path by not setting ECN-Capable Transport (ECT)
//# codepoints in subsequent packets sent on that path [RFC3168].
#[test]
fn validate_ce_suppression_remarked_to_ect0() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    // We sent one ECT0 and one CE
    let newly_acked_ecn_counts = helper_ecn_counts(1, 0, 1);
    let sent_packet_ecn_counts = helper_ecn_counts(1, 0, 1);
    // The peer remarked the CE as ECT0
    let ack_frame_ecn_counts = helper_ecn_counts(2, 0, 0);
    let outcome = controller.validate(
        newly_acked_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(ack_frame_ecn_counts),
        now,
        Duration::default(),
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Failed, outcome);
    assert!(matches!(controller.state, State::Failed(_)));
}

#[test]
fn validate_capable() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Unknown,
        ..Default::default()
    };
    let now = time::now();
    let expected_ecn_counts = helper_ecn_counts(2, 0, 0);
    let ack_frame_ecn_counts = helper_ecn_counts(2, 0, 0);
    let sent_packet_ecn_counts = helper_ecn_counts(2, 0, 0);
    let rtt = Duration::from_millis(50);
    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(ack_frame_ecn_counts),
        now,
        rtt,
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Passed, outcome);
    assert!(controller.is_capable());
    if let State::Capable(timer) = controller.state {
        assert_eq!(
            Some(now + *CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() as u32 * rtt),
            timer.next_expiration()
        );
    }
}

#[test]
fn validate_capable_congestion_experienced() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Unknown,
        ..Default::default()
    };
    let now = time::now();
    let expected_ecn_counts = helper_ecn_counts(2, 0, 5);
    let ack_frame_ecn_counts = helper_ecn_counts(1, 0, 12);
    let sent_packet_ecn_counts = helper_ecn_counts(2, 0, 5);
    let rtt = Duration::from_millis(50);
    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(ack_frame_ecn_counts),
        now,
        rtt,
        Path::test(),
        &mut publisher,
    );

    assert_eq!(
        ValidationOutcome::CongestionExperienced(7_u8.into()),
        outcome
    );
    assert!(controller.is_capable());
    if let State::Capable(timer) = controller.state {
        assert_eq!(
            Some(now + *CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() as u32 * rtt),
            timer.next_expiration()
        );
    }
}

#[test]
fn validate_capable_ce_suppression_test() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Unknown,
        ..Default::default()
    };
    let now = time::now();
    let expected_ecn_counts = helper_ecn_counts(2, 0, 1);
    let ack_frame_ecn_counts = helper_ecn_counts(2, 0, 1);
    let sent_packet_ecn_counts = helper_ecn_counts(2, 0, 1);
    let rtt = Duration::from_millis(50);
    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(ack_frame_ecn_counts),
        now,
        rtt,
        Path::test(),
        &mut publisher,
    );

    // The outcome should not be `CongestionExperienced` despite the increase in CE-count,
    // because the CE-count was coming from a packet we had marked as ECN-CE
    assert_eq!(ValidationOutcome::Passed, outcome);
    assert!(controller.is_capable());
    if let State::Capable(timer) = controller.state {
        assert_eq!(
            Some(now + *CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() as u32 * rtt),
            timer.next_expiration()
        );
    }
}

/// Successful validation when not in the Unknown state does not change the state
#[test]
fn validate_capable_not_in_unknown_state() {
    let mut publisher = Publisher::snapshot();
    for state in vec![
        State::Testing(0),
        State::Capable(Timer::default()),
        State::Failed(Timer::default()),
    ] {
        let mut controller = Controller {
            state,
            ..Default::default()
        };
        let now = time::now();
        let expected_ecn_counts = helper_ecn_counts(1, 0, 0);
        let ack_frame_ecn_counts = helper_ecn_counts(1, 0, 0);
        let sent_packet_ecn_counts = helper_ecn_counts(1, 0, 0);
        let rtt = Duration::from_millis(50);
        let expected_state = controller.state.clone();
        controller.validate(
            expected_ecn_counts,
            sent_packet_ecn_counts,
            EcnCounts::default(),
            Some(ack_frame_ecn_counts),
            now,
            rtt,
            Path::test(),
            &mut publisher,
        );

        assert_eq!(expected_state, controller.state);
    }
}

#[test]
fn validate_capable_lost_ack_frame() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Unknown,
        ..Default::default()
    };
    let now = time::now();

    // We sent three ECT0 packets
    let sent_packet_ecn_counts = helper_ecn_counts(3, 0, 0);

    // The peer is acknowledging 2 of them, the third was acknowledge in an ack frame
    // that was lost
    let ack_frame_ecn_counts = helper_ecn_counts(3, 0, 0);

    let expected_ecn_counts = helper_ecn_counts(2, 0, 0);
    let rtt = Duration::from_millis(50);

    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        EcnCounts::default(),
        Some(ack_frame_ecn_counts),
        now,
        rtt,
        Path::test(),
        &mut publisher,
    );

    assert_eq!(ValidationOutcome::Passed, outcome);
    assert!(controller.is_capable());
    if let State::Capable(timer) = controller.state {
        assert_eq!(
            Some(now + *CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() as u32 * rtt),
            timer.next_expiration()
        );
    }
}

#[test]
fn validate_capable_after_restart() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller {
        state: State::Unknown,
        ..Default::default()
    };
    let now = time::now();
    let sent_packet_ecn_counts = helper_ecn_counts(2, 0, 0);
    let expected_ecn_counts = helper_ecn_counts(2, 0, 0);
    // The Ect1 markings would normally fail validation, but since they are included
    // in the baseline ecn counts below, that means we've already accounted for them.
    let ack_frame_ecn_counts = helper_ecn_counts(1, 2, 1);
    let baseline_ecn_counts = helper_ecn_counts(0, 2, 0);
    let rtt = Duration::from_millis(50);
    let outcome = controller.validate(
        expected_ecn_counts,
        sent_packet_ecn_counts,
        baseline_ecn_counts,
        Some(ack_frame_ecn_counts),
        now,
        rtt,
        Path::test(),
        &mut publisher,
    );

    assert_eq!(
        ValidationOutcome::CongestionExperienced(1_u8.into()),
        outcome
    );
    assert!(controller.is_capable());
    if let State::Capable(timer) = controller.state {
        assert_eq!(
            Some(now + *CE_SUPPRESSION_TESTING_RTT_MULTIPLIER.start() as u32 * rtt),
            timer.next_expiration()
        );
    }
}

#[test]
fn on_packet_sent() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();

    for i in 0..TESTING_PACKET_THRESHOLD {
        assert_eq!(State::Testing(i), controller.state);
        controller.on_packet_sent(
            ExplicitCongestionNotification::Ect0,
            Path::test(),
            &mut publisher,
        );
    }

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

#[test]
fn on_packet_loss() {
    let mut publisher = Publisher::snapshot();
    for state in vec![
        State::Testing(0),
        State::Capable(Timer::default()),
        State::Unknown,
    ] {
        let mut controller = Controller {
            state,
            ..Default::default()
        };
        let now = time::now();
        let time_sent = now + Duration::from_secs(1);

        controller.last_acked_ecn_packet_timestamp = Some(now);

        for i in 0..TESTING_PACKET_THRESHOLD + 1 {
            assert_eq!(i, *controller.black_hole_counter.deref());
            assert!(!matches!(controller.state, State::Failed(_)));
            controller.on_packet_loss(
                time_sent,
                ExplicitCongestionNotification::Ect0,
                time_sent,
                Path::test(),
                &mut publisher,
            );
        }

        if let State::Failed(timer) = &controller.state {
            assert!(timer.is_armed());
            assert_eq!(
                Some(time_sent + RETEST_COOL_OFF_DURATION),
                controller.next_expiration()
            );
        } else {
            panic!("State should be Failed");
        }

        assert_eq!(0, *controller.black_hole_counter.deref());
    }
}

#[test]
fn on_packet_loss_already_failed() {
    let mut publisher = Publisher::snapshot();
    let mut controller = Controller::default();
    let now = time::now();
    let time_sent = now + Duration::from_secs(1);

    controller.last_acked_ecn_packet_timestamp = Some(now);
    controller.fail(now, Path::test(), &mut publisher);

    for _i in 0..TESTING_PACKET_THRESHOLD + 1 {
        assert_eq!(0, *controller.black_hole_counter.deref());
        assert!(matches!(controller.state, State::Failed(_)));
        controller.on_packet_loss(
            time_sent,
            ExplicitCongestionNotification::Ect0,
            time_sent,
            Path::test(),
            &mut publisher,
        );
    }

    if let State::Failed(timer) = &controller.state {
        assert!(timer.is_armed());
        assert_eq!(
            Some(now + RETEST_COOL_OFF_DURATION),
            controller.next_expiration()
        );
    } else {
        panic!("State should be Failed");
    }

    assert_eq!(0, *controller.black_hole_counter.deref());
}

#[test]
fn fuzz_validate() {
    let now = time::now();

    bolero::check!()
        .with_type::<(EcnCounts, EcnCounts, EcnCounts, Option<EcnCounts>, Duration)>()
        .cloned()
        .for_each(
            |(
                newly_acked_ecn_counts,
                sent_packet_ecn_counts,
                baseline_ecn_counts,
                ack_frame_ecn_counts,
                rtt,
            )| {
                let mut controller = Controller::default();
                let outcome = controller.validate(
                    newly_acked_ecn_counts,
                    sent_packet_ecn_counts,
                    baseline_ecn_counts,
                    ack_frame_ecn_counts,
                    now,
                    rtt,
                    Path::test(),
                    &mut Publisher::no_snapshot(),
                );

                if outcome == ValidationOutcome::Failed {
                    assert!(!controller.is_capable());
                    assert_eq!(
                        ExplicitCongestionNotification::NotEct,
                        controller.ecn(transmission::Mode::Normal, now)
                    );
                    assert!(matches!(controller.state, State::Failed(_)));
                }
            },
        );
}


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

use super::{tests::helper_path, *};
use crate::{
    connection::{ConnectionIdMapper, InternalConnectionIdGenerator},
    endpoint::testing::Server as Config,
    path,
};
use bolero::{check, generator::*};
use core::time::Duration;
use s2n_quic_core::{
    connection::LocalId,
    event::testing::Publisher,
    frame::path_validation,
    inet::{DatagramInfo, ExplicitCongestionNotification},
    random,
    random::testing::Generator,
    time::{testing::Clock, Clock as _},
};
use std::collections::HashMap;

type Manager = super::Manager<Config>;
type Handle = path::RemoteAddress;

/// Path information stored for modeling the path::Manager
#[derive(Debug)]
struct PathInfo {
    state: path::State,
    handle: Handle,
}

impl PathInfo {
    pub fn new(path: Path<Config>) -> PathInfo {
        PathInfo {
            handle: path.handle,
            state: path.state,
        }
    }

    fn remote_address(&self) -> Handle {
        self.handle
    }

    fn is_validated(&self) -> bool {
        self.state == path::State::Validated
    }
}

#[derive(Debug)]
struct Oracle {
    paths: HashMap<Handle, PathInfo>,
    active: Handle,
    last_known_active_validated_path: Option<u8>,
    prev_state_amplification_limited: bool,
}

impl Oracle {
    fn new(handle: Handle, path: PathInfo) -> Oracle {
        let mut paths = HashMap::new();
        paths.insert(handle, path);

        Oracle {
            paths,
            active: handle,
            last_known_active_validated_path: None,
            prev_state_amplification_limited: false,
        }
    }

    fn get_active_path(&self) -> &PathInfo {
        self.paths.get(&self.active).unwrap()
    }

    fn path(&mut self, handle: &Handle) -> Option<(&Handle, &mut PathInfo)> {
        self.paths
            .iter_mut()
            .find(|(path_handle, _path)| path_handle.eq(&handle))
    }

    /// Insert a new path while adhering to limits
    fn insert_new_path(&mut self, new_handle: Handle, new_path: PathInfo) {
        // We only store max of 4 paths because of limits imposed by Active cid limits
        if self.paths.len() < 4 {
            self.paths.insert(new_handle, new_path);
        }
    }

    /// Insert new path if the remote address doesnt match any existing path
    fn on_datagram_received(
        &mut self,
        handle: &Handle,
        payload_len: u16,
        still_amplification_limited: bool,
    ) {
        match self.path(handle) {
            Some(_path) => (),
            None => {
                let new_path = PathInfo {
                    handle: *handle,
                    state: path::State::Validated,
                };

                self.insert_new_path(*handle, new_path);
            }
        }

        // Verify receiving bytes unblocks amplification limited.
        //
        // Amplification is calculated in terms of packets rather than bytes. Therefore any
        // payload len that is > 0 will unblock amplification.
        if self.prev_state_amplification_limited && payload_len > 0 {
            assert!(!still_amplification_limited);
        }
    }

    fn on_timeout(&self, _millis: &u16) {
        // TODO implement
        // call timeout on each Path
        //
        // if active path is not validated and validation failed
        //   set the last_known_active_validated_path as active path
    }

    fn on_processed_packet(&mut self, handle: &Handle, probe: &path_validation::Probe) {
        if !probe.is_probing() {
            self.active = *handle;
        }
    }
}

#[derive(Debug, TypeGenerator)]
enum Operation {
    // on_datagram_received
    // on_processed_packet
    OnDatagramReceived {
        /// Used for calling on_processed_packet
        ///
        /// `Some` implies that the packet was successfully processed
        /// and that `on_processed_packet` should be called.
        probing: Option<path_validation::Probe>,

        /// The remote address from which the datagram was sent
        handle: Handle,

        /// The peer may switch the cid on the connection
        local_id: LocalId,

        // 9000 captures up to jumbo frames
        #[generator(0..9000)]
        payload_len: u16,
    },

    // path.on_bytes_transmitted
    OnBytesTransmit {
        /// Used to calculate path_id
        ///
        /// path_id = path_id_generator % paths.len()
        #[generator(1..100)]
        path_id_generator: u8,

        /// Bytes to be transmitted
        // 9000 captures up to jumbo frames
        #[generator(0..9000)]
        bytes: u16,
    },

    // on_timeout
    IncrementTime {
        /// The milli-second value by which to increase the timestamp
        millis: u16,
    },
}

#[derive(Debug)]
struct Model {
    oracle: Oracle,
    subject: Manager,

    /// A monotonically increasing timestamp
    timestamp: Timestamp,
}

impl Model {
    pub fn new() -> Model {
        let zero_conn_id = connection::PeerId::try_from_bytes(&[0]).unwrap();
        let new_addr = Handle::default();
        let oracle = {
            let zero_path = PathInfo::new(helper_path(zero_conn_id));
            Oracle::new(new_addr, zero_path)
        };

        let manager = {
            let zero_path = helper_path(zero_conn_id);
            let mut random_generator = random::testing::Generator(123);
            let mut peer_id_registry =
                ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server)
                    .create_server_peer_id_registry(
                        InternalConnectionIdGenerator::new().generate_id(),
                        zero_path.peer_connection_id,
                    );

            // register 3 more ids which means a total of 4 paths. cid 0 is retired as part of
            // on_new_connection_id call
            for i in 1..=3 {
                let cid = connection::PeerId::try_from_bytes(&[i]).unwrap();
                let token = &[i; 16].into();
                peer_id_registry
                    .on_new_connection_id(&cid, i.into(), 0, token)
                    .unwrap();
            }

            Manager::new(zero_path, peer_id_registry)
        };

        let clock = Clock::default();
        Model {
            oracle,
            subject: manager,
            timestamp: clock.get_time(),
        }
    }

    pub fn apply(&mut self, operation: &Operation) {
        match operation {
            Operation::IncrementTime { millis } => self.on_timeout(millis),
            Operation::OnDatagramReceived {
                probing,
                handle,
                local_id,
                payload_len,
            } => self.on_datagram_received(handle, probing, *local_id, *payload_len),
            Operation::OnBytesTransmit {
                path_id_generator,
                bytes,
            } => self.on_bytes_transmitted(*path_id_generator, *bytes),
        }
    }

    fn on_timeout(&mut self, millis: &u16) {
        // timestamp should be monotonically increasing
        self.timestamp += Duration::from_millis(*millis as u64);

        self.oracle.on_timeout(millis);

        self.subject
            .on_timeout(
                self.timestamp,
                &mut Generator::default(),
                &mut Publisher::no_snapshot(),
            )
            .unwrap();
    }

    fn on_datagram_received(
        &mut self,
        handle: &Handle,
        probing: &Option<path_validation::Probe>,
        local_id: LocalId,
        payload_len: u16,
    ) {
        // Handle is an alias to RemoteAddress so does not inherit the PathHandle
        // eq implementation which unmaps an ipv6 address into a ipv4 address
        let handle = path::RemoteAddress(handle.unmap());
        let datagram = DatagramInfo {
            timestamp: self.timestamp,
            payload_len: payload_len as usize,
            ecn: ExplicitCongestionNotification::NotEct,
            destination_connection_id: local_id,
            source_connection_id: None,
        };
        let mut migration_validator = path::migration::default::Validator;
        let mut random_generator = Generator::default();
        let mut publisher = Publisher::no_snapshot();

        self.oracle.prev_state_amplification_limited = self
            .subject
            .path(&handle)
            .map_or(false, |(_id, path)| path.at_amplification_limit());

        match self.subject.on_datagram_received(
            &handle,
            &datagram,
            true,
            &mut Default::default(),
            &mut migration_validator,
            MaxMtu::default(),
            &mut publisher,
        ) {
            Ok((id, _)) => {
                // Only call oracle if the subject can process on_datagram_received without errors
                self.oracle.on_datagram_received(
                    &handle,
                    payload_len,
                    self.subject[id].at_amplification_limit(),
                );

                if let Some(probe) = probing {
                    self.oracle.on_processed_packet(&handle, probe);

                    let (path_id, _path) = self.subject.path(&handle).unwrap();
                    self.subject
                        .on_processed_packet(
                            path_id,
                            None,
                            *probe,
                            &mut random_generator,
                            &mut publisher,
                        )
                        .unwrap();
                }
            }
            Err(datagram_drop_reason) => {
                match datagram_drop_reason {
                    // Ignore errors emitted by the migration::validator and peer_id_registry
                    DatagramDropReason::InsufficientConnectionIds => {}
                    DatagramDropReason::RejectedConnectionMigration => {}
                    datagram_drop_reason => panic!("{:?}", datagram_drop_reason),
                };
            }
        }
    }

    fn on_bytes_transmitted(&mut self, path_id_generator: u8, bytes: u16) {
        let id = path_id_generator as usize % self.subject.paths.len();

        let path = &mut self.subject[path_id(id as u8)];
        if !path.at_amplification_limit() {
            path.on_bytes_transmitted(bytes as usize);
        }
    }

    /// Check that the subject and oracle match.
    pub fn invariants(&self) {
        // compare total paths
        assert_eq!(self.oracle.paths.len(), self.subject.paths.len());

        // compare active path
        assert_eq!(
            self.oracle.get_active_path().remote_address(),
            self.subject.active_path().remote_address()
        );

        // compare last known valid path
        assert_eq!(
            self.oracle.last_known_active_validated_path,
            self.subject.last_known_active_validated_path
        );

        // compare path properties
        for (path_id, s_path) in self.subject.paths.iter().enumerate() {
            let o_path = self.oracle.paths.get(&s_path.remote_address()).unwrap();

            assert_eq!(
                o_path.remote_address(),
                s_path.remote_address(),
                "path_id: {}",
                path_id
            );
            assert_eq!(
                o_path.is_validated(),
                s_path.is_validated(),
                "path_id: {}",
                path_id
            );
        }
    }
}

#[test]
fn cm_model_test() {
    check!()
        .with_type::<Vec<Operation>>()
        .for_each(|operations| {
            let mut model = Model::new();
            for operation in operations.iter() {
                model.apply(operation);
            }

            model.invariants();
        })
}


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

use super::*;
use crate::{
    connection::{ConnectionIdMapper, InternalConnectionIdGenerator},
    contexts::testing::{MockWriteContext, OutgoingFrameBuffer},
    endpoint::{
        self,
        testing::{Client, Server},
    },
    path,
    path::DEFAULT_MAX_MTU,
};
use core::time::Duration;
use s2n_quic_core::{
    event::testing::Publisher,
    inet::{DatagramInfo, ExplicitCongestionNotification, SocketAddress},
    path::{migration, RemoteAddress},
    random::{self, Generator},
    recovery::RttEstimator,
    stateless_reset::token::testing::*,
    time::{Clock, NoopClock},
};
use std::net::SocketAddr;

type ServerManager = super::Manager<Server>;
type ServerPath = super::Path<Server>;
type ClientManager = super::Manager<Client>;
type ClientPath = super::Path<Client>;

// Helper function to easily create a PathManager as a Server
fn manager_server(first_path: ServerPath) -> ServerManager {
    let mut random_generator = random::testing::Generator(123);
    let peer_id_registry = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server)
        .create_server_peer_id_registry(
            InternalConnectionIdGenerator::new().generate_id(),
            first_path.peer_connection_id,
        );
    ServerManager::new(first_path, peer_id_registry)
}

// Helper function to easily create a PathManager as a Client
fn manager_client(first_path: ClientPath) -> ClientManager {
    let mut random_generator = random::testing::Generator(123);
    let peer_id_registry = ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Client)
        .create_client_peer_id_registry(InternalConnectionIdGenerator::new().generate_id());
    ClientManager::new(first_path, peer_id_registry)
}

#[test]
fn get_path_by_address_test() {
    let first_conn_id = connection::PeerId::try_from_bytes(&[0, 1, 2, 3, 4, 5]).unwrap();
    let first_local_conn_id = connection::LocalId::TEST_ID;
    let first_path = ServerPath::new(
        Default::default(),
        first_conn_id,
        first_local_conn_id,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );

    let second_conn_id = connection::PeerId::try_from_bytes(&[5, 4, 3, 2, 1]).unwrap();
    let second_path = ServerPath::new(
        Default::default(),
        second_conn_id,
        first_local_conn_id,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );

    let mut manager = manager_server(first_path.clone());
    manager.paths.push(second_path);
    assert_eq!(manager.paths.len(), 2);

    let (_id, matched_path) = manager.path(&RemoteAddress::default()).unwrap();
    assert_eq!(
        matched_path.peer_connection_id,
        first_path.peer_connection_id
    );
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-9.3.2
//= type=test
//# To protect the connection from failing due to such a spurious
//# migration, an endpoint MUST revert to using the last validated peer
//# address when validation of a new peer address fails.
#[test]
fn test_invalid_path_fallback() {
    let mut publisher = Publisher::snapshot();
    let first_conn_id = connection::PeerId::try_from_bytes(&[0, 1, 2, 3, 4, 5]).unwrap();
    let first_local_conn_id = connection::LocalId::TEST_ID;
    let mut first_path = ServerPath::new(
        Default::default(),
        first_conn_id,
        first_local_conn_id,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    // simulate receiving a handshake packet to force path validation
    first_path.on_handshake_packet();

    // Create a challenge that will expire in 100ms
    let now = NoopClock {}.get_time();
    let expiration = Duration::from_millis(1000);
    let challenge = challenge::Challenge::new(expiration, [0; 8]);
    let mut second_path = ServerPath::new(
        Default::default(),
        first_conn_id,
        first_local_conn_id,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    second_path.set_challenge(challenge);

    let mut manager = manager_server(first_path);
    manager.paths.push(second_path);
    assert_eq!(manager.last_known_active_validated_path, None);
    assert_eq!(manager.active, 0);
    assert!(manager.paths[0].is_validated());

    manager
        .update_active_path(
            path_id(1),
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();
    assert_eq!(manager.active, 1);
    assert_eq!(manager.last_known_active_validated_path, Some(0));

    // send challenge and arm abandon timer
    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut context = MockWriteContext::new(
        now,
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Client,
    );
    manager[path_id(1)].on_transmit(&mut context);

    // After a validation times out, the path should revert to the previous
    manager
        .on_timeout(
            now + expiration + Duration::from_millis(100),
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();
    assert_eq!(manager.active, 0);
    assert!(manager.last_known_active_validated_path.is_none());
}

#[test]
// a validated path should be assigned to last_known_active_validated_path
fn promote_validated_path_to_last_known_validated_path() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert!(!helper.manager.paths[helper.first_path_id.as_u8() as usize].is_validated());

    // Trigger:
    // simulate receiving a handshake packet to force path validation
    helper.manager.paths[helper.first_path_id.as_u8() as usize].on_handshake_packet();
    assert!(helper.manager.paths[helper.first_path_id.as_u8() as usize].is_validated());
    helper
        .manager
        .update_active_path(
            helper.second_path_id,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert_eq!(helper.manager.last_known_active_validated_path, Some(1));
}

#[test]
// a NOT validated path should NOT be assigned to last_known_active_validated_path
fn dont_promote_non_validated_path_to_last_known_validated_path() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert!(!helper.manager.paths[helper.first_path_id.as_u8() as usize].is_validated());

    // Trigger:
    helper
        .manager
        .update_active_path(
            helper.second_path_id,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert_eq!(helper.manager.last_known_active_validated_path, Some(0));
}

#[test]
// update path to the new active path
fn update_path_to_active_path() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());

    // Trigger:
    helper
        .manager
        .update_active_path(
            helper.second_path_id,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert_eq!(helper.manager.active, helper.second_path_id.as_u8());
}

#[test]
// Don't update path to the new active path if insufficient connection ids
fn dont_update_path_to_active_path_if_no_connection_id_available() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths_base(false, true, &mut publisher);
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());

    // Trigger:
    assert_eq!(
        helper.manager.update_active_path(
            helper.second_path_id,
            &mut random::testing::Generator(123),
            &mut publisher,
        ),
        Err(transport::Error::INTERNAL_ERROR)
    );

    // Expectation:
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());
}

#[test]
fn set_path_challenge_on_active_path_on_connection_migration() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    helper.manager[helper.zero_path_id].abandon_challenge(&mut publisher, 0);
    assert!(!helper.manager[helper.zero_path_id].is_challenge_pending());
    assert_eq!(
        helper.manager.last_known_active_validated_path.unwrap(),
        helper.zero_path_id.as_u8()
    );

    // Trigger:
    helper
        .manager
        .update_active_path(
            helper.second_path_id,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert!(helper.manager[helper.first_path_id].is_challenge_pending());
}

#[test]
// A path should be validated if a PATH_RESPONSE contains data sent
// via PATH_CHALLENGE
//
// Setup:
// - create manager with 2 paths
// - first path is active
// - second path is pending challenge/validation
//
// Trigger 1:
// - call on_timeout just BEFORE challenge should expire
//
// Expectation 1:
// - verify second path is pending challenge
// - verify second path is NOT validated
//
// Trigger 2:
// - call on_path_response with expected data for second path
//
// Expectation 2:
// - verify second path is validated
fn validate_path_before_challenge_expiration() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());

    // send challenge and arm abandon timer
    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut context = MockWriteContext::new(
        helper.now,
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Client,
    );
    helper.manager[helper.second_path_id].on_transmit(&mut context);
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());

    // Trigger 1:
    // A response 100ms before the challenge is abandoned
    helper
        .manager
        .on_timeout(
            helper.now + helper.challenge_expiration - Duration::from_millis(100),
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation 1:
    assert!(helper.manager[helper.second_path_id].is_challenge_pending(),);
    assert!(!helper.manager[helper.second_path_id].is_validated());

    // Trigger 2:
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.2
    //= type=test
    //# This requirement MUST NOT be enforced by the endpoint that initiates
    //# path validation, as that would enable an attack on migration; see
    //# Section 9.3.3.

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.3
    //= type=test
    //# A PATH_RESPONSE frame received on any network path validates the path
    //# on which the PATH_CHALLENGE was sent.
    //
    // The above requirements are satisfied because on_path_response is a path
    // agnostic function
    let frame = s2n_quic_core::frame::PathResponse {
        data: &helper.second_expected_data,
    };
    helper.manager.on_path_response(&frame, &mut publisher);

    // Expectation 2:
    assert!(helper.manager[helper.second_path_id].is_validated());
}

#[test]
// A path should NOT be validated if the challenge has been abandoned
//
// Setup:
// - create manager with 2 paths
// - first path is active
// - second path is pending challenge/validation
//
// Trigger 1:
// - call on_timeout just AFTER challenge should expire
//
// Expectation 1:
// - verify second path is NOT pending challenge
// - verify second path is NOT validated
//
//
// Trigger 2:
// - call on_path_response with expected data for second path
//
// Expectation 2:
// - verify second path is NOT validated
fn dont_validate_path_if_path_challenge_is_abandoned() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());

    // send challenge and arm abandon timer
    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut context = MockWriteContext::new(
        helper.now,
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Client,
    );
    helper.manager[helper.second_path_id].on_transmit(&mut context);
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());

    // Trigger 1:
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.4
    //= type=test
    //# Endpoints SHOULD abandon path validation based on a timer.
    // A response 100ms after the challenge should fail
    helper
        .manager
        .on_timeout(
            helper.now + helper.challenge_expiration + Duration::from_millis(100),
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation 1:
    assert!(!helper.manager[helper.second_path_id].is_challenge_pending());
    assert!(!helper.manager[helper.second_path_id].is_validated());

    // Trigger 2:
    let frame = s2n_quic_core::frame::PathResponse {
        data: &helper.second_expected_data,
    };
    helper.manager.on_path_response(&frame, &mut publisher);

    // Expectation 2:
    assert!(!helper.manager[helper.second_path_id].is_validated());
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9000#section-9.3
//# If the recipient permits the migration, it MUST send subsequent
//# packets to the new peer address and MUST initiate path validation
//# (Section 8.2) to verify the peer's ownership of the address if
//# validation is not already underway.
fn initiate_path_challenge_if_new_path_is_not_validated() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert!(!helper.manager[helper.first_path_id].is_validated());
    assert!(helper.manager[helper.first_path_id].is_challenge_pending());

    assert!(!helper.manager[helper.second_path_id].is_validated());
    helper.manager[helper.second_path_id].abandon_challenge(&mut publisher, 0);
    assert!(!helper.manager[helper.second_path_id].is_challenge_pending());
    assert_eq!(helper.manager.active_path_id(), helper.first_path_id);

    // Trigger:
    helper
        .manager
        .on_processed_packet(
            helper.second_path_id,
            None,
            path_validation::Probe::NonProbing,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert!(!helper.manager[helper.second_path_id].is_validated());
    assert_eq!(helper.manager.active_path_id(), helper.second_path_id);
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9000#section-9
//= type=test
//# When an endpoint has no validated path on which to send packets, it
//# MAY discard connection state.

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

//= https://www.rfc-editor.org/rfc/rfc9000#section-10
//= type=test
//# An endpoint MAY discard connection state if it does not have a
//# validated path on which it can send packets; see Section 8.2
//
// If there is no last_known_active_validated_path after a on_timeout then return a
// NoValidPath error
fn silently_return_when_there_is_no_valid_path() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let now = NoopClock {}.get_time();
    let expiration = Duration::from_millis(1000);
    let challenge = challenge::Challenge::new(expiration, [0; 8]);
    let mut first_path = ServerPath::new(
        Default::default(),
        connection::PeerId::try_from_bytes(&[1]).unwrap(),
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    first_path.set_challenge(challenge);
    let mut manager = manager_server(first_path);
    let first_path_id = path_id(0);

    assert!(!manager[first_path_id].is_validated());
    assert!(manager[first_path_id].is_challenge_pending());
    assert_eq!(manager.last_known_active_validated_path, None);

    // Trigger:
    // send challenge and arm abandon timer
    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut context = MockWriteContext::new(
        now,
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Client,
    );
    manager[first_path_id].on_transmit(&mut context);
    let res = manager.on_timeout(
        now + expiration + Duration::from_millis(100),
        &mut random::testing::Generator(123),
        &mut publisher,
    );

    // Expectation:
    assert!(!manager[first_path_id].is_challenge_pending());
    assert!(matches!(
        res.unwrap_err(),
        connection::Error::NoValidPath { .. }
    ));
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9000#section-9.3
//= type=test
//# After changing the address to which it sends non-probing packets, an
//# endpoint can abandon any path validation for other addresses.
//
// A non-probing (path validation probing) packet will cause the path to become an active
// path but the path is still not validated.
fn dont_abandon_path_challenge_if_new_path_is_not_validated() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert!(!helper.manager[helper.first_path_id].is_validated());
    assert!(helper.manager[helper.first_path_id].is_challenge_pending());

    assert!(!helper.manager[helper.second_path_id].is_validated());
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());
    assert_eq!(helper.manager.active_path_id(), helper.first_path_id);

    // Trigger:
    helper
        .manager
        .on_processed_packet(
            helper.second_path_id,
            None,
            path_validation::Probe::NonProbing,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert!(!helper.manager[helper.second_path_id].is_validated());
    assert_eq!(helper.manager.active_path_id(), helper.second_path_id);
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());
}

#[test]
fn abandon_path_challenges_if_new_path_is_validated() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert!(helper.manager[helper.first_path_id].is_challenge_pending());
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());
    assert_eq!(helper.manager.active_path_id(), helper.first_path_id);

    // simulate receiving a handshake packet to force path validation
    helper.manager[helper.second_path_id].on_handshake_packet();
    assert!(helper.manager[helper.second_path_id].is_validated());

    // Trigger:
    helper
        .manager
        .on_processed_packet(
            helper.second_path_id,
            None,
            path_validation::Probe::NonProbing,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert_eq!(helper.manager.active_path_id(), helper.second_path_id);
    assert!(!helper.manager[helper.first_path_id].is_challenge_pending());
    assert!(!helper.manager[helper.second_path_id].is_challenge_pending());
}

#[test]
fn abandon_all_path_challenges() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert!(helper.manager[helper.zero_path_id].is_challenge_pending());
    assert!(helper.manager[helper.first_path_id].is_challenge_pending());
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());

    // Trigger:
    helper.manager.abandon_all_path_challenges(&mut publisher);

    // Expectation:
    assert!(!helper.manager[helper.zero_path_id].is_challenge_pending());
    assert!(!helper.manager[helper.first_path_id].is_challenge_pending());
    assert!(!helper.manager[helper.second_path_id].is_challenge_pending());
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9000#section-9.2
//= type=test
//# An endpoint can migrate a connection to a new local address by
//# sending packets containing non-probing frames from that address.
//
// receiving a path_validation::Probing::NonProbing should update path to active path
fn non_probing_should_update_path_to_active_path() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());

    // Trigger:
    helper
        .manager
        .on_processed_packet(
            helper.second_path_id,
            None,
            path_validation::Probe::NonProbing,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert_eq!(helper.manager.active, helper.second_path_id.as_u8());
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9000#section-9.2
//= type=test
//# An endpoint can migrate a connection to a new local address by
//# sending packets containing non-probing frames from that address.
//
// receiving a path_validation::Probing::Probing should NOT update path to active path
fn probing_should_not_update_path_to_active_path() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());

    // Trigger:
    helper
        .manager
        .on_processed_packet(
            helper.second_path_id,
            None,
            path_validation::Probe::Probing,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert_eq!(helper.manager.active, helper.first_path_id.as_u8());
}

#[test]
// add new path when receiving a datagram on different remote address
// Setup:
// - create path manger with one path
//
// Trigger:
// - call on_datagram_received with new remote address
//
// Expectation:
// - assert we have two paths
fn test_adding_new_path() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let first_conn_id = connection::PeerId::try_from_bytes(&[1]).unwrap();
    let new_addr: SocketAddr = "127.0.0.1:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let first_path = ServerPath::new(
        new_addr,
        first_conn_id,
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);

    // verify we have one path
    assert!(manager.path(&new_addr).is_some());
    let new_addr: SocketAddr = "127.0.0.2:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    assert!(manager.path(&new_addr).is_none());
    assert_eq!(manager.paths.len(), 1);

    // Trigger:
    let datagram = DatagramInfo {
        timestamp: NoopClock {}.get_time(),
        payload_len: 0,
        ecn: ExplicitCongestionNotification::default(),
        destination_connection_id: connection::LocalId::TEST_ID,
        source_connection_id: None,
    };
    let (path_id, unblocked) = manager
        .on_datagram_received(
            &new_addr,
            &datagram,
            true,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        )
        .unwrap();

    // Expectation:
    assert_eq!(path_id.as_u8(), 1);
    assert!(!unblocked);
    assert!(manager.path(&new_addr).is_some());
    assert_eq!(manager.paths.len(), 2);
}

#[test]
// do NOT add new path if handshake is not confirmed
// Setup:
// - create path manger with one path
//
// Trigger:
// - call on_datagram_received with new remote address bit handshake_confirmed false
//
// Expectation:
// - asset on_datagram_received errors
// - assert we have one paths
fn do_not_add_new_path_if_handshake_not_confirmed() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let first_conn_id = connection::PeerId::try_from_bytes(&[1]).unwrap();
    let first_path = ServerPath::new(
        Default::default(),
        first_conn_id,
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);

    // verify we have one path
    let new_addr: SocketAddr = "127.0.0.1:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    assert_eq!(manager.paths.len(), 1);

    // Trigger:
    let datagram = DatagramInfo {
        timestamp: NoopClock {}.get_time(),
        payload_len: 0,
        ecn: ExplicitCongestionNotification::default(),
        destination_connection_id: connection::LocalId::TEST_ID,
        source_connection_id: None,
    };
    let handshake_confirmed = false;
    let on_datagram_result = manager.on_datagram_received(
        &new_addr,
        &datagram,
        handshake_confirmed,
        &mut Default::default(),
        &mut migration::default::Validator::default(),
        DEFAULT_MAX_MTU,
        &mut publisher,
    );

    // Expectation:
    assert!(on_datagram_result.is_err());
    assert!(manager.path(&new_addr).is_none());
    assert_eq!(manager.paths.len(), 1);
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9000#section-9
//= type=test
//# If a client receives packets from an unknown server address,
//# the client MUST discard these packets.
//
// Setup:
// - create path manager with one path as a client
//
// Trigger:
// - call on_datagram_received with new remote address bit
//
// Expectation:
// - asset on_datagram_received errors
// - assert we have one path
fn do_not_add_new_path_if_client() {
    // Setup:
    let first_conn_id = connection::PeerId::try_from_bytes(&[1]).unwrap();
    let first_path = ClientPath::new(
        Default::default(),
        first_conn_id,
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_client(first_path);
    let mut publisher = Publisher::snapshot();

    // verify we have one path
    let new_addr: SocketAddr = "127.0.0.1:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    assert_eq!(manager.paths.len(), 1);

    // Trigger:
    let datagram = DatagramInfo {
        timestamp: NoopClock {}.get_time(),
        payload_len: 0,
        ecn: ExplicitCongestionNotification::default(),
        destination_connection_id: connection::LocalId::TEST_ID,
        source_connection_id: None,
    };
    let on_datagram_result = manager.on_datagram_received(
        &new_addr,
        &datagram,
        true,
        &mut Default::default(),
        &mut migration::default::Validator::default(),
        DEFAULT_MAX_MTU,
        &mut publisher,
    );

    // Expectation:
    assert!(on_datagram_result.is_err());
    assert!(manager.path(&new_addr).is_none());
    assert_eq!(manager.paths.len(), 1);
}

#[test]
//= https://www.rfc-editor.org/rfc/rfc9000#section-7.2
//= type=test
//# Until a packet is received from the server, the client MUST
//# use the same Destination Connection ID value on all packets in this
//# connection.
fn switch_destination_connection_id_after_first_server_response() {
    // Setup:
    let initial_cid = connection::PeerId::try_from_bytes(&[0, 0]).unwrap();
    let zero_path_id = path_id(0);
    let path_handle = Default::default();
    let zero_path = ClientPath::new(
        path_handle,
        initial_cid,
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_client(zero_path);
    assert_eq!(manager[zero_path_id].peer_connection_id, initial_cid);

    // Trigger:
    let server_destination_cid = connection::PeerId::try_from_bytes(&[1, 1]).unwrap();
    let mut publisher = Publisher::snapshot();
    let processed_packet = manager.on_processed_packet(
        zero_path_id,
        Some(server_destination_cid),
        path_validation::Probe::NonProbing,
        &mut random::testing::Generator(123),
        &mut publisher,
    );

    // Expectation:
    assert!(processed_packet.is_ok());
    assert_eq!(
        manager[zero_path_id].peer_connection_id,
        server_destination_cid
    );
}

#[test]
fn limit_number_of_connection_migrations() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let new_addr: SocketAddr = "127.0.0.1:1".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let first_path = ServerPath::new(
        new_addr,
        connection::PeerId::try_from_bytes(&[1]).unwrap(),
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);
    let mut total_paths = 1;

    for i in 1..std::u8::MAX {
        let new_addr: SocketAddr = format!("127.0.0.2:{}", i).parse().unwrap();
        let new_addr = SocketAddress::from(new_addr);
        let new_addr = RemoteAddress::from(new_addr);
        let now = NoopClock {}.get_time();
        let datagram = DatagramInfo {
            timestamp: now,
            payload_len: 0,
            ecn: ExplicitCongestionNotification::default(),
            destination_connection_id: connection::LocalId::TEST_ID,
            source_connection_id: None,
        };

        let res = manager.handle_connection_migration(
            &new_addr,
            &datagram,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        );
        match res {
            Ok((id, _)) => {
                let _ = manager.on_processed_packet(
                    id,
                    None,
                    path_validation::Probe::NonProbing,
                    &mut random::testing::Generator(123),
                    &mut publisher,
                );
                total_paths += 1
            }
            Err(_) => break,
        }
    }
    assert_eq!(total_paths, MAX_ALLOWED_PATHS);
}

#[test]
fn connection_migration_challenge_behavior() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let first_conn_id = connection::PeerId::try_from_bytes(&[1]).unwrap();
    let new_addr: SocketAddr = "127.0.0.1:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let first_path = ServerPath::new(
        new_addr,
        first_conn_id,
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);

    let new_addr: SocketAddr = "127.0.0.2:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let now = NoopClock {}.get_time();
    let datagram = DatagramInfo {
        timestamp: now,
        payload_len: 0,
        ecn: ExplicitCongestionNotification::default(),
        destination_connection_id: connection::LocalId::TEST_ID,
        source_connection_id: None,
    };

    let (path_id, _unblocked) = manager
        .handle_connection_migration(
            &new_addr,
            &datagram,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        )
        .unwrap();

    // verify we have two paths
    assert!(manager.path(&new_addr).is_some());
    assert_eq!(manager.paths.len(), 2);

    assert!(!manager[path_id].is_challenge_pending());

    // notify the manager that the datagram was authenticated - the path should now issue a challenge
    let _ = manager.on_processed_packet(
        path_id,
        None,
        path_validation::Probe::NonProbing,
        &mut random::testing::Generator(123),
        &mut publisher,
    );

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

    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.1
    //= type=test
    //# The endpoint MUST use unpredictable data in every PATH_CHALLENGE
    //# frame so that it can associate the peer's response with the
    //# corresponding PATH_CHALLENGE.
    // Verify that the data stored in the challenge is taken from the random generator
    // TODO does the below actually work?? investigate
    let mut test_rnd_generator = random::testing::Generator(123);
    let mut expected_data: [u8; 8] = [0; 8];
    test_rnd_generator.public_random_fill(&mut expected_data);

    //= https://www.rfc-editor.org/rfc/rfc9000#section-9
    //= type=test
    //# An endpoint MUST
    //# perform path validation (Section 8.2) if it detects any change to a
    //# peer's address, unless it has previously validated that address.
    assert!(manager[path_id].on_path_response(&expected_data));
    assert!(manager[path_id].is_validated());
}

#[test]
// Abandon timer should use max PTO of active and new path(new path uses kInitialRtt)
// Setup 1:
// - create manager with path
// - create datagram for packet on second path
//
// Trigger 1:
// - call handle_connection_migration with packet for second path
//
// Expectation 1:
// - assert that new path uses max_ack_delay from the active path
fn connection_migration_use_max_ack_delay_from_active_path() {
    // Setup 1:
    let mut publisher = Publisher::snapshot();
    let new_addr: SocketAddr = "127.0.0.1:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let first_path = ServerPath::new(
        new_addr,
        connection::PeerId::try_from_bytes(&[1]).unwrap(),
        connection::LocalId::TEST_ID,
        RttEstimator::new(Duration::from_millis(30)),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);

    let new_addr: SocketAddr = "127.0.0.2:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let now = NoopClock {}.get_time();
    let datagram = DatagramInfo {
        timestamp: now,
        payload_len: 0,
        ecn: ExplicitCongestionNotification::default(),
        destination_connection_id: connection::LocalId::TEST_ID,
        source_connection_id: None,
    };

    // Trigger 1:
    let (second_path_id, _unblocked) = manager
        .handle_connection_migration(
            &new_addr,
            &datagram,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        )
        .unwrap();
    let first_path_id = path_id(0);

    // Expectation 1:
    // inherit max_ack_delay from the active path
    assert_eq!(manager.active, first_path_id.as_u8());
    assert_eq!(
        &manager[first_path_id].rtt_estimator.max_ack_delay(),
        &manager[second_path_id].rtt_estimator.max_ack_delay()
    );
}

#[test]
// Abandon timer should use max PTO of active and new path(new path uses kInitialRtt)
// Setup 1:
// - create manager with path
// - create datagram for packet on second path
// - call handle_connection_migration with packet for second path
//
// Trigger 1:
// - modify rtt for fist path to detect difference in PTO
//
// Expectation 1:
// - verify PTO of second path > PTO of first path
//
// Setup 2:
// - call on_transmit for second path to send challenge and arm abandon timer
//
// Trigger 2:
// - call second_path.on_timeout with abandon_time - 10ms
//
// Expectation 2:
// - verify challenge is NOT abandoned
//
// Trigger 3:
// - call second_path.on_timeout with abandon_time + 10ms
//
// Expectation 3:
// - verify challenge is abandoned
fn connection_migration_new_path_abandon_timer() {
    // Setup 1:
    let mut publisher = Publisher::snapshot();
    let new_addr: SocketAddr = "127.0.0.1:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let first_path = ServerPath::new(
        new_addr,
        connection::PeerId::try_from_bytes(&[1]).unwrap(),
        connection::LocalId::TEST_ID,
        RttEstimator::new(Duration::from_millis(30)),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);

    let new_addr: SocketAddr = "127.0.0.2:8001".parse().unwrap();
    let new_addr = SocketAddress::from(new_addr);
    let new_addr = RemoteAddress::from(new_addr);
    let now = NoopClock {}.get_time();
    let datagram = DatagramInfo {
        timestamp: now,
        payload_len: 0,
        ecn: ExplicitCongestionNotification::default(),
        destination_connection_id: connection::LocalId::TEST_ID,
        source_connection_id: None,
    };

    let (second_path_id, _unblocked) = manager
        .handle_connection_migration(
            &new_addr,
            &datagram,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        )
        .unwrap();
    let first_path_id = path_id(0);

    // notify the manager that the datagram was authenticated - the path should now issue a challenge
    let _ = manager.on_processed_packet(
        second_path_id,
        None,
        path_validation::Probe::NonProbing,
        &mut random::testing::Generator(123),
        &mut publisher,
    );

    // Trigger 1:
    // modify rtt for first path so we can detect differences
    manager[first_path_id].rtt_estimator.update_rtt(
        Duration::from_millis(0),
        Duration::from_millis(100),
        now,
        true,
        PacketNumberSpace::ApplicationData,
    );

    // Expectation 1:
    // verify the pto_period of the first path is less than the second path
    let first_path_pto = manager[first_path_id].pto_period(PacketNumberSpace::ApplicationData);
    let second_path_pto = manager[second_path_id].pto_period(PacketNumberSpace::ApplicationData);

    assert_eq!(first_path_pto, Duration::from_millis(330));
    assert_eq!(second_path_pto, Duration::from_millis(1_029));
    assert!(second_path_pto > first_path_pto);

    // Setup 2:
    // send challenge and arm abandon timer
    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut context = MockWriteContext::new(
        now,
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Client,
    );
    manager[second_path_id].on_transmit(&mut context);

    // Trigger 2:
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.4
    //= type=test
    //# Endpoints SHOULD abandon path validation based on a timer.
    //
    //= https://www.rfc-editor.org/rfc/rfc9000#section-8.2.4
    //= type=test
    //# Endpoints SHOULD abandon path validation based on a timer.  When
    //# setting this timer, implementations are cautioned that the new path
    //# could have a longer round-trip time than the original.  A value of
    //# three times the larger of the current PTO or the PTO for the new path
    //# (using kInitialRtt, as defined in [QUIC-RECOVERY]) is RECOMMENDED.
    // abandon_duration should use max pto_period: second path
    let abandon_time = now + (second_path_pto * 3);
    manager[second_path_id].on_timeout(
        abandon_time - Duration::from_millis(10),
        path::Id::test_id(),
        &mut random::testing::Generator(123),
        &mut publisher,
    );

    // Expectation 2:
    assert!(manager[second_path_id].is_challenge_pending());

    // Trigger 3:
    manager[second_path_id].on_timeout(
        abandon_time + Duration::from_millis(10),
        path::Id::test_id(),
        &mut random::testing::Generator(123),
        &mut publisher,
    );
    // Expectation 3:
    assert!(!manager[second_path_id].is_challenge_pending());
}

//= https://www.rfc-editor.org/rfc/rfc9000#section-5.1.2
//= type=test
//# Upon receipt of an increased Retire Prior To field, the peer MUST
//# stop using the corresponding connection IDs and retire them with
//# RETIRE_CONNECTION_ID frames before adding the newly provided
//# connection ID to the set of active connection IDs.
#[test]
fn stop_using_a_retired_connection_id() {
    let mut publisher = Publisher::snapshot();
    let id_1 = connection::PeerId::try_from_bytes(b"id01").unwrap();
    let first_path = ServerPath::new(
        Default::default(),
        id_1,
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);

    let id_2 = connection::PeerId::try_from_bytes(b"id02").unwrap();
    assert!(manager
        .on_new_connection_id(&id_2, 1, 1, &TEST_TOKEN_1, &mut publisher)
        .is_ok());

    assert_eq!(id_2, manager.paths[0].peer_connection_id);
}

#[test]
fn amplification_limited_true_if_all_paths_amplificaiton_limited() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let helper = helper_manager_with_paths_base(true, false, &mut publisher);
    let zp = &helper.manager[helper.zero_path_id];
    assert!(zp.at_amplification_limit());
    let fp = &helper.manager[helper.first_path_id];
    assert!(fp.at_amplification_limit());
    let sp = &helper.manager[helper.second_path_id];
    assert!(sp.at_amplification_limit());

    // Expectation:
    assert!(helper.manager.is_amplification_limited());
}

#[test]
fn amplification_limited_false_if_any_paths_amplificaiton_limited() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    let fp = &helper.manager[helper.first_path_id];
    assert!(fp.at_amplification_limit());
    let sp = &mut helper.manager[helper.second_path_id];
    sp.on_bytes_received(1200);
    assert!(!sp.at_amplification_limit());

    // Expectation:
    assert!(!helper.manager.is_amplification_limited());
}

#[test]
fn can_transmit_false_if_no_path_can_transmit() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let helper = helper_manager_with_paths_base(true, false, &mut publisher);

    let zp = &helper.manager[helper.zero_path_id];
    assert!(!transmission::Interest::None.can_transmit(zp.transmission_constraint()));

    let interest = transmission::Interest::Forced;
    let fp = &helper.manager[helper.first_path_id];
    assert!(!interest.can_transmit(fp.transmission_constraint()));
    let sp = &helper.manager[helper.second_path_id];
    assert!(!interest.can_transmit(sp.transmission_constraint()));

    // Expectation:
    assert!(!helper.manager.can_transmit(interest));
}

#[test]
fn can_transmit_true_if_any_path_can_transmit() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    let interest = transmission::Interest::Forced;
    let fp = &helper.manager[helper.first_path_id];
    assert!(!interest.can_transmit(fp.transmission_constraint()));

    let sp = &mut helper.manager[helper.second_path_id];
    sp.on_bytes_received(1200);
    assert!(interest.can_transmit(sp.transmission_constraint()));

    // Expectation:
    assert!(helper.manager.can_transmit(interest));
}

#[test]
// Return all paths that are pending challenge or response.
fn pending_paths_should_return_paths_pending_validation() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let mut helper = helper_manager_with_paths(&mut publisher);
    let third_path_id = path_id(3);
    let third_conn_id = connection::PeerId::try_from_bytes(&[3]).unwrap();
    let mut third_path = ServerPath::new(
        Default::default(),
        third_conn_id,
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let expected_response_data = [0; 8];
    third_path.on_path_challenge(&expected_response_data);
    helper.manager.paths.push(third_path);

    // not pending challenge or response
    helper.manager[helper.zero_path_id].abandon_challenge(&mut publisher, 0);
    assert!(!helper.manager[helper.zero_path_id].is_challenge_pending());
    assert!(!helper.manager[helper.zero_path_id].is_response_pending());

    // pending challenge
    assert!(helper.manager[helper.first_path_id].is_challenge_pending());
    assert!(!helper.manager[helper.first_path_id].is_response_pending());
    assert!(helper.manager[helper.second_path_id].is_challenge_pending());
    assert!(!helper.manager[helper.second_path_id].is_response_pending());

    // pending response
    assert!(!helper.manager[third_path_id].is_challenge_pending());
    assert!(helper.manager[third_path_id].is_response_pending());

    let mut pending_paths = helper.manager.paths_pending_validation();

    // inclusive range from 1 to 3
    for i in 1..=3 {
        // Trigger:
        let next = pending_paths.next_path();

        // Expectation:
        let (id, _path_manager) = next.unwrap();
        assert_eq!(id, path_id(i));
    }

    // Trigger:
    let next = pending_paths.next_path();

    // Expectation:
    assert!(next.is_none());
}

#[test]
// Ensure paths are temporary until after authenticating a packet on the path
fn temporary_until_authenticated() {
    let mut publisher = Publisher::snapshot();
    let now = NoopClock {}.get_time();
    let datagram = DatagramInfo {
        timestamp: now,
        payload_len: 0,
        ecn: ExplicitCongestionNotification::default(),
        destination_connection_id: connection::LocalId::TEST_ID,
        source_connection_id: None,
    };

    // create an initial path
    let first_addr: SocketAddr = "127.0.0.1:8001".parse().unwrap();
    let first_addr = SocketAddress::from(first_addr);
    let first_addr = RemoteAddress::from(first_addr);
    let first_path = ServerPath::new(
        first_addr,
        connection::PeerId::try_from_bytes(&[1]).unwrap(),
        connection::LocalId::TEST_ID,
        RttEstimator::default(),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    );
    let mut manager = manager_server(first_path);

    let second_addr: SocketAddr = "127.0.0.2:8001".parse().unwrap();
    let second_addr = SocketAddress::from(second_addr);
    let second_addr = RemoteAddress::from(second_addr);

    // create a second path without calling on_processed_packet
    let (second_path_id, _unblocked) = manager
        .on_datagram_received(
            &second_addr,
            &datagram,
            true,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        )
        .unwrap();

    assert!(
        !manager[second_path_id].is_challenge_pending(),
        "pending paths should not issue a challenge"
    );

    let third_addr: SocketAddr = "127.0.0.3:8001".parse().unwrap();
    let third_addr = SocketAddress::from(third_addr);
    let third_addr = RemoteAddress::from(third_addr);

    // create a third path
    let (third_path_id, _unblocked) = manager
        .on_datagram_received(
            &third_addr,
            &datagram,
            true,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        )
        .unwrap();

    assert_eq!(
        second_path_id, third_path_id,
        "third path should replace the second"
    );

    assert!(
        !manager[third_path_id].is_challenge_pending(),
        "pending paths should not issue a challenge"
    );

    // notify the manager that the packet was processed
    let _ = manager.on_processed_packet(
        third_path_id,
        None,
        path_validation::Probe::NonProbing,
        &mut random::testing::Generator(123),
        &mut publisher,
    );

    assert!(
        manager[third_path_id].is_challenge_pending(),
        "after processing a packet the path should issue a challenge"
    );

    // receive another datagram with the second_addr
    let (fourth_path_id, _unblocked) = manager
        .on_datagram_received(
            &second_addr,
            &datagram,
            true,
            &mut Default::default(),
            &mut migration::default::Validator::default(),
            DEFAULT_MAX_MTU,
            &mut publisher,
        )
        .unwrap();

    assert_ne!(
        fourth_path_id, third_path_id,
        "a new path should be created"
    );
}

// The last_known_active_validated_path needs to be both validated and also
// activated (the active path at some point in the connection).
//
// This test specifically checks that a currently non-active path can become
// the last_known_active_validated_path if it is activated and receives a
// PATH_RESPONSE which matches its PATH_CHALLENGE.
//
// Setup:
// - path 0 validated and active

// Trigger Setup 1:
// - path 1 non-probing packet
// Expectation Setup 1:
// - path 1 active + not valid + challenge pending
// - last_known_active_validated_path = path 0

// Trigger Setup 2:
// - path 2 non-probing packet
// Expectation Setup 2:
// - path 2 active + not valid + challenge pending
// - path 1 not valid + challenge pending
// - last_known_active_validated_path = path 0

// Trigger 1:
// - path response for path 1
// Expectation 1:
// - path 1 valid + no challenge pending
// - last_known_active_validated_path = path 1
//
// - path 2 active + not valid + challenge pending

// Trigger 2:
// - timeout for path 2 challenge
// Expectation 2:
// - path 1 active
// - path 2 not valid + no challenge pending
// - last_known_active_validated_path = None
#[test]
fn last_known_validated_path_should_update_on_path_response() {
    // Setup:
    let mut publisher = Publisher::snapshot();
    let zero_conn_id = connection::PeerId::try_from_bytes(&[0]).unwrap();
    let first_conn_id = connection::PeerId::try_from_bytes(&[1]).unwrap();
    let second_conn_id = connection::PeerId::try_from_bytes(&[2]).unwrap();

    // path zero
    let zero_path_id = path_id(0);
    let mut zero_path = helper_path(zero_conn_id);
    zero_path.on_handshake_packet();

    let mut random_generator = random::testing::Generator(123);
    let mut peer_id_registry =
        ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server)
            .create_server_peer_id_registry(
                InternalConnectionIdGenerator::new().generate_id(),
                zero_path.peer_connection_id,
            );
    assert!(peer_id_registry
        .on_new_connection_id(&first_conn_id, 1, 0, &TEST_TOKEN_1)
        .is_ok());

    assert!(peer_id_registry
        .on_new_connection_id(&second_conn_id, 2, 0, &TEST_TOKEN_2)
        .is_ok());

    let mut manager = Manager::new(zero_path, peer_id_registry);

    assert!(!manager[zero_path_id].is_challenge_pending());
    assert!(manager[zero_path_id].is_validated());
    assert_eq!(manager.active_path_id(), zero_path_id);

    // Trigger Setup 1:
    let first_path_id = path_id(1);
    let mut first_path = helper_path(first_conn_id);
    let now = NoopClock {}.get_time();
    let challenge_expiration = Duration::from_millis(10_000);
    let first_expected_data = [0; 8];
    let challenge = challenge::Challenge::new(challenge_expiration, first_expected_data);
    first_path.set_challenge(challenge);
    manager.paths.push(first_path);
    assert!(manager
        .update_active_path(
            first_path_id,
            &mut random::testing::Generator(123),
            &mut publisher
        )
        .is_ok());

    // Expectation Setup 1:
    assert_eq!(manager.active_path_id(), first_path_id);
    // first
    assert!(manager[first_path_id].is_challenge_pending());
    assert!(!manager[first_path_id].is_validated());
    // last_known_active_validated_path
    assert_eq!(
        manager.last_known_active_validated_path.unwrap(),
        zero_path_id.as_u8()
    );

    // Trigger Setup 2:
    let second_path_id = path_id(2);
    let second_expected_data = [1; 8];
    let challenge = challenge::Challenge::new(challenge_expiration, second_expected_data);
    let mut second_path = helper_path(second_conn_id);
    second_path.set_challenge(challenge);
    let mut frame_buffer = OutgoingFrameBuffer::new();
    let mut context = MockWriteContext::new(
        now,
        &mut frame_buffer,
        transmission::Constraint::None,
        transmission::Mode::Normal,
        endpoint::Type::Client,
    );
    second_path.on_transmit(&mut context); // send challenge and arm timer

    manager.paths.push(second_path);
    assert!(manager
        .update_active_path(
            second_path_id,
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .is_ok());

    // Expectation Setup 2:
    assert_eq!(manager.active_path_id(), second_path_id);
    // second
    assert!(manager[second_path_id].is_challenge_pending());
    assert!(!manager[second_path_id].is_validated());
    // first
    assert!(manager[first_path_id].is_challenge_pending());
    assert!(!manager[first_path_id].is_validated());
    // last_known_active_validated_path
    assert_eq!(
        manager.last_known_active_validated_path.unwrap(),
        zero_path_id.as_u8()
    );

    // Trigger 1:
    // - path response for path 1
    let frame = s2n_quic_core::frame::PathResponse {
        data: &first_expected_data,
    };
    manager.on_path_response(&frame, &mut publisher);
    // Expectation 1:
    assert_eq!(manager.active_path_id(), second_path_id);
    // second
    assert!(manager[second_path_id].is_challenge_pending());
    assert!(!manager[second_path_id].is_validated());
    // first
    assert!(!manager[first_path_id].is_challenge_pending());
    assert!(manager[first_path_id].is_validated());
    // last_known_active_validated_path
    assert_eq!(
        manager.last_known_active_validated_path.unwrap(),
        first_path_id.as_u8()
    );

    // Trigger 2:
    // - timeout for path 2 challenge
    manager
        .on_timeout(
            now + challenge_expiration + Duration::from_millis(100),
            &mut random::testing::Generator(123),
            &mut publisher,
        )
        .unwrap();

    // Expectation 2:
    assert_eq!(manager.active_path_id(), first_path_id);
    // second
    assert!(!manager[second_path_id].is_challenge_pending());
    assert!(!manager[second_path_id].is_validated());
    // first
    assert!(!manager[first_path_id].is_challenge_pending());
    assert!(manager[first_path_id].is_validated());
    // last_known_active_validated_path
    assert_eq!(manager.last_known_active_validated_path, None);
}

// creates a test path_manager. also check out `helper_manager_with_paths`
// which calls this helper with preset options
pub fn helper_manager_with_paths_base(
    register_second_conn_id: bool,
    validate_path_zero: bool,
    publisher: &mut Publisher,
) -> Helper {
    let zero_conn_id = connection::PeerId::try_from_bytes(&[0]).unwrap();
    let first_conn_id = connection::PeerId::try_from_bytes(&[1]).unwrap();
    let second_conn_id = connection::PeerId::try_from_bytes(&[2]).unwrap();
    let zero_path_id = path_id(0);
    let first_path_id = path_id(1);
    let second_path_id = path_id(2);
    let mut zero_path = helper_path(zero_conn_id);
    if validate_path_zero {
        // simulate receiving a handshake packet to force path validation
        zero_path.on_handshake_packet();
    }
    assert!(!zero_path.is_challenge_pending());

    let now = NoopClock {}.get_time();
    let challenge_expiration = Duration::from_millis(10_000);
    let first_expected_data = [0; 8];
    let challenge = challenge::Challenge::new(challenge_expiration, first_expected_data);

    let mut first_path = helper_path(first_conn_id);
    first_path.set_challenge(challenge);

    // Create a challenge that will expire in 100ms
    let second_expected_data = [1; 8];
    let challenge = challenge::Challenge::new(challenge_expiration, second_expected_data);
    let mut second_path = helper_path(second_conn_id);
    second_path.set_challenge(challenge);

    let mut random_generator = random::testing::Generator(123);
    let mut peer_id_registry =
        ConnectionIdMapper::new(&mut random_generator, endpoint::Type::Server)
            .create_server_peer_id_registry(
                InternalConnectionIdGenerator::new().generate_id(),
                zero_path.peer_connection_id,
            );
    assert!(peer_id_registry
        .on_new_connection_id(&first_conn_id, 1, 0, &TEST_TOKEN_1)
        .is_ok());

    if register_second_conn_id {
        assert!(peer_id_registry
            .on_new_connection_id(&second_conn_id, 2, 0, &TEST_TOKEN_2)
            .is_ok());
    }

    let mut manager = ServerManager::new(zero_path, peer_id_registry);
    assert!(manager.peer_id_registry.is_active(&first_conn_id));
    manager.paths.push(first_path);
    manager.paths.push(second_path);
    assert_eq!(manager.paths.len(), 3);

    // update active path to first_path
    assert_eq!(manager.active, zero_path_id.as_u8());
    if validate_path_zero {
        assert!(manager.active_path().is_validated());
    }

    assert!(manager
        .update_active_path(
            first_path_id,
            &mut random::testing::Generator(123),
            publisher,
        )
        .is_ok());
    if validate_path_zero {
        assert!(manager[zero_path_id].is_challenge_pending());
    }

    assert!(manager
        .peer_id_registry
        .consume_new_id_for_existing_path(path_id(0), zero_conn_id, publisher)
        .is_some());

    // assert first_path is active and last_known_active_validated_path
    assert!(manager.peer_id_registry.is_active(&first_conn_id));
    assert_eq!(manager.active, first_path_id.as_u8());

    if validate_path_zero {
        assert_eq!(
            manager.last_known_active_validated_path,
            Some(zero_path_id.as_u8())
        );
    } else {
        assert_eq!(manager.last_known_active_validated_path, None);
    }

    Helper {
        now,
        first_expected_data,
        second_expected_data,
        challenge_expiration,
        zero_path_id,
        first_path_id,
        second_path_id,
        manager,
    }
}

pub fn helper_path(peer_id: connection::PeerId) -> ServerPath {
    let local_conn_id = connection::LocalId::TEST_ID;
    ServerPath::new(
        Default::default(),
        peer_id,
        local_conn_id,
        RttEstimator::new(Duration::from_millis(30)),
        Default::default(),
        false,
        DEFAULT_MAX_MTU,
    )
}

fn helper_manager_with_paths(publisher: &mut Publisher) -> Helper {
    helper_manager_with_paths_base(true, true, publisher)
}

pub struct Helper {
    pub now: Timestamp,
    pub first_expected_data: challenge::Data,
    pub second_expected_data: challenge::Data,
    pub challenge_expiration: Duration,
    pub zero_path_id: Id,
    pub first_path_id: Id,
    pub second_path_id: Id,
    pub manager: ServerManager,
}