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

//! Defines time related datatypes and functions

#![allow(dead_code)]

use cfg_if::cfg_if;

cfg_if! {
    if #[cfg(any(feature = "testing", test))] {
        pub use if_testing::*;
    } else if #[cfg(feature = "std")] {
        pub use if_std::*;
    } else {
        pub use if_no_std::*;
    }
}

#[cfg(any(feature = "std", test))]
mod if_std {
    //! This module implements the clock functionality for the "std" feature
    //! and environment. In this environment we are directly using `Instant`
    //! types from the Rust standard library as Timestamps.
    use lazy_static::lazy_static;
    use s2n_quic_core::time::{Clock, StdClock, Timestamp};

    lazy_static! {
        static ref GLOBAL_CLOCK: StdClock = StdClock::default();
    }

    /// Returns the current [`Timestamp`] according to the system clock
    pub fn now() -> Timestamp {
        GLOBAL_CLOCK.get_time()
    }

    /// Returns a reference to the clock.
    pub fn clock() -> &'static dyn Clock {
        &*GLOBAL_CLOCK
    }
}

// The no-std version allows to set a global clock manually.
// It is not possible to overwrite the clock at runtime.
#[cfg(any(not(feature = "std"), test))]
mod if_no_std {
    //! This module implements the clock functionality for the "no-std"
    //! environments. In those environments we allow to configure a global clock
    //! via a trait object.
    //!
    //! The global clock has to be initialized via a call to
    //! `init_global_clock`.
    use core::sync::atomic::{AtomicUsize, Ordering};
    use s2n_quic_core::time::{Clock, NoopClock, Timestamp};

    /// The configured global clock
    static mut GLOBAL_CLOCK: &'static dyn Clock = &NoopClock {};

    const CLOCK_UNINITIALIZED: usize = 0;
    const CLOCK_INITIALIZING: usize = 1;
    const CLOCK_INITIALIZED: usize = 2;

    /// Tracks whether the global clock had already been initialized
    static CLOCK_STATE: AtomicUsize = AtomicUsize::new(CLOCK_UNINITIALIZED);

    /// Initialize the global Clock for use in `no-std` mode.
    /// The configured clock will be queried on all `crate::time::now()` calls.
    /// The global clock can only be set once.
    pub fn init_global_clock(clock: &'static dyn Clock) -> Result<(), ()> {
        unsafe {
            match CLOCK_STATE.compare_exchange(
                CLOCK_UNINITIALIZED,
                CLOCK_INITIALIZING,
                Ordering::SeqCst,
                Ordering::SeqCst,
            ) {
                Ok(_) => {
                    GLOBAL_CLOCK = clock;
                    CLOCK_STATE.store(CLOCK_INITIALIZED, Ordering::SeqCst);
                    Ok(())
                }
                Err(err) if err == CLOCK_INITIALIZING => {
                    // Wait until a different thread has initialized the clock
                    while CLOCK_STATE.load(Ordering::SeqCst) != CLOCK_INITIALIZED {}
                    Err(())
                }
                _ => Err(()),
            }
        }
    }

    /// Returns the current [`Timestamp`] according to the system clock
    pub fn now() -> Timestamp {
        clock().get_time()
    }

    /// Returns a reference to the clock.
    ///
    /// If a clock has not been set, a no-op implementation is returned.
    pub fn clock() -> &'static dyn Clock {
        unsafe {
            if CLOCK_STATE.load(Ordering::SeqCst) != CLOCK_INITIALIZED {
                static NOP: NoopClock = NoopClock {};
                &NOP
            } else {
                GLOBAL_CLOCK
            }
        }
    }
}

/// Clock implementation if the "testing" feature is enabled
#[cfg(any(feature = "testing", test))]
mod if_testing {
    use cfg_if::cfg_if;
    use core::cell::RefCell;
    use s2n_quic_core::time::{Clock, Duration, Timestamp};
    use std::sync::Arc;

    cfg_if! {
        if #[cfg(feature = "std")] {
            use super::if_std::now as inner_now;
        } else {
            use super::if_no_std::now as inner_now;
        }
    }

    thread_local! {
        static LOCAL_CLOCK: ClockHolder = ClockHolder {
            inner: RefCell::new(None),
        }
    }

    struct ClockHolder {
        inner: RefCell<Option<Arc<dyn Clock>>>,
    }

    impl Clock for ClockHolder {
        fn get_time(&self) -> Timestamp {
            match &*self.inner.borrow() {
                Some(clock) => clock.get_time(),
                None => inner_now(),
            }
        }
    }

    /// Returns the current [`Timestamp`] according to the system clock
    pub fn now() -> Timestamp {
        (&LOCAL_CLOCK).get_time()
    }

    /// Returns a reference to the clock.
    pub fn clock() -> &'static dyn Clock {
        &&LOCAL_CLOCK
    }

    pub mod testing {
        use super::*;
        use std::sync::Mutex;

        /// Configures a [`Clock`] which will be utilized for the following
        /// calls to `crate::time::now()` on the current thread.
        ///
        /// Example:
        ///
        /// ```ignore
        /// # use core::time::Duration;
        /// use std::sync::Arc;
        /// use s2n_quic_platform::time::{self, testing};
        /// let clock = Arc::new(testing::MockClock::new());
        /// testing::set_local_clock(clock.clone());
        ///
        /// let before = time::now();
        /// clock.adjust_by(Duration::from_millis(333));
        /// let after = time::now();
        /// assert_eq!(after - before, Duration::from_millis(333));
        /// ```
        pub fn set_local_clock(clock: Arc<dyn Clock>) {
            LOCAL_CLOCK.with(|current_local_clock| {
                *current_local_clock.inner.borrow_mut() = Some(clock);
            });
        }

        /// Resets the local clock.
        ///
        /// Following invocations to [`crate::time::now()`]
        /// will return the system time again.
        pub fn reset_local_clock() {
            LOCAL_CLOCK.with(|current_local_clock| {
                *current_local_clock.inner.borrow_mut() = None;
            });
        }

        /// A [`Clock`] for testing purposes.
        ///
        /// The timestamp stored by the clock can be adjusted through the
        /// `adjust_by` and `set_time` functions.
        /// Following calls to `get_time` return the adjusted timestamp.
        pub struct MockClock {
            timestamp: Mutex<Timestamp>,
        }

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

        impl MockClock {
            /// Creates a new clock instance for testing purposes.
            ///
            /// The Clock will default to a default [`Timestamp`], which
            /// represents the lowest possible [`Timestamp`] which can ever be
            /// returned by this Clock. It is not allowed to adjust the Clock
            /// to a [`Timestamp`] before its initial time.
            pub fn new() -> MockClock {
                MockClock {
                    timestamp: Mutex::new(inner_now()),
                }
            }

            /// Sets the current time to the given [`Timestamp`].
            /// Follow-up calls to [`Self::get_time`] will return this [`Timestamp`]
            /// until the time had been adjusted again.
            pub fn set_time(&self, timestamp: Timestamp) {
                let mut guard = self.timestamp.lock().unwrap();
                *guard = timestamp;
            }

            /// Adjusts the time which is returned by the clock by the given
            /// Duration. This will not perform any overflow or underflow checks.
            /// If the clock went backwards more than to the initial zero
            /// timestamp, the call will lead to a panic.
            pub fn adjust_by(&self, duration: Duration) {
                let mut guard = self.timestamp.lock().unwrap();
                *guard += duration;
            }
        }

        impl Clock for MockClock {
            fn get_time(&self) -> Timestamp {
                *self.timestamp.lock().unwrap()
            }
        }
    }

    #[test]
    fn use_mocked_clock() {
        use std::sync::Arc;

        let original_time = now();

        // Switch to a mock clock
        let clock = Arc::new(testing::MockClock::new());
        testing::set_local_clock(clock.clone());

        let ts1 = now();
        clock.adjust_by(Duration::from_millis(333));
        let ts2 = now();
        assert_eq!(ts2 - ts1, Duration::from_millis(333));
        clock.adjust_by(Duration::from_millis(111));
        let ts3 = now();
        assert_eq!(ts3 - ts1, Duration::from_millis(444));

        clock.set_time(ts1);
        assert_eq!(ts1, now());

        // Switch back to the original clock
        testing::reset_local_clock();
        let restored_time = now();
        assert!(restored_time - original_time >= Duration::from_millis(0));
        assert!(restored_time - original_time <= Duration::from_millis(100));
    }
}


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

mod h09;
mod h3;
pub mod interop;
pub mod perf;

pub use interop::Interop;
pub use perf::Perf;


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

use crate::Result;
use bytes::{BufMut, Bytes, BytesMut};
use core::mem::MaybeUninit;
use futures::{ready, stream::Stream};
use std::{
    path::{Path, PathBuf},
    pin::Pin,
    task::{Context, Poll},
};
use tokio::{
    fs,
    io::{AsyncRead, ReadBuf},
};

const CAPACITY: usize = 4096;

pub struct File {
    file: fs::File,
    buf: BytesMut,
}

impl File {
    pub async fn open<P: AsRef<Path>>(p: P) -> Result<Self> {
        let file = fs::File::open(p).await?;
        Ok(Self {
            file,
            buf: BytesMut::with_capacity(CAPACITY),
        })
    }
}

impl Stream for File {
    type Item = std::io::Result<Bytes>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        if self.buf.capacity() == 0 {
            self.buf.reserve(CAPACITY);
        }

        let n = {
            let bytes = self.buf.chunk_mut();
            let dst = unsafe { &mut *(bytes as *mut _ as *mut [MaybeUninit<u8>]) };
            let mut buf = ReadBuf::uninit(dst);
            ready!(AsyncRead::poll_read(
                Pin::new(&mut self.file),
                cx,
                &mut buf
            )?);
            buf.filled().len()
        };

        if n == 0 {
            return Poll::Ready(None);
        }

        // Safety: This is guaranteed to be the number of initialized (and read)
        // bytes due to the invariants provided by `ReadBuf::filled`.
        unsafe {
            self.buf.advance_mut(n);
        }

        let chunk = self.buf.split();
        Poll::Ready(Some(Ok(chunk.freeze())))
    }
}

pub(crate) fn abs_path(path: &str, www_dir: &Path) -> PathBuf {
    let mut abs_path = www_dir.to_path_buf();
    abs_path.extend(
        path.split('/')
            .filter(|segment| !segment.starts_with('.'))
            .map(std::path::Path::new),
    );
    abs_path
}


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

use crate::Result;
use core::str::FromStr;
use std::io;

/// From <https://github.com/marten-seemann/quic-interop-runner#test-cases>
#[derive(Clone, Copy, Debug)]
pub enum Testcase {
    /// Tests that a server sends a Version Negotiation packet in response to an unknown QUIC version number.
    ///
    /// The client should start a connection using an unsupported version number (it can use a reserved version number to do so),
    /// and should abort the connection attempt when receiving the Version Negotiation packet. Currently disabled due to #20.
    VersionNegotiation,

    /// Tests the successful completion of the handshake.
    ///
    /// The client is expected to establish a single QUIC connection to
    /// the server and download one or multiple small files. Servers should not send a Retry packet in this test case.
    Handshake,

    /// Tests both flow control and stream multiplexing.
    ///
    /// The client should use small initial flow control windows for both
    /// stream- and connection-level flow control, such that the during the transfer of files on the order of 1 MB the flow
    /// control window needs to be increased. The client is expected to establish a single QUIC connection, and use multiple
    /// streams to concurrently download the files.
    Transfer,

    /// Tests support for ChaCha20.
    ///
    /// In this test, client and server are expected to offer only ChaCha20 as a ciphersuite. The client then downloads the files.
    ChaCha20,

    /// Tests support for key updates (client only)
    ///
    /// The client is expected to make sure that a key update happens early in the connection (during the first MB transferred).
    /// It doesn't matter which peer actually initiated the update.
    KeyUpdate,

    /// Tests that the server can generate a Retry, and that the client can act upon it.
    ///
    /// The client should use the Token provided in the Retry packet in the Initial packet.
    Retry,

    /// Tests QUIC session resumption (without 0-RTT).
    ///
    /// The client is expected to establish a connection and download the first file. The server is expected to
    /// provide the client with a session ticket that allows it to resume the connection.
    /// After downloading the first file, the client has to close the connection, establish a resumed connection using the
    /// session ticket, and use this connection to download the remaining file(s).
    Resumption,

    /// Tests QUIC 0-RTT.
    ///
    /// The client is expected to establish a connection and download the first file. The server is expected to provide the
    /// client with a session ticket that allows it establish a 0-RTT connection on the next connection attempt. After downloading
    ///  the first file, the client has to close the connection, establish and request the remaining file(s) in 0-RTT.
    ZeroRtt,

    /// Tests a simple HTTP/3 connection.
    ///
    /// The client is expected to download multiple files using HTTP/3. Files should be requested and transferred in parallel.
    Http3,

    /// Tests resilience of the handshake to high loss.
    ///
    ///  The client is expected to establish multiple connections, sequential or in parallel, and use each connection to download a single file.
    Multiconnect,

    /// Tests support for ECN markings
    Ecn,

    /// Tests an active connection migration
    ///
    /// A transfer succeeded during which the client performed an active migration.
    ConnectionMigration,
}

impl Testcase {
    pub const TESTCASES: &'static [Self] = &[
        Self::VersionNegotiation,
        Self::Handshake,
        Self::Transfer,
        Self::ChaCha20,
        Self::KeyUpdate,
        Self::Retry,
        Self::Resumption,
        Self::ZeroRtt,
        Self::Http3,
        Self::Multiconnect,
        Self::Ecn,
        Self::ConnectionMigration,
    ];

    pub const fn as_str(self) -> &'static str {
        use Testcase::*;
        match self {
            VersionNegotiation => "versionnegotiation",
            Handshake => "handshake",
            Transfer => "transfer",
            ChaCha20 => "chacha20",
            KeyUpdate => "keyupdate",
            Retry => "retry",
            Resumption => "resumption",
            ZeroRtt => "zerortt",
            Http3 => "http3",
            Multiconnect => "multiconnect",
            Ecn => "ecn",
            ConnectionMigration => "connectionmigration",
        }
    }

    pub fn supported(f: impl Fn(Self) -> bool) -> Vec<&'static str> {
        let mut results = vec![];

        for testcase in Self::TESTCASES.iter().copied() {
            if f(testcase) {
                results.push(testcase.as_str());
            }
        }

        results
    }
}

impl FromStr for Testcase {
    type Err = crate::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        use Testcase::*;

        Ok(match s {
            "versionnegotiation" => VersionNegotiation,
            "handshake" => Handshake,
            "transfer" => Transfer,
            "chacha20" => ChaCha20,
            "keyupdate" => KeyUpdate,
            "retry" => Retry,
            "resumption" => Resumption,
            "zerortt" => ZeroRtt,
            "http3" => Http3,
            "multiconnect" => Multiconnect,
            "ecn" => Ecn,
            "connectionmigration" => ConnectionMigration,
            _ => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidInput,
                    format!("Unsupported test case: {s}"),
                )
                .into())
            }
        })
    }
}


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

use crate::Result;
use s2n_quic::provider::io;
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub struct Server {
    #[structopt(short, long, default_value = "::")]
    pub ip: std::net::IpAddr,

    #[structopt(short, long, default_value = "443")]
    pub port: u16,

    #[structopt(long)]
    pub disable_gso: bool,

    #[structopt(long, default_value = "9000")]
    pub max_mtu: u16,

    #[cfg(feature = "xdp")]
    #[structopt(flatten)]
    xdp: crate::xdp::Xdp,
}

impl Server {
    #[cfg(feature = "xdp")]
    pub fn build(&self) -> Result<impl io::Provider> {
        // GSO isn't currently supported for XDP so just read it to avoid a dead_code warning
        let _ = self.disable_gso;
        let _ = self.max_mtu;

        let addr = (self.ip, self.port).into();

        self.xdp.server(addr)
    }

    #[cfg(not(feature = "xdp"))]
    pub fn build(&self) -> Result<impl io::Provider> {
        let mut io_builder = io::Default::builder()
            .with_receive_address((self.ip, self.port).into())?
            .with_max_mtu(self.max_mtu)?;

        if self.disable_gso {
            io_builder = io_builder.with_gso_disabled()?;
        }

        Ok(io_builder.build()?)
    }
}

#[derive(Debug, StructOpt)]
pub struct Client {
    #[structopt(long)]
    pub disable_gso: bool,

    #[structopt(long, default_value = "9000")]
    pub max_mtu: u16,

    #[structopt(short, long, default_value = "::")]
    pub local_ip: std::net::IpAddr,

    #[cfg(feature = "xdp")]
    #[structopt(flatten)]
    xdp: crate::xdp::Xdp,
}

impl Client {
    #[cfg(feature = "xdp")]
    pub fn build(&self) -> Result<impl io::Provider> {
        // GSO isn't currently supported for XDP so just read it to avoid a dead_code warning
        let _ = self.disable_gso;
        let _ = self.max_mtu;

        let addr = (self.local_ip, 0u16).into();

        self.xdp.client(addr)
    }

    #[cfg(not(feature = "xdp"))]
    pub fn build(&self) -> Result<impl io::Provider> {
        let mut io_builder = io::Default::builder()
            .with_receive_address((self.local_ip, 0u16).into())?
            .with_max_mtu(self.max_mtu)?;

        if self.disable_gso {
            io_builder = io_builder.with_gso_disabled()?;
        }

        Ok(io_builder.build()?)
    }
}


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

use structopt::StructOpt;

pub type Error = Box<dyn 'static + std::error::Error + Send + Sync>;
pub type Result<T, E = Error> = core::result::Result<T, E>;

mod client;
mod file;
mod interop;
mod io;
mod perf;
mod server;
mod tls;
#[cfg(feature = "xdp")]
mod xdp;

/// This message is searched in interop logs to ensure the application doesn't panic
///
/// Do not change it without updating it elsewhere
const CRASH_ERROR_MESSAGE: &str = "The s2n-quic-qns application shut down unexpectedly";

#[cfg(not(target_os = "android"))]
#[global_allocator]
static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc;

fn main() {
    let format = tracing_subscriber::fmt::format()
        .with_level(false) // don't include levels in formatted output
        .with_timer(tracing_subscriber::fmt::time::uptime())
        .with_ansi(false)
        .compact(); // Use a less verbose output format.

    tracing_subscriber::fmt()
        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
        .event_format(format)
        .init();

    match Arguments::from_args_safe() {
        Ok(args) => {
            if let Err(error) = args.run() {
                eprintln!("Error: {error:?}");
            }
        }
        Err(error) => {
            if error.use_stderr() {
                eprintln!("{error}");

                // https://github.com/marten-seemann/quic-interop-runner/blob/cd223804bf3f102c3567758ea100577febe486ff/interop.py#L102
                // The interop runner wants us to exit with code 127 when an invalid argument is passed
                std::process::exit(127);
            } else {
                println!("{error}");
            }
        }
    };
}

#[derive(Debug, StructOpt)]
struct Arguments {
    #[structopt(long, env = "S2N_QUIC_MULTI_THREAD")]
    multithread: bool,

    #[structopt(flatten)]
    program: Program,
}

impl Arguments {
    pub fn run(&self) -> Result<()> {
        let runtime = if self.multithread {
            tokio::runtime::Builder::new_multi_thread()
        } else {
            tokio::runtime::Builder::new_current_thread()
        }
        .enable_all()
        .build()?;

        runtime.block_on(self.program.run())?;

        Ok(())
    }
}

#[derive(Debug, StructOpt)]
enum Program {
    Interop(Interop),
    Perf(Perf),
}

impl Program {
    pub async fn run(&self) -> Result<()> {
        match self {
            Self::Interop(subject) => subject.run().await,
            Self::Perf(subject) => subject.run().await,
        }
    }
}

#[derive(Debug, StructOpt)]
enum Interop {
    Server(server::Interop),
    Client(client::Interop),
}

impl Interop {
    pub async fn run(&self) -> Result<()> {
        match self {
            Self::Server(subject) => subject.run().await,
            Self::Client(subject) => subject.run().await,
        }
    }
}

#[derive(Debug, StructOpt)]
enum Perf {
    Server(server::Perf),
    Client(client::Perf),
}

impl Perf {
    pub async fn run(&self) -> Result<()> {
        match self {
            Self::Server(subject) => subject.run().await,
            Self::Client(subject) => subject.run().await,
        }
    }
}


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

use crate::Result;
use bytes::Bytes;
use core::time::Duration;
use s2n_quic::{
    provider::event,
    stream::{ReceiveStream, SendStream},
};
use std::sync::{
    atomic::{AtomicU64, Ordering},
    Arc,
};

/// Drains a receive stream
pub async fn handle_receive_stream(mut stream: ReceiveStream) -> Result<()> {
    let mut chunks: [_; 64] = core::array::from_fn(|_| Bytes::new());

    loop {
        let (len, is_open) = stream.receive_vectored(&mut chunks).await?;

        if !is_open {
            break;
        }

        for chunk in chunks[..len].iter_mut() {
            // discard chunks
            *chunk = Bytes::new();
        }
    }

    Ok(())
}

/// Sends a specified amount of data on a send stream
pub async fn handle_send_stream(mut stream: SendStream, len: u64) -> Result<()> {
    let mut chunks: [_; 64] = core::array::from_fn(|_| Bytes::new());

    //= https://tools.ietf.org/id/draft-banks-quic-performance-00#4.1
    //# Since the goal here is to measure the efficiency of the QUIC
    //# implementation and not any application protocol, the performance
    //# application layer should be as light-weight as possible.  To this
    //# end, the client and server application layer may use a single
    //# preallocated and initialized buffer that it queues to send when any
    //# payload needs to be sent out.
    let mut data = s2n_quic_core::stream::testing::Data::new(len);

    loop {
        match data.send(usize::MAX, &mut chunks) {
            Some(count) => {
                stream.send_vectored(&mut chunks[..count]).await?;
            }
            None => {
                stream.finish()?;
                break;
            }
        }
    }

    Ok(())
}

//= https://tools.ietf.org/id/draft-banks-quic-performance-00#2.3.1
//# Every stream opened by the client uses the first 8 bytes of the
//# stream data to encode a 64-bit unsigned integer in network byte order
//# to indicate the length of data the client wishes the server to
//# respond with.
pub async fn write_stream_size(stream: &mut SendStream, len: u64) -> Result<()> {
    let size = len.to_be_bytes();
    let chunk = Bytes::copy_from_slice(&size);
    stream.send(chunk).await?;
    Ok(())
}

pub async fn read_stream_size(stream: &mut ReceiveStream) -> Result<(u64, Bytes)> {
    let mut chunk = Bytes::new();
    let mut offset = 0;
    let mut id = [0u8; core::mem::size_of::<u64>()];

    while offset < id.len() {
        chunk = stream
            .receive()
            .await?
            .expect("every stream should be prefixed with the scenario ID");

        let needed_len = id.len() - offset;
        let len = chunk.len().min(needed_len);

        id[offset..offset + len].copy_from_slice(&chunk[..len]);
        offset += len;
        bytes::Buf::advance(&mut chunk, len);
    }

    let id = u64::from_be_bytes(id);

    Ok((id, chunk))
}

#[derive(Debug, structopt::StructOpt)]
pub struct Limits {
    /// The maximum bits/sec for each connection
    #[structopt(long, default_value = "150")]
    pub max_throughput: u64,

    /// The expected RTT in milliseconds
    #[structopt(long, default_value = "100")]
    pub expected_rtt: u64,
}

impl Limits {
    pub fn limits(&self) -> s2n_quic::provider::limits::Limits {
        let data_window = self.data_window();

        s2n_quic::provider::limits::Limits::default()
            .with_data_window(data_window)
            .unwrap()
            .with_max_send_buffer_size(data_window.min(u32::MAX as _) as _)
            .unwrap()
            .with_bidirectional_local_data_window(data_window)
            .unwrap()
            .with_bidirectional_remote_data_window(data_window)
            .unwrap()
            .with_unidirectional_data_window(data_window)
            .unwrap()
    }

    fn data_window(&self) -> u64 {
        s2n_quic_core::transport::parameters::compute_data_window(
            self.max_throughput,
            core::time::Duration::from_millis(self.expected_rtt),
            2,
        )
        .as_u64()
    }
}

#[derive(Debug, Default)]
pub struct Subscriber {
    pub counters: Arc<Counters>,
}

impl Subscriber {
    pub fn spawn(&self, frequency: Duration) {
        let counters = self.counters.clone();
        counters.print_header();
        tokio::spawn(async move {
            loop {
                tokio::time::sleep(frequency).await;
                counters.print(frequency);
            }
        });
    }
}

impl s2n_quic::provider::event::Subscriber for Subscriber {
    type ConnectionContext = ();

    #[inline]
    fn create_connection_context(
        &mut self,
        _meta: &event::ConnectionMeta,
        _info: &event::ConnectionInfo,
    ) -> Self::ConnectionContext {
    }

    #[inline]
    fn on_tx_stream_progress(
        &mut self,
        _context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::TxStreamProgress,
    ) {
        self.counters
            .send_progress
            .fetch_add(event.bytes as _, Ordering::Relaxed);
    }

    #[inline]
    fn on_rx_stream_progress(
        &mut self,
        _context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::RxStreamProgress,
    ) {
        self.counters
            .receive_progress
            .fetch_add(event.bytes as _, Ordering::Relaxed);
    }

    #[inline]
    fn on_recovery_metrics(
        &mut self,
        _context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::RecoveryMetrics,
    ) {
        self.counters
            .max_cwnd
            .fetch_max(event.congestion_window as _, Ordering::Relaxed);
        self.counters
            .max_bytes_in_flight
            .fetch_max(event.bytes_in_flight as _, Ordering::Relaxed);
        self.counters
            .max_rtt
            .fetch_max(event.latest_rtt.as_nanos() as _, Ordering::Relaxed);
        self.counters
            .max_smoothed_rtt
            .fetch_max(event.smoothed_rtt.as_nanos() as _, Ordering::Relaxed);
        self.counters
            .pto_count
            .fetch_max(event.pto_count as _, Ordering::Relaxed);
    }

    #[inline]
    fn on_packet_lost(
        &mut self,
        _context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        _event: &event::events::PacketLost,
    ) {
        self.counters.lost_packets.fetch_add(1, Ordering::Relaxed);
    }

    #[inline]
    fn on_platform_event_loop_wakeup(
        &mut self,
        _meta: &event::events::EndpointMeta,
        _event: &event::events::PlatformEventLoopWakeup,
    ) {
        self.counters
            .event_loop_wakeup
            .fetch_add(1, Ordering::Relaxed);
    }

    #[inline]
    fn on_platform_event_loop_sleep(
        &mut self,
        _meta: &event::events::EndpointMeta,
        event: &event::events::PlatformEventLoopSleep,
    ) {
        self.counters.timeout.store(
            event.timeout.unwrap_or_default().as_nanos() as _,
            Ordering::Relaxed,
        );
    }

    #[inline]
    fn on_pacing_rate_updated(
        &mut self,
        _context: &mut Self::ConnectionContext,
        _meta: &event::events::ConnectionMeta,
        event: &event::events::PacingRateUpdated,
    ) {
        self.counters
            .max_pacing_rate
            .fetch_max(event.bytes_per_second, Ordering::Relaxed);
    }

    #[inline]
    fn on_delivery_rate_sampled(
        &mut self,
        _context: &mut Self::ConnectionContext,
        _meta: &event::events::ConnectionMeta,
        event: &event::events::DeliveryRateSampled,
    ) {
        self.counters.max_delivery_rate.fetch_max(
            event.rate_sample.delivery_rate_bytes_per_second,
            Ordering::Relaxed,
        );
    }
}

#[derive(Debug, Default)]
pub struct Counters {
    send_progress: AtomicU64,
    receive_progress: AtomicU64,
    max_cwnd: AtomicU64,
    max_bytes_in_flight: AtomicU64,
    max_rtt: AtomicU64,
    max_smoothed_rtt: AtomicU64,
    lost_packets: AtomicU64,
    event_loop_wakeup: AtomicU64,
    timeout: AtomicU64,
    pto_count: AtomicU64,
    max_pacing_rate: AtomicU64,
    max_delivery_rate: AtomicU64,
}

impl Counters {
    pub fn print_header(&self) {
        println!(
            "Tx Rate\t\
            Rx Rate\t\
            Max Cwnd\t\
            Max Inflight\t\
            Lost Packets\t\
            Wakeups\t\
            Duration\t\
            Max RTT\t\
            Max SRTT\t\
            PTO Count\t\
            Max Pacing Rate\t\
            Max Delivery Rate"
        );
    }

    pub fn print(&self, duration: Duration) {
        // The goodput of data transmitted to the peer
        let send_progress = self.send_progress.swap(0, Ordering::Relaxed);
        let send_rate = rate(send_progress, duration);
        // The goodput of data received from the peer
        let receive_progress = self.receive_progress.swap(0, Ordering::Relaxed);
        let receive_rate = rate(receive_progress, duration);
        // The maximum congestion window observed during the interval
        let max_cwnd = self.max_cwnd.swap(0, Ordering::Relaxed);
        let max_cwnd = bytes(max_cwnd);
        // The maximum amount of unacknowledged data in flight during the interval
        let max_bytes_in_flight = self.max_bytes_in_flight.swap(0, Ordering::Relaxed);
        let max_bytes_in_flight = bytes(max_bytes_in_flight);
        // The maximum round trip time observed during the interval
        let max_rtt = self.max_rtt.swap(0, Ordering::Relaxed);
        let max_rtt = Duration::from_nanos(max_rtt);
        // The maximum smoothed (weighted average) round trip time observed during the interval
        let max_smoothed_rtt = self.max_smoothed_rtt.swap(0, Ordering::Relaxed);
        let max_smoothed_rtt = Duration::from_nanos(max_smoothed_rtt);
        // The number of packets recorded lost during the interval
        let lost_packets = self.lost_packets.swap(0, Ordering::Relaxed);
        // The number of event loop wakeups recorded during the interval
        let wakeups = self.event_loop_wakeup.swap(0, Ordering::Relaxed);
        // The duration of the latest event loop wakeup
        let duration = self.timeout.swap(0, Ordering::Relaxed);
        let duration = Duration::from_nanos(duration);
        // The number of packet time out events observed during the interval
        let pto_count = self.pto_count.swap(0, Ordering::Relaxed);
        // The maximum rate at which packets are paced out observed during the interval
        let max_pacing_rate = self.max_pacing_rate.swap(0, Ordering::Relaxed);
        let max_pacing_rate = rate(max_pacing_rate, Duration::from_secs(1));
        // The maximum estimate of bandwidth observed during the interval. Only output for BBRv2
        let max_delivery_rate = self.max_delivery_rate.swap(0, Ordering::Relaxed);
        let max_delivery_rate = rate(max_delivery_rate, Duration::from_secs(1));
        println!(
            "{send_rate}\t\
            {receive_rate}\t\
            {max_cwnd}\t\
            {max_bytes_in_flight}\t\
            {lost_packets}\t\
            {wakeups}\t\
            {duration:?}\t\
            {max_rtt:?}\t\
            {max_smoothed_rtt:?}\t\
            {pto_count}\t\
            {max_pacing_rate}\t\
            {max_delivery_rate}",
        );
    }
}

fn rate(bytes: u64, duration: Duration) -> String {
    use humansize::{format_size, DECIMAL};

    let bits = bytes * 8;
    let value = format_size(bits, DECIMAL.space_after_value(false));
    let value = value.trim_end_matches('B');

    if duration == Duration::from_secs(1) {
        format!("{value}bps")
    } else {
        format!("{value}/{duration:?}")
    }
}

fn bytes(value: u64) -> String {
    use humansize::{format_size, DECIMAL};
    format_size(value, DECIMAL.space_after_value(false))
}


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

mod h09;
mod h3;
pub mod interop;
pub mod perf;
#[cfg(all(s2n_quic_unstable, feature = "unstable_client_hello"))]
mod unstable;

pub use interop::Interop;
pub use perf::Perf;


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

use crate::Result;
use s2n_quic::{
    client::{self, ClientProviders},
    server::{self, ServerProviders},
};
use std::{path::PathBuf, str::FromStr};
use structopt::StructOpt;

#[derive(Debug, StructOpt)]
pub struct Server {
    #[structopt(long)]
    pub certificate: Option<PathBuf>,

    #[structopt(long)]
    pub private_key: Option<PathBuf>,

    #[structopt(long, default_value)]
    pub tls: TlsProviders,
}

impl Server {
    pub fn build(
        &self,
        builder: server::Builder<impl ServerProviders>,
        alpns: &[String],
    ) -> Result<s2n_quic::Server> {
        Ok(match self.tls {
            #[cfg(unix)]
            TlsProviders::S2N => {
                let tls = self.build_s2n_tls(alpns)?;

                builder.with_tls(tls)?.start().unwrap()
            }
            TlsProviders::Rustls => {
                let tls = self.build_rustls(alpns)?;

                builder.with_tls(tls)?.start().unwrap()
            }
            TlsProviders::Null => {
                let tls = self.build_null()?;

                builder.with_tls(tls)?.start().unwrap()
            }
        })
    }

    #[cfg(unix)]
    fn build_s2n_tls(&self, alpns: &[String]) -> Result<s2n_tls::Server> {
        // The server builder defaults to a chain because this allows certs to just work, whether
        // the PEM contains a single cert or a chain

        let tls = s2n_tls::Server::builder()
            .with_certificate(
                s2n_tls::ca(self.certificate.as_ref())?,
                s2n_tls::private_key(self.private_key.as_ref())?,
            )?
            .with_application_protocols(alpns.iter().map(String::as_bytes))?
            .with_key_logging()?;

        cfg_if::cfg_if! {
            if #[cfg(all(
                s2n_quic_unstable,
                feature = "unstable_client_hello"
            ))] {
                use super::unstable::MyClientHelloHandler;
                let tls = tls.with_client_hello_handler(MyClientHelloHandler {})?;
            }
        }

        Ok(tls.build()?)
    }

    fn build_rustls(&self, alpns: &[String]) -> Result<rustls::Server> {
        // The server builder defaults to a chain because this allows certs to just work, whether
        // the PEM contains a single cert or a chain
        let tls = rustls::Server::builder()
            .with_certificate(
                rustls::ca(self.certificate.as_ref())?,
                rustls::private_key(self.private_key.as_ref())?,
            )?
            .with_application_protocols(alpns.iter().map(String::as_bytes))?
            .with_key_logging()?
            .build()?;

        Ok(tls)
    }

    fn build_null(&self) -> Result<null::Provider> {
        Ok(null::Provider::default())
    }
}

#[derive(Debug, StructOpt)]
pub struct Client {
    #[structopt(long)]
    pub ca: Option<PathBuf>,

    #[structopt(long, default_value)]
    pub tls: TlsProviders,
}

impl Client {
    pub fn build(
        &self,
        builder: client::Builder<impl ClientProviders>,
        alpns: &[String],
    ) -> Result<s2n_quic::Client> {
        Ok(match self.tls {
            #[cfg(unix)]
            TlsProviders::S2N => {
                let tls = self.build_s2n_tls(alpns)?;

                builder.with_tls(tls)?.start().unwrap()
            }
            TlsProviders::Rustls => {
                let tls = self.build_rustls(alpns)?;

                builder.with_tls(tls)?.start().unwrap()
            }
            TlsProviders::Null => {
                let tls = self.build_null()?;

                builder.with_tls(tls)?.start().unwrap()
            }
        })
    }

    #[cfg(unix)]
    fn build_s2n_tls(&self, alpns: &[String]) -> Result<s2n_tls::Client> {
        let tls = s2n_tls::Client::builder()
            .with_certificate(s2n_tls::ca(self.ca.as_ref())?)?
            // the "amplificationlimit" tests generates a very large chain so bump the limit
            .with_max_cert_chain_depth(10)?
            .with_application_protocols(alpns.iter().map(String::as_bytes))?
            .with_key_logging()?
            .build()?;

        Ok(tls)
    }

    fn build_rustls(&self, alpns: &[String]) -> Result<rustls::Client> {
        let tls = rustls::Client::builder()
            .with_certificate(rustls::ca(self.ca.as_ref())?)?
            // the "amplificationlimit" tests generates a very large chain so bump the limit
            .with_max_cert_chain_depth(10)?
            .with_application_protocols(alpns.iter().map(String::as_bytes))?
            .with_key_logging()?
            .build()?;

        Ok(tls)
    }

    fn build_null(&self) -> Result<null::Provider> {
        Ok(null::Provider::default())
    }
}

#[derive(Clone, Copy, Debug)]
pub enum TlsProviders {
    /// Use s2n-tls as the tls provider
    #[cfg(unix)]
    S2N,
    /// Use rustls as the tls provider
    Rustls,
    /// Use the null tls provider
    Null,
}

impl Default for TlsProviders {
    fn default() -> Self {
        #[cfg(unix)]
        return Self::S2N;
        #[cfg(not(unix))]
        return Self::Rustls;
    }
}

impl ToString for TlsProviders {
    fn to_string(&self) -> String {
        match self {
            #[cfg(unix)]
            TlsProviders::S2N => String::from("s2n-tls"),
            TlsProviders::Rustls => String::from("rustls"),
            TlsProviders::Null => String::from("null"),
        }
    }
}

impl FromStr for TlsProviders {
    type Err = crate::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        Ok(match s {
            "rustls" => Self::Rustls,
            "null" => Self::Null,
            #[cfg(unix)]
            "s2n-tls" => Self::S2N,
            _ => {
                return Err(std::io::Error::new(
                    std::io::ErrorKind::InvalidInput,
                    format!("Unsupported tls provider: {s}"),
                )
                .into())
            }
        })
    }
}

pub mod default {
    #[cfg(not(unix))]
    pub use super::rustls::*;
    #[cfg(unix)]
    pub use super::s2n_tls::*;
}

#[cfg(unix)]
pub mod s2n_tls {
    use super::*;
    pub use s2n_quic::provider::tls::s2n_tls::{
        certificate::{Certificate, IntoCertificate, IntoPrivateKey, PrivateKey},
        Client, Server,
    };

    pub fn ca(ca: Option<&PathBuf>) -> Result<Certificate> {
        Ok(if let Some(pathbuf) = ca.as_ref() {
            pathbuf.into_certificate()?
        } else {
            s2n_quic_core::crypto::tls::testing::certificates::CERT_PEM.into_certificate()?
        })
    }

    pub fn private_key(private_key: Option<&PathBuf>) -> Result<PrivateKey> {
        Ok(if let Some(pathbuf) = private_key.as_ref() {
            pathbuf.into_private_key()?
        } else {
            s2n_quic_core::crypto::tls::testing::certificates::KEY_PEM.into_private_key()?
        })
    }
}

pub mod rustls {
    use super::*;
    pub use s2n_quic::provider::tls::rustls::{
        certificate::{Certificate, IntoCertificate, IntoPrivateKey, PrivateKey},
        Client, Server,
    };

    pub fn ca(ca: Option<&PathBuf>) -> Result<Certificate> {
        Ok(if let Some(pathbuf) = ca.as_ref() {
            pathbuf.into_certificate()?
        } else {
            s2n_quic_core::crypto::tls::testing::certificates::CERT_PEM.into_certificate()?
        })
    }

    pub fn private_key(private_key: Option<&PathBuf>) -> Result<PrivateKey> {
        Ok(if let Some(pathbuf) = private_key.as_ref() {
            pathbuf.into_private_key()?
        } else {
            s2n_quic_core::crypto::tls::testing::certificates::KEY_PEM.into_private_key()?
        })
    }
}

pub mod null {
    use super::*;
    use s2n_quic_core::crypto::tls::null::*;

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

    impl s2n_quic::provider::tls::Provider for Provider {
        type Server = Endpoint;
        type Client = Endpoint;
        type Error = String;

        fn start_server(self) -> Result<Self::Server, Self::Error> {
            Ok(Self::Server::default())
        }

        fn start_client(self) -> Result<Self::Client, Self::Error> {
            Ok(Self::Client::default())
        }
    }
}


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

use crate::Result;
use aya::{
    maps::{HashMap, MapData, XskMap},
    programs, Bpf,
};
use s2n_quic::provider::io::{
    self,
    xdp::{
        bpf, encoder,
        if_xdp::{self, XdpFlags},
        io::{self as xdp_io},
        ring, socket, syscall,
        tx::{self, TxExt as _},
        umem, Provider,
    },
};
use std::{ffi::CString, net::SocketAddr, os::unix::io::AsRawFd, sync::Arc};
use structopt::StructOpt;
use tokio::{io::unix::AsyncFd, net::UdpSocket};

#[derive(Debug, StructOpt)]
pub struct Xdp {
    #[structopt(long, default_value = "lo")]
    interface: String,

    // Default values come from https://elixir.bootlin.com/linux/v6.3.9/source/tools/testing/selftests/bpf/xsk.h#L185
    #[structopt(long, default_value = "2048")]
    tx_queue_len: u32,

    #[structopt(long, default_value = "2048")]
    rx_queue_len: u32,

    #[structopt(long, default_value = "4096")]
    frame_size: u32,

    #[structopt(long)]
    xdp_stats: bool,

    #[structopt(long)]
    bpf_trace: bool,

    #[structopt(long, default_value = "auto")]
    xdp_mode: XdpMode,

    #[structopt(long)]
    no_checksum: bool,
}

#[derive(Clone, Copy, Debug)]
enum XdpMode {
    /// Automatically selects an XDP mode based on the capabilities of the NIC
    Auto,
    /// Uses the software SKB (socket buffer) mode - usually requires no NIC support
    Skb,
    /// Uses the driver mode, which integrates with XDP directly in the kernel driver
    Drv,
    /// Uses the hardware mode, which integrates with XDP directly in the actual NIC hardware
    Hw,
}

impl core::str::FromStr for XdpMode {
    type Err = crate::Error;

    fn from_str(v: &str) -> Result<Self, Self::Err> {
        Ok(match v {
            "auto" => Self::Auto,
            "skb" => Self::Skb,
            "drv" | "driver" => Self::Drv,
            "hw" | "hardware" => Self::Hw,
            _ => return Err(format!("invalid xdp-mode: {v:?}").into()),
        })
    }
}

impl From<XdpMode> for programs::xdp::XdpFlags {
    fn from(mode: XdpMode) -> Self {
        match mode {
            XdpMode::Auto => Self::default(),
            XdpMode::Skb => Self::SKB_MODE,
            XdpMode::Drv => Self::DRV_MODE,
            XdpMode::Hw => Self::HW_MODE,
        }
    }
}

type SetupResult = Result<(
    umem::Umem,
    Vec<xdp_io::rx::Channel<Arc<AsyncFd<socket::Fd>>>>,
    Vec<(u32, socket::Fd)>,
    Vec<xdp_io::tx::Channel<xdp_io::tx::BusyPoll>>,
)>;

impl Xdp {
    fn setup(&self) -> SetupResult {
        let frame_size = self.frame_size;
        let rx_queue_len = self.rx_queue_len;
        let tx_queue_len = self.tx_queue_len;
        let fill_ring_len = rx_queue_len * 2;
        let completion_ring_len = tx_queue_len;

        let max_queues = syscall::max_queues(&self.interface);
        let umem_size = (rx_queue_len + tx_queue_len) * max_queues;

        // create a UMEM
        let umem = umem::Builder {
            frame_count: umem_size,
            frame_size,
            ..Default::default()
        }
        .build()?;

        // setup the address we're going to bind to
        let mut address = if_xdp::Address {
            flags: XdpFlags::USE_NEED_WAKEUP,
            ..Default::default()
        };
        address.set_if_name(&CString::new(self.interface.clone())?)?;

        let mut shared_umem_fd = None;
        let mut tx_channels = vec![];
        let mut rx_channels = vec![];
        let mut rx_fds = vec![];

        let mut desc = umem.frames();

        // iterate over all of the queues and create sockets for each one
        for queue_id in 0..max_queues {
            let socket = socket::Fd::open()?;

            // if we've already attached a socket to the UMEM, then reuse the first FD
            if let Some(fd) = shared_umem_fd {
                address.set_shared_umem(&fd);
            } else {
                socket.attach_umem(&umem)?;
                shared_umem_fd = Some(socket.as_raw_fd());
            }

            // set the queue id to the current value
            address.queue_id = queue_id;

            // file descriptors can only be added once so wrap it in an Arc
            let async_fd = Arc::new(AsyncFd::new(socket.clone())?);

            // get the offsets for each of the rings
            let offsets = syscall::offsets(&socket)?;

            {
                // create a pair of rings for receiving packets
                let mut fill = ring::Fill::new(socket.clone(), &offsets, fill_ring_len)?;
                let rx = ring::Rx::new(socket.clone(), &offsets, rx_queue_len)?;

                // remember the FD so we can add it to the XSK map later
                rx_fds.push((queue_id, socket.clone()));

                // put descriptors in the Fill queue
                fill.init((&mut desc).take(rx_queue_len as _));

                rx_channels.push(xdp_io::rx::Channel {
                    rx,
                    fill,
                    driver: async_fd.clone(),
                });
            };

            {
                // create a pair of rings for transmitting packets
                let mut completion =
                    ring::Completion::new(socket.clone(), &offsets, completion_ring_len)?;
                let tx = ring::Tx::new(socket.clone(), &offsets, tx_queue_len)?;

                // put descriptors in the completion queue
                completion.init((&mut desc).take(tx_queue_len as _));

                tx_channels.push(xdp_io::tx::Channel {
                    tx,
                    completion,
                    driver: xdp_io::tx::BusyPoll,
                });
            };

            // finally bind the socket to the configured address
            syscall::bind(&socket, &mut address)?;
        }

        // make sure we've allocated all descriptors from the UMEM to a queue
        assert_eq!(desc.count(), 0, "descriptors have been leaked");

        Ok((umem, rx_channels, rx_fds, tx_channels))
    }

    fn bpf_task(&self, port: u16, rx_fds: Vec<(u32, socket::Fd)>) -> Result<()> {
        // load the default BPF program from s2n-quic-xdp
        let mut bpf = if self.bpf_trace {
            let mut bpf = Bpf::load(bpf::DEFAULT_PROGRAM_TRACE)?;

            if let Err(err) = aya_log::BpfLogger::init(&mut bpf) {
                eprint!("error initializing BPF trace: {err:?}");
            }

            bpf
        } else {
            Bpf::load(bpf::DEFAULT_PROGRAM)?
        };

        let interface = self.interface.clone();
        let xdp_stats = self.xdp_stats;
        let xdp_mode = self.xdp_mode.into();

        let program: &mut programs::Xdp = bpf
            .program_mut(bpf::PROGRAM_NAME)
            .expect("missing default program")
            .try_into()?;
        program.load()?;

        // attach the BPF program to the interface
        let link_id = program.attach(&interface, xdp_mode)?;

        let bpf_task = async move {
            // register the port as active
            let mut ports: HashMap<_, _, _> = bpf
                .map_mut(bpf::PORT_MAP_NAME)
                .expect("missing port map")
                .try_into()?;

            // the BPF program just needs to have a non-zero value for the port
            let enabled = 1u8;
            // no flags are needed
            let flags = 0;
            ports.insert(port, enabled, flags)?;

            // register all of the RX sockets on each of the queues
            let mut xskmap: XskMap<&mut MapData> = bpf
                .map_mut(bpf::XSK_MAP_NAME)
                .expect("missing socket map")
                .try_into()?;

            for (queue_id, socket) in &rx_fds {
                xskmap.set(*queue_id, socket.as_raw_fd(), 0)?;
            }

            // print xdp stats every second, if configured
            if xdp_stats {
                loop {
                    tokio::time::sleep(core::time::Duration::from_secs(1)).await;
                    for (queue_id, socket) in &rx_fds {
                        if let Ok(stats) = syscall::statistics(socket) {
                            println!("stats[{queue_id}]: {stats:?}");
                        }
                    }
                }
            }

            // we want this future to go until the end of the program so we can keep the BPF
            // program active on the the NIC.
            core::future::pending::<()>().await;

            // retain the bpf program for the duration of execution
            let _ = bpf;
            let _ = link_id;
            let _ = rx_fds;

            Result::<(), crate::Error>::Ok(())
        };

        tokio::spawn(async move {
            if let Err(error) = bpf_task.await {
                panic!("BPF ERROR: {error}");
            }
        });

        Ok(())
    }

    fn udp_socket(&self, addr: SocketAddr) -> Result<std::net::UdpSocket> {
        let interface = CString::new(self.interface.clone())?;

        let udp_socket = socket::bind_udp(&interface, addr)?;

        Ok(udp_socket)
    }

    fn tx_encoder(&self) -> encoder::Config {
        let mut encoder = encoder::Config::default();
        if self.no_checksum {
            encoder.set_checksum(false);
        }
        encoder
    }

    fn spawn_udp_rx(&self, udp_socket: UdpSocket) {
        // Set up a task to read from the bound UDP socket
        //
        // If everything is working properly, this won't ever get a packet, since the AF_XDP socket
        // is intercepting all packets on this port. If it does start logging, something has gone
        // wrong with the BPF setup.
        let mut recv_buffer = vec![0; self.frame_size as usize];
        tokio::spawn(async move {
            let result = udp_socket.recv_from(&mut recv_buffer).await;
            eprintln!("ERROR: received packet on regular UDP socket: {:?}", result);

            std::process::exit(1);
        });
    }

    pub fn server(&self, addr: SocketAddr) -> Result<impl io::Provider> {
        let udp_socket = self.udp_socket(addr)?;
        let udp_socket = UdpSocket::from_std(udp_socket)?;

        let (umem, rx, rx_fds, tx) = self.setup()?;

        self.bpf_task(addr.port(), rx_fds)?;

        let io_rx = xdp_io::rx::Rx::new(rx, umem.clone());
        let io_tx = xdp_io::tx::Tx::new(tx, umem, self.tx_encoder());

        let provider = Provider::builder()
            .with_rx(io_rx)
            .with_tx(io_tx)
            .with_frame_size(self.frame_size as _)?
            .build();

        self.spawn_udp_rx(udp_socket);

        Ok(provider)
    }

    pub fn client(&self, addr: SocketAddr) -> Result<impl io::Provider> {
        let udp_socket = self.udp_socket(addr)?;
        // query the actual port that was selected for the socket in case it was `0`
        let addr = udp_socket.local_addr()?;

        let recv_udp_socket = udp_socket.try_clone();

        let (umem, rx, rx_fds, tx) = self.setup()?;

        self.bpf_task(addr.port(), rx_fds)?;

        let io_rx = xdp_io::rx::Rx::new(rx, umem.clone());

        let io_tx = {
            let tx = xdp_io::tx::Tx::new(tx, umem, self.tx_encoder());

            let udp_tx = {
                let (udp_tx, udp_task) = tx::channel(udp_socket);

                tokio::spawn(udp_task);

                udp_tx.with_handle_map(|handle: &io::xdp::PathHandle| {
                    // convert the XDP handle into the regular UDP handle
                    handle.into()
                })
            };

            // route any initial packets to the UDP socket so we can offload address resolution on
            // the OS
            tx.with_router(xdp_io::router::Router::default(), udp_tx)
        };

        let provider = Provider::builder()
            .with_rx(io_rx)
            .with_tx(io_tx)
            .with_frame_size(self.frame_size as _)?
            .build();

        if let Ok(udp_socket) = recv_udp_socket {
            let udp_socket = UdpSocket::from_std(udp_socket)?;
            self.spawn_udp_rx(udp_socket);
        }

        Ok(provider)
    }
}


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

use crate::Result;
use bytes::Bytes;
use futures::future::try_join_all;
use s2n_quic::{client::Connect, connection::Handle, stream::SendStream, Client};
use std::{
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};
use tokio::{fs::File, io::AsyncWriteExt, spawn};
use url::Url;

pub(crate) async fn create_connection<'a, R: IntoIterator<Item = &'a Url>>(
    client: Client,
    connect: Connect,
    requests: R,
    download_dir: Arc<Option<PathBuf>>,
    keep_alive: Option<Duration>,
) -> Result<()> {
    eprintln!("connecting to {connect:#}");
    let mut connection = client.connect(connect).await?;

    if keep_alive.is_some() {
        connection.keep_alive(true)?;
    }

    let mut streams = vec![];
    for request in requests {
        streams.push(spawn(create_stream(
            connection.handle(),
            request.path().to_string(),
            download_dir.clone(),
        )));
    }

    for result in try_join_all(streams).await? {
        // `try_join_all` should be returning an Err if any stream fails, but it
        // seems to just include the Err in the Vec of results. This will force
        // any Error to bubble up so it can be printed in the output.
        result?;
    }

    if let Some(keep_alive) = keep_alive {
        tokio::time::sleep(keep_alive).await;
        connection.keep_alive(false)?;
    }

    Ok(())
}

async fn create_stream(
    connection: Handle,
    request: String,
    download_dir: Arc<Option<PathBuf>>,
) -> Result<()> {
    eprintln!("GET {request}");

    match create_stream_inner(connection, &request, download_dir).await {
        Ok(()) => {
            eprintln!("Request {request} completed successfully");
            Ok(())
        }
        Err(error) => {
            eprintln!("Request {request} failed: {error:?}");
            Err(error)
        }
    }
}

async fn create_stream_inner(
    mut connection: Handle,
    request: &str,
    download_dir: Arc<Option<PathBuf>>,
) -> Result<()> {
    let stream = connection.open_bidirectional_stream().await?;
    let (mut rx_stream, tx_stream) = stream.split();

    write_request(tx_stream, request).await?;

    if let Some(download_dir) = download_dir.as_ref() {
        if download_dir == Path::new("/dev/null") {
            crate::perf::handle_receive_stream(rx_stream).await?;
        } else {
            let mut abs_path = download_dir.to_path_buf();
            abs_path.push(Path::new(request.trim_start_matches('/')));
            let mut file = File::create(&abs_path).await?;
            tokio::io::copy(&mut rx_stream, &mut file).await?;
            file.flush().await?;
        }
    } else {
        let mut stdout = tokio::io::stdout();
        tokio::io::copy(&mut rx_stream, &mut stdout).await?;
        stdout.flush().await?;
    };

    Ok(())
}

async fn write_request(mut stream: SendStream, request: &str) -> Result<()> {
    static GET: Bytes = Bytes::from_static(b"GET ");
    static END_OF_REQUEST: Bytes = Bytes::from_static(b"\r\n");

    stream
        .send_vectored(&mut [
            GET.clone(),
            Bytes::copy_from_slice(request.as_bytes()),
            END_OF_REQUEST.clone(),
        ])
        .await?;

    stream.finish()?;

    Ok(())
}


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

use crate::Result;
use bytes::Buf;
use futures::future::try_join_all;
use http::Uri;
use s2n_quic::{client::Connect, Client};
use s2n_quic_h3::h3;
use std::{
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};
use tokio::{fs::File, io::AsyncWriteExt, spawn};
use url::Url;

pub async fn create_connection<'a, R: IntoIterator<Item = &'a Url>>(
    client: Client,
    connect: Connect,
    requests: R,
    download_dir: Arc<Option<PathBuf>>,
    keep_alive: Option<Duration>,
) -> Result<()> {
    eprintln!("connecting to {connect:#}");
    let mut connection = client.connect(connect).await?;

    if keep_alive.is_some() {
        connection.keep_alive(true)?;
    }

    let (mut driver, send_request) =
        h3::client::new(s2n_quic_h3::Connection::new(connection)).await?;

    let mut streams = vec![];
    for request in requests {
        streams.push(spawn(create_stream(
            send_request.clone(),
            request.clone(),
            download_dir.clone(),
        )));
    }

    for result in try_join_all(streams).await? {
        // `try_join_all` should be returning an Err if any stream fails, but it
        // seems to just include the Err in the Vec of results. This will force
        // any Error to bubble up so it can be printed in the output.
        result?;
    }

    if let Some(keep_alive) = keep_alive {
        tokio::time::sleep(keep_alive).await;
    }

    driver.shutdown(0).await?;

    Ok(())
}

async fn create_stream<B: Buf, T: h3::quic::OpenStreams<B>>(
    send_request: h3::client::SendRequest<T, B>,
    request: Url,
    download_dir: Arc<Option<PathBuf>>,
) -> Result<()> {
    eprintln!("GET {request}");

    match create_stream_inner(send_request, request.clone(), download_dir).await {
        Ok(()) => {
            eprintln!("Request {request} completed successfully");
            Ok(())
        }
        Err(error) => {
            eprintln!("Request {request} failed: {error:?}");
            Err(error)
        }
    }
}

async fn create_stream_inner<B: Buf, T: h3::quic::OpenStreams<B>>(
    mut send_request: h3::client::SendRequest<T, B>,
    request: Url,
    download_dir: Arc<Option<PathBuf>>,
) -> Result<()> {
    let uri = request.to_string().parse::<Uri>().unwrap();
    let req = http::Request::builder().uri(uri).body(())?;
    let mut stream = send_request.send_request(req).await?;
    stream.finish().await?;

    let resp = stream.recv_response().await?;

    eprintln!("Response: {:?} {}", resp.version(), resp.status());
    eprintln!("Headers: {:#?}", resp.headers());

    if let Some(download_dir) = download_dir.as_ref() {
        if download_dir == Path::new("/dev/null") {
            while stream.recv_data().await?.is_some() {}
        } else {
            let mut abs_path = download_dir.to_path_buf();
            abs_path.push(Path::new(request.path().trim_start_matches('/')));
            let mut file = File::create(&abs_path).await?;
            while let Some(mut chunk) = stream.recv_data().await? {
                file.write_all_buf(&mut chunk).await?;
            }
            file.flush().await?;
        }
    };

    Ok(())
}


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

use crate::{
    client::{h09, h3},
    interop::Testcase,
    tls, Result,
};
use core::time::Duration;
use s2n_quic::{client::Connect, provider::event, Client};
use std::{
    collections::{hash_map::Entry, HashMap},
    path::PathBuf,
    sync::Arc,
};
use structopt::StructOpt;
use tokio::net::lookup_host;
use url::{Host, Url};

#[derive(Debug, StructOpt)]
pub struct Interop {
    #[structopt(short, long)]
    ip: Option<std::net::IpAddr>,

    #[structopt(short, long, default_value = "443")]
    port: u16,

    #[structopt(long, default_value = "hq-interop")]
    application_protocols: Vec<String>,

    #[structopt(long)]
    download_dir: Option<PathBuf>,

    #[structopt(long, parse(try_from_str = parse_duration))]
    keep_alive: Option<Duration>,

    #[structopt(long, env = "TESTCASE", possible_values = &Testcase::supported(is_supported_testcase))]
    testcase: Option<Testcase>,

    #[structopt(min_values = 1, required = true)]
    requests: Vec<Url>,

    #[structopt(flatten)]
    io: crate::io::Client,

    #[structopt(flatten)]
    tls: tls::Client,
}

impl Interop {
    pub async fn run(&self) -> Result<()> {
        let mut client = self.client()?;

        let download_dir = Arc::new(self.download_dir.clone());
        if self.requests.len() > 1 && download_dir.is_none() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "`--download-dir` must be specified if there is more than one request",
            )
            .into());
        }

        let endpoints = self.endpoints().await?;

        // https://github.com/marten-seemann/quic-interop-runner#test-cases
        // Handshake Loss (multiconnect): Tests resilience of the handshake to high loss.
        // The client is expected to establish multiple connections, sequential or in parallel,
        // and use each connection to download a single file.
        if let Some(Testcase::Multiconnect) = self.testcase {
            for request in &self.requests {
                let connect = endpoints.get(&request.host().unwrap()).unwrap().clone();
                let requests = core::iter::once(request);
                h09::create_connection(
                    client.clone(),
                    connect,
                    requests,
                    download_dir.clone(),
                    self.keep_alive,
                )
                .await?;
            }
        } else {
            // establish a connection per endpoint rather than per request
            for (host, connect) in endpoints.iter() {
                let host = host.clone();
                let connect = connect.clone();
                let requests = self
                    .requests
                    .iter()
                    .filter(|req| req.host().as_ref() == Some(&host))
                    .collect::<Vec<_>>();

                if let Some(Testcase::Http3) = self.testcase {
                    h3::create_connection(
                        client.clone(),
                        connect,
                        requests,
                        download_dir.clone(),
                        self.keep_alive,
                    )
                    .await?;
                } else {
                    h09::create_connection(
                        client.clone(),
                        connect,
                        requests,
                        download_dir.clone(),
                        self.keep_alive,
                    )
                    .await?;
                }
            }
        }

        client.wait_idle().await?;

        Ok(())
    }

    fn client(&self) -> Result<Client> {
        let io = self.io.build()?;

        let client = Client::builder()
            .with_io(io)?
            .with_event(event::tracing::Subscriber::default())?;

        let client = self.tls.build(client, &self.application_protocols)?;

        Ok(client)
    }

    async fn endpoints(&self) -> Result<HashMap<Host<&str>, Connect>> {
        let mut endpoints = HashMap::new();

        for req in &self.requests {
            if let Some(host) = req.host() {
                if let Entry::Vacant(entry) = endpoints.entry(host.clone()) {
                    let (ip, server_name) = match host {
                        Host::Domain(domain) => {
                            let ip = if let Some(ip) = self.ip {
                                ip
                            } else {
                                lookup_host((domain, self.port))
                                    .await?
                                    .next()
                                    .unwrap_or_else(|| {
                                        panic!("host {domain:?} did not resolve to any addresses")
                                    })
                                    .ip()
                            };

                            (ip, Some(domain))
                        }
                        Host::Ipv4(ip) => (ip.into(), None),
                        Host::Ipv6(ip) => (ip.into(), None),
                    };

                    let port = req.port().unwrap_or(self.port);

                    let connect = Connect::new((ip, port));

                    let connect = if let Some(server_name) = server_name {
                        connect.with_server_name(server_name)
                    } else {
                        // TODO make it optional
                        connect.with_server_name("localhost")
                    };

                    entry.insert(connect);
                }
            }
        }

        Ok(endpoints)
    }
}

fn is_supported_testcase(testcase: Testcase) -> bool {
    use Testcase::*;
    match testcase {
        // TODO add the ability to override the QUIC version
        VersionNegotiation => false,
        Handshake => true,
        Transfer => true,
        // TODO enable _only_ chacha20 on supported ciphersuites
        ChaCha20 => false,
        // TODO add the ability to trigger a key update from the application
        KeyUpdate => false,
        Retry => true,
        // TODO support storing tickets
        Resumption => false,
        // TODO implement 0rtt
        ZeroRtt => false,
        Http3 => true,
        Multiconnect => true,
        Ecn => true,
        // TODO support the ability to actively migrate on the client
        ConnectionMigration => false,
    }
}

fn parse_duration(duration: &str) -> Result<Duration> {
    let seconds = duration.parse()?;
    Ok(Duration::from_secs(seconds))
}


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

use crate::{perf, tls, Result};
use s2n_quic::{client, Client, Connection};
use structopt::StructOpt;
use tokio::task::JoinSet;

#[derive(Debug, StructOpt)]
pub struct Perf {
    #[structopt(short, long, default_value = "127.0.0.1")]
    ip: std::net::IpAddr,

    #[structopt(short, long, default_value = "443")]
    port: u16,

    #[structopt(short, long)]
    server_name: Option<String>,

    //= https://tools.ietf.org/id/draft-banks-quic-performance-00#2.1
    //# The ALPN used by the QUIC performance protocol is "perf".
    #[structopt(long, default_value = "perf")]
    application_protocols: Vec<String>,

    #[structopt(long)]
    connections: Option<usize>,

    #[structopt(long, default_value)]
    send: u64,

    #[structopt(long, default_value)]
    receive: u64,

    #[structopt(long, default_value = "1")]
    streams: u64,

    #[structopt(flatten)]
    limits: perf::Limits,

    /// Logs statistics for the endpoint
    #[structopt(long)]
    stats: bool,

    #[structopt(flatten)]
    io: crate::io::Client,

    #[structopt(flatten)]
    tls: tls::Client,
}

impl Perf {
    pub async fn run(&self) -> Result<()> {
        let mut client = self.client()?;

        let mut requests = JoinSet::new();

        // TODO support a richer connection strategy
        for _ in 0..self.connections.unwrap_or(1) {
            let mut connect = client::Connect::new((self.ip, self.port));
            if let Some(server_name) = self.server_name.as_deref() {
                connect = connect.with_server_name(server_name);
            } else {
                // TODO allow skipping setting the server_name
                connect = connect.with_server_name("localhost");
            }
            let connection = client.connect(connect).await?;

            requests.spawn(handle_connection(
                connection,
                self.streams,
                self.send,
                self.receive,
            ));
        }

        // wait until all of the connections finish before closing
        while requests.join_next().await.is_some() {}

        client.wait_idle().await?;

        return Ok(());

        async fn handle_connection(
            mut connection: Connection,
            streams: u64,
            send: u64,
            receive: u64,
        ) -> Result<()> {
            if send == 0 && receive == 0 {
                return Ok(());
            }

            for _ in 0..streams {
                let stream = connection.open_bidirectional_stream().await?;
                let (receive_stream, mut send_stream) = stream.split();

                let s = async move {
                    perf::write_stream_size(&mut send_stream, receive).await?;
                    perf::handle_send_stream(send_stream, send).await?;
                    <Result<()>>::Ok(())
                };

                let r = perf::handle_receive_stream(receive_stream);

                let _ = tokio::try_join!(s, r)?;
            }

            Ok(())
        }
    }

    fn client(&self) -> Result<Client> {
        let io = self.io.build()?;

        let subscriber = perf::Subscriber::default();

        if self.stats {
            subscriber.spawn(core::time::Duration::from_secs(1));
        }

        let subscriber = (
            subscriber,
            s2n_quic::provider::event::tracing::Subscriber::default(),
        );

        let client = Client::builder()
            .with_limits(self.limits.limits())?
            .with_io(io)?
            .with_event(subscriber)?;

        let client = self.tls.build(client, &self.application_protocols)?;

        Ok(client)
    }
}


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

use crate::{
    file::{abs_path, File},
    server::interop::MyConnectionContext,
    Result,
};
use bytes::Bytes;
use futures::StreamExt;
use s2n_quic::{
    stream::{BidirectionalStream, ReceiveStream, SendStream},
    Connection,
};
use std::{convert::TryInto, path::Path, sync::Arc, time::Duration};
use tokio::time::timeout;
use tracing::debug;

pub(crate) async fn handle_connection(mut connection: Connection, www_dir: Arc<Path>) {
    loop {
        match connection.accept_bidirectional_stream().await {
            Ok(Some(stream)) => {
                let _ = connection.query_event_context_mut(|context: &mut MyConnectionContext| {
                    context.stream_requests += 1
                });

                let www_dir = www_dir.clone();
                // spawn a task per stream
                tokio::spawn(async move {
                    if let Err(err) = handle_stream(stream, www_dir).await {
                        eprintln!("Stream error: {err:?}")
                    }
                });
            }
            Ok(None) => {
                // the connection was closed without an error
                let context = connection
                    .query_event_context(|context: &MyConnectionContext| *context)
                    .expect("query should execute");
                debug!("Final stats: {context:?}");
                return;
            }
            Err(err) => {
                eprintln!("error while accepting stream: {err}");
                let context = connection
                    .query_event_context(|context: &MyConnectionContext| *context)
                    .expect("query should execute");
                debug!("Final stats: {context:?}");
                return;
            }
        }
    }
}

async fn handle_stream(stream: BidirectionalStream, www_dir: Arc<Path>) -> Result<()> {
    let (rx_stream, mut tx_stream) = stream.split();
    let path = read_request(rx_stream).await?;

    if let Some(amount) = path.strip_prefix("_perf/").and_then(|v| v.parse().ok()) {
        return handle_perf_stream(amount, tx_stream).await;
    }

    let abs_path = abs_path(&path, &www_dir);
    let mut file = File::open(&abs_path).await?;
    loop {
        match timeout(Duration::from_secs(1), file.next()).await {
            Ok(Some(Ok(chunk))) => {
                let len = chunk.len();
                debug!(
                    "{:?} bytes ready to send on Stream({:?})",
                    len,
                    tx_stream.id()
                );
                tx_stream.send(chunk).await?;
                debug!("{:?} bytes sent on Stream({:?})", len, tx_stream.id());
            }
            Ok(Some(Err(err))) => {
                eprintln!("error opening {abs_path:?}");
                tx_stream.reset(1u32.try_into()?)?;
                return Err(err.into());
            }
            Ok(None) => {
                tx_stream.finish()?;
                return Ok(());
            }
            Err(_) => {
                eprintln!("timeout opening {abs_path:?}");
            }
        }
    }
}

async fn handle_perf_stream(amount: u64, mut stream: SendStream) -> Result<()> {
    let mut data = s2n_quic_core::stream::testing::Data::new(amount);

    while let Some(chunk) = data.send_one(usize::MAX) {
        stream.send(chunk).await?;
    }

    stream.finish()?;

    Ok(())
}

async fn read_request(mut stream: ReceiveStream) -> Result<String> {
    let mut path = String::new();
    let mut chunks = vec![Bytes::new(), Bytes::new()];
    let mut total_chunks = 0;
    loop {
        // grow the chunks
        if chunks.len() == total_chunks {
            chunks.push(Bytes::new());
        }
        let (consumed, is_open) = stream.receive_vectored(&mut chunks[total_chunks..]).await?;
        total_chunks += consumed;
        if parse_h09_request(&chunks[..total_chunks], &mut path, is_open)? {
            return Ok(path);
        }
    }
}

fn parse_h09_request(chunks: &[Bytes], path: &mut String, is_open: bool) -> Result<bool> {
    let mut bytes = chunks.iter().flat_map(|chunk| chunk.iter().cloned());

    macro_rules! expect {
        ($char:literal) => {
            match bytes.next() {
                Some($char) => {}
                None if is_open => return Ok(false),
                _ => return Err("invalid request".into()),
            }
        };
    }

    expect!(b'G');
    expect!(b'E');
    expect!(b'T');
    expect!(b' ');
    expect!(b'/');

    // reset the copied path in case this isn't the first time a path is being parsed
    path.clear();

    loop {
        match bytes.next() {
            Some(c @ b'0'..=b'9') => path.push(c as char),
            Some(c @ b'a'..=b'z') => path.push(c as char),
            Some(c @ b'A'..=b'Z') => path.push(c as char),
            Some(b'.') => path.push('.'),
            Some(b'/') => path.push('/'),
            Some(b'-') => path.push('-'),
            Some(b'_') => path.push('_'),
            Some(b'\n' | b'\r') => return Ok(true),
            // https://www.w3.org/Protocols/HTTP/AsImplemented.html
            // > The document address will consist of a single word (ie no spaces).
            // > If any further words are found on the request line, they MUST either be ignored,
            // > or else treated according to the full HTTP spec.
            Some(b' ') => return Ok(true),
            Some(c) => return Err(format!("invalid request {}", c as char).into()),
            None => return Ok(!is_open),
        }
    }
}

#[test]
fn parse_h09_request_test() {
    macro_rules! test {
        ([$($chunk:expr),* $(,)?], $expected:pat) => {{
            let chunks = [$(Bytes::from_static($chunk.as_bytes())),*];
            let mut path = String::new();

            for idx in 0..chunks.len() {
                let _ = parse_h09_request(&chunks[..idx], &mut path, true);
            }

            let result = parse_h09_request(&chunks, &mut path, false);
            let result = result.map(|has_request| if has_request { Some(path) } else { None });
            let result = result.as_ref().map(|v| v.as_deref());

            assert!(matches!(result, $expected), "{:?}", result);
        }}
    }

    test!([], Err(_));
    test!(["GET /"], Ok(Some("")));
    test!(["GET /abc"], Ok(Some("abc")));
    test!(["GET /abc/123"], Ok(Some("abc/123")));
    test!(["GET /CAPS/lower"], Ok(Some("CAPS/lower")));
    test!(["GET /abc\rextra stuff"], Ok(Some("abc")));
    test!(
        ["G", "E", "T", " ", "/", "t", "E", "s", "T"],
        Ok(Some("tEsT"))
    );
}


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

use crate::{
    file::{abs_path, File},
    Result,
};
use bytes::Bytes;
use futures::StreamExt;
use h3::{quic::BidiStream, server::RequestStream};
use http::StatusCode;
use s2n_quic::Connection;
use s2n_quic_h3::h3;
use std::{path::Path, sync::Arc, time::Duration};
use tokio::time::timeout;

pub async fn handle_connection(connection: Connection, www_dir: Arc<Path>) {
    let mut conn = h3::server::Connection::new(s2n_quic_h3::Connection::new(connection))
        .await
        .unwrap();

    while let Ok(Some((req, stream))) = conn.accept().await {
        if let Some(amount) = req
            .uri()
            .path()
            .strip_prefix("/_perf/")
            .and_then(|v| v.parse().ok())
        {
            tokio::spawn(async move {
                if let Err(err) = handle_perf_stream(amount, stream).await {
                    eprintln!("Stream error: {err:?}");
                }
            });
            continue;
        }

        let www_dir = www_dir.clone();
        tokio::spawn(async {
            if let Err(err) = handle_stream(req, stream, www_dir).await {
                eprintln!("Stream error: {err:?}")
            }
        });
    }
}

async fn handle_stream<T>(
    req: http::Request<()>,
    mut stream: RequestStream<T, Bytes>,
    www_dir: Arc<Path>,
) -> Result<()>
where
    T: BidiStream<Bytes>,
{
    let abs_path = abs_path(req.uri().path(), &www_dir);
    let mut file = File::open(&abs_path).await?;
    let resp = http::Response::builder().status(StatusCode::OK).body(())?;

    stream.send_response(resp).await?;

    loop {
        match timeout(Duration::from_secs(1), file.next()).await {
            Ok(Some(Ok(chunk))) => {
                stream.send_data(chunk).await?;
            }
            Ok(Some(Err(err))) => {
                eprintln!("error opening {abs_path:?}");
                stream
                    .send_response(
                        http::Response::builder()
                            .status(StatusCode::INTERNAL_SERVER_ERROR)
                            .body(())?,
                    )
                    .await?;
                return Err(err.into());
            }
            Ok(None) => {
                stream.finish().await?;
                return Ok(());
            }
            Err(_) => {
                eprintln!("timeout opening {abs_path:?}");
            }
        }
    }
}

async fn handle_perf_stream<T>(amount: u64, mut stream: RequestStream<T, Bytes>) -> Result<()>
where
    T: BidiStream<Bytes>,
{
    let mut data = s2n_quic_core::stream::testing::Data::new(amount);
    let resp = http::Response::builder().status(StatusCode::OK).body(())?;

    stream.send_response(resp).await?;

    while let Some(chunk) = data.send_one(usize::MAX) {
        stream.send_data(chunk).await?;
    }

    stream.finish().await?;

    Ok(())
}


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

use crate::{
    interop::Testcase,
    server::{h09, h3},
    tls, Result,
};
use s2n_quic::{
    provider::{
        endpoint_limits,
        event::{events, Subscriber},
    },
    Server,
};
use std::{
    path::{Path, PathBuf},
    sync::Arc,
};
use structopt::StructOpt;
use tokio::spawn;

#[derive(Debug, StructOpt)]
pub struct Interop {
    #[structopt(long, default_value = "hq-interop")]
    application_protocols: Vec<String>,

    #[structopt(long, default_value = ".")]
    www_dir: PathBuf,

    #[structopt(long, env = "TESTCASE", possible_values = &Testcase::supported(is_supported_testcase))]
    testcase: Option<Testcase>,

    #[structopt(flatten)]
    tls: tls::Server,

    #[structopt(flatten)]
    io: crate::io::Server,
}

impl Interop {
    pub async fn run(&self) -> Result<()> {
        let mut server = self.server()?;

        let www_dir: Arc<Path> = Arc::from(self.www_dir.as_path());

        while let Some(connection) = server.accept().await {
            let unspecified: std::net::SocketAddr = ([0, 0, 0, 0], 0).into();
            println!(
                "Accepted a QUIC connection from {} on {}",
                connection.remote_addr().unwrap_or(unspecified),
                connection.local_addr().unwrap_or(unspecified)
            );

            // spawn a task per connection
            match &(connection.application_protocol()?)[..] {
                b"h3" => spawn(h3::handle_connection(connection, www_dir.clone())),
                b"hq-interop" => spawn(h09::handle_connection(connection, www_dir.clone())),
                _ => spawn(async move {
                    eprintln!(
                        "Unsupported application protocol: {:?}",
                        connection.application_protocol()
                    );
                }),
            };
        }

        Err(crate::CRASH_ERROR_MESSAGE.into())
    }

    fn server(&self) -> Result<Server> {
        let mut max_handshakes = 100;
        if let Some(Testcase::Retry) = self.testcase {
            max_handshakes = 0;
        }

        let limits = endpoint_limits::Default::builder()
            .with_inflight_handshake_limit(max_handshakes)?
            .build()?;

        let io = self.io.build()?;

        let server = Server::builder()
            .with_io(io)?
            .with_endpoint_limits(limits)?
            .with_event((
                EventSubscriber(1),
                s2n_quic::provider::event::tracing::Subscriber::default(),
            ))?;
        let server = self.tls.build(server, &self.application_protocols)?;

        eprintln!("Server listening on port {}", self.io.port);

        Ok(server)
    }
}

fn is_supported_testcase(testcase: Testcase) -> bool {
    use Testcase::*;
    match testcase {
        VersionNegotiation => true,
        Handshake => true,
        Transfer => true,
        ChaCha20 => true,
        // KeyUpdate is client only
        KeyUpdate => false,
        Retry => true,
        // TODO support issuing tickets
        Resumption => false,
        // TODO implement 0rtt
        ZeroRtt => false,
        Http3 => true,
        Multiconnect => true,
        Ecn => true,
        ConnectionMigration => true,
    }
}

#[derive(Debug, Clone, Copy)]
pub struct MyConnectionContext {
    packet_sent: u64,
    pub(crate) stream_requests: u64,
}

pub struct EventSubscriber(usize);

impl Subscriber for EventSubscriber {
    type ConnectionContext = MyConnectionContext;

    fn create_connection_context(
        &mut self,
        _meta: &events::ConnectionMeta,
        _info: &events::ConnectionInfo,
    ) -> Self::ConnectionContext {
        MyConnectionContext {
            packet_sent: 0,
            stream_requests: 0,
        }
    }

    fn on_packet_sent(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &events::ConnectionMeta,
        _event: &events::PacketSent,
    ) {
        context.packet_sent += 1;
    }
}


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

use crate::{perf, tls, Result};
use futures::future::try_join_all;
use s2n_quic::{
    stream::{BidirectionalStream, ReceiveStream, SendStream},
    Connection, Server,
};
use structopt::StructOpt;
use tokio::spawn;

#[derive(Debug, StructOpt)]
pub struct Perf {
    //= https://tools.ietf.org/id/draft-banks-quic-performance-00#2.1
    //# The ALPN used by the QUIC performance protocol is "perf".
    #[structopt(long, default_value = "perf")]
    application_protocols: Vec<String>,

    #[structopt(long)]
    connections: Option<usize>,

    #[structopt(flatten)]
    limits: perf::Limits,

    /// Logs statistics for the endpoint
    #[structopt(long)]
    stats: bool,

    #[structopt(flatten)]
    tls: tls::Server,

    #[structopt(flatten)]
    io: crate::io::Server,
}

impl Perf {
    pub async fn run(&self) -> Result<()> {
        let mut server = self.server()?;

        if let Some(limit) = self.connections {
            let mut connections = vec![];

            while connections.len() < limit {
                if let Some(connection) = server.accept().await {
                    // spawn a task per connection
                    connections.push(spawn(handle_connection(connection)));
                } else {
                    break;
                }
            }

            let did_panic = connections.len() != limit;

            try_join_all(connections).await?;

            println!("closing server after {limit} connections");

            if did_panic {
                return Err(crate::CRASH_ERROR_MESSAGE.into());
            }

            return Ok(());
        } else {
            while let Some(connection) = server.accept().await {
                // spawn a task per connection
                spawn(handle_connection(connection));
            }

            return Err(crate::CRASH_ERROR_MESSAGE.into());
        }

        async fn handle_connection(connection: Connection) {
            let (mut handle, acceptor) = connection.split();
            let (mut bidi, mut uni) = acceptor.split();

            let bidi = tokio::spawn(async move {
                loop {
                    match bidi.accept_bidirectional_stream().await? {
                        Some(stream) => {
                            // spawn a task per stream
                            tokio::spawn(async move {
                                //= https://tools.ietf.org/id/draft-banks-quic-performance-00#2.3.1
                                //# On the server side, any stream that is closed before all 8 bytes are
                                //# received should just be ignored, and gracefully closed on its end (if
                                //# applicable).
                                let _ = handle_bidi_stream(stream).await;
                            });
                        }
                        None => {
                            // the connection was closed without an error
                            return <Result<()>>::Ok(());
                        }
                    }
                }
            });

            let uni = tokio::spawn(async move {
                loop {
                    match uni.accept_receive_stream().await? {
                        Some(receiver) => {
                            let sender = handle.open_send_stream().await?;
                            // spawn a task per stream
                            tokio::spawn(async move {
                                //= https://tools.ietf.org/id/draft-banks-quic-performance-00#2.3.1
                                //# On the server side, any stream that is closed before all 8 bytes are
                                //# received should just be ignored, and gracefully closed on its end (if
                                //# applicable).
                                let _ = handle_uni_stream(receiver, sender).await;
                            });
                        }
                        None => {
                            // the connection was closed without an error
                            return <Result<()>>::Ok(());
                        }
                    }
                }
            });

            let _ = futures::try_join!(bidi, uni);
        }

        //= https://tools.ietf.org/id/draft-banks-quic-performance-00#2.3.2
        //# When a client uses a bidirectional stream to request a response
        //# payload from the server, the server sends the requested data on the
        //# same stream.  If no data is requested by the client, the server
        //# merely closes its side of the stream.
        async fn handle_bidi_stream(stream: BidirectionalStream) -> Result<()> {
            let (mut receiver, sender) = stream.split();
            let (size, _prelude) = perf::read_stream_size(&mut receiver).await?;

            let receiver = tokio::spawn(async move { perf::handle_receive_stream(receiver).await });
            let sender = tokio::spawn(async move { perf::handle_send_stream(sender, size).await });

            let _ = futures::try_join!(receiver, sender);

            Ok(())
        }

        //= https://tools.ietf.org/id/draft-banks-quic-performance-00#2.3.2
        //# When a client uses a unidirectional stream to request a response
        //# payload from the server, the server opens a new unidirectional stream
        //# to send the requested data.  If no data is requested by the client,
        //# the server need take no action.
        async fn handle_uni_stream(mut receiver: ReceiveStream, sender: SendStream) -> Result<()> {
            let (size, _prelude) = perf::read_stream_size(&mut receiver).await?;

            let receiver = tokio::spawn(async move { perf::handle_receive_stream(receiver).await });
            let sender = tokio::spawn(async move { perf::handle_send_stream(sender, size).await });

            let _ = futures::try_join!(receiver, sender);

            Ok(())
        }
    }

    fn server(&self) -> Result<Server> {
        let io = self.io.build()?;

        let subscriber = perf::Subscriber::default();

        if self.stats {
            subscriber.spawn(core::time::Duration::from_secs(1));
        }

        let subscriber = (
            subscriber,
            s2n_quic::provider::event::tracing::Subscriber::default(),
        );

        let server = Server::builder()
            .with_limits(self.limits.limits())?
            .with_io(io)?
            .with_event(subscriber)?;

        let server = self.tls.build(server, &self.application_protocols)?;

        eprintln!("Server listening on port {}", self.io.port);

        Ok(server)
    }
}


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

use core::task::Poll;
use s2n_quic::provider::tls::s2n_tls::{ClientHelloHandler, Connection};

pub struct MyClientHelloHandler {}

impl ClientHelloHandler for MyClientHelloHandler {
    fn poll_client_hello(&self, connection: &mut Connection) -> core::task::Poll<Result<(), ()>> {
        Poll::Ready(Ok(()))
    }
}


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

#![allow(dead_code)]

use rustls::Error;

macro_rules! cert_type {
    ($name:ident, $trait:ident, $method:ident, $inner:ty) => {
        pub struct $name(pub(crate) $inner);

        pub trait $trait {
            fn $method(self) -> Result<$name, Error>;
        }

        impl $trait for $name {
            fn $method(self) -> Result<$name, Error> {
                Ok(self)
            }
        }

        impl $trait for String {
            fn $method(self) -> Result<$name, Error> {
                let cert = pem::$method(self.as_bytes())?;
                Ok($name(cert))
            }
        }

        impl $trait for &String {
            fn $method(self) -> Result<$name, Error> {
                let cert = pem::$method(self.as_bytes())?;
                Ok($name(cert))
            }
        }

        impl $trait for &str {
            fn $method(self) -> Result<$name, Error> {
                let cert = pem::$method(self.as_bytes())?;
                Ok($name(cert))
            }
        }

        impl $trait for Vec<u8> {
            fn $method(self) -> Result<$name, Error> {
                let cert = der::$method(self)?;
                Ok($name(cert))
            }
        }

        impl $trait for &[u8] {
            fn $method(self) -> Result<$name, Error> {
                self.to_vec().$method()
            }
        }

        impl $trait for &std::path::Path {
            fn $method(self) -> Result<$name, Error> {
                match self.extension() {
                    Some(ext) if ext == "der" => {
                        let pem =
                            std::fs::read(self).map_err(|err| Error::General(err.to_string()))?;
                        pem.$method()
                    }
                    _ => {
                        let pem = std::fs::read_to_string(self)
                            .map_err(|err| Error::General(err.to_string()))?;
                        pem.$method()
                    }
                }
            }
        }
    };
}

cert_type!(
    PrivateKey,
    IntoPrivateKey,
    into_private_key,
    rustls::PrivateKey
);
cert_type!(
    Certificate,
    IntoCertificate,
    into_certificate,
    Vec<rustls::Certificate>
);

mod pem {
    use super::*;

    pub fn into_certificate(contents: &[u8]) -> Result<Vec<rustls::Certificate>, Error> {
        let mut cursor = std::io::Cursor::new(contents);
        let certs = rustls_pemfile::certs(&mut cursor)
            .map(|certs| certs.into_iter().map(rustls::Certificate).collect())
            .map_err(|_| Error::General("Could not read certificate".to_string()))?;
        Ok(certs)
    }

    pub fn into_private_key(contents: &[u8]) -> Result<rustls::PrivateKey, Error> {
        let mut cursor = std::io::Cursor::new(contents);

        let parsers = [
            rustls_pemfile::rsa_private_keys,
            rustls_pemfile::pkcs8_private_keys,
        ];

        for parser in parsers.iter() {
            cursor.set_position(0);

            match parser(&mut cursor) {
                Ok(keys) if keys.is_empty() => continue,
                Ok(mut keys) if keys.len() == 1 => {
                    return Ok(rustls::PrivateKey(keys.pop().unwrap()))
                }
                Ok(keys) => {
                    return Err(Error::General(format!(
                        "Unexpected number of keys: {} (only 1 supported)",
                        keys.len()
                    )));
                }
                // try the next parser
                Err(_) => continue,
            }
        }

        Err(Error::General(
            "could not load any valid private keys".to_string(),
        ))
    }
}

mod der {
    use super::*;

    pub fn into_certificate(contents: Vec<u8>) -> Result<Vec<rustls::Certificate>, Error> {
        // der files only have a single cert
        Ok(vec![rustls::Certificate(contents)])
    }

    pub fn into_private_key(contents: Vec<u8>) -> Result<rustls::PrivateKey, Error> {
        Ok(rustls::PrivateKey(contents))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use s2n_quic_core::crypto::tls::testing::certificates::*;

    #[test]
    fn load() {
        let _ = CERT_PEM.into_certificate().unwrap();
        let _ = CERT_DER.into_certificate().unwrap();

        let _ = KEY_PEM.into_private_key().unwrap();
        let _ = KEY_DER.into_private_key().unwrap();
    }
}