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

use super::{events, CliRange};
use s2n_quic::{
    client::Connect,
    provider::{
        event::tracing::Subscriber as Tracing,
        io::testing::{primary, rand, spawn, time, Handle, Result},
    },
    Client, Server,
};
use s2n_quic_core::{crypto::tls::testing::certificates, stream::testing::Data};
use std::net::SocketAddr;

pub fn server(handle: &Handle, events: events::Events) -> Result<SocketAddr> {
    let mut server = Server::builder()
        .with_io(handle.builder().build().unwrap())?
        .with_tls((certificates::CERT_PEM, certificates::KEY_PEM))?
        .with_event((events, Tracing::default()))?
        .start()?;
    let server_addr = server.local_addr()?;

    // accept connections and echo back
    spawn(async move {
        while let Some(mut connection) = server.accept().await {
            primary::spawn(async move {
                while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                    primary::spawn(async move {
                        while let Ok(Some(chunk)) = stream.receive().await {
                            let _ = chunk;
                        }
                    });
                }
            });
        }
    });

    Ok(server_addr)
}

pub fn client(
    handle: &Handle,
    events: events::Events,
    servers: &[SocketAddr],
    count: usize,
    delay: CliRange<humantime::Duration>,
    streams: CliRange<u32>,
    stream_data: CliRange<u64>,
) -> Result {
    let client = Client::builder()
        .with_io(handle.builder().build().unwrap())?
        .with_tls(certificates::CERT_PEM)?
        .with_event((events, Tracing::default()))?
        .start()?;

    let mut total_delay = core::time::Duration::ZERO;

    for _ in 0..count {
        total_delay += delay.gen_duration();

        // pick a random server to connect to
        let server_addr = *rand::one_of(servers);
        let delay = total_delay;

        let client = client.clone();
        primary::spawn(async move {
            if !delay.is_zero() {
                time::delay(delay).await;
            }

            let connect = Connect::new(server_addr).with_server_name("localhost");
            let mut connection = client.connect(connect).await?;

            for _ in 0..streams.gen() {
                let stream = connection.open_bidirectional_stream().await?;
                primary::spawn(async move {
                    let (mut recv, mut send) = stream.split();

                    let mut send_data = Data::new(stream_data.gen());

                    let mut recv_data = send_data;
                    primary::spawn(async move {
                        while let Some(chunk) = recv.receive().await? {
                            recv_data.receive(&[chunk]);
                        }

                        <s2n_quic::stream::Result<()>>::Ok(())
                    });

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

                    <s2n_quic::stream::Result<()>>::Ok(())
                })
                .await?;
            }

            <s2n_quic::stream::Result<()>>::Ok(())
        });
    }

    Ok(())
}


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

use crate::stats::{Connection, Parameters};
use core::time::Duration;
use once_cell::sync::Lazy;
use s2n_quic::{
    connection,
    provider::{
        event,
        io::testing::{primary, time},
    },
};
use std::sync::{
    atomic::{AtomicBool, AtomicU64, Ordering},
    Arc,
};

static IDS: Lazy<Arc<AtomicU64>> = Lazy::new(Default::default);

static IS_OPEN: AtomicBool = AtomicBool::new(true);

pub fn is_open() -> bool {
    IS_OPEN.load(Ordering::Relaxed)
}

pub fn close() {
    IS_OPEN.store(false, Ordering::Relaxed);
}

#[derive(Clone, Debug)]
pub struct Events {
    params: Arc<DumpOnDrop<Parameters>>,
}

fn now() -> Duration {
    unsafe { time::now().as_duration() }
}

pub struct PrimaryContext<Inner> {
    #[allow(dead_code)]
    guard: primary::Guard,
    inner: Inner,
}

impl<Inner> core::ops::Deref for PrimaryContext<Inner> {
    type Target = Inner;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl<Inner> core::ops::DerefMut for PrimaryContext<Inner> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

trait OrDefault<T> {
    fn or_default(&mut self) -> &mut T;
}

impl<T: Default> OrDefault<T> for Option<T> {
    fn or_default(&mut self) -> &mut T {
        if self.is_none() {
            *self = Some(T::default());
        }

        self.as_mut().unwrap()
    }
}

impl event::Subscriber for Events {
    type ConnectionContext = PrimaryContext<DumpOnDrop<Connection>>;

    fn create_connection_context(
        &mut self,
        meta: &event::ConnectionMeta,
        _info: &event::ConnectionInfo,
    ) -> Self::ConnectionContext {
        let seed = self.params.seed;
        let id = IDS.fetch_add(1, Ordering::Relaxed);

        let mut conn = Connection {
            seed,
            start_time: Some(now().into()),
            ..Default::default()
        };

        if matches!(
            meta.endpoint_type,
            event::events::EndpointType::Server { .. }
        ) {
            conn.server_id = Some(id);
        } else {
            conn.client_id = Some(id);
        }

        let inner = DumpOnDrop(conn);
        PrimaryContext {
            guard: primary::guard(),
            inner,
        }
    }

    #[inline]
    fn on_packet_sent(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::PacketSent,
    ) {
        context.tx.or_default().inc_packet(&event.packet_header);
    }

    #[inline]
    fn on_packet_received(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::PacketReceived,
    ) {
        context.rx.or_default().inc_packet(&event.packet_header);
    }

    #[inline]
    fn on_packet_lost(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::PacketLost,
    ) {
        context.loss.or_default().inc_packet(&event.packet_header);
    }

    #[inline]
    fn on_congestion(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        _event: &event::events::Congestion,
    ) {
        context.congestion += 1;
    }

    #[inline]
    fn on_recovery_metrics(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::RecoveryMetrics,
    ) {
        context.max_cwin = context.max_cwin.max(event.congestion_window as _);
        context.max_bytes_in_flight = context.max_bytes_in_flight.max(event.bytes_in_flight as _);
        context.max_rtt = Some(
            context
                .max_rtt
                .unwrap_or_default()
                .as_duration()
                .max(event.smoothed_rtt)
                .into(),
        );
        context.min_rtt = Some(
            context
                .min_rtt
                .map_or(event.min_rtt, |prev| prev.as_duration().min(event.min_rtt))
                .into(),
        );
        context.smoothed_rtt = Some(event.smoothed_rtt.into());
    }

    #[inline]
    fn on_handshake_status_updated(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::HandshakeStatusUpdated,
    ) {
        match event.status {
            event::events::HandshakeStatus::Complete { .. } => {
                context.handshake.or_default().complete = Some(now().into())
            }
            event::events::HandshakeStatus::Confirmed { .. } => {
                context.handshake.or_default().confirmed = Some(now().into())
            }
            _ => {}
        }
    }

    #[inline]
    fn on_frame_sent(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::FrameSent,
    ) {
        context.tx.or_default().inc_frame(&event.frame);
    }

    #[inline]
    fn on_frame_received(
        &mut self,
        context: &mut Self::ConnectionContext,
        _meta: &event::ConnectionMeta,
        event: &event::events::FrameReceived,
    ) {
        context.rx.or_default().inc_frame(&event.frame);
    }

    #[inline]
    fn on_rx_stream_progress(
        &mut self,
        context: &mut Self::ConnectionContext,
        meta: &event::ConnectionMeta,
        event: &event::events::RxStreamProgress,
    ) {
        context
            .rx
            .or_default()
            .stream_progress(meta.timestamp, event.bytes);
    }

    #[inline]
    fn on_tx_stream_progress(
        &mut self,
        context: &mut Self::ConnectionContext,
        meta: &event::ConnectionMeta,
        event: &event::events::TxStreamProgress,
    ) {
        context
            .tx
            .or_default()
            .stream_progress(meta.timestamp, event.bytes);
    }

    #[inline]
    fn on_connection_closed(
        &mut self,
        context: &mut Self::ConnectionContext,
        meta: &event::ConnectionMeta,
        event: &event::events::ConnectionClosed,
    ) {
        context.end_time = Some(meta.timestamp.duration_since_start().into());

        match event.error {
            connection::Error::Closed { .. } => {}
            connection::Error::Transport { code, .. } => {
                context.transport_error = Some(code.as_u64());
            }
            connection::Error::Application { error, .. } => {
                let error = *error;
                // if the application closed it with `0` then it's no error
                if error > 0 {
                    context.application_error = Some(error);
                }
            }
            connection::Error::IdleTimerExpired { .. } => context.idle_timer_error = true,
            connection::Error::MaxHandshakeDurationExceeded { .. } => {
                context.handshake_duration_exceeded_error = true
            }
            _ => context.unspecified_error = true,
        }
    }
}

impl Dump for Connection {
    fn dump(&mut self) {
        if self.seed == 0 {
            return;
        }

        if self.end_time.is_none() {
            self.end_time = Some(now().into());
        }

        dump(|io| self.write(io));
    }
}

impl From<Parameters> for Events {
    fn from(s: Parameters) -> Self {
        Self {
            params: Arc::new(DumpOnDrop(s)),
        }
    }
}

impl Dump for Parameters {
    fn dump(&mut self) {
        if self.seed == 0 {
            return;
        }

        if self.end_time.is_none() {
            self.end_time = Some(now().into());
        }

        dump(|io| self.write(io));
    }
}

pub trait Dump {
    fn dump(&mut self);
}

#[derive(Debug)]
pub struct DumpOnDrop<T: Dump>(T);

impl<T: Dump> core::ops::Deref for DumpOnDrop<T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T: Dump> core::ops::DerefMut for DumpOnDrop<T> {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<T: Dump> Drop for DumpOnDrop<T> {
    fn drop(&mut self) {
        self.0.dump();
    }
}

pub fn dump<F: FnOnce(&mut std::io::StdoutLock) -> std::io::Result<()>>(f: F) {
    if !IS_OPEN.load(Ordering::Relaxed) {
        return;
    }

    let stdout = std::io::stdout();
    let mut stdout = stdout.lock();

    let res = f(&mut stdout);

    if res.is_err() {
        // close the process as the reader is no longer interested
        IS_OPEN.store(false, Ordering::Relaxed);
    }
}


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

use core::{fmt, marker::PhantomData, str::FromStr, time::Duration};
use s2n_quic::provider::io::testing::rand;
use serde::Deserialize;

#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct CliRange<T> {
    pub start: T,
    pub end: T,
}

impl Default for CliRange<f64> {
    fn default() -> Self {
        Self {
            start: 0.0,
            end: 0.0,
        }
    }
}

impl Default for CliRange<u64> {
    fn default() -> Self {
        Self { start: 0, end: 0 }
    }
}

impl Default for CliRange<humantime::Duration> {
    fn default() -> Self {
        Self {
            start: Duration::ZERO.into(),
            end: Duration::ZERO.into(),
        }
    }
}

impl<T: PartialEq + fmt::Display> fmt::Display for CliRange<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.start == self.end {
            self.start.fmt(f)
        } else {
            write!(f, "{}..{}", self.start, self.end)
        }
    }
}

impl<T> CliRange<T>
where
    T: Copy + PartialOrd + ::rand::distributions::uniform::SampleUniform,
{
    pub fn gen(&self) -> T {
        if self.start == self.end {
            return self.start;
        }

        rand::gen_range(self.start..self.end)
    }
}

impl CliRange<humantime::Duration> {
    pub fn gen_duration(&self) -> Duration {
        let start = self.start.as_nanos();
        let end = self.end.as_nanos();

        if start == end {
            return Duration::from_nanos(start as _);
        }

        let nanos = rand::gen_range(start..end);
        Duration::from_nanos(nanos as _)
    }
}

impl<T: Copy + FromStr> FromStr for CliRange<T> {
    type Err = T::Err;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Some((start, end)) = s.split_once("..") {
            let start = start.parse()?;
            let end = end.parse()?;
            Ok(Self { start, end })
        } else {
            let start = s.parse()?;
            let end = start;
            Ok(Self { start, end })
        }
    }
}

impl<'de, T> Deserialize<'de> for CliRange<T>
where
    T: Copy + FromStr,
    <T as FromStr>::Err: core::fmt::Display,
{
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        deserializer.deserialize_any(Visitor::<T>(PhantomData))
    }
}

struct Visitor<T>(PhantomData<T>);

impl<'de, T> serde::de::Visitor<'de> for Visitor<T>
where
    T: Copy + FromStr,
    <T as FromStr>::Err: core::fmt::Display,
{
    type Value = CliRange<T>;

    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(formatter, "a range or individual value")
    }

    fn visit_i64<E>(self, v: i64) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        v.to_string().parse().map_err(E::custom)
    }

    fn visit_u64<E>(self, v: u64) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        v.to_string().parse().map_err(E::custom)
    }

    fn visit_f64<E>(self, v: f64) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        v.to_string().parse().map_err(E::custom)
    }

    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
    where
        E: serde::de::Error,
    {
        v.parse().map_err(E::custom)
    }
}


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

use bytes::{Bytes, BytesMut};
use core::{ffi::c_void, marker::PhantomData};
use s2n_quic_core::{
    application::ServerName,
    crypto::{tls, CryptoError, CryptoSuite},
    endpoint, transport,
};
use s2n_quic_crypto::{
    handshake::HandshakeKey, hkdf, one_rtt::OneRttKey, ring_aead as aead, Prk, SecretPair, Suite,
};
use s2n_tls::{connection::Connection, error::Fallible, ffi::*};

/// The preallocated size of the outgoing buffer
///
/// By allocating a larger buffer, we only do a few allocations, even when
/// s2n-tls sends small chunks
const SEND_BUFFER_CAPACITY: usize = 2048;

/// Handles all callback contexts for each session
pub struct Callback<'a, T, C> {
    pub context: &'a mut T,
    pub endpoint: endpoint::Type,
    pub state: &'a mut State,
    pub suite: PhantomData<C>,
    pub err: Option<transport::Error>,
    pub send_buffer: &'a mut BytesMut,
    pub emitted_server_name: &'a mut bool,
    pub server_name: &'a Option<ServerName>,
}

impl<'a, T, C> Callback<'a, T, C>
where
    T: 'a + tls::Context<C>,
    C: CryptoSuite<
        HandshakeKey = <Suite as CryptoSuite>::HandshakeKey,
        HandshakeHeaderKey = <Suite as CryptoSuite>::HandshakeHeaderKey,
        InitialKey = <Suite as CryptoSuite>::InitialKey,
        InitialHeaderKey = <Suite as CryptoSuite>::InitialHeaderKey,
        OneRttKey = <Suite as CryptoSuite>::OneRttKey,
        OneRttHeaderKey = <Suite as CryptoSuite>::OneRttHeaderKey,
        ZeroRttKey = <Suite as CryptoSuite>::ZeroRttKey,
        ZeroRttHeaderKey = <Suite as CryptoSuite>::ZeroRttHeaderKey,
        RetryKey = <Suite as CryptoSuite>::RetryKey,
    >,
{
    /// Initializes the s2n-tls connection with all of the contexts and callbacks
    ///
    /// # Safety
    ///
    /// * The Callback struct must live at least as long as the connection
    /// * or the `unset` method should be called if it doesn't
    pub unsafe fn set(&mut self, connection: &mut Connection) {
        let context = self as *mut Self as *mut c_void;

        // We use unwrap here since s2n-tls will just check if connection is not null
        connection
            .set_secret_callback(Some(Self::secret_cb), context)
            .unwrap();
        connection.set_send_callback(Some(Self::send_cb)).unwrap();
        connection.set_send_context(context).unwrap();
        connection
            .set_receive_callback(Some(Self::recv_cb))
            .unwrap();
        connection.set_receive_context(context).unwrap();
        // A Waker is provided for use with the client hello callback.
        connection.set_waker(Some(self.context.waker())).unwrap();
    }

    /// Removes all of the callback and context pointers from the connection
    pub fn unset(mut self, connection: &mut Connection) -> Result<(), transport::Error> {
        unsafe {
            unsafe extern "C" fn secret_cb(
                _context: *mut c_void,
                _conn: *mut s2n_connection,
                _secret_type: s2n_secret_type_t::Type,
                _secret: *mut u8,
                _secret_size: u8,
            ) -> s2n_status_code::Type {
                -1
            }

            unsafe extern "C" fn send_cb(
                _context: *mut c_void,
                _data: *const u8,
                _len: u32,
            ) -> s2n_status_code::Type {
                -1
            }

            unsafe extern "C" fn recv_cb(
                _context: *mut c_void,
                _data: *mut u8,
                _len: u32,
            ) -> s2n_status_code::Type {
                -1
            }

            // We use unwrap here since s2n-tls will just check if connection is not null
            connection
                .set_secret_callback(Some(secret_cb), core::ptr::null_mut())
                .unwrap();
            connection.set_send_callback(Some(send_cb)).unwrap();
            connection.set_send_context(core::ptr::null_mut()).unwrap();
            connection.set_receive_callback(Some(recv_cb)).unwrap();
            connection
                .set_receive_context(core::ptr::null_mut())
                .unwrap();
            connection.set_waker(None).unwrap();

            // Flush the send buffer before returning to the connection
            self.flush();
            // attempt to emit server name after making progress and prior to error handling
            if !*self.emitted_server_name {
                if let Some(server_name) = self.server_name.clone().or_else(|| {
                    connection
                        .server_name()
                        .map(|server_name| server_name.into())
                }) {
                    self.context.on_server_name(server_name)?;
                    *self.emitted_server_name = true;
                }
            }

            if let Some(err) = self.err {
                return Err(err);
            }

            Ok(())
        }
    }

    /// The function s2n-tls calls when it emits secrets
    unsafe extern "C" fn secret_cb(
        context: *mut c_void,
        conn: *mut s2n_connection,
        secret_type: s2n_secret_type_t::Type,
        secret: *mut u8,
        secret_size: u8,
    ) -> s2n_status_code::Type {
        let context = &mut *(context as *mut Self);
        let secret = core::slice::from_raw_parts_mut(secret, secret_size as _);
        match context.on_secret(conn, secret_type, secret) {
            Ok(()) => 0,
            Err(err) => {
                context.err = Some(err);
                -1
            }
        }
    }

    /// Handles secrets from the s2n-tls connection
    fn on_secret(
        &mut self,
        conn: *mut s2n_connection,
        id: s2n_secret_type_t::Type,
        secret: &mut [u8],
    ) -> Result<(), transport::Error> {
        match core::mem::replace(&mut self.state.secrets, Secrets::Waiting) {
            Secrets::Waiting => {
                if id == s2n_secret_type_t::CLIENT_EARLY_TRAFFIC_SECRET {
                    // TODO enable with 0rtt
                    // I couldn't find a good citation but here's the issue: https://github.com/aws/s2n-quic/issues/301
                    return Ok(());
                }

                let (prk_algo, _aead) = get_algo_type(conn).ok_or(CryptoError::INTERNAL_ERROR)?;
                let secret = Prk::new_less_safe(prk_algo, secret);
                self.state.secrets = Secrets::Half { secret, id };

                Ok(())
            }
            Secrets::Half {
                id: other_id,
                secret: other_secret,
            } => {
                let (prk_algo, aead_algo) =
                    get_algo_type(conn).ok_or(CryptoError::INTERNAL_ERROR)?;
                let secret = Prk::new_less_safe(prk_algo, secret);
                let pair = match (id, other_id) {
                    (
                        s2n_secret_type_t::CLIENT_HANDSHAKE_TRAFFIC_SECRET,
                        s2n_secret_type_t::SERVER_HANDSHAKE_TRAFFIC_SECRET,
                    )
                    | (
                        s2n_secret_type_t::CLIENT_APPLICATION_TRAFFIC_SECRET,
                        s2n_secret_type_t::SERVER_APPLICATION_TRAFFIC_SECRET,
                    ) => SecretPair {
                        client: secret,
                        server: other_secret,
                    },
                    (
                        s2n_secret_type_t::SERVER_HANDSHAKE_TRAFFIC_SECRET,
                        s2n_secret_type_t::CLIENT_HANDSHAKE_TRAFFIC_SECRET,
                    )
                    | (
                        s2n_secret_type_t::SERVER_APPLICATION_TRAFFIC_SECRET,
                        s2n_secret_type_t::CLIENT_APPLICATION_TRAFFIC_SECRET,
                    ) => SecretPair {
                        server: secret,
                        client: other_secret,
                    },
                    _ => {
                        debug_assert!(false, "invalid key phase");
                        return Err(transport::Error::INTERNAL_ERROR);
                    }
                };

                // Flush the send buffer before transitioning to the next phase
                self.flush();

                match self.state.tx_phase {
                    HandshakePhase::Initial => {
                        let (key, header_key) = HandshakeKey::new(self.endpoint, aead_algo, pair)
                            .expect("invalid cipher");

                        self.context.on_handshake_keys(key, header_key)?;
                        self.state.tx_phase.transition();
                        self.state.rx_phase.transition();
                    }
                    _ => {
                        let (key, header_key) =
                            OneRttKey::new(self.endpoint, aead_algo, pair).expect("invalid cipher");

                        let params = unsafe {
                            // Safety: conn needs to outlive params
                            //
                            // TODO use interning for these values
                            // issue: https://github.com/aws/s2n-quic/issues/248
                            //
                            // Move this event to where `on_server_name` is emitted once we expose
                            // the functionality in s2n_tls bindings
                            let application_protocol =
                                Bytes::copy_from_slice(get_application_protocol(conn)?);
                            self.context.on_application_protocol(application_protocol)?;
                            get_application_params(conn)?
                        };

                        self.context.on_one_rtt_keys(key, header_key, params)?;
                    }
                }

                Ok(())
            }
        }
    }

    /// The function s2n-tls calls when it wants to send data
    unsafe extern "C" fn send_cb(
        context: *mut c_void,
        data: *const u8,
        len: u32,
    ) -> s2n_status_code::Type {
        let context = &mut *(context as *mut Self);
        let data = core::slice::from_raw_parts(data, len as _);
        context.on_write(data) as _
    }

    /// Called when sending data
    fn on_write(&mut self, data: &[u8]) -> usize {
        // If this write would cause the current send buffer to reallocate,
        // we should flush and create a new send buffer.
        let remaining_capacity = self.send_buffer.capacity() - self.send_buffer.len();

        if remaining_capacity < data.len() {
            // Flush the send buffer before reallocating it
            self.flush();

            // ensure we only do one allocation for this write
            let len = SEND_BUFFER_CAPACITY.max(data.len());

            debug_assert!(
                self.send_buffer.is_empty(),
                "dropping a send buffer with data will result in data loss"
            );
            *self.send_buffer = BytesMut::with_capacity(len);
        }

        // Write the current data to the send buffer
        //
        // NOTE: we don't immediately flush to the packet space since s2n-tls may do
        //       several small writes in a row.
        self.send_buffer.extend_from_slice(data);

        data.len()
    }

    /// Flushes the send buffer into the current TX space
    fn flush(&mut self) {
        if !self.send_buffer.is_empty() {
            let chunk = self.send_buffer.split().freeze();

            match self.state.tx_phase {
                HandshakePhase::Initial => self.context.send_initial(chunk),
                HandshakePhase::Handshake => self.context.send_handshake(chunk),
                HandshakePhase::Application => self.context.send_application(chunk),
            }
        }
    }

    /// The function s2n-tls calls when it wants to receive data
    unsafe extern "C" fn recv_cb(
        context: *mut c_void,
        data: *mut u8,
        len: u32,
    ) -> s2n_status_code::Type {
        let context = &mut *(context as *mut Self);
        let data = core::slice::from_raw_parts_mut(data, len as _);
        match context.on_read(data) {
            0 => {
                // https://github.com/aws/s2n-tls/blob/main/docs/USAGE-GUIDE.md#s2n_connection_set_send_cb
                // s2n-tls wants us to set the global errno to signal blocked
                errno::set_errno(errno::Errno(libc::EWOULDBLOCK));
                -1
            }
            len => len as _,
        }
    }

    /// Called when receiving data
    fn on_read(&mut self, data: &mut [u8]) -> usize {
        let max_len = Some(data.len());

        let chunk = match self.state.rx_phase {
            HandshakePhase::Initial => self.context.receive_initial(max_len),
            HandshakePhase::Handshake => self.context.receive_handshake(max_len),
            HandshakePhase::Application => self.context.receive_application(max_len),
        };

        if let Some(chunk) = chunk {
            let len = chunk.len();
            data[..len].copy_from_slice(&chunk);
            len
        } else {
            0
        }
    }
}

#[derive(Default, Debug)]
pub struct State {
    rx_phase: HandshakePhase,
    tx_phase: HandshakePhase,
    secrets: Secrets,
}

impl State {
    /// Complete the handshake
    pub fn on_handshake_complete(&mut self) {
        debug_assert_eq!(self.tx_phase, HandshakePhase::Handshake);
        debug_assert_eq!(self.rx_phase, HandshakePhase::Handshake);
        self.tx_phase.transition();
        self.rx_phase.transition();
        debug_assert_eq!(self.tx_phase, HandshakePhase::Application);
        debug_assert_eq!(self.rx_phase, HandshakePhase::Application);
    }
}

#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
enum HandshakePhase {
    Initial,
    Handshake,
    Application,
}

impl HandshakePhase {
    fn transition(&mut self) {
        *self = match self {
            Self::Initial => Self::Handshake,
            _ => Self::Application,
        };
    }
}

impl Default for HandshakePhase {
    fn default() -> Self {
        Self::Initial
    }
}

#[derive(Debug)]
enum Secrets {
    Waiting,
    Half {
        secret: Prk,
        id: s2n_secret_type_t::Type,
    },
}

impl Default for Secrets {
    fn default() -> Self {
        Self::Waiting
    }
}

fn get_algo_type(
    connection: *mut s2n_connection,
) -> Option<(hkdf::Algorithm, &'static aead::Algorithm)> {
    let mut cipher = [0, 0];
    unsafe {
        s2n_connection_get_cipher_iana_value(connection, &mut cipher[0], &mut cipher[1])
            .into_result()
            .ok()?;
    }

    //= https://www.rfc-editor.org/rfc/rfc8446#appendix-B.4
    //# This specification defines the following cipher suites for use with
    //# TLS 1.3.
    //#
    //#              +------------------------------+-------------+
    //#              | Description                  | Value       |
    //#              +------------------------------+-------------+
    //#              | TLS_AES_128_GCM_SHA256       | {0x13,0x01} |
    //#              |                              |             |
    //#              | TLS_AES_256_GCM_SHA384       | {0x13,0x02} |
    //#              |                              |             |
    //#              | TLS_CHACHA20_POLY1305_SHA256 | {0x13,0x03} |
    //#              |                              |             |
    //#              | TLS_AES_128_CCM_SHA256       | {0x13,0x04} |
    //#              |                              |             |
    //#              | TLS_AES_128_CCM_8_SHA256     | {0x13,0x05} |
    //#              +------------------------------+-------------+
    const TLS_AES_128_GCM_SHA256: [u8; 2] = [0x13, 0x01];
    const TLS_AES_256_GCM_SHA384: [u8; 2] = [0x13, 0x02];
    const TLS_CHACHA20_POLY1305_SHA256: [u8; 2] = [0x13, 0x03];

    // NOTE: we don't have CCM support implemented currently

    // NOTE: CCM_8 is not allowed by QUIC
    //= https://www.rfc-editor.org/rfc/rfc9001#section-5.3
    //# QUIC can use any of the cipher suites defined in [TLS13] with the
    //# exception of TLS_AES_128_CCM_8_SHA256.

    match cipher {
        TLS_AES_128_GCM_SHA256 => Some((hkdf::HKDF_SHA256, &aead::AES_128_GCM)),
        TLS_AES_256_GCM_SHA384 => Some((hkdf::HKDF_SHA384, &aead::AES_256_GCM)),
        TLS_CHACHA20_POLY1305_SHA256 => Some((hkdf::HKDF_SHA256, &aead::CHACHA20_POLY1305)),
        _ => None,
    }
}

unsafe fn get_application_params<'a>(
    connection: *mut s2n_connection,
) -> Result<tls::ApplicationParameters<'a>, CryptoError> {
    //= https://www.rfc-editor.org/rfc/rfc9001#section-8.1
    //# When using ALPN, endpoints MUST immediately close a connection (see
    //# Section 10.2 of [QUIC-TRANSPORT]) with a no_application_protocol TLS
    //# alert (QUIC error code 0x178; see Section 4.8) if an application
    //# protocol is not negotiated.

    //= https://www.rfc-editor.org/rfc/rfc9001#section-8.1
    //# While [ALPN] only specifies that servers
    //# use this alert, QUIC clients MUST use error 0x178 to terminate a
    //# connection when ALPN negotiation fails.
    let transport_parameters =
        get_transport_parameters(connection).ok_or(CryptoError::MISSING_EXTENSION)?;

    Ok(tls::ApplicationParameters {
        transport_parameters,
    })
}

//= https://www.rfc-editor.org/rfc/rfc9001#section-8.1
//# Unless
//# another mechanism is used for agreeing on an application protocol,
//# endpoints MUST use ALPN for this purpose.
//
//= https://www.rfc-editor.org/rfc/rfc7301#section-3.1
//# Client                                              Server
//#
//#    ClientHello                     -------->       ServerHello
//#      (ALPN extension &                               (ALPN extension &
//#       list of protocols)                              selected protocol)
//#                                                    [ChangeCipherSpec]
//#                                    <--------       Finished
//#    [ChangeCipherSpec]
//#    Finished                        -------->
//#    Application Data                <------->       Application Data
unsafe fn get_application_protocol<'a>(
    connection: *mut s2n_connection,
) -> Result<&'a [u8], CryptoError> {
    let ptr = s2n_get_application_protocol(connection).into_result().ok();
    ptr.and_then(|ptr| get_cstr_slice(ptr))
        .ok_or(CryptoError::MISSING_EXTENSION)
}

unsafe fn get_transport_parameters<'a>(connection: *mut s2n_connection) -> Option<&'a [u8]> {
    let mut ptr = core::ptr::null();
    let mut len = 0u16;

    s2n_connection_get_quic_transport_parameters(connection, &mut ptr, &mut len)
        .into_result()
        .ok()?;

    get_slice(ptr, len as _)
}

unsafe fn get_cstr_slice<'a>(ptr: *const libc::c_char) -> Option<&'a [u8]> {
    let len = libc::strlen(ptr);
    get_slice(ptr as *const _, len)
}

unsafe fn get_slice<'a>(ptr: *const u8, len: usize) -> Option<&'a [u8]> {
    if ptr.is_null() || len == 0 {
        return None;
    }

    Some(core::slice::from_raw_parts(ptr, len as _))
}


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

use bytes::Bytes;
use s2n_tls::error::Error;

impl Format {
    pub fn as_pem(&self) -> Option<&[u8]> {
        if let Format::Pem(bytes) = &self {
            Some(bytes.as_ref())
        } else {
            None
        }
    }

    #[allow(dead_code)] // remove if s2n-tls ever starts supporting DER certs
    pub fn as_der(&self) -> Option<&[u8]> {
        if let Format::Der(bytes) = &self {
            Some(bytes.as_ref())
        } else {
            None
        }
    }
}

pub(crate) enum Format {
    Pem(Bytes),
    Der(Bytes),
    #[allow(dead_code)] // Only used if private key offloading supported
    None,
}

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

        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 bytes = self.into_bytes();
                let bytes = Bytes::from(bytes);
                let bytes = Format::Pem(bytes);
                Ok($name(bytes))
            }
        }

        impl $trait for &String {
            fn $method(self) -> Result<$name, Error> {
                let bytes = self.as_bytes();
                let bytes = Bytes::copy_from_slice(bytes);
                let bytes = Format::Pem(bytes);
                Ok($name(bytes))
            }
        }

        impl $trait for &str {
            fn $method(self) -> Result<$name, Error> {
                let bytes = self.as_bytes();
                let bytes = Bytes::copy_from_slice(bytes);
                let bytes = Format::Pem(bytes);
                Ok($name(bytes))
            }
        }

        impl $trait for Vec<u8> {
            fn $method(self) -> Result<$name, Error> {
                let bytes = Bytes::from(self);
                let bytes = Format::Der(bytes);
                Ok($name(bytes))
            }
        }

        impl $trait for &[u8] {
            fn $method(self) -> Result<$name, Error> {
                let bytes = Bytes::copy_from_slice(self);
                let bytes = Format::Der(bytes);
                Ok($name(bytes))
            }
        }

        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::io_error(err))?;
                        pem.$method()
                    }
                    // assume it's in pem format
                    _ => {
                        let pem =
                            std::fs::read_to_string(self).map_err(|err| Error::io_error(err))?;
                        pem.$method()
                    }
                }
            }
        }
    };
}

cert_type!(PrivateKey, IntoPrivateKey, into_private_key);
cert_type!(Certificate, IntoCertificate, into_certificate);

#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_private_key")))]
pub const OFFLOAD_PRIVATE_KEY: PrivateKey = PrivateKey(Format::None);


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

use crate::{
    certificate::{IntoCertificate, IntoPrivateKey},
    keylog::KeyLogHandle,
    params::Params,
    session::Session,
    ConfigLoader,
};
use s2n_codec::EncoderValue;
use s2n_quic_core::{application::ServerName, crypto::tls, endpoint};
use s2n_tls::{
    callbacks::VerifyHostNameCallback,
    config::{self, Config},
    enums::ClientAuthType,
    error::Error,
};
use std::sync::Arc;

pub struct Client<L: ConfigLoader = Config> {
    loader: L,
    #[allow(dead_code)] // we need to hold on to the handle to ensure it is cleaned up correctly
    keylog: Option<KeyLogHandle>,
    params: Params,
}

impl Client {
    pub fn builder() -> Builder {
        Builder::default()
    }
}

impl<L: ConfigLoader> Client<L> {
    /// Creates a [`Client`] from a [`ConfigLoader`]
    ///
    /// The caller is responsible for building the `Config`
    /// correctly for QUIC settings. This includes:
    /// * setting a security policy that supports TLS 1.3
    /// * enabling QUIC support
    /// * setting at least one application protocol
    pub fn from_loader(loader: L) -> Self {
        Self {
            loader,
            keylog: None,
            params: Default::default(),
        }
    }
}

impl Default for Client {
    fn default() -> Self {
        Self::builder()
            .build()
            .expect("could not create a default client")
    }
}

impl<L: ConfigLoader> ConfigLoader for Client<L> {
    #[inline]
    fn load(&mut self, cx: crate::ConnectionContext) -> s2n_tls::config::Config {
        self.loader.load(cx)
    }
}

pub struct Builder {
    config: config::Builder,
    keylog: Option<KeyLogHandle>,
}

impl Default for Builder {
    fn default() -> Self {
        let mut config = config::Builder::default();
        config.enable_quic().unwrap();
        // https://github.com/aws/s2n-tls/blob/main/docs/USAGE-GUIDE.md#s2n_config_set_cipher_preferences
        config.set_security_policy(crate::DEFAULT_POLICY).unwrap();
        config.set_application_protocol_preference([b"h3"]).unwrap();

        Self {
            config,
            keylog: None,
        }
    }
}

impl Builder {
    pub fn with_application_protocols<P: IntoIterator<Item = I>, I: AsRef<[u8]>>(
        mut self,
        protocols: P,
    ) -> Result<Self, Error> {
        self.config.set_application_protocol_preference(protocols)?;
        Ok(self)
    }

    pub fn with_certificate<C: IntoCertificate>(mut self, certificate: C) -> Result<Self, Error> {
        let certificate = certificate.into_certificate()?;
        let certificate = certificate
            .0
            .as_pem()
            .expect("pem is currently the only certificate format supported");
        self.config.trust_pem(certificate)?;
        Ok(self)
    }

    /// Clears the default trust store for this client
    ///
    /// By default, the trust store is initialized with common
    /// trust store locations for the host operating system.
    /// By invoking this method, the trust store will be cleared.
    ///
    /// Note that call ordering matters. The caller should call this
    /// method before making any calls to `with_trust_client_certificate_signed_by()`.
    /// Calling this method after a method that modifies the trust store will clear it.
    pub fn with_empty_trust_store(mut self) -> Result<Self, Error> {
        self.config.wipe_trust_store()?;
        Ok(self)
    }

    /// Add the cert and key to the key store.
    ///
    /// This must be set when the server requires client authentication (mutual TLS).
    /// The client will offer the certificate to the server when it is requested
    /// as part of the TLS handshake.
    pub fn with_client_identity<C: IntoCertificate, PK: IntoPrivateKey>(
        mut self,
        certificate: C,
        private_key: PK,
    ) -> Result<Self, Error> {
        let certificate = certificate.into_certificate()?;
        let private_key = private_key.into_private_key()?;
        self.config.load_pem(
            certificate
                .0
                .as_pem()
                .expect("pem is currently the only certificate format supported"),
            private_key
                .0
                .as_pem()
                .expect("pem is currently the only certificate format supported"),
        )?;
        self.config.set_client_auth_type(ClientAuthType::Required)?;
        Ok(self)
    }

    /// Set the host name verification callback.
    ///
    /// This will be invoked when a server certificate is presented during a TLS
    /// handshake. If this function is invoked, the default server name validation
    /// logic is disabled; this should only be used in very specific cases where normal
    /// TLS hostname validation is not appropriate.
    pub fn with_verify_host_name_callback<T: 'static + VerifyHostNameCallback>(
        mut self,
        handler: T,
    ) -> Result<Self, Error> {
        self.config.set_verify_host_callback(handler)?;
        Ok(self)
    }

    pub fn with_max_cert_chain_depth(mut self, len: u16) -> Result<Self, Error> {
        self.config.set_max_cert_chain_depth(len)?;
        Ok(self)
    }

    pub fn with_key_logging(mut self) -> Result<Self, Error> {
        use crate::keylog::KeyLog;

        self.keylog = KeyLog::try_open();

        unsafe {
            // Safety: the KeyLog is stored on `self` to ensure it outlives `config`
            if let Some(keylog) = self.keylog.as_ref() {
                self.config
                    .set_key_log_callback(Some(KeyLog::callback), Arc::as_ptr(keylog) as *mut _)?;
            } else {
                // disable key logging if it failed to create a file
                self.config
                    .set_key_log_callback(None, core::ptr::null_mut())?;
            }
        }

        Ok(self)
    }

    pub fn build(self) -> Result<Client, Error> {
        Ok(Client {
            loader: self.config.build()?,
            keylog: self.keylog,
            params: Default::default(),
        })
    }
}

impl<L: ConfigLoader> tls::Endpoint for Client<L> {
    type Session = Session;

    fn new_server_session<Params: EncoderValue>(&mut self, _params: &Params) -> Self::Session {
        panic!("cannot create a server session from a client config");
    }

    fn new_client_session<Params: EncoderValue>(
        &mut self,
        params: &Params,
        server_name: ServerName,
    ) -> Self::Session {
        let config = self.loader.load(crate::ConnectionContext {
            server_name: Some(&server_name),
        });
        self.params.with(params, |params| {
            Session::new(endpoint::Type::Client, config, params, Some(server_name)).unwrap()
        })
    }

    fn max_tag_length(&self) -> usize {
        s2n_quic_crypto::MAX_TAG_LEN
    }
}


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

use libc::{c_int, c_void};
use s2n_tls::ffi::*;
use std::{
    fs::{File, OpenOptions},
    io::{BufWriter, Write},
    sync::{Arc, Mutex},
};

pub type KeyLogHandle = Arc<KeyLog>;

pub struct KeyLog(Mutex<BufWriter<File>>);

impl KeyLog {
    pub fn try_open() -> Option<KeyLogHandle> {
        let path = std::env::var("SSLKEYLOGFILE").ok()?;
        let file = OpenOptions::new()
            .append(true)
            .create(true)
            .open(path)
            .ok()?;
        let file = BufWriter::new(file);
        let file = Mutex::new(file);
        let keylog = Self(file);
        let keylog = Arc::new(keylog);
        Some(keylog)
    }

    pub unsafe extern "C" fn callback(
        ctx: *mut c_void,
        _conn: *mut s2n_connection,
        logline: *mut u8,
        len: usize,
    ) -> c_int {
        let handle = &mut *(ctx as *mut Self);
        let logline = core::slice::from_raw_parts(logline, len);

        // ignore any errors
        let _ = handle.on_logline(logline);

        0
    }

    fn on_logline(&mut self, logline: &[u8]) -> Option<()> {
        let mut file = self.0.lock().ok()?;
        file.write_all(logline).ok()?;
        file.write_all(b"\n").ok()?;

        // ensure keys are immediately written so tools can use them
        file.flush().ok()?;

        Some(())
    }
}


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

use s2n_quic_core::application::ServerName;
use s2n_tls::config::Config;

/// Ensure memory is correctly managed in tests
#[cfg(test)]
#[global_allocator]
static ALLOCATOR: checkers::Allocator = checkers::Allocator::system();

#[cfg(all(s2n_quic_unstable, s2n_quic_enable_pq_tls))]
static DEFAULT_POLICY: &s2n_tls::security::Policy = &s2n_tls::security::TESTING_PQ;
#[cfg(not(all(s2n_quic_unstable, s2n_quic_enable_pq_tls)))]
static DEFAULT_POLICY: &s2n_tls::security::Policy = &s2n_tls::security::DEFAULT_TLS13;

#[non_exhaustive]
pub struct ConnectionContext<'a> {
    pub server_name: Option<&'a ServerName>,
}

/// Loads a config for a given connection
///
/// This trait can be implemented to override the default config loading for a QUIC endpoint
pub trait ConfigLoader: 'static + Send {
    fn load(&mut self, cx: ConnectionContext) -> Config;
}

impl ConfigLoader for Config {
    #[inline]
    fn load(&mut self, _cx: ConnectionContext) -> Config {
        self.clone()
    }
}

impl<T: FnMut(ConnectionContext) -> Config + Send + 'static> ConfigLoader for T {
    #[inline]
    fn load(&mut self, cx: ConnectionContext) -> Config {
        (self)(cx)
    }
}

impl ConfigLoader for Box<dyn ConfigLoader> {
    #[inline]
    fn load(&mut self, cx: ConnectionContext) -> Config {
        (**self).load(cx)
    }
}

mod callback;
mod keylog;
mod params;
mod session;

pub mod certificate;
pub mod client;
pub mod server;

pub use client::Client;
pub use server::Server;

// Re-export the `ClientHelloHandler` and `Connection` to make it easier for users
// to consume. This depends on experimental behavior in s2n-tls.
#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_client_hello")))]
pub use s2n_tls::{self, callbacks::ClientHelloCallback, connection::Connection};

#[cfg(test)]
mod tests;


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

use s2n_codec::{EncoderBuffer, EncoderValue};

/// A buffer used by an endpoint to encode transport parameters
///
/// We do a single allocation at the endpoint since s2n-tls will just
/// copy the buffer anyway
#[derive(Debug, Default)]
pub struct Params {
    buffer: Vec<u8>,
}

impl Params {
    pub fn with<P, F, R>(&mut self, params: &P, f: F) -> R
    where
        P: EncoderValue,
        F: FnOnce(&[u8]) -> R,
    {
        let len = params.encoding_size();
        self.buffer.resize(len, 0);
        params.encode(&mut EncoderBuffer::new(&mut self.buffer));
        f(&self.buffer)
    }
}


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

use crate::{
    certificate::{Format, IntoCertificate, IntoPrivateKey},
    keylog::KeyLogHandle,
    params::Params,
    session::Session,
    ConfigLoader,
};
use s2n_codec::EncoderValue;
use s2n_quic_core::{application::ServerName, crypto::tls, endpoint};
#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_client_hello")))]
use s2n_tls::callbacks::ClientHelloCallback;
#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_private_key")))]
use s2n_tls::callbacks::PrivateKeyCallback;
use s2n_tls::{
    callbacks::VerifyHostNameCallback,
    config::{self, Config},
    enums::ClientAuthType,
    error::Error,
};
use std::sync::Arc;

pub struct Server<L: ConfigLoader = Config> {
    loader: L,
    #[allow(dead_code)] // we need to hold on to the handle to ensure it is cleaned up correctly
    keylog: Option<KeyLogHandle>,
    params: Params,
}

impl Server {
    pub fn builder() -> Builder {
        Builder::default()
    }
}

impl<L: ConfigLoader> Server<L> {
    /// Creates a [`Server`] from a [`ConfigLoader`]
    ///
    /// The caller is responsible for building the `Config`
    /// correctly for QUIC settings. This includes:
    /// * setting a security policy that supports TLS 1.3
    /// * enabling QUIC support
    /// * setting at least one application protocol
    pub fn from_loader(loader: L) -> Self {
        Self {
            loader,
            keylog: None,
            params: Default::default(),
        }
    }
}

impl Default for Server {
    fn default() -> Self {
        Self::builder()
            .build()
            .expect("could not create a default server")
    }
}

impl<L: ConfigLoader> ConfigLoader for Server<L> {
    #[inline]
    fn load(&mut self, cx: crate::ConnectionContext) -> s2n_tls::config::Config {
        self.loader.load(cx)
    }
}

impl<L: ConfigLoader> From<Server<L>> for Config {
    fn from(mut server: Server<L>) -> Self {
        server.load(crate::ConnectionContext { server_name: None })
    }
}

pub struct Builder {
    config: config::Builder,
    keylog: Option<KeyLogHandle>,
}

impl Default for Builder {
    fn default() -> Self {
        let mut config = config::Builder::default();
        config.enable_quic().unwrap();
        // https://github.com/aws/s2n-tls/blob/main/docs/USAGE-GUIDE.md#s2n_config_set_cipher_preferences
        config.set_security_policy(crate::DEFAULT_POLICY).unwrap();
        config.set_application_protocol_preference([b"h3"]).unwrap();

        Self {
            config,
            keylog: None,
        }
    }
}

impl Builder {
    #[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_client_hello")))]
    pub fn with_client_hello_handler<T: 'static + ClientHelloCallback>(
        mut self,
        handler: T,
    ) -> Result<Self, Error> {
        self.config.set_client_hello_callback(handler)?;
        Ok(self)
    }

    #[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_private_key")))]
    pub fn with_private_key_handler<T: 'static + PrivateKeyCallback>(
        mut self,
        handler: T,
    ) -> Result<Self, Error> {
        self.config.set_private_key_callback(handler)?;
        Ok(self)
    }

    pub fn with_application_protocols<P: IntoIterator<Item = I>, I: AsRef<[u8]>>(
        mut self,
        protocols: P,
    ) -> Result<Self, Error> {
        self.config.set_application_protocol_preference(protocols)?;
        Ok(self)
    }

    pub fn with_certificate<C: IntoCertificate, PK: IntoPrivateKey>(
        mut self,
        certificate: C,
        private_key: PK,
    ) -> Result<Self, Error> {
        let private_key = private_key.into_private_key()?.0;
        let certificate = certificate.into_certificate()?.0;
        let certificate = certificate
            .as_pem()
            .expect("pem is currently the only certificate format supported");
        match private_key {
            Format::Pem(bytes) => self.config.load_pem(certificate, bytes.as_ref())?,
            Format::None => self.config.load_public_pem(certificate)?,
            Format::Der(_) => panic!("der private keys not supported"),
        };
        Ok(self)
    }

    pub fn with_trusted_certificate<C: IntoCertificate>(
        mut self,
        certificate: C,
    ) -> Result<Self, Error> {
        let certificate = certificate.into_certificate()?;
        let certificate = certificate
            .0
            .as_pem()
            .expect("pem is currently the only certificate format supported");
        self.config.trust_pem(certificate)?;
        Ok(self)
    }

    /// Clears the default trust store for this client.
    ///
    /// By default, the trust store is initialized with common
    /// trust store locations for the host operating system.
    /// By invoking this method, the trust store will be cleared.
    ///
    /// Note that call ordering matters. The caller should call this
    /// method before making any calls to `with_trusted_certificate()`.
    /// Calling this method after a method that modifies the trust store will clear it.
    pub fn with_empty_trust_store(mut self) -> Result<Self, Error> {
        self.config.wipe_trust_store()?;
        Ok(self)
    }

    /// Configures this server instance to require client authentication (mutual TLS).
    pub fn with_client_authentication(mut self) -> Result<Self, Error> {
        self.config.set_client_auth_type(ClientAuthType::Required)?;
        Ok(self)
    }

    /// Set the application level certificate verification handler which will be invoked on this
    /// server instance when a client certificate is presented during the mutual TLS handshake.
    #[deprecated(note = "use `with_verify_host_name_callback` instead")]
    pub fn with_verify_client_certificate_handler<T: 'static + VerifyHostNameCallback>(
        mut self,
        handler: T,
    ) -> Result<Self, Error> {
        self.config.set_verify_host_callback(handler)?;
        Ok(self)
    }

    /// Set the host name verification callback.
    ///
    /// This will be invoked when a client certificate is presented during a mutual TLS
    /// handshake.
    pub fn with_verify_host_name_callback<T: 'static + VerifyHostNameCallback>(
        mut self,
        handler: T,
    ) -> Result<Self, Error> {
        self.config.set_verify_host_callback(handler)?;
        Ok(self)
    }

    pub fn with_key_logging(mut self) -> Result<Self, Error> {
        use crate::keylog::KeyLog;

        self.keylog = KeyLog::try_open();

        unsafe {
            // Safety: the KeyLog is stored on `self` to ensure it outlives `config`
            if let Some(keylog) = self.keylog.as_ref() {
                self.config
                    .set_key_log_callback(Some(KeyLog::callback), Arc::as_ptr(keylog) as *mut _)?;
            } else {
                // disable key logging if it failed to create a file
                self.config
                    .set_key_log_callback(None, core::ptr::null_mut())?;
            }
        }

        Ok(self)
    }

    pub fn build(self) -> Result<Server, Error> {
        Ok(Server {
            loader: self.config.build()?,
            keylog: self.keylog,
            params: Default::default(),
        })
    }
}

impl<L: ConfigLoader> tls::Endpoint for Server<L> {
    type Session = Session;

    fn new_server_session<Params: EncoderValue>(&mut self, params: &Params) -> Self::Session {
        let config = self
            .loader
            .load(crate::ConnectionContext { server_name: None });
        self.params.with(params, |params| {
            Session::new(endpoint::Type::Server, config, params, None).unwrap()
        })
    }

    fn new_client_session<Params: EncoderValue>(
        &mut self,
        _transport_parameters: &Params,
        _erver_name: ServerName,
    ) -> Self::Session {
        panic!("cannot create a client session from a server config");
    }

    fn max_tag_length(&self) -> usize {
        s2n_quic_crypto::MAX_TAG_LEN
    }
}


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

use crate::callback::{self, Callback};
use bytes::BytesMut;
use core::{marker::PhantomData, task::Poll};
use s2n_quic_core::{
    application::ServerName,
    crypto::{tls, CryptoError, CryptoSuite},
    endpoint, transport,
};
use s2n_quic_crypto::Suite;
use s2n_tls::{
    config::Config,
    connection::Connection,
    enums::{Blinding, Mode},
    error::Error,
};

#[derive(Debug)]
pub struct Session {
    endpoint: endpoint::Type,
    pub(crate) connection: Connection,
    state: callback::State,
    handshake_complete: bool,
    send_buffer: BytesMut,
    emitted_server_name: bool,
    // This is only set for the client to avoid an extra allocation
    server_name: Option<ServerName>,
}

impl Session {
    pub fn new(
        endpoint: endpoint::Type,
        config: Config,
        params: &[u8],
        server_name: Option<ServerName>,
    ) -> Result<Self, Error> {
        let mut connection = Connection::new(match endpoint {
            endpoint::Type::Server => Mode::Server,
            endpoint::Type::Client => Mode::Client,
        });

        connection.set_config(config)?;
        connection.enable_quic()?;
        connection.set_quic_transport_parameters(params)?;
        // QUIC handles sending alerts, so no need to apply TLS blinding
        connection.set_blinding(Blinding::SelfService)?;

        if let Some(server_name) = server_name.as_ref() {
            connection
                .set_server_name(server_name)
                .expect("invalid server name value");
        }

        Ok(Self {
            endpoint,
            connection,
            state: Default::default(),
            handshake_complete: false,
            send_buffer: BytesMut::new(),
            emitted_server_name: false,
            server_name,
        })
    }
}

impl CryptoSuite for Session {
    type HandshakeKey = <Suite as CryptoSuite>::HandshakeKey;
    type HandshakeHeaderKey = <Suite as CryptoSuite>::HandshakeHeaderKey;
    type InitialKey = <Suite as CryptoSuite>::InitialKey;
    type InitialHeaderKey = <Suite as CryptoSuite>::InitialHeaderKey;
    type OneRttKey = <Suite as CryptoSuite>::OneRttKey;
    type OneRttHeaderKey = <Suite as CryptoSuite>::OneRttHeaderKey;
    type ZeroRttKey = <Suite as CryptoSuite>::ZeroRttKey;
    type ZeroRttHeaderKey = <Suite as CryptoSuite>::ZeroRttHeaderKey;
    type RetryKey = <Suite as CryptoSuite>::RetryKey;
}

impl tls::Session for Session {
    fn poll<W>(&mut self, context: &mut W) -> Poll<Result<(), transport::Error>>
    where
        W: tls::Context<Self>,
    {
        let mut callback: Callback<W, Self> = Callback {
            context,
            endpoint: self.endpoint,
            state: &mut self.state,
            suite: PhantomData,
            err: None,
            send_buffer: &mut self.send_buffer,
            emitted_server_name: &mut self.emitted_server_name,
            server_name: &self.server_name,
        };

        unsafe {
            // Safety: the callback struct must live as long as the callbacks are
            // set on on the connection
            callback.set(&mut self.connection);
        }

        let result = self.connection.poll_negotiate().map_ok(|_| ());

        callback.unset(&mut self.connection)?;

        match result {
            Poll::Ready(Ok(())) => {
                // s2n-tls has indicated that the handshake is complete
                if !self.handshake_complete {
                    self.state.on_handshake_complete();
                    context.on_handshake_complete()?;
                    self.handshake_complete = true;
                }
                Poll::Ready(Ok(()))
            }
            Poll::Ready(Err(e)) => Poll::Ready(Err(e
                .alert()
                .map(CryptoError::new)
                .unwrap_or(CryptoError::HANDSHAKE_FAILURE)
                .into())),
            Poll::Pending => Poll::Pending,
        }
    }
}


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

use crate::{certificate, client, server};
use core::{
    sync::atomic::{AtomicBool, AtomicU8, Ordering},
    task::Poll,
};
use openssl::{ec::EcKey, ecdsa::EcdsaSig};
use pin_project::pin_project;
use s2n_quic_core::{
    crypto::tls::{
        self,
        testing::certificates::{CERT_PEM, KEY_PEM, UNTRUSTED_CERT_PEM, UNTRUSTED_KEY_PEM},
        Endpoint,
    },
    transport,
};
#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_client_hello")))]
use s2n_tls::callbacks::ClientHelloCallback;
#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_client_hello")))]
use s2n_tls::callbacks::{PrivateKeyCallback, PrivateKeyOperation};
use s2n_tls::{
    callbacks::{ConnectionFuture, VerifyHostNameCallback},
    connection::Connection,
    error::Error,
};
use std::sync::Arc;

pub struct MyCallbackHandler {
    done: Arc<AtomicBool>,
    wait_counter: Arc<AtomicU8>,
}

impl MyCallbackHandler {
    fn new(wait_counter: u8) -> Self {
        MyCallbackHandler {
            done: Arc::new(AtomicBool::new(false)),
            wait_counter: Arc::new(AtomicU8::new(wait_counter)),
        }
    }
}

#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_client_hello")))]
impl ClientHelloCallback for MyCallbackHandler {
    fn on_client_hello(
        &self,
        _connection: &mut Connection,
    ) -> Result<Option<std::pin::Pin<Box<dyn s2n_tls::callbacks::ConnectionFuture>>>, Error> {
        let fut = MyConnectionFuture {
            done: self.done.clone(),
            wait_counter: self.wait_counter.clone(),
        };
        Ok(Some(Box::pin(fut)))
    }
}

struct MyConnectionFuture {
    done: Arc<AtomicBool>,
    wait_counter: Arc<AtomicU8>,
}

impl ConnectionFuture for MyConnectionFuture {
    fn poll(
        self: std::pin::Pin<&mut Self>,
        _connection: &mut Connection,
        _ctx: &mut core::task::Context,
    ) -> Poll<Result<(), Error>> {
        if self.wait_counter.fetch_sub(1, Ordering::SeqCst) == 0 {
            self.done.store(true, Ordering::SeqCst);
            return Poll::Ready(Ok(()));
        }

        Poll::Pending
    }
}

#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_private_key")))]
impl PrivateKeyCallback for MyCallbackHandler {
    fn handle_operation(
        &self,
        _connection: &mut Connection,
        op: PrivateKeyOperation,
    ) -> Result<Option<std::pin::Pin<Box<dyn s2n_tls::callbacks::ConnectionFuture>>>, Error> {
        let future = MyConnectionFuture {
            done: self.done.clone(),
            wait_counter: self.wait_counter.clone(),
        };
        let op = Some(op);
        let pkey_future = MyPrivateKeyFuture { future, op };
        Ok(Some(Box::pin(pkey_future)))
    }
}

#[pin_project]
struct MyPrivateKeyFuture {
    #[pin]
    future: MyConnectionFuture,
    #[pin]
    op: Option<PrivateKeyOperation>,
}

impl ConnectionFuture for MyPrivateKeyFuture {
    fn poll(
        self: std::pin::Pin<&mut Self>,
        conn: &mut Connection,
        ctx: &mut core::task::Context,
    ) -> Poll<Result<(), Error>> {
        let mut this = self.project();
        if this.future.poll(conn, ctx).is_pending() {
            return Poll::Pending;
        }

        let op = this.op.take().expect("Missing pkey operation");
        let in_buf_size = op.input_size()?;
        let mut in_buf = vec![0; in_buf_size];
        op.input(&mut in_buf)?;

        let key = EcKey::private_key_from_pem(KEY_PEM.as_bytes())
            .expect("Failed to create EcKey from pem");
        let sig = EcdsaSig::sign(&in_buf, &key).expect("Failed to sign input");
        let out = sig.to_der().expect("Failed to convert signature to der");

        op.set_output(conn, &out)?;
        Poll::Ready(Ok(()))
    }
}

pub struct VerifyHostNameClientCertVerifier {
    host_name: String,
}

impl VerifyHostNameCallback for VerifyHostNameClientCertVerifier {
    fn verify_host_name(&self, host_name: &str) -> bool {
        self.host_name == host_name
    }
}

impl VerifyHostNameClientCertVerifier {
    pub fn new(host_name: impl ToString) -> VerifyHostNameClientCertVerifier {
        VerifyHostNameClientCertVerifier {
            host_name: host_name.to_string(),
        }
    }
}

#[derive(Default)]
pub struct RejectAllClientCertificatesHandler {}
impl VerifyHostNameCallback for RejectAllClientCertificatesHandler {
    fn verify_host_name(&self, _host_name: &str) -> bool {
        false
    }
}

fn s2n_client() -> client::Client {
    client::Builder::default()
        .with_certificate(CERT_PEM)
        .unwrap()
        .build()
        .unwrap()
}

fn s2n_client_with_client_auth() -> Result<client::Client, Error> {
    client::Builder::default()
        .with_empty_trust_store()?
        .with_certificate(CERT_PEM)?
        .with_client_identity(CERT_PEM, KEY_PEM)?
        .build()
}

fn s2n_client_with_untrusted_client_auth() -> Result<client::Client, Error> {
    client::Builder::default()
        .with_empty_trust_store()?
        .with_certificate(CERT_PEM)?
        .with_client_identity(UNTRUSTED_CERT_PEM, UNTRUSTED_KEY_PEM)?
        .build()
}

fn s2n_client_with_fixed_hostname_auth(host_name: &str) -> Result<client::Client, Error> {
    client::Builder::default()
        .with_certificate(CERT_PEM)
        .unwrap()
        .with_verify_host_name_callback(VerifyHostNameClientCertVerifier::new(host_name))
        .unwrap()
        .build()
}

fn s2n_server() -> server::Server {
    server::Builder::default()
        .with_certificate(CERT_PEM, KEY_PEM)
        .unwrap()
        .build()
        .unwrap()
}

fn s2n_server_with_client_auth() -> Result<server::Server, Error> {
    server::Builder::default()
        .with_empty_trust_store()?
        .with_client_authentication()?
        .with_verify_host_name_callback(VerifyHostNameClientCertVerifier::new("qlaws.qlaws"))?
        .with_certificate(CERT_PEM, KEY_PEM)?
        .with_trusted_certificate(CERT_PEM)?
        .build()
}

fn s2n_server_with_client_auth_verifier_rejects_client_certs() -> Result<server::Server, Error> {
    server::Builder::default()
        .with_empty_trust_store()?
        .with_client_authentication()?
        .with_verify_host_name_callback(RejectAllClientCertificatesHandler::default())?
        .with_certificate(CERT_PEM, KEY_PEM)?
        .with_trusted_certificate(CERT_PEM)?
        .build()
}

#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_client_hello")))]
fn s2n_server_with_client_hello_callback(wait_counter: u8) -> (server::Server, Arc<AtomicBool>) {
    let handle = MyCallbackHandler::new(wait_counter);
    let done = handle.done.clone();
    let tls = server::Builder::default()
        .with_certificate(CERT_PEM, KEY_PEM)
        .unwrap()
        .with_client_hello_handler(handle)
        .unwrap()
        .build()
        .unwrap();
    (tls, done)
}

#[cfg(any(test, all(s2n_quic_unstable, feature = "unstable_private_key")))]
fn s2n_server_with_private_key_callback(wait_counter: u8) -> (server::Server, Arc<AtomicBool>) {
    let handle = MyCallbackHandler::new(wait_counter);
    let done = handle.done.clone();
    let tls = server::Builder::default()
        .with_certificate(CERT_PEM, certificate::OFFLOAD_PRIVATE_KEY)
        .unwrap()
        .with_private_key_handler(handle)
        .unwrap()
        .build()
        .unwrap();
    (tls, done)
}

fn rustls_server() -> s2n_quic_rustls::server::Server {
    s2n_quic_rustls::server::Builder::default()
        .with_certificate(CERT_PEM, KEY_PEM)
        .unwrap()
        .build()
        .unwrap()
}

fn rustls_client() -> s2n_quic_rustls::client::Client {
    s2n_quic_rustls::client::Builder::default()
        .with_certificate(CERT_PEM)
        .unwrap()
        .build()
        .unwrap()
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_s2n_server_ch_callback_test() {
    for wait_counter in 0..=10 {
        let mut client_endpoint = s2n_client();
        let (mut server_endpoint, done) = s2n_server_with_client_hello_callback(wait_counter);

        run(&mut server_endpoint, &mut client_endpoint, Some(done));
    }
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_s2n_server_pkey_callback_test() {
    for wait_counter in 0..=10 {
        let mut client_endpoint = s2n_client();
        let (mut server_endpoint, done) = s2n_server_with_private_key_callback(wait_counter);

        run(&mut server_endpoint, &mut client_endpoint, Some(done));
    }
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_s2n_server_test() {
    let mut client_endpoint = s2n_client();
    let mut server_endpoint = s2n_server();

    run(&mut server_endpoint, &mut client_endpoint, None);
}

#[test]
#[cfg_attr(miri, ignore)]
fn rustls_client_s2n_server_test() {
    let mut client_endpoint = rustls_client();
    let mut server_endpoint = s2n_server();

    run(&mut server_endpoint, &mut client_endpoint, None);
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_rustls_server_test() {
    let mut client_endpoint = s2n_client();
    let mut server_endpoint = rustls_server();

    run(&mut server_endpoint, &mut client_endpoint, None);
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_s2n_server_client_auth_test() {
    let mut client_endpoint = s2n_client_with_client_auth().unwrap();
    let mut server_endpoint = s2n_server_with_client_auth().unwrap();

    run(&mut server_endpoint, &mut client_endpoint, None);
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_no_client_auth_s2n_server_requires_client_auth_test() {
    let mut client_endpoint = s2n_client();
    let mut server_endpoint = s2n_server_with_client_auth().unwrap();

    let test_result = run_result(&mut server_endpoint, &mut client_endpoint, None);

    // The handshake should fail because the server requires client auth,
    // but the client does not support it.
    assert!(test_result.is_err());
    let e = test_result.unwrap_err();
    assert_eq!(e.description().unwrap(), "UNEXPECTED_MESSAGE");
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_with_client_auth_s2n_server_does_not_require_client_auth_test() {
    let mut client_endpoint = s2n_client_with_client_auth().unwrap();
    let mut server_endpoint = s2n_server();

    let test_result = run_result(&mut server_endpoint, &mut client_endpoint, None);

    // The handshake should fail because the client requires client auth,
    // but the server does not support it.
    assert!(test_result.is_err());
    let e = test_result.unwrap_err();
    assert_eq!(e.description().unwrap(), "UNEXPECTED_MESSAGE");
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_with_client_auth_s2n_server_does_not_trust_client_certificate() {
    let mut client_endpoint = s2n_client_with_client_auth().unwrap();
    let mut server_endpoint = s2n_server_with_client_auth_verifier_rejects_client_certs().unwrap();

    let test_result = run_result(&mut server_endpoint, &mut client_endpoint, None);

    // The handshake should fail because the certificate presented by the client is rejected by the
    // application level host verification check on the cert.
    assert!(test_result.is_err());
    let e = test_result.unwrap_err();
    assert_eq!(e.description().unwrap(), "HANDSHAKE_FAILURE");
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_with_client_auth_s2n_server_does_not_trust_issuer() {
    let mut client_endpoint = s2n_client_with_untrusted_client_auth().unwrap();
    let mut server_endpoint = s2n_server_with_client_auth().unwrap();

    let test_result = run_result(&mut server_endpoint, &mut client_endpoint, None);

    // The handshake should fail because the certificate presented by the client is issued
    // by a CA that is not in the server trust store, even though the host name is validated.
    assert!(test_result.is_err());
    let e = test_result.unwrap_err();
    assert_eq!(e.description().unwrap(), "HANDSHAKE_FAILURE");
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_with_custom_hostname_auth_rejects_server_name() {
    let mut client_endpoint = s2n_client_with_fixed_hostname_auth("not-localhost").unwrap();
    let mut server_endpoint = s2n_server();

    let test_result = run_result(&mut server_endpoint, &mut client_endpoint, None);

    // The handshake should fail because the hostname ("localhost") is not validated
    assert!(test_result.is_err());
    let e = test_result.unwrap_err();
    assert_eq!(e.description().unwrap(), "HANDSHAKE_FAILURE");
}

#[test]
#[cfg_attr(miri, ignore)]
fn s2n_client_with_custom_hostname_auth_accepts_server_name() {
    let mut client_endpoint = s2n_client_with_fixed_hostname_auth("localhost").unwrap();
    let mut server_endpoint = s2n_server();

    run_result(&mut server_endpoint, &mut client_endpoint, None).unwrap();
}

/// Executes the handshake to completion
fn run_result<S: Endpoint, C: Endpoint>(
    server: &mut S,
    client: &mut C,
    client_hello_cb_done: Option<Arc<AtomicBool>>,
) -> Result<(), transport::Error> {
    let mut pair = tls::testing::Pair::new(server, client, "localhost".into());

    while pair.is_handshaking() {
        pair.poll(client_hello_cb_done.clone())?;
    }

    pair.finish();
    Ok(())
}

/// Executes the handshake to completion
fn run<S: Endpoint, C: Endpoint>(
    server: &mut S,
    client: &mut C,
    client_hello_cb_done: Option<Arc<AtomicBool>>,
) {
    run_result(server, client, client_hello_cb_done).unwrap();
}

#[test]
fn config_loader() {
    use crate::{ConfigLoader, Server};

    let server = Server::default();

    // make sure the loader can be a static type
    let server: Server<Server> = Server::from_loader(server);

    // make sure the loader can be a dynamic type
    let server: Box<dyn ConfigLoader> = Box::new(server);
    let mut server: Server<Box<dyn ConfigLoader>> = Server::from_loader(server);

    // make sure the server can actually create a session
    let _ = server.new_server_session(&1);
}


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

#[cfg(not(unix))]
pub use s2n_quic_rustls::*;
#[cfg(unix)]
pub use s2n_quic_tls::*;


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

//! This module contains all main runtime components for receiving and sending
//! data via the QUIC protocol.

#![deny(unused_must_use)]
extern crate alloc;

mod ack;
mod contexts;
mod processed_packet;
mod space;
mod sync;
mod transmission;
mod wakeup_queue;

pub mod connection;
pub mod endpoint;
pub mod path;
pub mod recovery;
pub mod stream;


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

use s2n_quic_core::{
    frame::{ack_elicitation::AckElicitation, path_validation, FrameTrait},
    inet::DatagramInfo,
    packet::number::PacketNumber,
};

/// Tracks information about a packet that has been processed
#[derive(Clone, Copy, Debug)]
pub struct ProcessedPacket<'a> {
    pub(crate) packet_number: PacketNumber,
    pub(crate) datagram: &'a DatagramInfo,
    pub(crate) ack_elicitation: AckElicitation,
    pub(crate) path_challenge_on_active_path: bool,
    pub(crate) frames: usize,
    pub(crate) path_validation_probing: path_validation::Probe,
    pub(crate) bytes_progressed: usize,
}

impl<'a> ProcessedPacket<'a> {
    /// Creates a processed packet tracker
    pub fn new(packet_number: PacketNumber, datagram: &'a DatagramInfo) -> Self {
        Self {
            packet_number,
            datagram,
            ack_elicitation: AckElicitation::default(),
            path_challenge_on_active_path: false,
            frames: 0,
            path_validation_probing: path_validation::Probe::default(),
            bytes_progressed: 0,
        }
    }

    /// Records information about a processed frame
    pub fn on_processed_frame<F: FrameTrait>(&mut self, frame: &F) {
        self.ack_elicitation |= frame.ack_elicitation();
        self.frames += 1;
        self.path_validation_probing |= frame.path_validation();
    }

    /// Returns `true` if any of the processed frames are ack eliciting
    pub fn is_ack_eliciting(&self) -> bool {
        self.ack_elicitation.is_ack_eliciting()
    }
}


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

//! A queue which allows to wake up a QUIC endpoint which is blocked on packet
//! reception or timers. This queue is used in case connections inside the endpoint
//! change their readiness state (e.g. they get ready to write).

use alloc::{collections::VecDeque, sync::Arc};
use core::{
    sync::atomic::{AtomicBool, Ordering},
    task::{Context, Waker},
};
use std::{sync::Mutex, task::Wake};

/// The shared state of the [`WakeupQueue`].
#[derive(Debug)]
struct QueueState<T> {
    /// The IDs of connections which have been woken
    woken_connections: VecDeque<T>,
    /// The waker which should be used to wake up the connection
    waker: Option<Waker>,
    /// Whether a wakeup is already in progress
    wakeup_in_progress: bool,
}

impl<T: Copy> QueueState<T> {
    fn new() -> Self {
        Self {
            woken_connections: VecDeque::new(),
            waker: None,
            wakeup_in_progress: false,
        }
    }

    fn queue_wakeup(&mut self, wakeup_handle_id: T) -> Option<Waker> {
        self.woken_connections.push_back(wakeup_handle_id);
        // If pushing another handle already notified the processing thread that it should dequeue
        // pending notifications there is no need to do this again.
        if self.wakeup_in_progress {
            return None;
        }

        self.wakeup_in_progress = true;
        self.waker.clone()
    }

    /// Polls for queued wakeup events.
    ///
    /// The method gets passed a queued which is used to store further wakeup events.
    /// It will returns a queue of occurred events.
    /// If no wakeup occurred, the method will store the passed [`Waker`] and notify it as soon as
    /// a wakeup occurred.
    fn poll_pending_wakeups(&mut self, swap_queue: &mut VecDeque<T>, context: &Context) {
        debug_assert!(swap_queue.is_empty());
        self.wakeup_in_progress = false;

        if self.woken_connections.is_empty() {
            // If no wakeup was pending, store or update the `Waker`
            match &self.waker {
                Some(w) => {
                    if !w.will_wake(context.waker()) {
                        self.waker = Some(context.waker().clone());
                    }
                }
                None => self.waker = Some(context.waker().clone()),
            }

            return;
        }

        core::mem::swap(&mut self.woken_connections, swap_queue);
    }

    #[cfg(any(feature = "testing", test))]
    fn test_inspect(&self) -> (usize, bool) {
        (self.woken_connections.len(), self.wakeup_in_progress)
    }
}

/// A queue which allows individual components to wakeups to a common blocked thread.
///
/// Multiple components can notify the thread to unblocked and to dequeue handles of components.
/// Each component is identified by a handle of type `T`.
///
/// A single thread is expected to deque the handles of blocked components and to inform those.
#[derive(Debug)]
pub struct WakeupQueue<T> {
    state: Arc<Mutex<QueueState<T>>>,
}

impl<T: Copy> WakeupQueue<T> {
    /// Creates a new `WakeupQueue`.
    ///
    /// If a wakeup is triggered, the given [`Waker`] will be notified.
    pub fn new() -> Self {
        Self {
            state: Arc::new(Mutex::new(QueueState::new())),
        }
    }

    /// Creates a new [`WakeupHandle`] which will wake up this [`WakeupQueue`] if
    /// [`WakeupHandle::wakeup`] is called.
    pub fn create_wakeup_handle(&self, wakeup_handle_id: T) -> WakeupHandle<T> {
        WakeupHandle::new(self.state.clone(), wakeup_handle_id)
    }

    /// Returns the list of component handles which need to get woken.
    /// Those component handles are retrieved inside a `VecDeque`. In order to avoid
    /// memory allocations, the caller is expected to pass in a new `VecDequeue` which will
    /// by utilized for further queueing. Thereby a double-buffering approach for wakeups is
    /// enabled.
    pub fn poll_pending_wakeups(&mut self, swap_queue: &mut VecDeque<T>, context: &Context) {
        let mut guard = self
            .state
            .lock()
            .expect("Locking can only fail if locks are poisoned");
        guard.poll_pending_wakeups(swap_queue, context)
    }

    #[cfg(any(feature = "testing", test))]
    fn test_state(&self) -> Arc<Mutex<QueueState<T>>> {
        self.state.clone()
    }
}

/// A handle which refers to a wakeup queue. The handles allows to notify the
/// queue that a wakeup is required, and that after the wakeup the owner of the handle
/// wants to be notified.
#[derive(Debug)]
pub struct WakeupHandle<T> {
    /// The queue this handle is referring to
    queue: Arc<Mutex<QueueState<T>>>,
    /// The internal ID of this wakeup handle. This can be used to distinguish which
    /// handle had woken up the [`WakeupQueue`].
    wakeup_handle_id: T,
    /// Whether a wakeup for this handle had already been queued since the last time
    /// the wakeup handler was called
    wakeup_queued: AtomicBool,
}

impl<T: Copy> WakeupHandle<T> {
    /// Creates a new [`WakeupHandle`] which delegates wakeups to the given `queue`.
    fn new(queue: Arc<Mutex<QueueState<T>>>, wakeup_handle_id: T) -> Self {
        Self {
            queue,
            wakeup_handle_id,
            wakeup_queued: AtomicBool::new(false),
        }
    }

    /// Notifies the queue to wake up. If a `wakeup()` had been issued for the same
    /// [`WakeupHandle`] without having been handled yet, the new [`wakeup()`] request will be
    /// ignored, since the wakeup will already be pending.
    pub fn wakeup(&self) {
        // Check if a wakeup had been queued earlier
        if self.wakeup_queued.swap(true, Ordering::SeqCst) {
            return;
        }

        // Enqueue the wakeup request
        let maybe_waker = {
            let mut guard = self
                .queue
                .lock()
                .expect("Locking can only fail if locks are poisoned");
            guard.queue_wakeup(self.wakeup_handle_id)
        };

        // If the queue handling thread wasn't notified earlier by another thread,
        // notify it now.
        if let Some(waker) = maybe_waker {
            waker.wake();
        }
    }

    /// Notifies the `WakeupHandle` that a wakeup for this handle had been processed.
    ///
    /// Further calls to [`wakeup`] will be queued again.
    pub fn wakeup_handled(&self) {
        self.wakeup_queued.store(false, Ordering::SeqCst);
    }
}

impl<T: Copy> Wake for WakeupHandle<T> {
    fn wake(self: Arc<Self>) {
        self.wakeup()
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.wakeup()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures_test::task::new_count_waker;

    macro_rules! vec_deque {
        ($($value:expr),*) => {{
            #[allow(unused_mut)]
            let mut value = VecDeque::new();
            $(
                value.push_back($value);
            )*
            value
        }}
    }

    #[cfg(any(feature = "testing", test))]
    #[test]
    fn waker() {
        // check state len and if pending wakeup
        macro_rules! check_state {
            ($state:ident, $len:literal, $wakeup_in_progress:literal) => {{
                let state = $state
                    .lock()
                    .expect("Locking can only fail if locks are poisoned");
                assert_eq!(state.test_inspect(), ($len, $wakeup_in_progress));
            }};
        }

        // check queued wakeups and reset state (wakeup_handled)
        macro_rules! reset_state {
            ($state:ident, $len:literal, $handle1:ident, $handle2:ident) => {{
                let mut state = $state
                    .lock()
                    .expect("Locking can only fail if locks are poisoned");
                let (waker, _counter) = new_count_waker();
                let mut pending = VecDeque::new();
                state.poll_pending_wakeups(&mut pending, &Context::from_waker(&waker));
                assert_eq!(pending.len(), $len);

                $handle1.wakeup_handled();
                $handle2.wakeup_handled();
            }};
        }

        let queue = WakeupQueue::new();
        let state = queue.test_state();
        check_state!(state, 0, false);

        let handle1 = Arc::new(queue.create_wakeup_handle(1u32));
        let handle2 = Arc::new(queue.create_wakeup_handle(2u32));
        let waker1 = Waker::from(handle1.clone());
        let waker2 = Waker::from(handle2.clone());

        // wake_by_ref
        waker1.wake_by_ref();
        check_state!(state, 1, true);
        // calling wake on same handle should not queue a wakeup
        waker1.wake_by_ref();
        check_state!(state, 1, true);
        waker2.wake_by_ref();
        check_state!(state, 2, true);

        reset_state!(state, 2, handle1, handle2);
        check_state!(state, 0, false);

        // clone and call wake
        waker1.clone().wake();
        check_state!(state, 1, true);
        waker2.clone().wake();
        check_state!(state, 2, true);

        // drop handle
        drop(waker1.clone());
        drop(waker2.clone());

        reset_state!(state, 2, handle1, handle2);
        check_state!(state, 0, false);

        // call wake on original handle
        waker1.wake();
        check_state!(state, 1, true);
        waker2.wake();
        check_state!(state, 2, true);
    }

    #[test]
    fn queue_wakeups() {
        let (waker, counter) = new_count_waker();
        let mut queue = WakeupQueue::new();
        let mut pending = VecDeque::new();

        let handle1 = queue.create_wakeup_handle(1u32);
        let handle2 = queue.create_wakeup_handle(2u32);
        assert_eq!(counter, 0);

        // Initially no wakeup should be signalled - but the Waker should be stored
        queue.poll_pending_wakeups(&mut pending, &Context::from_waker(&waker));
        assert!(pending.is_empty());

        // After a wakeup the waker should be notified
        handle1.wakeup();
        assert_eq!(counter, 1);
        // A second wakeup on the same handle should not lead to another global wakeup
        handle1.wakeup();
        assert_eq!(counter, 1);

        // Even a second wakeup on the other handle should not lead to a global wakeup
        handle2.wakeup();
        assert_eq!(counter, 1);

        // The pending wakeups should be signaled
        queue.poll_pending_wakeups(&mut pending, &Context::from_waker(&waker));
        assert_eq!(vec_deque![1u32, 2u32], pending);
        pending.clear();

        // In the next query no wakeups should be signaled
        queue.poll_pending_wakeups(&mut pending, &Context::from_waker(&waker));
        assert!(pending.is_empty());

        // As long as wakeups are not handled, no new ones are enqueued
        handle2.wakeup();
        assert_eq!(counter, 1);
        queue.poll_pending_wakeups(&mut pending, &Context::from_waker(&waker));
        assert!(pending.is_empty());

        // If wakeups are handled, wakeups are forwarded again
        handle1.wakeup_handled();
        handle2.wakeup_handled();

        handle2.wakeup();
        assert_eq!(counter, 2);
        queue.poll_pending_wakeups(&mut pending, &Context::from_waker(&waker));
        assert_eq!(vec_deque![2u32], pending);
        pending.clear();
    }
}


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

use core::ops::RangeInclusive;
use s2n_quic_core::{ack, packet::number::PacketNumber};

pub type AckRange = RangeInclusive<PacketNumber>;

#[derive(Clone, Debug, Default)]
pub struct AckElicitingTransmissionSet {
    /// A stable ack-eliciting transmission
    ///
    /// In this case, "stable" means the oldest transmission that
    /// hasn't been acked by the peer.
    ///
    /// This transmission is stored to ensure the packet number ranges
    /// are always either removed or declared lost. Without it,
    /// the TX packet number would be a moving target
    /// and packet number ranges would never be removed.
    stable: Option<AckElicitingTransmission>,

    /// The latest ack-eliciting transmission
    latest: Option<AckElicitingTransmission>,
}

impl AckElicitingTransmissionSet {
    /// Called when an ACK frame is bundled with an ack eliciting packet
    pub fn on_transmit(&mut self, transmission: AckElicitingTransmission) {
        self.latest = Some(transmission);

        // only set the stable transmission if it's not set
        if self.stable.is_none() {
            self.stable = Some(transmission);
        }
    }

    /// Called when a set of packets was acknowledged or lost
    pub fn on_update<A: ack::Set>(&mut self, ack_set: &A) -> Option<AckRange> {
        if let Some(ack_range) = self
            .latest
            .as_ref()
            .and_then(|transmission| transmission.ack_range(ack_set))
        {
            // the latest was either ACKed or lost so there is nothing pending
            self.stable = None;
            self.latest = None;
            return Some(ack_range);
        }

        if let Some(ack_range) = self
            .stable
            .as_ref()
            .and_then(|transmission| transmission.ack_range(ack_set))
        {
            // assign the latest transmission to stable
            self.stable = self.latest;
            return Some(ack_range);
        }

        None
    }

    /// Returns `true` if there are no pending ack-eliciting transmissions
    /// Only used by tests
    #[cfg(test)]
    pub fn is_empty(&self) -> bool {
        self.stable.is_none() && self.latest.is_none()
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd)]
pub struct AckElicitingTransmission {
    pub sent_in_packet: PacketNumber,
    pub largest_received_packet_number_acked: PacketNumber,
}

impl AckElicitingTransmission {
    /// Called when a set of packets was acknowledged or lost
    pub fn ack_range<A: ack::Set>(&self, ack_set: &A) -> Option<AckRange> {
        //= https://www.rfc-editor.org/rfc/rfc9000#section-13.2.4
        //# When a packet containing an ACK frame is acknowledged, the receiver can stop
        //# acknowledging packets less than or equal to the Largest Acknowledged
        //# field in the sent ACK frame.
        if ack_set.contains(self.sent_in_packet) {
            // create a range from 0..=largest
            let pn_zero = self
                .largest_received_packet_number_acked
                .space()
                .new_packet_number(Default::default());

            Some(pn_zero..=self.largest_received_packet_number_acked)
        } else {
            None
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{super::tests::transmissions_iter, *};

    /// This test is meant to simulate an immediate ACK rate from a peer
    #[test]
    fn latest_ack_test() {
        let mut set = AckElicitingTransmissionSet::default();
        assert!(set.is_empty());

        let mut transmissions = transmissions_iter();

        for transmission_count in 0..10 {
            let stable = transmissions.next().unwrap();
            set.on_transmit(stable);

            // simulate a few transmissions before ACKing
            let mut latest = stable;
            for _ in 0..transmission_count {
                latest = transmissions.next().unwrap();
                set.on_transmit(latest);
            }

            // ACKing `latest` simulates an immediate ACK from peer
            assert_eq!(
                set.on_update(&latest.sent_in_packet).unwrap().end(),
                &latest.largest_received_packet_number_acked
            );

            // ACKing the latest removes all items
            assert!(set.is_empty());
        }
    }

    /// This test is meant to simulate a delayed ACK rate from a peer
    #[test]
    fn stable_ack_test() {
        let mut set = AckElicitingTransmissionSet::default();
        assert!(set.is_empty());

        let mut transmissions = transmissions_iter();

        // create an initial transmission
        let mut stable = transmissions.next().unwrap();
        set.on_transmit(stable);

        for transmission_count in 0..10 {
            // simulate a few transmissions before ACKing
            let mut latest = stable;
            for _ in 0..transmission_count {
                latest = transmissions.next().unwrap();
                set.on_transmit(latest);
            }

            // ACKing `stable` simulates a delayed ACK from peer
            assert_eq!(
                set.on_update(&stable.sent_in_packet).unwrap().end(),
                &stable.largest_received_packet_number_acked
            );

            // on update latest is moved into stable
            stable = latest;

            // if there weren't any transmissions between, the set
            // should be empty
            if transmission_count == 0 {
                assert!(set.is_empty());
                set.on_transmit(stable);
            }
        }
    }

    #[test]
    fn size_of_snapshots() {
        use core::mem::size_of;
        use insta::assert_debug_snapshot;

        assert_debug_snapshot!(
            "AckElicitingTransmission",
            size_of::<AckElicitingTransmission>()
        );
        assert_debug_snapshot!(
            "AckElicitingTransmissionSet",
            size_of::<AckElicitingTransmissionSet>()
        );
    }
}