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

use s2n_quic_transport::stream::Stream;

/// A QUIC stream that may send and receive data.
#[derive(Debug)]
pub struct BidirectionalStream(Stream);

impl BidirectionalStream {
    pub(crate) const fn new(stream: Stream) -> Self {
        Self(stream)
    }

    /// Returns the stream's identifier
    ///
    /// This value is unique to a particular connection. The format follows the same as what is
    /// defined in the
    /// [QUIC Transport RFC](https://www.rfc-editor.org/rfc/rfc9000.html#name-stream-types-and-identifier).
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # async fn test() -> s2n_quic::stream::Result<()> {
    /// #   let connection: s2n_quic::connection::Connection = todo!();
    /// #
    /// let stream = connection.open_send_stream().await?;
    /// println!("New stream's id: {}", stream.id());
    /// #
    /// #   Ok(())
    /// # }
    /// ```
    pub fn id(&self) -> u64 {
        self.0.id().into()
    }

    impl_connection_api!(|stream| crate::connection::Handle(stream.0.connection().clone()));

    impl_receive_stream_api!(|stream, call| call!(stream.0));

    impl_send_stream_api!(|stream, call| call!(stream.0));

    /// Splits the stream into [`ReceiveStream`](crate::stream::ReceiveStream) and
    /// [`SendStream`](crate::stream::SendStream) halves.
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use bytes::Bytes;
    /// # async fn test() -> s2n_quic::stream::Result<()> {
    /// #   let connection: s2n_quic::connection::Connection = todo!();
    /// #
    /// let stream = connection.open_bidirectional_stream().await?;
    /// let (recv, mut send) = stream.split();
    ///
    /// tokio::spawn(async move {
    ///     let _ = send.send(Bytes::from_static(&[1, 2, 3])).await;
    /// });
    ///
    /// while let Some(chunk) = recv.receive().await? {
    ///     println!("received: {:?}", chunk);
    /// }
    /// #
    /// #   Ok(())
    /// # }
    /// ```
    pub fn split(self) -> (crate::stream::ReceiveStream, crate::stream::SendStream) {
        let (recv, send) = self.0.split();
        (
            crate::stream::ReceiveStream::new(recv),
            crate::stream::SendStream::new(send),
        )
    }
}

impl_receive_stream_trait!(BidirectionalStream, |stream, call| call!(stream.0));
impl_send_stream_trait!(BidirectionalStream, |stream, call| call!(stream.0));
impl_splittable_stream_trait!(BidirectionalStream, |stream| {
    let (recv, send) = Self::split(stream);
    (Some(recv), Some(send))
});


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

macro_rules! impl_connection_api {
    (| $self:ident | $convert:expr) => {
        /// Returns the [`connection::Handle`](crate::connection::Handle) associated with the stream.
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # let stream: s2n_quic::stream::Stream = todo!();
        /// #
        /// let connection = stream.connection();
        ///
        /// println!("The stream's connection id is: {}", connection.id());
        /// ```
        #[inline]
        pub fn connection(&self) -> $crate::connection::Handle {
            let $self = self;
            $convert
        }
    };
}


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

use super::{BidirectionalStream, SendStream};

/// An enum of all the possible types of QUIC streams that may be opened by a local endpoint.
#[derive(Debug)]
pub enum LocalStream {
    Bidirectional(BidirectionalStream),
    Send(SendStream),
}

impl LocalStream {
    /// Returns the stream's identifier
    ///
    /// This value is unique to a particular connection. The format follows the same as what is
    /// defined in the
    /// [QUIC Transport RFC](https://www.rfc-editor.org/rfc/rfc9000.html#name-stream-types-and-identifier).
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # async fn test() -> s2n_quic::stream::Result<()> {
    /// #   let connection: s2n_quic::connection::Connection = todo!();
    /// #
    /// use s2n_quic::stream::{Type, LocalStream};
    ///
    /// let stream: LocalStream = connection.open_stream(Type::Bidirectional).await?;
    /// println!("New stream's id: {}", stream.id());
    /// #
    /// #   Ok(())
    /// # }
    /// ```
    pub fn id(&self) -> u64 {
        match self {
            Self::Bidirectional(stream) => stream.id(),
            Self::Send(stream) => stream.id(),
        }
    }

    impl_connection_api!(|stream| match stream {
        LocalStream::Bidirectional(stream) => stream.connection(),
        LocalStream::Send(stream) => stream.connection(),
    });

    impl_receive_stream_api!(|stream, dispatch| match stream {
        LocalStream::Bidirectional(stream) => dispatch!(stream),
        LocalStream::Send(_stream) => dispatch!(),
    });

    impl_send_stream_api!(|stream, dispatch| match stream {
        LocalStream::Bidirectional(stream) => dispatch!(stream),
        LocalStream::Send(stream) => dispatch!(stream),
    });

    impl_splittable_stream_api!();
}

impl_receive_stream_trait!(LocalStream, |stream, dispatch| match stream {
    LocalStream::Bidirectional(stream) => dispatch!(stream),
    LocalStream::Send(_stream) => dispatch!(),
});
impl_send_stream_trait!(LocalStream, |stream, dispatch| match stream {
    LocalStream::Bidirectional(stream) => dispatch!(stream),
    LocalStream::Send(stream) => dispatch!(stream),
});
impl_splittable_stream_trait!(LocalStream, |stream| match stream {
    LocalStream::Bidirectional(stream) => super::SplittableStream::split(stream),
    LocalStream::Send(stream) => super::SplittableStream::split(stream),
});

impl From<SendStream> for LocalStream {
    #[inline]
    fn from(stream: SendStream) -> Self {
        Self::Send(stream)
    }
}

impl From<BidirectionalStream> for LocalStream {
    #[inline]
    fn from(stream: BidirectionalStream) -> Self {
        Self::Bidirectional(stream)
    }
}


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

use super::{BidirectionalStream, ReceiveStream};

/// An enum of all the possible types of QUIC streams that may be opened by a peer.
#[derive(Debug)]
pub enum PeerStream {
    Bidirectional(BidirectionalStream),
    Receive(ReceiveStream),
}

impl PeerStream {
    /// Returns the stream's identifier
    ///
    /// This value is unique to a particular connection. The format follows the same as what is
    /// defined in the
    /// [QUIC Transport RFC](https://www.rfc-editor.org/rfc/rfc9000.html#name-stream-types-and-identifier).
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # async fn test() -> s2n_quic::stream::Result<()> {
    /// #   let connection: s2n_quic::connection::Connection = todo!();
    /// #
    /// use s2n_quic::stream::Type;
    ///
    /// while let Some(stream) = connection.accept().await? {
    ///     println!("New stream's id: {}", stream.id());
    /// }
    /// #
    /// #   Ok(())
    /// # }
    /// ```
    pub fn id(&self) -> u64 {
        match self {
            Self::Bidirectional(stream) => stream.id(),
            Self::Receive(stream) => stream.id(),
        }
    }

    impl_connection_api!(|stream| match stream {
        PeerStream::Bidirectional(stream) => stream.connection(),
        PeerStream::Receive(stream) => stream.connection(),
    });

    impl_receive_stream_api!(|stream, dispatch| match stream {
        PeerStream::Bidirectional(stream) => dispatch!(stream),
        PeerStream::Receive(stream) => dispatch!(stream),
    });

    impl_send_stream_api!(|stream, dispatch| match stream {
        PeerStream::Bidirectional(stream) => dispatch!(stream),
        PeerStream::Receive(_stream) => dispatch!(),
    });

    impl_splittable_stream_api!();
}

impl_receive_stream_trait!(PeerStream, |stream, dispatch| match stream {
    PeerStream::Bidirectional(stream) => dispatch!(stream),
    PeerStream::Receive(stream) => dispatch!(stream),
});
impl_send_stream_trait!(PeerStream, |stream, dispatch| match stream {
    PeerStream::Bidirectional(stream) => dispatch!(stream),
    PeerStream::Receive(_stream) => dispatch!(),
});
impl_splittable_stream_trait!(PeerStream, |stream| match stream {
    PeerStream::Bidirectional(stream) => super::SplittableStream::split(stream),
    PeerStream::Receive(stream) => super::SplittableStream::split(stream),
});

impl From<ReceiveStream> for PeerStream {
    #[inline]
    fn from(stream: ReceiveStream) -> Self {
        Self::Receive(stream)
    }
}

impl From<BidirectionalStream> for PeerStream {
    #[inline]
    fn from(stream: BidirectionalStream) -> Self {
        Self::Bidirectional(stream)
    }
}


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

use s2n_quic_transport::stream;

/// A QUIC stream that is only allowed to receive data.
#[derive(Debug)]
pub struct ReceiveStream(stream::ReceiveStream);

macro_rules! impl_receive_stream_api {
    (| $stream:ident, $dispatch:ident | $dispatch_body:expr) => {
        /// Receives a chunk of data from the stream.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Ok(Some(chunk))` if the stream is open and data was available.
        /// - `Ok(None)` if the stream was finished and all of the data was consumed.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// #   let mut stream: s2n_quic::stream::ReceiveStream = todo!();
        /// #
        /// while let Some(chunk) = stream.receive().await? {
        ///     println!("received: {:?}", chunk);
        /// }
        ///
        /// println!("finished");
        /// #
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub async fn receive(&mut self) -> $crate::stream::Result<Option<bytes::Bytes>> {
            ::futures::future::poll_fn(|cx| self.poll_receive(cx)).await
        }

        /// Poll for a chunk of data from the stream.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Poll::Pending` if the stream is waiting to receive data from the peer. In this case,
        ///   the caller should retry receiving after the [`Waker`](core::task::Waker) on the provided
        ///   [`Context`](core::task::Context) is notified.
        /// - `Poll::Ready(Ok(Some(chunk)))` if the stream is open and data was available.
        /// - `Poll::Ready(Ok(None))` if the stream was finished and all of the data was consumed.
        /// - `Poll::Ready(Err(e))` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn poll_receive(
            &mut self,
            cx: &mut core::task::Context,
        ) -> core::task::Poll<$crate::stream::Result<Option<bytes::Bytes>>> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_readable()).into()
                };
                ($variant: expr) => {
                    $variant.poll_receive(cx)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Receives a slice of chunks of data from the stream.
        ///
        /// This can be more efficient than calling [`receive`](Self::receive) for each chunk,
        /// especially when receiving large amounts of data.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Ok((count, is_open))` if the stream received data into the slice,
        ///   where `count` was the number of chunks received, and `is_open` indicating if the stream is
        ///   still open. If `is_open == true`, `count` will be at least `1`. If `is_open == false`, future calls to
        ///   [`receive_vectored`](Self::receive_vectored) will always return
        ///   `Ok((0, false))`.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// #   let mut stream: s2n_quic::stream::ReceiveStream = todo!();
        /// #
        /// # use bytes::Bytes;
        /// #
        /// loop {
        ///     let mut chunks = [Bytes::new(), Bytes::new(), Bytes::new()];
        ///     let (count, is_open) = stream.receive_vectored(&mut chunks).await?;
        ///
        ///     for chunk in &chunks[..count] {
        ///         println!("received: {:?}", chunk);
        ///     }
        ///
        ///     if !is_open {
        ///         break;
        ///     }
        /// }
        ///
        /// println!("finished");
        /// #
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub async fn receive_vectored(
            &mut self,
            chunks: &mut [bytes::Bytes],
        ) -> $crate::stream::Result<(usize, bool)> {
            ::futures::future::poll_fn(|cx| self.poll_receive_vectored(chunks, cx)).await
        }

        /// Polls for receiving a slice of chunks of data from the stream.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Poll::Pending` if the stream is waiting to receive data from the peer. In this case,
        ///   the caller should retry receiving after the [`Waker`](core::task::Waker) on the provided
        ///   [`Context`](core::task::Context) is notified.
        /// - `Poll::Ready(Ok((count, is_open)))` if the stream received data into the slice,
        ///   where `count` was the number of chunks received, and `is_open` indicating if the stream is
        ///   still open. If `is_open == true`, `count` will be at least `1`. If `is_open == false`, future calls to
        ///   [`poll_receive_vectored`](Self::poll_receive_vectored) will always return
        ///   `Poll::Ready(Ok((0, false)))`.
        /// - `Poll::Ready(Err(e))` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn poll_receive_vectored(
            &mut self,
            chunks: &mut [bytes::Bytes],
            cx: &mut core::task::Context,
        ) -> core::task::Poll<$crate::stream::Result<(usize, bool)>> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_readable()).into()
                };
                ($variant: expr) => {
                    $variant.poll_receive_vectored(chunks, cx)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Notifies the peer to stop sending data on the stream.
        ///
        /// This requests the peer to finish the stream as soon as possible
        /// by issuing a [`reset`](crate::stream::SendStream::reset) with the
        /// provided [`error_code`](crate::application::Error).
        ///
        /// Since this is merely a request for the peer to reset the stream, the
        /// stream will not immediately be in a reset state after issuing this
        /// call.
        ///
        /// If the stream has already been reset by the peer or if all data has
        /// been received, the call will not trigger any action.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Ok(())` if the stop sending message was enqueued for the peer.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// #   let mut connection: s2n_quic::connection::Connection = todo!();
        /// #
        /// while let Some(stream) = connection.accept_receive_stream().await? {
        ///     stream.stop_sending(123u8.into());
        /// }
        /// #
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub fn stop_sending(
            &mut self,
            error_code: $crate::application::Error,
        ) -> $crate::stream::Result<()> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_readable())
                };
                ($variant: expr) => {
                    $variant.stop_sending(error_code)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Create a batch request for receiving data
        #[inline]
        pub(crate) fn rx_request(
            &mut self,
        ) -> $crate::stream::Result<s2n_quic_transport::stream::RxRequest> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_readable())
                };
                ($variant: expr) => {
                    $variant.rx_request()
                };
            }

            let $stream = self;
            $dispatch_body
        }
    };
}

macro_rules! impl_receive_stream_trait {
    ($name:ident, | $stream:ident, $dispatch:ident | $dispatch_body:expr) => {
        impl futures::stream::Stream for $name {
            type Item = $crate::stream::Result<bytes::Bytes>;

            #[inline]
            fn poll_next(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<Option<Self::Item>> {
                match futures::ready!(self.poll_receive(cx)) {
                    Ok(Some(v)) => Some(Ok(v)),
                    Ok(None) => None,
                    Err(err) => Some(Err(err)),
                }
                .into()
            }
        }

        impl futures::io::AsyncRead for $name {
            fn poll_read(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
                buf: &mut [u8],
            ) -> core::task::Poll<std::io::Result<usize>> {
                use bytes::Bytes;

                if buf.is_empty() {
                    return Ok(0).into();
                }

                // create some chunks on the stack to receive into
                // TODO investigate a better default number
                let mut chunks = [
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                ];

                let high_watermark = buf.len();

                let response = futures::ready!(self
                    .rx_request()?
                    .receive(&mut chunks)
                    // don't receive more than we're capable of storing
                    .with_high_watermark(high_watermark)
                    .poll(Some(cx))?
                    .into_poll());

                let chunks = &chunks[..response.chunks.consumed];
                let mut bufs = [buf];
                let copied_len = s2n_quic_core::slice::vectored_copy(chunks, &mut bufs);

                debug_assert_eq!(
                    copied_len, response.bytes.consumed,
                    "the consumed bytes should always have enough capacity in bufs"
                );

                Ok(response.bytes.consumed).into()
            }

            fn poll_read_vectored(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
                bufs: &mut [futures::io::IoSliceMut],
            ) -> core::task::Poll<std::io::Result<usize>> {
                use bytes::Bytes;

                if bufs.is_empty() {
                    return Ok(0).into();
                }

                // create some chunks on the stack to receive into
                // TODO investigate a better default number
                let mut chunks = [
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                ];

                let high_watermark = bufs.iter().map(|buf| buf.len()).sum();

                let response = futures::ready!(self
                    .rx_request()?
                    .receive(&mut chunks)
                    // don't receive more than we're capable of storing
                    .with_high_watermark(high_watermark)
                    .poll(Some(cx))?
                    .into_poll());

                let chunks = &chunks[..response.chunks.consumed];
                let copied_len = s2n_quic_core::slice::vectored_copy(chunks, bufs);

                debug_assert_eq!(
                    copied_len, response.bytes.consumed,
                    "the consumed bytes should always have enough capacity in bufs"
                );

                Ok(copied_len).into()
            }
        }

        impl tokio::io::AsyncRead for $name {
            fn poll_read(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
                buf: &mut tokio::io::ReadBuf,
            ) -> core::task::Poll<std::io::Result<()>> {
                use bytes::Bytes;

                if buf.remaining() == 0 {
                    return Ok(()).into();
                }

                // create some chunks on the stack to receive into
                // TODO investigate a better default number
                let mut chunks = [
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                    Bytes::new(),
                ];

                let high_watermark = buf.remaining();

                let response = futures::ready!(self
                    .rx_request()?
                    .receive(&mut chunks)
                    // don't receive more than we're capable of storing
                    .with_high_watermark(high_watermark)
                    .poll(Some(cx))?
                    .into_poll());

                for chunk in &chunks[..response.chunks.consumed] {
                    buf.put_slice(chunk);
                }

                Ok(()).into()
            }
        }
    };
}

impl ReceiveStream {
    #[inline]
    pub(crate) const fn new(stream: stream::ReceiveStream) -> Self {
        Self(stream)
    }

    /// Returns the stream's identifier
    ///
    /// This value is unique to a particular connection. The format follows the same as what is
    /// defined in the
    /// [QUIC Transport RFC](https://www.rfc-editor.org/rfc/rfc9000.html#name-stream-types-and-identifier).
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # async fn test() -> s2n_quic::stream::Result<()> {
    /// #   let connection: s2n_quic::connection::Connection = todo!();
    /// #
    /// while let Some(stream) = connection.accept_receive_stream().await? {
    ///     println!("New stream's id: {}", stream.id());
    /// }
    /// #
    /// #   Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn id(&self) -> u64 {
        self.0.id().into()
    }

    impl_connection_api!(|stream| crate::connection::Handle(stream.0.connection().clone()));

    impl_receive_stream_api!(|stream, dispatch| dispatch!(stream.0));
}

impl_splittable_stream_trait!(ReceiveStream, |stream| (Some(stream), None));
impl_receive_stream_trait!(ReceiveStream, |stream, dispatch| dispatch!(stream.0));


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

use s2n_quic_transport::stream;

/// A QUIC stream that is only allowed to send data.
#[derive(Debug)]
pub struct SendStream(stream::SendStream);

macro_rules! impl_send_stream_api {
    (| $stream:ident, $dispatch:ident | $dispatch_body:expr) => {
        /// Enqueues a chunk of data for sending it towards the peer.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Ok(())` if the data was enqueued for sending.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// #   let stream: s2n_quic::stream::SendStream = todo!();
        /// #
        /// let data = bytes::Bytes::from_static(&[1, 2, 3, 4]);
        /// stream.send(data).await?;
        /// #
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub async fn send(&mut self, mut data: bytes::Bytes) -> $crate::stream::Result<()> {
            ::futures::future::poll_fn(|cx| self.poll_send(&mut data, cx)).await
        }

        /// Enqueues a chunk of data for sending it towards the peer.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Poll::Pending` if the stream's send buffer capacity is currently exhausted. In this case,
        ///   the caller should retry sending after the [`Waker`](core::task::Waker) on the provided
        ///   [`Context`](core::task::Context) is notified.
        /// - `Poll::Ready(Ok(()))` if the data was enqueued for sending. The provided `chunk` will
        ///   be replaced with an empty [`Bytes`](bytes::Bytes).
        /// - `Poll::Ready(Err(e))` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn poll_send(
            &mut self,
            chunk: &mut bytes::Bytes,
            cx: &mut core::task::Context,
        ) -> core::task::Poll<$crate::stream::Result<()>> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable()).into()
                };
                ($variant: expr) => {
                    $variant.poll_send(chunk, cx)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Enqueues a slice of chunks of data for sending it towards the peer.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Ok(())` if all of the chunks of data were enqueued for sending. Each of the
        ///   consumed [`Bytes`](bytes::Bytes) will be replaced with an empty [`Bytes`](bytes::Bytes).
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// #   let stream: s2n_quic::stream::SendStream = todo!();
        /// #
        /// let mut data1 = bytes::Bytes::from_static(&[1, 2, 3]);
        /// let mut data2 = bytes::Bytes::from_static(&[4, 5, 6]);
        /// let mut data3 = bytes::Bytes::from_static(&[7, 8, 9]);
        /// let chunks = [data1, data2, data3];
        /// stream.send_vectored(&mut chunks).await?;
        /// #
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub async fn send_vectored(
            &mut self,
            chunks: &mut [bytes::Bytes],
        ) -> $crate::stream::Result<()> {
            let mut sent_chunks = 0;

            ::futures::future::poll_fn(|cx| {
                sent_chunks +=
                    ::futures::ready!(self.poll_send_vectored(&mut chunks[sent_chunks..], cx))?;
                if sent_chunks == chunks.len() {
                    return Ok(()).into();
                }
                core::task::Poll::Pending
            })
            .await
        }

        /// Polls enqueueing a slice of chunks of data for sending it towards the peer.
        ///
        /// # Return value
        ///
        /// The function returns:
        ///
        /// - `Poll::Pending` if the stream's send buffer capacity is currently exhausted. In this case,
        ///   the caller should retry sending after the [`Waker`](core::task::Waker) on the provided
        ///   [`Context`](core::task::Context) is notified.
        /// - `Poll::Ready(Ok(count))` if one or more chunks of data were enqueued for sending. Any of the
        ///   consumed [`Bytes`](bytes::Bytes) will be replaced with an empty [`Bytes`](bytes::Bytes).
        ///   If `count` does not equal the total number of chunks, the stream will store the
        ///   [Waker](core::task::Waker) and notify the task once more capacity is available.
        /// - `Poll::Ready(Err(e))` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn poll_send_vectored(
            &mut self,
            chunks: &mut [bytes::Bytes],
            cx: &mut core::task::Context,
        ) -> core::task::Poll<$crate::stream::Result<usize>> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable()).into()
                };
                ($variant: expr) => {
                    $variant.poll_send_vectored(chunks, cx)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Polls send readiness for the given stream.
        ///
        /// This method _must_ be called before calling [`send_data`](Self::send_data).
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Poll::Pending` if the stream's send buffer capacity is currently exhausted. In this case,
        ///   the caller should retry sending after the [`Waker`](core::task::Waker) on the provided
        ///   [`Context`](core::task::Context) is notified.
        /// - `Poll::Ready(Ok(available_bytes))` if the stream is ready to send data, where
        ///   `available_bytes` is how many bytes the stream can currently accept.
        /// - `Poll::Ready(Err(e))` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn poll_send_ready(
            &mut self,
            cx: &mut core::task::Context,
        ) -> core::task::Poll<$crate::stream::Result<usize>> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable()).into()
                };
                ($variant: expr) => {
                    $variant.poll_send_ready(cx)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Sends data on the stream without blocking the task.
        ///
        /// [`poll_send_ready`](Self::poll_send_ready) _must_ be called before calling this method.
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Ok(())` if the data was enqueued for sending.
        /// - `Err(SendingBlocked)` if the stream did not have enough capacity to enqueue the
        ///   `chunk`.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn send_data(&mut self, chunk: bytes::Bytes) -> $crate::stream::Result<()> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable())
                };
                ($variant: expr) => {
                    $variant.send_data(chunk)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Flushes the stream and waits for the peer to receive all outstanding data.
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Ok(())` if the send buffer was completely flushed and acknowledged by
        ///   the peer.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// #   let stream: s2n_quic::stream::SendStream = todo!();
        /// #
        /// let data = bytes::Bytes::from_static(&[1, 2, 3, 4]);
        /// stream.send(data).await?;
        /// stream.flush().await?;
        /// // at this point, the peer has received all of the `data`
        /// #
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub async fn flush(&mut self) -> $crate::stream::Result<()> {
            ::futures::future::poll_fn(|cx| self.poll_flush(cx)).await
        }

        /// Polls flushing the stream and waits for the peer to receive all outstanding data.
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Poll::Pending` if the stream's send buffer is still being sent. In this case,
        ///   the caller should retry sending after the [`Waker`](core::task::Waker) on the provided
        ///   [`Context`](core::task::Context) is notified.
        /// - `Poll::Ready(Ok(()))` if the send buffer was completely flushed and acknowledged by
        ///   the peer.
        /// - `Poll::Ready(Err(e))` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn poll_flush(
            &mut self,
            cx: &mut core::task::Context,
        ) -> core::task::Poll<$crate::stream::Result<()>> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable()).into()
                };
                ($variant: expr) => {
                    $variant.poll_flush(cx)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Marks the stream as finished.
        ///
        /// This method returns immediately without notifying the caller that all of the outstanding
        /// data has been received by the peer. An application wanting to both [`finish`](Self::finish)
        /// and [`flush`](Self::flush) the outstanding data can use [`close`](Self::close) to accomplish
        /// this.
        ///
        /// __NOTE__: This method will be called when the [`stream`](Self) is dropped.
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Ok(())` if the stream was finished successfully.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn finish(&mut self) -> $crate::stream::Result<()> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable()).into()
                };
                ($variant: expr) => {
                    $variant.finish()
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Marks the stream as finished and waits for all outstanding data to be acknowledged.
        ///
        /// This method is equivalent to calling [`finish`](Self::finish) and [`flush`](Self::flush).
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Ok(())` if the send buffer was completely flushed and acknowledged by
        ///   the peer.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error).
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// #   let stream: s2n_quic::stream::SendStream = todo!();
        /// #
        /// let data = bytes::Bytes::from_static(&[1, 2, 3, 4]);
        /// stream.send(data).await?;
        /// stream.close().await?;
        /// // at this point, the peer has received all of the `data` and has acknowledged the
        /// // stream being finished.
        /// #
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub async fn close(&mut self) -> $crate::stream::Result<()> {
            ::futures::future::poll_fn(|cx| self.poll_close(cx)).await
        }

        /// Marks the stream as finished and polls for all outstanding data to be acknowledged.
        ///
        /// This method is equivalent to calling [`finish`](Self::finish) and [`flush`](Self::flush).
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Poll::Pending` if the stream's send buffer is still being sent. In this case,
        ///   the caller should retry sending after the [`Waker`](core::task::Waker) on the provided
        ///   [`Context`](core::task::Context) is notified.
        /// - `Poll::Ready(Ok(()))` if the send buffer was completely flushed and acknowledged by
        ///   the peer.
        /// - `Poll::Ready(Err(e))` if the stream encountered a [`stream::Error`](crate::stream::Error).
        #[inline]
        pub fn poll_close(
            &mut self,
            cx: &mut core::task::Context,
        ) -> core::task::Poll<$crate::stream::Result<()>> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable()).into()
                };
                ($variant: expr) => {
                    $variant.poll_close(cx)
                };
            }

            let $stream = self;
            $dispatch_body
        }

        /// Closes the stream with an [error code](crate::application::Error).
        ///
        /// After calling this, the stream is closed and will not accept any additional data to be
        /// sent to the peer. The peer will also be notified of the [error
        /// code](crate::application::Error).
        ///
        /// # Return value
        ///
        /// The function returns:
        /// - `Ok(())` if the stream was reset successfully.
        /// - `Err(e)` if the stream encountered a [`stream::Error`](crate::stream::Error). The
        ///   stream may have been reset previously, or the connection itself was closed.
        #[inline]
        pub fn reset(
            &mut self,
            error_code: $crate::application::Error,
        ) -> $crate::stream::Result<()> {
            macro_rules! $dispatch {
                () => {
                    Err($crate::stream::Error::non_writable())
                };
                ($variant: expr) => {
                    $variant.reset(error_code)
                };
            }

            let $stream = self;
            $dispatch_body
        }
    };
}

macro_rules! impl_send_stream_trait {
    ($name:ident, | $stream:ident, $dispatch:ident | $dispatch_body:expr) => {
        impl futures::sink::Sink<bytes::Bytes> for $name {
            type Error = $crate::stream::Error;

            #[inline]
            fn poll_ready(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<$crate::stream::Result<()>> {
                futures::ready!(self.poll_send_ready(cx))?;
                Ok(()).into()
            }

            #[inline]
            fn start_send(
                mut self: core::pin::Pin<&mut Self>,
                data: bytes::Bytes,
            ) -> $crate::stream::Result<()> {
                self.send_data(data)
            }

            #[inline]
            fn poll_flush(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<$crate::stream::Result<()>> {
                Self::poll_flush(&mut self, cx)
            }

            #[inline]
            fn poll_close(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<$crate::stream::Result<()>> {
                Self::poll_close(&mut self, cx)
            }
        }

        impl futures::io::AsyncWrite for $name {
            fn poll_write(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
                buf: &[u8],
            ) -> core::task::Poll<std::io::Result<usize>> {
                if buf.is_empty() {
                    return Ok(0).into();
                }

                let len = futures::ready!(self.poll_send_ready(cx))?.min(buf.len());
                let data = bytes::Bytes::copy_from_slice(&buf[..len]);
                self.send_data(data)?;
                Ok(len).into()
            }

            fn poll_write_vectored(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
                bufs: &[futures::io::IoSlice],
            ) -> core::task::Poll<std::io::Result<usize>> {
                if bufs.is_empty() {
                    return Ok(0).into();
                }

                let len = futures::ready!(self.poll_send_ready(cx))?;
                let capacity = bufs.iter().map(|buf| buf.len()).sum();
                let len = len.min(capacity);

                let mut data = bytes::BytesMut::with_capacity(len);
                for buf in bufs {
                    // only copy what the window will allow
                    let to_copy = buf.len().min(len - data.len());
                    data.extend_from_slice(&buf[..to_copy]);

                    // we're done filling the buffer
                    if data.len() == len {
                        break;
                    }
                }

                self.send_data(data.freeze())?;

                Ok(len).into()
            }

            #[inline]
            fn poll_flush(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<std::io::Result<()>> {
                futures::ready!($name::poll_flush(&mut self, cx))?;
                Ok(()).into()
            }

            #[inline]
            fn poll_close(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<std::io::Result<()>> {
                futures::ready!($name::poll_close(&mut self, cx))?;
                Ok(()).into()
            }
        }

        impl tokio::io::AsyncWrite for $name {
            #[inline]
            fn poll_write(
                self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
                buf: &[u8],
            ) -> core::task::Poll<std::io::Result<usize>> {
                futures::io::AsyncWrite::poll_write(self, cx, buf)
            }

            fn poll_write_vectored(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
                bufs: &[std::io::IoSlice],
            ) -> core::task::Poll<std::io::Result<usize>> {
                if bufs.is_empty() {
                    return Ok(0).into();
                }

                let len = futures::ready!(self.poll_send_ready(cx))?;
                let capacity = bufs.iter().map(|buf| buf.len()).sum();
                let len = len.min(capacity);

                let mut data = bytes::BytesMut::with_capacity(len);
                for buf in bufs {
                    // only copy what the window will allow
                    let to_copy = buf.len().min(len - data.len());
                    data.extend_from_slice(&buf[..to_copy]);

                    // we're done filling the buffer
                    if data.len() == len {
                        break;
                    }
                }

                self.send_data(data.freeze())?;

                Ok(len).into()
            }

            #[inline]
            fn is_write_vectored(&self) -> bool {
                true
            }

            #[inline]
            fn poll_flush(
                self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<std::io::Result<()>> {
                futures::io::AsyncWrite::poll_flush(self, cx)
            }

            #[inline]
            fn poll_shutdown(
                mut self: core::pin::Pin<&mut Self>,
                cx: &mut core::task::Context<'_>,
            ) -> core::task::Poll<std::io::Result<()>> {
                futures::ready!(self.poll_close(cx))?;
                Ok(()).into()
            }
        }
    };
}

impl SendStream {
    #[inline]
    pub(crate) const fn new(stream: stream::SendStream) -> Self {
        Self(stream)
    }

    /// Returns the stream's identifier
    ///
    /// This value is unique to a particular connection. The format follows the same as what is
    /// defined in the
    /// [QUIC Transport RFC](https://www.rfc-editor.org/rfc/rfc9000.html#name-stream-types-and-identifier).
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # async fn test() -> s2n_quic::stream::Result<()> {
    /// #   let connection: s2n_quic::connection::Connection = todo!();
    /// #
    /// let stream = connection.open_send_stream().await?;
    /// println!("New stream's id: {}", stream.id());
    /// #
    /// #   Ok(())
    /// # }
    /// ```
    #[inline]
    pub fn id(&self) -> u64 {
        self.0.id().into()
    }

    impl_connection_api!(|stream| crate::connection::Handle(stream.0.connection().clone()));

    impl_send_stream_api!(|stream, dispatch| dispatch!(stream.0));
}

impl_splittable_stream_trait!(SendStream, |stream| (None, Some(stream)));
impl_send_stream_trait!(SendStream, |stream, dispatch| dispatch!(stream.0));


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

use super::{ReceiveStream, SendStream};

/// A trait that enables a stream to split into a [`ReceiveStream`] and [`SendStream`].
///
/// Note that if a stream is only allowed to send, then the receiving side will be [`None`].
/// The same is true for streams that are only allowed to receive: the sending side will be [`None`].
pub trait SplittableStream {
    /// Splits the stream into [`ReceiveStream`] and [`SendStream`] halves
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// # use bytes::Bytes;
    /// # async fn test() -> s2n_quic::stream::Result<()> {
    /// #   let connection: s2n_quic::connection::Connection = todo!();
    /// #
    /// let stream = connection.open_bidirectional_stream().await?;
    /// let (recv, send) = s2n_quic::stream::SplittableStream::split(stream);
    /// let mut recv = recv.expect("bidirectional streams have receiving sides");
    /// let mut send = send.expect("bidirectional streams have sending sides");
    ///
    /// tokio::spawn(async move {
    ///     let _ = send.send(Bytes::from_static(&[1, 2, 3])).await;
    /// });
    ///
    /// while let Some(chunk) = recv.receive().await? {
    ///     println!("received: {:?}", chunk);
    /// }
    /// #
    /// #   Ok(())
    /// # }
    /// ```
    fn split(self) -> (Option<ReceiveStream>, Option<SendStream>);
}

macro_rules! impl_splittable_stream_api {
    () => {
        /// Splits the stream into [`ReceiveStream`](crate::stream::ReceiveStream) and
        /// [`SendStream`](crate::stream::SendStream) halves
        ///
        /// # Examples
        ///
        /// ```rust,no_run
        /// # async fn test() -> s2n_quic::stream::Result<()> {
        /// # let stream: s2n_quic::stream::Stream = todo!();
        /// #
        /// let (recv, send) = stream.split();
        ///
        /// if let Some(recv) = recv {
        ///    // the stream has a receive half
        /// }
        ///
        /// if let Some(send) = send {
        ///    // the stream has a send half
        /// }
        /// #   Ok(())
        /// # }
        /// ```
        #[inline]
        pub fn split(
            self,
        ) -> (
            Option<$crate::stream::ReceiveStream>,
            Option<$crate::stream::SendStream>,
        ) {
            $crate::stream::SplittableStream::split(self)
        }
    };
}

macro_rules! impl_splittable_stream_trait {
    ($name:ident, | $self:ident | $convert:expr) => {
        impl $crate::stream::SplittableStream for $name {
            #[inline]
            fn split(
                self,
            ) -> (
                Option<$crate::stream::ReceiveStream>,
                Option<$crate::stream::SendStream>,
            ) {
                let $self = self;
                $convert
            }
        }
    };
}


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

//! This module shows an example of an event provider that blocks the client from opening any
//! streams until the handshake is completely confirmed by the server.
//!
//! When using mTLS, the client is technically allowed to start sending application data before the
//! server has actually authenticated the client. For clients that want to wait before getting a
//! confirmation of the server accepting the client's certificate, they can use the [`ClientConfirm`]
//! struct as an event provider on the [`Client`].
//!
//! Note that waiting for handshake confirmation will add another round trip to the handshake so
//! keep that in mind when deciding if this functionality is useful.
//!
//! This module includes 3 distinct tests to show the behavior:
//!
//! * The `mtls_happy_case` creates a successful connection and sends some data
//! * The `mtls_failure_no_wait` shows the case when the `wait_ready` functionality is not used.
//!   The client opens a connection and stream, sends some data, only to find that the server
//!   rejected the client's certificate when trying to receive the response.
//! * The `mtls_failure_with_wait` shows that when using the `wait_ready` functionality the client
//!   is unable to open a stream, as the task was blocked until the server either confirmed or
//!   rejected the connection attempt.

use super::*;
use crate::provider::event::events;
use core::task::{Context, Poll, Waker};

struct ClientConfirm;

impl ClientConfirm {
    /// Blocks the task until the provided connection has either confirmed the handshake or closed
    /// with an error
    pub async fn wait_ready(conn: &mut crate::Connection) {
        futures::future::poll_fn(|cx| {
            conn.query_event_context_mut(|context: &mut ClientConfirmContext| {
                context.poll_ready(cx)
            })
            .unwrap_or(Poll::Ready(()))
        })
        .await;
    }
}

#[derive(Default)]
struct ClientConfirmContext {
    waker: Option<Waker>,
    state: State,
}

impl ClientConfirmContext {
    /// Updates the state on the context
    fn update(&mut self, state: State) {
        self.state = state;

        // notify the application that the state was updated
        self.wake();
    }

    /// Polls the context for handshake confirmation
    fn poll_ready(&mut self, cx: &mut Context) -> Poll<()> {
        // if we're ready then let the application know
        if matches!(self.state, State::Ready) {
            return Poll::Ready(());
        }

        // store the waker so we can notify the application of state updates
        self.waker = Some(cx.waker().clone());

        Poll::Pending
    }

    /// notify the application of a state update
    fn wake(&mut self) {
        if let Some(waker) = self.waker.take() {
            waker.wake();
        }
    }
}

impl Drop for ClientConfirmContext {
    // make sure the application is notified that we're closing the connection
    fn drop(&mut self) {
        self.wake();
    }
}

#[derive(Default)]
enum State {
    #[default]
    Waiting,
    Ready,
}

impl Subscriber for ClientConfirm {
    type ConnectionContext = ClientConfirmContext;

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

    #[inline]
    fn on_handshake_status_updated(
        &mut self,
        context: &mut Self::ConnectionContext,
        _: &ConnectionMeta,
        event: &events::HandshakeStatusUpdated,
    ) {
        if let events::HandshakeStatus::Confirmed { .. } = event.status {
            // notify the application that the handshake has been confirmed by the server
            context.update(State::Ready);
        }
    }

    #[inline]
    fn on_connection_closed(
        &mut self,
        context: &mut Self::ConnectionContext,
        _: &ConnectionMeta,
        _event: &events::ConnectionClosed,
    ) {
        // notify the application if we close the connection
        context.update(State::Ready);
    }
}

fn mtl_test<C>(server_cert: &str, f: fn(crate::Connection) -> C)
where
    C: 'static + core::future::Future<Output = ()> + Send,
{
    let model = Model::default();
    model.set_delay(Duration::from_millis(50));

    test(model, |handle| {
        let server_tls = build_server_mtls_provider(server_cert)?;
        let server = Server::builder()
            .with_io(handle.builder().build()?)?
            .with_tls(server_tls)?
            .with_event(events())?
            .start()?;

        let addr = start_server(server)?;

        let client_tls = build_client_mtls_provider(certificates::MTLS_CA_CERT)?;
        let client = Client::builder()
            .with_io(handle.builder().build().unwrap())?
            .with_tls(client_tls)?
            .with_event((ClientConfirm, events()))?
            .start()?;

        // show it working for several connections
        for _ in 0..10 {
            let client = client.clone();
            primary::spawn(async move {
                let connect = Connect::new(addr).with_server_name("localhost");
                let conn = client.connect(connect).await.unwrap();
                f(conn).await;
            });
        }

        Ok(addr)
    })
    .unwrap();
}

#[test]
fn mtls_happy_case() {
    const LEN: usize = 1000;

    mtl_test(certificates::MTLS_CA_CERT, |mut conn| {
        async move {
            // make sure we get confirmation of the handshake before opening a stream and sending
            // data
            ClientConfirm::wait_ready(&mut conn).await;

            let mut stream = conn.open_bidirectional_stream().await.unwrap();

            stream.send(vec![42; LEN].into()).await.unwrap();
            stream.finish().unwrap();

            let mut recv_len = 0;
            while let Some(chunk) = stream.receive().await.unwrap() {
                recv_len += chunk.len();
            }
            assert_eq!(LEN, recv_len);
        }
    });
}

#[test]
fn mtls_failure_no_wait() {
    const LEN: usize = 1000;

    mtl_test(certificates::UNTRUSTED_CERT_PEM, |mut conn| {
        async move {
            // We don't use `wait_ready` here to show that we can open a stream that will fail
            // later

            let mut stream = conn.open_bidirectional_stream().await.unwrap();

            stream.send(vec![42; LEN].into()).await.unwrap();
            stream.finish().unwrap();

            // we get an error only when receiving from the peer
            stream.receive().await.unwrap_err();
        }
    });
}

#[test]
fn mtls_failure_with_wait() {
    mtl_test(certificates::UNTRUSTED_CERT_PEM, |mut conn| {
        async move {
            // make sure we get confirmation of the handshake before opening a stream and sending
            // data
            ClientConfirm::wait_ready(&mut conn).await;

            // the open should fail since we waited for handshake confirmation
            conn.open_bidirectional_stream().await.unwrap_err();
        }
    });
}


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

use super::*;
use crate::provider::tls::Provider;
use s2n_quic_core::crypto::tls::null;

#[derive(Default)]
pub struct NoTlsProvider {}

impl Provider for NoTlsProvider {
    type Server = null::Endpoint;
    type Client = null::Endpoint;
    type Error = String;

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

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

#[test]
fn no_tls_test() {
    let model = Model::default();

    test(model, |handle| {
        let server = Server::builder()
            .with_io(handle.builder().build()?)?
            .with_tls(NoTlsProvider::default())?
            .with_event(events())?
            .start()?;
        let client = Client::builder()
            .with_io(handle.builder().build()?)?
            .with_tls(NoTlsProvider::default())?
            .with_event(events())?
            .start()?;
        let addr = start_server(server)?;
        start_client(client, addr, Data::new(1000))?;
        Ok(addr)
    })
    .unwrap();
}


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

use crate::{
    client::Connect,
    provider::{
        event,
        io::testing::{primary, spawn, Handle, Result},
    },
    Client, Server,
};
use rand::{Rng, RngCore};
use s2n_quic_core::{crypto::tls::testing::certificates, havoc, stream::testing::Data};
use std::net::SocketAddr;

pub static SERVER_CERTS: (&str, &str) = (certificates::CERT_PEM, certificates::KEY_PEM);

pub fn events() -> event::tracing::Subscriber {
    use std::sync::Once;

    static TRACING: Once = Once::new();

    // make sure this only gets initialized once
    TRACING.call_once(|| {
        let format = tracing_subscriber::fmt::format()
            .with_level(false) // don't include levels in formatted output
            .with_timer(Uptime)
            .with_ansi(false)
            .compact(); // Use a less verbose output format.

        struct Uptime;

        // Generate the timestamp from the testing IO provider rather than wall clock.
        impl tracing_subscriber::fmt::time::FormatTime for Uptime {
            fn format_time(
                &self,
                w: &mut tracing_subscriber::fmt::format::Writer<'_>,
            ) -> std::fmt::Result {
                write!(w, "{}", crate::provider::io::testing::now())
            }
        }

        tracing_subscriber::fmt()
            .with_max_level(tracing_subscriber::filter::LevelFilter::DEBUG)
            .event_format(format)
            .with_test_writer()
            .init();
    });

    event::tracing::Subscriber::default()
}

pub fn start_server(mut server: Server) -> Result<SocketAddr> {
    let server_addr = server.local_addr()?;

    // accept connections and echo back
    spawn(async move {
        while let Some(mut connection) = server.accept().await {
            tracing::debug!("accepted server connection: {}", connection.id());
            spawn(async move {
                while let Ok(Some(mut stream)) = connection.accept_bidirectional_stream().await {
                    tracing::debug!("accepted server stream: {}", stream.id());
                    spawn(async move {
                        while let Ok(Some(chunk)) = stream.receive().await {
                            let _ = stream.send(chunk).await;
                        }
                    });
                }
            });
        }
    });

    Ok(server_addr)
}

pub fn server(handle: &Handle) -> Result<SocketAddr> {
    let server = build_server(handle)?;
    start_server(server)
}

pub fn build_server(handle: &Handle) -> Result<Server> {
    Ok(Server::builder()
        .with_io(handle.builder().build().unwrap())?
        .with_tls(SERVER_CERTS)?
        .with_event(events())?
        .start()?)
}

pub fn client(handle: &Handle, server_addr: SocketAddr) -> Result {
    let client = build_client(handle)?;
    start_client(client, server_addr, Data::new(10_000))
}

pub fn start_client(client: Client, server_addr: SocketAddr, data: Data) -> Result {
    primary::spawn(async move {
        let connect = Connect::new(server_addr).with_server_name("localhost");
        let mut connection = client.connect(connect).await.unwrap();

        tracing::debug!("connected with client connection: {}", connection.id());

        let stream = connection.open_bidirectional_stream().await.unwrap();
        tracing::debug!("opened client stream: {}", stream.id());

        let (mut recv, mut send) = stream.split();

        let mut send_data = data;
        let mut recv_data = data;

        primary::spawn(async move {
            while let Some(chunk) = recv.receive().await.unwrap() {
                recv_data.receive(&[chunk]);
            }
            assert!(recv_data.is_finished());
        });

        while let Some(chunk) = send_data.send_one(usize::MAX) {
            tracing::debug!("client sending {} chunk", chunk.len());
            send.send(chunk).await.unwrap();
        }
    });

    Ok(())
}

pub fn build_client(handle: &Handle) -> Result<Client> {
    Ok(Client::builder()
        .with_io(handle.builder().build().unwrap())?
        .with_tls(certificates::CERT_PEM)?
        .with_event(events())?
        .start()?)
}

pub fn client_server(handle: &Handle) -> Result<SocketAddr> {
    let addr = server(handle)?;
    client(handle, addr)?;
    Ok(addr)
}

pub struct Random {
    inner: rand_chacha::ChaCha8Rng,
}

impl Random {
    pub fn with_seed(seed: u64) -> Self {
        use rand::SeedableRng;
        Self {
            inner: rand_chacha::ChaCha8Rng::seed_from_u64(seed),
        }
    }
}

impl havoc::Random for Random {
    fn fill(&mut self, bytes: &mut [u8]) {
        self.fill_bytes(bytes);
    }

    fn gen_range(&mut self, range: std::ops::Range<u64>) -> u64 {
        self.inner.gen_range(range)
    }
}

impl RngCore for Random {
    fn try_fill_bytes(&mut self, dest: &mut [u8]) -> Result<(), rand::Error> {
        self.inner.try_fill_bytes(dest)
    }

    fn fill_bytes(&mut self, dest: &mut [u8]) {
        self.inner.fill_bytes(dest)
    }

    fn next_u32(&mut self) -> u32 {
        self.inner.next_u32()
    }

    fn next_u64(&mut self) -> u64 {
        self.inner.next_u64()
    }
}

#[cfg(not(target_os = "windows"))]
mod mtls {
    use super::*;
    use crate::provider::tls;

    pub fn build_client_mtls_provider(ca_cert: &str) -> Result<tls::default::Client> {
        let tls = tls::default::Client::builder()
            .with_certificate(ca_cert)?
            .with_client_identity(
                certificates::MTLS_CLIENT_CERT,
                certificates::MTLS_CLIENT_KEY,
            )?
            .build()?;
        Ok(tls)
    }

    pub fn build_server_mtls_provider(ca_cert: &str) -> Result<tls::default::Server> {
        let tls = tls::default::Server::builder()
            .with_certificate(
                certificates::MTLS_SERVER_CERT,
                certificates::MTLS_SERVER_KEY,
            )?
            .with_client_authentication()?
            .with_trusted_certificate(ca_cert)?
            .build()?;
        Ok(tls)
    }
}

#[cfg(not(target_os = "windows"))]
pub use mtls::*;


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

use criterion::{criterion_group, criterion_main};

criterion_group!(benches, s2n_quic_bench::benchmarks);
criterion_main!(benches);


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

use criterion::{BenchmarkId, Criterion, Throughput};
use s2n_quic_core::{buffer::ReceiveBuffer, varint::VarInt};

pub fn benchmarks(c: &mut Criterion) {
    let mut group = c.benchmark_group("buffer");

    for size in [1, 100, 1000, 1450] {
        let input = vec![42u8; size];

        group.throughput(Throughput::Bytes(input.len() as _));
        group.bench_with_input(BenchmarkId::new("write_at", size), &input, |b, input| {
            let mut buffer = ReceiveBuffer::new();
            let mut offset = VarInt::from_u8(0);
            let len = VarInt::new(input.len() as _).unwrap();
            b.iter(move || {
                buffer.write_at(offset, input).unwrap();
                offset += len;
            });
        });
    }
}


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

use criterion::Criterion;

mod aes;
mod aesgcm;
mod ghash;

pub fn benchmarks(c: &mut Criterion) {
    aes::benchmarks(c);
    aesgcm::benchmarks(c);
    ghash::benchmarks(c);
}


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

use criterion::{black_box, BenchmarkId, Criterion, Throughput};
use s2n_codec::{DecoderBufferMut, Encoder, EncoderBuffer, EncoderValue};
use s2n_quic_core::frame::FrameMut;

pub fn benchmarks(c: &mut Criterion) {
    codec(c);
}

struct Input {
    name: &'static str,
    buffer: &'static [u8],
}

macro_rules! input {
    ($name:expr) => {
        Input {
            name: $name,
            buffer: include_bytes!(concat!(
                env!("CARGO_MANIFEST_DIR"),
                "/../s2n-quic-core/src/frame/test_samples/",
                $name,
                ".bin"
            )),
        }
    };
}

fn codec(c: &mut Criterion) {
    let mut group = c.benchmark_group("frame");

    let mut inputs = [
        input!("ack"),
        input!("connection_close"),
        input!("crypto"),
        input!("data_blocked"),
        input!("handshake_done"),
        input!("max_data"),
        input!("max_stream_data"),
        input!("max_streams"),
        input!("new_connection_id"),
        input!("new_token"),
        input!("padding"),
        input!("path_challenge"),
        input!("path_response"),
        input!("ping"),
        input!("reset_stream"),
        input!("retire_connection_id"),
        input!("stop_sending"),
        input!("stream"),
        input!("stream_data_blocked"),
        input!("streams_blocked"),
    ];

    // sort by length to make the graphs nicer
    inputs.sort_by_key(|input| input.buffer.len());

    for input in &inputs {
        group.throughput(Throughput::Bytes(input.buffer.len() as _));
        group.bench_with_input(
            BenchmarkId::new("decode", input.name),
            &input.buffer,
            |b, buffer| {
                let mut buffer = buffer.to_vec();
                b.iter(move || {
                    let buffer = DecoderBufferMut::new(&mut buffer);
                    let _ = black_box(buffer.decode::<FrameMut>());
                });
            },
        );
        group.bench_with_input(
            BenchmarkId::new("encode", input.name),
            &input.buffer,
            |b, buffer| {
                let mut buffer = buffer.to_vec();
                let buffer = DecoderBufferMut::new(&mut buffer);
                let (frame, _remaining) = buffer.decode::<FrameMut>().unwrap();
                let frame = black_box(frame);
                let mut buffer = vec![0; frame.encoding_size()];
                b.iter(move || {
                    EncoderBuffer::new(&mut buffer).encode(&frame);
                });
            },
        );
    }

    group.finish();
}


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

use core::hash::Hasher;
use criterion::{black_box, BenchmarkId, Criterion, Throughput};
use s2n_quic_core::inet::checksum::Checksum;

pub fn benchmarks(c: &mut Criterion) {
    let mut group = c.benchmark_group("inet");
    for block in [20, 1500, 9000, 1 << 16] {
        let data = vec![123u8; block];
        group.throughput(Throughput::Bytes(block as u64));
        group.bench_with_input(
            BenchmarkId::new("s2n/checksum", block),
            &data,
            |b, block| {
                let cs = Checksum::default();
                let input = black_box(&block[..]);
                b.iter(|| {
                    let mut checksum = cs;
                    checksum.write(input);
                    black_box(checksum.finish())
                })
            },
        );
        group.bench_with_input(
            BenchmarkId::new("fuchsia/checksum", block),
            &data,
            |b, block| {
                let input = black_box(&block[..]);
                b.iter(|| black_box(internet_checksum::checksum(input)))
            },
        );
    }
    group.finish();
}


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

use criterion::Criterion;

mod buffer;
mod crypto;
mod frame;
mod inet;
mod packet;
mod sync;
mod varint;
mod xdp;

pub fn benchmarks(c: &mut Criterion) {
    buffer::benchmarks(c);
    crypto::benchmarks(c);
    frame::benchmarks(c);
    inet::benchmarks(c);
    packet::benchmarks(c);
    sync::benchmarks(c);
    varint::benchmarks(c);
    xdp::benchmarks(c);
}


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

use criterion::{black_box, BenchmarkId, Criterion, Throughput};
use s2n_codec::DecoderBufferMut;
use s2n_quic_core::{connection::id::ConnectionInfo, inet::SocketAddress, packet::ProtectedPacket};

pub fn benchmarks(c: &mut Criterion) {
    codec(c);
}

struct Input {
    name: &'static str,
    buffer: &'static [u8],
}

macro_rules! input {
    ($name:expr) => {
        Input {
            name: $name,
            buffer: include_bytes!(concat!(
                env!("CARGO_MANIFEST_DIR"),
                "/../s2n-quic-core/src/packet/test_samples/",
                $name,
                ".bin"
            )),
        }
    };
}

fn codec(c: &mut Criterion) {
    let mut group = c.benchmark_group("packet");

    let mut inputs = [
        input!("handshake"),
        input!("initial"),
        input!("short"),
        input!("zero_rtt"),
        input!("version_negotiation"),
        input!("retry"),
    ];

    // sort by length to make the graphs nicer
    inputs.sort_by_key(|input| input.buffer.len());

    for input in &inputs {
        group.throughput(Throughput::Bytes(input.buffer.len() as _));
        group.bench_with_input(
            BenchmarkId::new("decode", input.name),
            &input.buffer,
            |b, buffer| {
                let mut buffer = buffer.to_vec();
                b.iter(move || {
                    let buffer = DecoderBufferMut::new(&mut buffer);
                    let remote_address = SocketAddress::default();
                    let connection_info = ConnectionInfo::new(&remote_address);
                    let _ = black_box(ProtectedPacket::decode(buffer, &connection_info, &20));
                })
            },
        );
    }

    group.finish();
}


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

use criterion::{BenchmarkId, Criterion, Throughput};
use crossbeam_channel::bounded;
use s2n_quic_core::sync::spsc;

pub fn benchmarks(c: &mut Criterion) {
    spsc_benches(c);
}

fn spsc_benches(c: &mut Criterion) {
    let mut group = c.benchmark_group("spsc");

    for i in [1, 64, 1024, 4096] {
        group.throughput(Throughput::Elements(i as _));
        group.bench_with_input(BenchmarkId::new("s2n/send_recv", i), &i, |b, input| {
            let (mut sender, mut receiver) = spsc::channel(*input);
            b.iter(|| {
                {
                    let mut slice = sender.try_slice().unwrap().unwrap();
                    while slice.push(123usize).is_ok() {}
                }

                {
                    let mut slice = receiver.try_slice().unwrap().unwrap();
                    while slice.pop().is_some() {}
                }
            });
        });
        group.bench_with_input(
            BenchmarkId::new("crossbeam/send_recv", i),
            &i,
            |b, input| {
                let (sender, receiver) = bounded(*input);
                b.iter(|| {
                    {
                        while sender.try_send(123usize).is_ok() {}
                    }

                    {
                        while receiver.try_recv().is_ok() {}
                    }
                });
            },
        );

        group.bench_with_input(BenchmarkId::new("s2n/send_recv_iter", i), &i, |b, input| {
            let (mut sender, mut receiver) = spsc::channel(*input);
            b.iter(|| {
                {
                    let mut slice = sender.try_slice().unwrap().unwrap();
                    let _ = slice.extend(&mut core::iter::repeat(123usize));
                }

                {
                    let mut slice = receiver.try_slice().unwrap().unwrap();
                    slice.clear();
                }
            });
        });
        group.bench_with_input(
            BenchmarkId::new("crossbeam/send_recv_iter", i),
            &i,
            |b, input| {
                let (sender, receiver) = bounded(*input);
                b.iter(|| {
                    {
                        while sender.try_send(123usize).is_ok() {}
                    }

                    {
                        for _ in receiver.try_iter() {}
                    }
                });
            },
        );
    }
    group.finish();
}


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

use criterion::{black_box, BenchmarkId, Criterion};
use s2n_codec::{DecoderBuffer, EncoderBuffer, EncoderValue};
use s2n_quic_core::varint::VarInt;

pub fn benchmarks(c: &mut Criterion) {
    round_trip(c);
}

fn round_trip(c: &mut Criterion) {
    let mut group = c.benchmark_group("varint");
    for i in [0, 1, 5, 6, 13, 14, 29, 30, 61] {
        let i = VarInt::new(2u64.pow(i)).unwrap();

        group.bench_with_input(BenchmarkId::new("round_trip", i), &i, |b, input| {
            let mut buffer = vec![0; 8];
            b.iter(|| {
                input.encode(&mut EncoderBuffer::new(&mut buffer));
                let (actual, _) = DecoderBuffer::new(&buffer).decode::<VarInt>().unwrap();
                black_box(actual);
            });
        });
    }
    group.finish();
}


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

use criterion::{black_box, BenchmarkId, Criterion, Throughput};
use s2n_codec::EncoderBuffer;
use s2n_quic_core::{
    inet::{ExplicitCongestionNotification, IpV4Address, IpV6Address},
    io::tx::{self, PayloadBuffer},
    xdp::{
        encoder::{encode_packet, State},
        path,
    },
};

pub fn benchmarks(c: &mut Criterion) {
    let mut group = c.benchmark_group("xdp/encoder");
    let overhead = 100;

    let paths = [
        (
            "ipv4",
            true,
            path::Tuple {
                remote_address: path::RemoteAddress {
                    mac: Default::default(),
                    ip: IpV4Address::default().into(),
                    port: 0,
                },
                local_address: path::LocalAddress {
                    mac: Default::default(),
                    ip: IpV4Address::default().into(),
                    port: 0,
                },
            },
        ),
        (
            "ipv4-no-checksum",
            false,
            path::Tuple {
                remote_address: path::RemoteAddress {
                    mac: Default::default(),
                    ip: IpV4Address::default().into(),
                    port: 0,
                },
                local_address: path::LocalAddress {
                    mac: Default::default(),
                    ip: IpV4Address::default().into(),
                    port: 0,
                },
            },
        ),
        (
            "ipv6",
            true,
            path::Tuple {
                remote_address: path::RemoteAddress {
                    mac: Default::default(),
                    ip: IpV6Address::default().into(),
                    port: 0,
                },
                local_address: path::LocalAddress {
                    mac: Default::default(),
                    ip: IpV6Address::default().into(),
                    port: 0,
                },
            },
        ),
    ];

    for (label, ipv4_checksum, path) in paths {
        for payload_len in [1500, 9000, 1 << 16] {
            let message = Message {
                path,
                ecn: Default::default(),
                ipv6_flow_label: 123,
                payload_len: payload_len - overhead,
            };

            group.throughput(Throughput::Elements(1));
            group.bench_with_input(
                BenchmarkId::new(label, payload_len),
                &message,
                |b, mut message| {
                    let mut buffer = vec![0u8; payload_len];
                    let mut state = State::default();
                    state.set_checksum(ipv4_checksum);

                    b.iter(|| {
                        let mut encoder = EncoderBuffer::new(&mut buffer);
                        let _ = black_box(encode_packet(
                            black_box(&mut encoder),
                            black_box(&mut message),
                            black_box(&mut state),
                        ));
                    })
                },
            );
        }
    }
    group.finish();
}

#[derive(Debug)]
struct Message {
    path: path::Tuple,
    ecn: ExplicitCongestionNotification,
    ipv6_flow_label: u32,
    payload_len: usize,
}

impl<'a> tx::Message for &'a Message {
    type Handle = path::Tuple;

    fn path_handle(&self) -> &Self::Handle {
        &self.path
    }

    fn ecn(&mut self) -> ExplicitCongestionNotification {
        self.ecn
    }

    fn delay(&mut self) -> core::time::Duration {
        Default::default()
    }

    fn ipv6_flow_label(&mut self) -> u32 {
        self.ipv6_flow_label
    }

    fn can_gso(&self, _: usize, _: usize) -> bool {
        true
    }

    fn write_payload(
        &mut self,
        _buffer: PayloadBuffer,
        _gso_offset: usize,
    ) -> Result<usize, tx::Error> {
        // skip copying the payload to just measure the header overhead
        Ok(self.payload_len)
    }
}


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

use criterion::{BenchmarkId, Criterion, Throughput};
use s2n_quic_crypto::testing;

macro_rules! aes_impl {
    ($name:ident) => {
        fn $name(c: &mut Criterion) {
            aes_impl!(c, $name, encrypt);
            aes_impl!(c, $name, decrypt);
        }
    };
    ($c:ident, $name:ident, $call:ident) => {{
        let mut group = $c.benchmark_group(concat!(
            "crypto/aes/",
            stringify!($name),
            "/",
            stringify!($call)
        ));
        for imp in testing::aes::$name::implementations() {
            for block in testing::BLOCK_SIZES.iter() {
                group.throughput(Throughput::Bytes(block.len() as _));
                group.bench_with_input(BenchmarkId::new(imp.name(), block), &block, |b, block| {
                    let key = imp.new([1; testing::aes::$name::KEY_LEN]);
                    let mut input = block.to_vec();
                    b.iter(|| key.$call(&mut input));
                });
            }
        }
        group.finish();
    }};
}

aes_impl!(aes128);
aes_impl!(aes256);

pub fn benchmarks(c: &mut Criterion) {
    aes128(c);
    aes256(c);
}


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

use core::convert::TryInto;
use criterion::{black_box, BenchmarkId, Criterion, Throughput};
use s2n_quic_crypto::testing::{
    self,
    aesgcm::{Key, NONCE_LEN, TAG_LEN},
};

struct Implementation {
    key: Key,
    name: &'static str,
}

macro_rules! impls {
    ($name:ident) => {{
        use s2n_quic_crypto::testing::aes::$name::KEY_LEN;

        let mut impls = vec![];

        for imp in testing::aesgcm::$name::implementations() {
            impls.push(Implementation {
                key: imp.new(black_box([1; KEY_LEN])),
                name: imp.name(),
            });
        }

        (stringify!($name), impls)
    }};
}

pub fn benchmarks(c: &mut Criterion) {
    let impls = [impls!(aes128), impls!(aes256)];

    // create some shared values
    let nonce = black_box([123u8; NONCE_LEN]);
    let aad = black_box([123u8; 20]);

    for (group, impls) in impls.iter() {
        let mut encrypt = c.benchmark_group(format!("crypto/aesgcm/{group}/encrypt"));
        for imp in impls.iter() {
            for block in testing::BLOCK_SIZES.iter() {
                encrypt.throughput(Throughput::Bytes(block.len() as _));
                encrypt.bench_with_input(
                    BenchmarkId::new(imp.name, block),
                    &block,
                    move |b, block| {
                        let key = &imp.key;

                        let mut input = block.to_vec();
                        let payload_len = input.len();
                        input.extend_from_slice(&[0; TAG_LEN]);

                        let (payload, tag) = input.split_at_mut(payload_len);
                        let tag: &mut [u8; TAG_LEN] = tag.try_into().unwrap();
                        b.iter(|| {
                            let _ = key.encrypt(&nonce, &aad, payload, tag);
                        });
                    },
                );
            }
        }
        encrypt.finish();

        let mut decrypt = c.benchmark_group(format!("crypto/aesgcm/{group}/decrypt"));
        for imp in impls.iter() {
            for block in testing::BLOCK_SIZES.iter() {
                decrypt.throughput(Throughput::Bytes(block.len() as _));
                decrypt.bench_with_input(BenchmarkId::new(imp.name, block), &block, |b, block| {
                    let key = &imp.key;

                    let mut input = block.to_vec();
                    let payload_len = input.len();
                    input.extend_from_slice(&[0; TAG_LEN]);

                    let (payload, tag) = input.split_at_mut(payload_len);
                    let tag: &mut [u8; TAG_LEN] = tag.try_into().unwrap();

                    // create a valid encrypted payload
                    key.encrypt(&nonce, &aad, payload, tag).unwrap();
                    let tag = &*tag;

                    b.iter_batched(
                        || payload.to_vec(),
                        |mut payload| {
                            let _ = black_box(key.decrypt(&nonce, &aad, &mut payload, tag));
                        },
                        criterion::BatchSize::LargeInput,
                    );
                });
            }
        }
        decrypt.finish();
    }
}


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

use criterion::{black_box, BenchmarkId, Criterion, Throughput};
use s2n_quic_crypto::testing;

pub fn benchmarks(c: &mut Criterion) {
    let mut group = c.benchmark_group("crypto/ghash");
    for imp in testing::ghash::implementations() {
        for block in testing::BLOCK_SIZES.iter() {
            group.throughput(Throughput::Bytes(block.len() as _));
            group.bench_with_input(BenchmarkId::new(imp.name(), block), &block, |b, block| {
                let key = imp.new([1; 16]);
                let input = &block[..];
                b.iter(|| black_box(key.hash(input)))
            });
        }
    }
    group.finish();
}


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

pub mod set;
pub mod settings;

pub use set::Set;
pub use settings::Settings;


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

use crate::number::{
    CheckedAddAssign, CheckedMulAssign, CheckedSubAssign, SaturatingAddAssign, SaturatingMulAssign,
    SaturatingSubAssign, UpcastFrom,
};
use core::{cmp::Ordering, convert::TryFrom, marker::PhantomData, ops};

/// A checked-overflow counter
///
/// Rather than silently wrapping, we want to ensure counting errors stay somewhat isolated so the
/// counter will saturate rather than wrap. The counter operates in 3 modes:
///
/// * If `debug_assertions` are enabled, the counter will panic on overflow
/// * If the `checked-counters` feature flag is defined, the counter will panic on overflow, even in
///   release builds.
/// * Otherwise, the counter will saturate
///
/// The counter can also be configured to always saturate by passing the `Saturating` behavior:
///
/// ```rust
/// use s2n_quic_core::counter::{Counter, Saturating};
///
/// let counter: Counter<u32, Saturating> = Default::default();
/// ```
#[derive(Clone, Copy, Debug, Default, Hash)]
pub struct Counter<T, Behavior = ()>(T, PhantomData<Behavior>);

/// Overrides the behavior of a counter to always saturate
#[derive(Clone, Copy, Debug, Default, Hash)]
pub struct Saturating;

impl<T, Behavior> Counter<T, Behavior> {
    /// Creates a new counter with an initial value
    #[inline]
    pub const fn new(value: T) -> Self {
        Self(value, PhantomData)
    }

    #[inline]
    pub fn set(&mut self, value: T) {
        self.0 = value;
    }

    /// Tries to convert V to T and add it to the current counter value
    #[inline]
    pub fn try_add<V>(&mut self, value: V) -> Result<(), T::Error>
    where
        T: TryFrom<V>,
        Self: ops::AddAssign<T>,
    {
        let value = T::try_from(value)?;
        *self += value;
        Ok(())
    }

    /// Tries to convert V to T and subtract it from the current counter value
    #[inline]
    pub fn try_sub<V>(&mut self, value: V) -> Result<(), T::Error>
    where
        T: TryFrom<V>,
        Self: ops::SubAssign<T>,
    {
        let value = T::try_from(value)?;
        *self -= value;
        Ok(())
    }
}

/// Generates an assign trait implementation for the Counter
macro_rules! assign_trait {
    (
        $op:ident,
        $method:ident,
        $saturating_trait:ident,
        $saturating_method:ident,
        $checked_trait:ident,
        $checked:ident
    ) => {
        impl<T, R> ops::$op<R> for Counter<T, ()>
        where
            T: $saturating_trait<R> + $checked_trait<R> + ops::$op + UpcastFrom<R>,
        {
            #[inline]
            fn $method(&mut self, rhs: R) {
                if cfg!(feature = "checked-counters") {
                    (self.0).$checked(rhs).expect("counter overflow");
                } else if cfg!(debug_assertions) {
                    (self.0).$method(T::upcast_from(rhs));
                } else {
                    (self.0).$saturating_method(rhs);
                }
            }
        }

        impl<T, R> ops::$op<R> for Counter<T, Saturating>
        where
            T: $saturating_trait<R>,
        {
            #[inline]
            fn $method(&mut self, rhs: R) {
                (self.0).$saturating_method(rhs);
            }
        }
    };
}

assign_trait!(
    AddAssign,
    add_assign,
    SaturatingAddAssign,
    saturating_add_assign,
    CheckedAddAssign,
    checked_add_assign
);

assign_trait!(
    SubAssign,
    sub_assign,
    SaturatingSubAssign,
    saturating_sub_assign,
    CheckedSubAssign,
    checked_sub_assign
);

assign_trait!(
    MulAssign,
    mul_assign,
    SaturatingMulAssign,
    saturating_mul_assign,
    CheckedMulAssign,
    checked_mul_assign
);

impl<T, B> UpcastFrom<Counter<T, B>> for T {
    #[inline]
    fn upcast_from(value: Counter<T, B>) -> Self {
        value.0
    }
}

impl<T, B> UpcastFrom<&Counter<T, B>> for T
where
    T: for<'a> UpcastFrom<&'a T>,
{
    #[inline]
    fn upcast_from(value: &Counter<T, B>) -> Self {
        T::upcast_from(&value.0)
    }
}

impl<T, B> ops::Deref for Counter<T, B> {
    type Target = T;

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

impl<T, B, R> PartialEq<R> for Counter<T, B>
where
    Self: PartialOrd<R>,
{
    #[inline]
    fn eq(&self, other: &R) -> bool {
        self.partial_cmp(other) == Some(Ordering::Equal)
    }
}

impl<T, B> PartialOrd<T> for Counter<T, B>
where
    T: PartialOrd<T>,
{
    #[inline]
    fn partial_cmp(&self, other: &T) -> Option<Ordering> {
        self.0.partial_cmp(other)
    }
}

impl<T, B> PartialOrd for Counter<T, B>
where
    T: PartialOrd<T>,
{
    #[inline]
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

impl<T, B> Eq for Counter<T, B> where Self: Ord {}

impl<T, B> Ord for Counter<T, B>
where
    T: Ord,
{
    #[inline]
    fn cmp(&self, other: &Self) -> Ordering {
        self.0.cmp(&other.0)
    }
}

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

    #[test]
    fn automatic_upcast() {
        let mut a: Counter<u32> = Counter::new(0);
        a += 1u8;
        a += 2u16;
        a += 3u32;

        assert_eq!(a, Counter::new(6));
        assert_eq!(a, 6u32);
    }

    #[test]
    fn saturating() {
        let mut a: Counter<u8, Saturating> = Counter::new(0);
        a += 250;
        a += 250;
        a += 123;

        assert_eq!(a, Counter::new(255));
    }
}


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

use core::ops;
use num_traits::ops::{
    checked::{CheckedDiv, CheckedRem, CheckedShl, CheckedShr},
    overflowing::{OverflowingAdd, OverflowingMul, OverflowingSub},
};
pub use subtle::{Choice, ConditionallySelectable, ConstantTimeEq, CtOption};

/// A best-effort constant-time number used for reducing branching
/// based on secret information
#[derive(Copy, Clone, Debug)]
pub struct Number<T>(CtOption<T>);

impl<T> Number<T> {
    pub fn new(value: T) -> Self {
        Self(CtOption::new(value, Choice::from(1u8)))
    }

    pub fn is_valid(&self) -> Choice {
        self.0.is_some()
    }

    pub fn unwrap_or_default(&self) -> T
    where
        T: ConditionallySelectable + Default,
    {
        self.0.unwrap_or_else(Default::default)
    }

    pub fn and_then<U, F, C>(self, f: F) -> Number<U>
    where
        T: ConditionallySelectable + Default,
        F: FnOnce(T) -> (U, C),
        C: Into<Choice>,
    {
        Number(self.0.and_then(|value| {
            let (next, is_valid) = f(value);
            CtOption::new(next, is_valid.into())
        }))
    }

    #[must_use]
    pub fn filter<F, C>(self, f: F) -> Self
    where
        T: ConditionallySelectable + Default,
        F: FnOnce(T) -> C,
        C: Into<Choice>,
    {
        Number(self.0.and_then(|value| {
            let is_valid = f(value);
            CtOption::new(value, is_valid.into())
        }))
    }

    pub fn ct_lt(self, rhs: Self) -> Choice
    where
        T: ConditionallySelectable + Default + OverflowingSub,
    {
        (self - rhs).0.is_none()
    }

    pub fn ct_le(self, rhs: Self) -> Choice
    where
        T: ConditionallySelectable + Default + OverflowingSub,
    {
        (rhs - self).0.is_some()
    }

    pub fn ct_ge(self, rhs: Self) -> Choice
    where
        T: ConditionallySelectable + Default + OverflowingSub,
    {
        (self - rhs).0.is_some()
    }

    pub fn ct_gt(self, rhs: Self) -> Choice
    where
        T: ConditionallySelectable + Default + OverflowingSub,
    {
        (rhs - self).0.is_none()
    }
}

impl<T> From<T> for Number<T> {
    fn from(value: T) -> Self {
        Self::new(value)
    }
}

impl<T> ConditionallySelectable for Number<T>
where
    T: ConditionallySelectable,
{
    fn conditional_select(a: &Self, b: &Self, choice: Choice) -> Self {
        Self(CtOption::conditional_select(&a.0, &b.0, choice))
    }
}

impl<T> ConstantTimeEq for Number<T>
where
    T: ConstantTimeEq,
{
    fn ct_eq(&self, other: &Self) -> Choice {
        self.0.ct_eq(&other.0)
    }
}

impl<T> ops::Add for Number<T>
where
    T: ConditionallySelectable + Default + OverflowingAdd,
{
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(rhs.0.and_then(|rhs| (self + rhs).0))
    }
}

impl<T> ops::Add<T> for Number<T>
where
    T: ConditionallySelectable + Default + OverflowingAdd,
{
    type Output = Self;

    fn add(self, rhs: T) -> Self::Output {
        Self(self.0.and_then(|prev| {
            let (next, overflowed) = prev.overflowing_add(&rhs);
            let is_valid = !overflowed as u8;
            CtOption::new(next, is_valid.into())
        }))
    }
}

impl<T> ops::Sub for Number<T>
where
    T: ConditionallySelectable + Default + OverflowingSub,
{
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(rhs.0.and_then(|rhs| (self - rhs).0))
    }
}

impl<T> ops::Sub<T> for Number<T>
where
    T: ConditionallySelectable + Default + OverflowingSub,
{
    type Output = Self;

    fn sub(self, rhs: T) -> Self::Output {
        Self(self.0.and_then(|prev| {
            let (next, overflowed) = prev.overflowing_sub(&rhs);
            let is_valid = !overflowed as u8;
            CtOption::new(next, is_valid.into())
        }))
    }
}

impl<T> ops::Mul for Number<T>
where
    T: ConditionallySelectable + Default + OverflowingMul,
{
    type Output = Self;

    fn mul(self, rhs: Self) -> Self::Output {
        Self(rhs.0.and_then(|rhs| (self * rhs).0))
    }
}

impl<T> ops::Mul<T> for Number<T>
where
    T: ConditionallySelectable + Default + OverflowingMul,
{
    type Output = Self;

    fn mul(self, rhs: T) -> Self::Output {
        Self(self.0.and_then(|prev| {
            let (next, overflowed) = prev.overflowing_mul(&rhs);
            let is_valid = !overflowed as u8;
            CtOption::new(next, is_valid.into())
        }))
    }
}

impl<T> ops::Div for Number<T>
where
    T: ConditionallySelectable + Default + CheckedDiv,
{
    type Output = Self;

    fn div(self, rhs: Self) -> Self::Output {
        Self(rhs.0.and_then(|rhs| (self / rhs).0))
    }
}

impl<T> ops::Div<T> for Number<T>
where
    T: ConditionallySelectable + Default + CheckedDiv,
{
    type Output = Self;

    fn div(self, rhs: T) -> Self::Output {
        Self(self.0.and_then(|prev| {
            let next = prev.checked_div(&rhs);
            let is_valid = next.is_some() as u8;
            let next = next.unwrap_or_default();
            CtOption::new(next, is_valid.into())
        }))
    }
}

impl<T> ops::Rem for Number<T>
where
    T: ConditionallySelectable + Default + CheckedRem,
{
    type Output = Self;

    fn rem(self, rhs: Self) -> Self::Output {
        Self(rhs.0.and_then(|rhs| (self % rhs).0))
    }
}

impl<T> ops::Rem<T> for Number<T>
where
    T: ConditionallySelectable + Default + CheckedRem,
{
    type Output = Self;

    fn rem(self, rhs: T) -> Self::Output {
        Self(self.0.and_then(|prev| {
            let next = prev.checked_rem(&rhs);
            let is_valid = next.is_some() as u8;
            let next = next.unwrap_or_default();
            CtOption::new(next, is_valid.into())
        }))
    }
}

impl<T> ops::Shl<Number<u32>> for Number<T>
where
    T: ConditionallySelectable + Default + CheckedShl,
{
    type Output = Self;

    fn shl(self, rhs: Number<u32>) -> Self::Output {
        Self(rhs.0.and_then(|rhs| (self << rhs).0))
    }
}

impl<T> ops::Shl<u32> for Number<T>
where
    T: ConditionallySelectable + Default + CheckedShl,
{
    type Output = Self;

    fn shl(self, rhs: u32) -> Self::Output {
        Self(self.0.and_then(|prev| {
            let next = prev.checked_shl(rhs);
            let is_valid = next.is_some() as u8;
            let next = next.unwrap_or_default();
            CtOption::new(next, is_valid.into())
        }))
    }
}

impl<T> ops::Shr<Number<u32>> for Number<T>
where
    T: ConditionallySelectable + Default + CheckedShr,
{
    type Output = Self;

    fn shr(self, rhs: Number<u32>) -> Self::Output {
        Self(rhs.0.and_then(|rhs| (self >> rhs).0))
    }
}

impl<T> ops::Shr<u32> for Number<T>
where
    T: ConditionallySelectable + Default + CheckedShr,
{
    type Output = Self;

    fn shr(self, rhs: u32) -> Self::Output {
        Self(self.0.and_then(|prev| {
            let next = prev.checked_shr(rhs);
            let is_valid = next.is_some() as u8;
            let next = next.unwrap_or_default();
            CtOption::new(next, is_valid.into())
        }))
    }
}

impl<T> ops::Not for Number<T>
where
    T: ConditionallySelectable + Default + ops::Not,
{
    type Output = Number<T::Output>;

    fn not(self) -> Self::Output {
        Number(self.0.map(|prev| prev.not()))
    }
}

impl<T> ops::BitAnd for Number<T>
where
    T: ConditionallySelectable + Default + ops::BitAnd,
{
    type Output = Number<T::Output>;

    fn bitand(self, rhs: Self) -> Self::Output {
        Number(self.0.and_then(|prev| rhs.0.map(|rhs| prev.bitand(rhs))))
    }
}

impl<T> ops::BitOr for Number<T>
where
    T: ConditionallySelectable + Default + ops::BitOr,
{
    type Output = Number<T::Output>;

    fn bitor(self, rhs: Self) -> Self::Output {
        Number(self.0.and_then(|prev| rhs.0.map(|rhs| prev.bitor(rhs))))
    }
}

impl<T> ops::BitXor for Number<T>
where
    T: ConditionallySelectable + Default + ops::BitXor,
{
    type Output = Number<T::Output>;

    fn bitxor(self, rhs: Self) -> Self::Output {
        Number(self.0.and_then(|prev| rhs.0.map(|rhs| prev.bitxor(rhs))))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bolero::check;
    use ops::*;

    macro_rules! binop_test {
        ($op:ident, $checked_op:ident) => {
            #[test]
            #[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
            fn $op() {
                check!()
                    .with_type::<(u8, u8)>()
                    .cloned()
                    .for_each(|(a, b)| {
                        let actual = Number::new(a).$op(Number::new(b)).unwrap_or_default();
                        if let Some(expected) = a.$checked_op(b) {
                            assert_eq!(actual, expected);
                        } else {
                            assert_eq!(actual, 0);
                        }
                    });
            }
        };
    }

    binop_test!(add, checked_add);
    binop_test!(sub, checked_sub);
    binop_test!(mul, checked_mul);
    binop_test!(div, checked_div);
    binop_test!(rem, checked_rem);

    macro_rules! cmp_test {
        ($op:ident, $core_op:ident) => {
            #[test]
            #[cfg_attr(kani, kani::proof, kani::unwind(5), kani::solver(kissat))]
            fn $op() {
                check!()
                    .with_type::<(u8, u8)>()
                    .cloned()
                    .for_each(|(a, b)| {
                        let actual: bool = Number::new(a).$op(Number::new(b)).into();
                        let expected = a.$core_op(&b);
                        assert_eq!(actual, expected);
                    });
            }
        };
    }

    cmp_test!(ct_lt, lt);
    cmp_test!(ct_le, le);
    cmp_test!(ct_gt, gt);
    cmp_test!(ct_ge, ge);
}