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

//! Defines the Stream objects that applications are interacting with

use crate::connection::Connection;
use bytes::Bytes;
use core::{
    fmt,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
pub use s2n_quic_core::{
    application,
    stream::{ops, StreamError, StreamId, StreamType},
};

#[derive(Clone)]
struct State {
    connection: Connection,
    stream_id: StreamId,
    rx: ops::Status,
    tx: ops::Status,
}

impl State {
    fn new(connection: Connection, stream_id: StreamId) -> Self {
        Self {
            connection,
            stream_id,
            rx: ops::Status::Open,
            tx: ops::Status::Open,
        }
    }

    fn poll_request(
        &mut self,
        request: &mut ops::Request,
        context: Option<&Context>,
    ) -> Result<ops::Response, StreamError> {
        let id = self.stream_id;
        self.connection.poll_request(id, request, context)
    }

    fn request(&mut self) -> Request {
        Request {
            state: self,
            request: ops::Request::default(),
        }
    }
}

impl Drop for State {
    fn drop(&mut self) {
        let is_rx_open = !self.rx.is_closed();
        let is_tx_open = !self.tx.is_closed();

        if is_rx_open || is_tx_open {
            let mut request = self.request();

            if is_tx_open {
                // Dropping a send stream will automatically finish the stream
                //
                // This is to stay consistent with std::net::TcpStream
                request.finish().detach_tx();
            }

            if is_rx_open {
                // Send a STOP_SENDING message on the receiving half of the `Stream`,
                // for the case the application did not consume all data.
                // If that already happened, this will be a noop.
                request
                    .stop_sending(application::Error::UNKNOWN)
                    .detach_rx();
            }

            let _ = request.poll(None);
        }
    }
}

macro_rules! ready {
    ($e:expr $(,)?) => {
        match $e {
            core::task::Poll::Ready(t) => t,
            core::task::Poll::Pending => return core::task::Poll::Pending,
        }
    };
}

macro_rules! tx_stream_apis {
    () => {
        /// Enqueues a chunk of data for sending it towards the peer.
        ///
        /// The method will return:
        /// - `Poll::Ready(Ok(()))` if the data was enqueued for sending. The provided `Bytes` will
        ///   be replaced with an empty `Bytes`, in order to reduce needless ref count increases.
        /// - `Poll::Ready(Err(stream_error))` if the data could not be sent, because the stream
        ///   had previously entered an error state.
        /// - `Poll::Pending` if the send buffer capacity is currently exhausted. In this case, the
        ///   caller should retry sending after the `Waker` on the provided `Context` is notified.
        pub fn poll_send(
            &mut self,
            chunk: &mut Bytes,
            cx: &mut Context,
        ) -> Poll<Result<(), StreamError>> {
            if chunk.is_empty() {
                return Poll::Ready(Ok(()));
            }

            self.tx_request()?
                .send(core::slice::from_mut(chunk))
                .poll(Some(cx))?
                .into()
        }

        /// Enqueues a slice of chunks of data for sending it towards the peer.
        ///
        /// The method will return:
        /// - `Poll::Ready(Ok(count))` if part of the slice was enqueued for sending. Any of the
        ///   consumed `Bytes` will be replaced with an empty `Bytes`, in order to reduce needless
        ///   ref count increases. If `count` does not equal the total number of chunks, the stream
        ///   will store the waker and wake the task once more capacity is available.
        /// - `Poll::Ready(Err(stream_error))` if the data could not be sent, because the stream
        ///   had previously entered an error state.
        /// - `Poll::Pending` if the send buffer capacity is currently exhausted. In this case, the
        ///   caller should retry sending after the `Waker` on the provided `Context` is notified.
        pub fn poll_send_vectored(
            &mut self,
            chunks: &mut [Bytes],
            cx: &mut Context,
        ) -> Poll<Result<usize, StreamError>> {
            if chunks.is_empty() {
                return Poll::Ready(Ok(0));
            }

            let response = self.tx_request()?.send(chunks).poll(Some(cx))?;

            if response.chunks.consumed == 0 {
                return Poll::Pending;
            }

            Ok(response.tx().expect("invalid response").chunks.consumed).into()
        }

        /// Polls send readiness for the given stream.
        ///
        /// The method will return:
        /// - `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(stream_error))` if the data could not be sent, because the stream
        ///   had previously entered an error state.
        /// - `Poll::Pending` if the send buffer capacity is currently exhausted. In this case, the
        ///   caller should retry sending after the `Waker` on the provided `Context` is notified.
        pub fn poll_send_ready(&mut self, cx: &mut Context) -> Poll<Result<usize, StreamError>> {
            let response = ready!(self
                .tx_request()?
                .send_readiness()
                .poll(Some(cx))?
                .into_poll());
            Ok(response.tx().expect("invalid response").bytes.available).into()
        }

        /// Enqueues a chunk of data for sending it towards the peer.
        ///
        /// This method should only be called after calling `poll_send_ready` first, as the stream
        /// may not have available send buffer capacity.
        ///
        /// The method will return:
        /// - `Ok(())` if the data was enqueued for sending.
        /// - `Err(stream_error)` if the data could not be sent, because the stream
        ///   had previously entered an error state, or the stream was not ready to send data.
        pub fn send_data(&mut self, chunk: Bytes) -> Result<(), StreamError> {
            if chunk.is_empty() {
                return Ok(());
            }

            match self.tx_request()?.send(&mut [chunk]).poll(None)? {
                response if response.tx().expect("invalid response").chunks.consumed == 1 => Ok(()),
                _ => Err(StreamError::sending_blocked()),
            }
        }

        /// Flushes the send buffer and waits for acknowledgement from the peer.
        ///
        /// The method will return:
        /// - `Poll::Ready(Ok(()))` if the send buffer was completely flushed and acknowledged.
        /// - `Poll::Ready(Err(stream_error))` if the stream could not be flushed, because the stream
        ///   had previously entered an error state.
        /// - `Poll::Pending` if the send buffer is still being flushed. In this case, the
        ///   caller should retry sending after the `Waker` on the provided `Context` is notified.
        pub fn poll_flush(&mut self, cx: &mut Context) -> Poll<Result<(), StreamError>> {
            self.tx_request()?.flush().poll(Some(cx))?.into()
        }

        /// Marks the stream as finished.
        ///
        /// The method will return:
        /// - `Ok(())` if the stream was finished successfully
        /// - `Err(stream_error)` if the stream could not be finished, because the stream
        ///   had previously entered an error state.
        pub fn finish(&mut self) -> Result<(), StreamError> {
            self.tx_request()?.finish().poll(None)?;
            Ok(())
        }

        /// Marks the stream as finished and waits for all outstanding data to be acknowledged
        ///
        /// The method will return:
        /// - `Poll::Ready(Ok(()))` if the stream was completely flushed and acknowledged.
        /// - `Poll::Ready(Err(stream_error))` if the stream could not be flushed, because the stream
        ///   had previously entered an error state.
        /// - `Poll::Pending` if the stream is still being flushed. In this case, the
        ///   caller should retry sending after the `Waker` on the provided `Context` is notified.
        pub fn poll_close(&mut self, cx: &mut Context) -> Poll<Result<(), StreamError>> {
            self.tx_request()?.finish().flush().poll(Some(cx))?.into()
        }

        /// Initiates a `RESET` on the stream.
        ///
        /// This will close the stream and notify the peer of the provided `error_code`.
        pub fn reset(&mut self, error_code: application::Error) -> Result<(), StreamError> {
            self.tx_request()?.reset(error_code).poll(None)?;
            Ok(())
        }
    };
}

macro_rules! rx_stream_apis {
    () => {
        /// Receives a chunk of data from the stream.
        ///
        /// The method will return:
        /// - `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(stream_error))` if the stream could not be read, because the stream
        ///   had previously entered an error state.
        /// - `Poll::Pending` if the stream is waiting to receive data from the peer. In this case, the
        ///   caller should retry sending after the `Waker` on the provided `Context` is notified.
        pub fn poll_receive(
            &mut self,
            cx: &mut Context,
        ) -> Poll<Result<Option<Bytes>, StreamError>> {
            let mut chunk = Bytes::new();
            let response =
                ready!(self.poll_receive_vectored(core::slice::from_mut(&mut chunk), cx))?;

            Ok(match response {
                // return the chunk if it was consumed
                (consumed, _) if consumed > 0 => Some(chunk),
                // any other response means the stream has ended
                _ => None,
            })
            .into()
        }

        /// Receives a slice of chunks of data from the stream.
        ///
        /// The method will return:
        /// - `Poll::Ready(Ok((len, is_open)))` if the stream received data into the slice,
        ///   where `len` was the number of chunks received, and `is_open` indicating if the stream is
        ///   still open. If `is_open == false`, future calls to `poll_receive_vectored` will
        ///   always return `Poll::Ready(Ok((0, false)))`.
        /// - `Poll::Ready(Err(stream_error))` if the stream could not be read, because the stream
        ///   had previously entered an error state.
        /// - `Poll::Pending` if the stream is waiting to receive data from the peer. In this case, the
        ///   caller should retry sending after the `Waker` on the provided `Context` is notified.
        pub fn poll_receive_vectored(
            &mut self,
            chunks: &mut [Bytes],
            cx: &mut Context,
        ) -> Poll<Result<(usize, bool), StreamError>> {
            let response = ready!(self
                .rx_request()?
                .receive(chunks)
                .poll(Some(cx))?
                .into_poll());

            let rx = response.rx().expect("invalid response");
            let consumed = rx.chunks.consumed;
            debug_assert!(
                consumed <= chunks.len(),
                "consumed exceeded the number of chunks provided"
            );
            // return if the stream is still open to receiving more data
            let is_open = rx.status.is_open() || rx.status.is_finishing();
            Poll::Ready(Ok((consumed, is_open)))
        }

        /// Sends a `STOP_SENDING` message to the peer. This requests the peer to
        /// finish the `Stream` as soon as possible by issuing a `RESET` with the
        /// provided `error_code`.
        ///
        /// Since this is merely a request to the peer to `RESET` the `Stream`, the
        /// `Stream` will not immediately be in a `RESET` state after issuing this
        /// API call.
        ///
        /// If the `Stream` had been previously reset by the peer or if all data had
        /// already been received the API call will not trigger any action.
        pub fn stop_sending(&mut self, error_code: application::Error) -> Result<(), StreamError> {
            self.rx_request()?.stop_sending(error_code).poll(None)?;
            Ok(())
        }
    };
}

/// A readable and writeable QUIC stream
pub struct Stream(State);

impl fmt::Debug for Stream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let is_alternate = f.alternate();

        let mut s = f.debug_struct("Stream");
        s.field("id", &self.id());

        // return additional information
        if is_alternate {
            // TODO query the shared state
        }

        s.finish()
    }
}

impl Stream {
    /// Creates a `Stream` instance, which represents a QUIC stream with the
    /// given ID. All interactions with the `Stream` will be performed through
    /// the provided [`SynchronizedSharedConnectionState`].
    pub(crate) fn new(connection: Connection, stream_id: StreamId) -> Self {
        Self(State::new(connection, stream_id))
    }

    pub fn id(&self) -> StreamId {
        self.0.stream_id
    }

    pub fn connection(&self) -> &Connection {
        &self.0.connection
    }

    pub fn request(&mut self) -> Request {
        self.0.request()
    }

    pub fn tx_request(&mut self) -> Result<TxRequest, StreamError> {
        Ok(TxRequest {
            state: &mut self.0,
            request: ops::Request::default(),
        })
    }

    pub fn rx_request(&mut self) -> Result<RxRequest, StreamError> {
        Ok(RxRequest {
            state: &mut self.0,
            request: ops::Request::default(),
        })
    }

    tx_stream_apis!();
    rx_stream_apis!();

    /// Splits a bidirectional QUIC Stream in two halves.
    ///
    /// One half can be used to read data from the Stream.
    /// The other half can be used to send data.
    pub fn split(self) -> (ReceiveStream, SendStream) {
        let mut rx_state = self.0;
        let mut tx_state = rx_state.clone();

        // close the opposite sides
        rx_state.tx = ops::Status::Finished;
        tx_state.rx = ops::Status::Finished;

        (ReceiveStream(rx_state), SendStream(tx_state))
    }
}

/// A writeable QUIC stream
pub struct SendStream(State);

impl fmt::Debug for SendStream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let is_alternate = f.alternate();

        let mut s = f.debug_struct("SendStream");
        s.field("id", &self.id());

        // return additional information
        if is_alternate {
            // TODO query the shared state
        }

        s.finish()
    }
}

impl SendStream {
    pub fn id(&self) -> StreamId {
        self.0.stream_id
    }

    pub fn connection(&self) -> &Connection {
        &self.0.connection
    }

    pub fn tx_request(&mut self) -> Result<TxRequest, StreamError> {
        Ok(TxRequest {
            state: &mut self.0,
            request: ops::Request::default(),
        })
    }

    tx_stream_apis!();
}

impl From<Stream> for SendStream {
    fn from(stream: Stream) -> Self {
        Self(stream.0)
    }
}

/// A readable QUIC stream
pub struct ReceiveStream(State);

impl fmt::Debug for ReceiveStream {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let is_alternate = f.alternate();

        let mut s = f.debug_struct("ReceiveStream");
        s.field("id", &self.id());

        // return additional information
        if is_alternate {
            // TODO query the shared state
        }

        s.finish()
    }
}

impl ReceiveStream {
    pub fn id(&self) -> StreamId {
        self.0.stream_id
    }

    pub fn connection(&self) -> &Connection {
        &self.0.connection
    }

    pub fn rx_request(&mut self) -> Result<RxRequest, StreamError> {
        Ok(RxRequest {
            state: &mut self.0,
            request: ops::Request::default(),
        })
    }

    rx_stream_apis!();
}

impl From<Stream> for ReceiveStream {
    fn from(stream: Stream) -> Self {
        Self(stream.0)
    }
}

macro_rules! tx_request_apis {
    () => {
        pub fn send(&mut self, chunks: &'chunks mut [Bytes]) -> &mut Self {
            self.request.send(chunks);
            self
        }

        pub fn send_readiness(&mut self) -> &mut Self {
            // express interest in tx
            if self.request.tx.is_none() {
                self.request.tx = Some(Default::default());
            }
            self
        }

        pub fn finish(&mut self) -> &mut Self {
            self.request.finish();
            self
        }

        pub fn reset(&mut self, error_code: application::Error) -> &mut Self {
            self.request.reset(error_code);
            self
        }

        pub fn flush(&mut self) -> &mut Self {
            self.request.flush();
            self
        }
    };
}

macro_rules! rx_request_apis {
    () => {
        pub fn receive(&mut self, chunks: &'chunks mut [Bytes]) -> &mut Self {
            self.request.receive(chunks);
            self
        }

        pub fn with_watermark(&mut self, low: usize, high: usize) -> &mut Self {
            self.request.with_watermark(low, high);
            self
        }

        pub fn with_low_watermark(&mut self, low: usize) -> &mut Self {
            self.request.with_low_watermark(low);
            self
        }

        pub fn with_high_watermark(&mut self, high: usize) -> &mut Self {
            self.request.with_high_watermark(high);
            self
        }

        pub fn stop_sending(&mut self, error_code: application::Error) -> &mut Self {
            self.request.stop_sending(error_code);
            self
        }
    };
}

pub struct Request<'state, 'chunks> {
    state: &'state mut State,
    request: ops::Request<'chunks>,
}

impl<'state, 'chunks> Request<'state, 'chunks> {
    tx_request_apis!();
    rx_request_apis!();

    fn detach_tx(&mut self) -> &mut Self {
        self.request.detach_tx();
        self
    }

    fn detach_rx(&mut self) -> &mut Self {
        self.request.detach_rx();
        self
    }

    pub fn poll(&mut self, context: Option<&Context>) -> Result<ops::Response, StreamError> {
        if self.state.rx.is_finished() && self.state.tx.is_finished() {
            // Tx and Rx are both finished, so return early to avoid sending a request
            // for a stream that has been removed from the stream container already
            return Ok(ops::Response {
                tx: Some(ops::tx::Response {
                    status: ops::Status::Finished,
                    ..Default::default()
                }),
                rx: Some(ops::rx::Response {
                    status: ops::Status::Finished,
                    ..Default::default()
                }),
            });
        }

        let response = self.state.poll_request(&mut self.request, context)?;

        if let Some(rx) = response.rx() {
            self.state.rx = rx.status;
        }

        if let Some(tx) = response.tx() {
            self.state.tx = tx.status;
        }

        Ok(response)
    }
}

impl<'state, 'chunks> Future for Request<'state, 'chunks> {
    type Output = Result<ops::Response, StreamError>;

    fn poll(
        mut self: Pin<&mut Self>,
        context: &mut Context,
    ) -> Poll<Result<ops::Response, StreamError>> {
        Self::poll(&mut self, Some(context))?.into()
    }
}

pub struct TxRequest<'state, 'chunks> {
    state: &'state mut State,
    request: ops::Request<'chunks>,
}

impl<'state, 'chunks> TxRequest<'state, 'chunks> {
    tx_request_apis!();

    pub fn poll(&mut self, context: Option<&Context>) -> Result<ops::tx::Response, StreamError> {
        if self.state.tx.is_finished() {
            // return early to avoid sending a request for a stream that has been
            // removed from the stream container already
            return Ok(ops::tx::Response {
                status: ops::Status::Finished,
                ..Default::default()
            });
        }

        let response = self
            .state
            .poll_request(&mut self.request, context)?
            .tx
            .expect("invalid response");

        self.state.tx = response.status;

        Ok(response)
    }
}

impl<'state, 'chunks> Future for TxRequest<'state, 'chunks> {
    type Output = Result<ops::tx::Response, StreamError>;

    fn poll(
        mut self: Pin<&mut Self>,
        context: &mut Context,
    ) -> Poll<Result<ops::tx::Response, StreamError>> {
        Self::poll(&mut self, Some(context))?.into()
    }
}

pub struct RxRequest<'state, 'chunks> {
    state: &'state mut State,
    request: ops::Request<'chunks>,
}

impl<'state, 'chunks> RxRequest<'state, 'chunks> {
    rx_request_apis!();

    pub fn poll(&mut self, context: Option<&Context>) -> Result<ops::rx::Response, StreamError> {
        if self.state.rx.is_finished() {
            // return early to avoid sending a request for a stream that has been
            // removed from the stream container already
            return Ok(ops::rx::Response {
                status: ops::Status::Finished,
                ..Default::default()
            });
        }

        let response = self
            .state
            .poll_request(&mut self.request, context)?
            .rx
            .expect("invalid response");

        self.state.rx = response.status;

        Ok(response)
    }
}

impl<'state, 'chunks> Future for RxRequest<'state, 'chunks> {
    type Output = Result<ops::rx::Response, StreamError>;

    fn poll(
        mut self: Pin<&mut Self>,
        context: &mut Context,
    ) -> Poll<Result<ops::rx::Response, StreamError>> {
        Self::poll(&mut self, Some(context))?.into()
    }
}


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

use bytes::Bytes;
use core::task::Context;
use s2n_quic_core::{
    application,
    stream::{ops, StreamError},
};

/// Request snapshot
pub struct Request {
    tx: Option<tx::Request>,
    rx: Option<rx::Request>,
}

impl<'a> From<&'a ops::Request<'a>> for Request {
    fn from(req: &'a ops::Request<'a>) -> Self {
        Self {
            tx: req.tx.as_ref().map(|req| req.into()),
            rx: req.rx.as_ref().map(|req| req.into()),
        }
    }
}

impl Request {
    pub fn validate_response(
        &self,
        request: &ops::Request,
        response: Result<&ops::Response, &StreamError>,
        context: Option<&Context>,
    ) {
        macro_rules! validate_response {
            ($ty:ident) => {
                if let Some(req) = self.$ty.as_ref() {
                    let res = response.map(|res| {
                        res.$ty
                            .as_ref()
                            .expect(concat!(stringify!($ty), " request should yield a response"))
                    });

                    let original = request
                        .$ty
                        .as_ref()
                        .expect("original request should always match snapshot");

                    req.validate_response(original, res, context);
                } else if let Ok(response) = response {
                    assert_eq!(response.$ty, None)
                }
            };
        }

        validate_response!(tx);
        validate_response!(rx);
    }
}

pub mod tx {
    use super::*;

    pub struct Request {
        chunks: Option<Vec<Bytes>>,
        reset: Option<application::Error>,
        flush: bool,
        finish: bool,
    }

    impl<'a> From<&'a ops::tx::Request<'a>> for Request {
        fn from(tx: &'a ops::tx::Request) -> Self {
            Self {
                chunks: tx.chunks.as_ref().map(|chunks| chunks.to_vec()),
                reset: tx.reset,
                flush: tx.flush,
                finish: tx.finish,
            }
        }
    }

    impl Request {
        #[allow(clippy::cognitive_complexity)]
        pub fn validate_response(
            &self,
            request: &ops::tx::Request,
            response: Result<&ops::tx::Response, &StreamError>,
            context: Option<&Context>,
        ) {
            // general response consistency checks
            if let Ok(response) = response {
                if response.will_wake {
                    assert!(
                        context.is_some(),
                        "will_wake should only be set when a context is provided"
                    );
                }

                if response.is_finished() {
                    assert_eq!(
                        response.bytes.available, 0,
                        "a finished stream should never report available bytes"
                    );
                    assert_eq!(
                        response.chunks.available, 0,
                        "a finished stream should never report available chunks"
                    );
                }
            }

            // resetting takes priority
            if self.reset.is_some() || response.map_or(false, |res| res.is_reset()) {
                let response = response.expect("reset should never fail");

                assert_eq!(
                    response.bytes,
                    Default::default(),
                    "resetting should always return empty bytes"
                );
                assert_eq!(
                    response.chunks,
                    Default::default(),
                    "resetting should always return empty chunks"
                );

                if response.will_wake {
                    assert!(self.flush, "resetting should only wake on flush");
                    assert!(
                        response.is_resetting(),
                        "resetting with flush should report a Resetting status"
                    );
                } else if self.flush && context.is_some() {
                    assert!(
                        response.is_reset(),
                        "resetting should always set the finish flag"
                    );
                } else {
                    assert!(
                        response.is_resetting() || response.is_reset(),
                        "resetting with flush should report a Resetting status"
                    );
                }

                // none of the other checks apply so return early
                return;
            }

            // the request is interested in push availability
            if self
                .chunks
                .as_ref()
                .map_or(true, |chunks| chunks.is_empty())
                && !self.finish
                && !self.flush
                && context.is_some()
            {
                if let Ok(response) = response {
                    assert!(
                        !response.is_closing() && !response.is_closed(),
                        "availability queries should not finish"
                    );
                    assert_eq!(
                        response.bytes.consumed, 0,
                        "availability queries should never consume bytes"
                    );
                    assert_eq!(
                        response.chunks.consumed, 0,
                        "availability queries should never consume chunks"
                    );

                    if response.will_wake {
                        assert_eq!(
                            response.bytes,
                            Default::default(),
                            "response should only wake when no bytes available"
                        );
                        assert_eq!(
                            response.chunks,
                            Default::default(),
                            "response should only wake when no chunks available"
                        );
                    } else {
                        assert_ne!(
                            response.bytes.available, 0,
                            "response should wake when no bytes available"
                        );
                        assert_ne!(
                            response.chunks.available, 0,
                            "response should wake when no chunks available"
                        );
                    }
                }

                // none of the other checks apply so return early
                return;
            }

            if let Some(chunks) = self.chunks.as_ref() {
                let original = request.chunks.as_ref().unwrap();

                if let Ok(response) = response {
                    // make sure the consumed chunks line up with the byte lengths
                    let mut actual_bytes = 0;
                    for chunk in chunks.iter().take(response.chunks.consumed) {
                        actual_bytes += chunk.len();
                    }
                    assert_eq!(
                        response.bytes.consumed, actual_bytes,
                        "consumed bytes should be reported accurately"
                    );

                    // make sure all of the consumed chunks are empty in the original slice
                    for chunk in original.iter().take(response.chunks.consumed) {
                        assert!(chunk.is_empty(), "consumed chunks should always be empty");
                    }

                    assert_eq!(
                        chunks.len(),
                        original.len(),
                        "snapshot should always have the same number of chunks as the original request"
                    );

                    // make sure the remaining chunks haven't been modified
                    for (original, snapshot) in chunks
                        .iter()
                        .zip(original.iter())
                        .skip(response.chunks.consumed)
                    {
                        assert_eq!(
                            original, snapshot,
                            "non-consumed chunks should not be modified"
                        );
                    }

                    if response.will_wake {
                        let mut should_wake = false;
                        // empty chunks means we're querying status
                        should_wake |= chunks.is_empty();
                        // not all of the chunks were consumed
                        should_wake |= response.chunks.consumed < chunks.len();
                        // the request wanted all of the chunks to be flushed
                        should_wake |= self.flush;
                        assert!(
                            should_wake,
                            concat!(
                                "waker should only wake when not all of the provided chunks were consumed ",
                                "or when a flush was requested",
                            )
                        );
                    } else if context.is_some() {
                        assert_eq!(
                            response.chunks.consumed,
                            chunks.len(),
                            concat!(
                                "if a context was provided and will_wake was not set, all of the ",
                                "provided chunks should be consumed"
                            )
                        );
                    }
                }
            }

            if self.finish {
                if let Ok(response) = response {
                    // finishing only happens after we don't consume anything
                    if response.chunks.consumed == 0 {
                        if response.will_wake {
                            assert!(
                            response.is_finishing(),
                            "finished responses that will_wake should have a Finishing status; actual: {:?}",
                            response.status
                        );
                            assert!(
                                response.chunks.consumed < self.chunks_len() || self.flush,
                                concat!(
                                    "waker should only wake when not all of the provided chunks were consumed ",
                                    "or when a flush was requested",
                                ),
                            );
                        } else if self.flush && context.is_some() {
                            assert!(response.is_finished());
                        } else {
                            assert!(
                                response.is_finishing() || response.is_finished(),
                                "finishing should transition to finish status; actual: {:?}",
                                response.status
                            );
                        }
                    }
                }
            }

            if self.flush {
                if let Ok(response) = response {
                    if !response.will_wake && context.is_some() && response.is_open() {
                        assert_ne!(
                            response.bytes.available, 0,
                            "flushing should result in available bytes"
                        );
                        assert_ne!(
                            response.chunks.available, 0,
                            "flushing should result in available chunks"
                        );
                    }
                }
            }
        }

        fn chunks_len(&self) -> usize {
            self.chunks.as_ref().map(|chunks| chunks.len()).unwrap_or(0)
        }
    }
}

pub mod rx {
    use super::*;

    pub struct Request {
        chunks: Option<Vec<Bytes>>,
        low_watermark: usize,
        high_watermark: usize,
        stop_sending: Option<application::Error>,
    }

    impl<'a> From<&'a ops::rx::Request<'a>> for Request {
        fn from(rx: &'a ops::rx::Request) -> Self {
            Self {
                chunks: rx.chunks.as_ref().map(|chunks| chunks.to_vec()),
                high_watermark: rx.high_watermark,
                low_watermark: rx.low_watermark,
                stop_sending: rx.stop_sending,
            }
        }
    }

    impl Request {
        #[allow(clippy::cognitive_complexity)]
        pub fn validate_response(
            &self,
            request: &ops::rx::Request,
            response: Result<&ops::rx::Response, &StreamError>,
            context: Option<&Context>,
        ) {
            // general response consistency checks
            if let Ok(response) = response {
                if response.will_wake {
                    assert!(
                        context.is_some(),
                        "will_wake should only be set when a context was provided"
                    );
                }

                if (response.is_open() || response.is_finishing()) && response.will_wake {
                    assert_eq!(
                        response.bytes.consumed, 0,
                        "will_wake should only be set when nothing was consumed"
                    );
                    assert_eq!(
                        response.chunks.consumed, 0,
                        "will_wake should only be set when nothing was consumed"
                    );
                }

                if response.is_finishing() && context.is_some() && !response.will_wake {
                    assert_ne!(
                        response.bytes.available, 0,
                        "a finishing stream should always report available bytes",
                    );
                    assert_ne!(
                        response.chunks.available, 0,
                        "a finishing stream should always report available chunks",
                    );
                }

                if response.is_resetting() || response.is_closed() {
                    assert!(!response.will_wake, "a closed stream should never wake");
                    assert_eq!(
                        response.bytes.available, 0,
                        "a closed stream should never report available bytes"
                    );
                    assert_eq!(
                        response.chunks.available, 0,
                        "a closed stream should never report available chunks"
                    );
                }

                assert_eq!(
                    request.low_watermark,
                    self.low_watermark.saturating_sub(response.bytes.consumed),
                    "the low watermark should be lowered as data is consumed"
                );

                assert_eq!(
                    request.high_watermark,
                    self.high_watermark.saturating_sub(response.bytes.consumed),
                    "the high watermark should be lowered as data is consumed"
                );

                assert!(
                    self.high_watermark >= response.bytes.consumed,
                    "the number of bytes consumed should not exceed the high watermark"
                );
            }

            if self.stop_sending.is_some() {
                let response = response.expect("stop_sending should never fail");
                assert!(
                    response.is_reset() || response.is_finished(),
                    "stop_sending should reset or finish the stream"
                );

                assert_eq!(
                    response.bytes,
                    Default::default(),
                    "resetting should always return empty bytes"
                );
                assert_eq!(
                    response.chunks,
                    Default::default(),
                    "resetting should always return empty chunks"
                );

                // none of the other checks apply so return early
                return;
            }

            if let Some(chunks) = self.chunks.as_ref() {
                // if any of the provided chunks are non-empty, it should return an error
                let should_error = chunks.iter().any(|chunk| !chunk.is_empty());

                if should_error {
                    assert!(matches!(response, Err(StreamError::NonEmptyOutput { .. })));
                    return;
                }

                if let Ok(response) = response {
                    let mut iter = request.chunks.as_ref().unwrap().iter();

                    // add up all of the actual bytes returned
                    let mut actual_bytes = 0;
                    for _ in 0..response.chunks.consumed {
                        let chunk = iter.next().unwrap();
                        actual_bytes += chunk.len();
                    }

                    assert_eq!(
                        actual_bytes, response.bytes.consumed,
                        "reported consumed bytes should reflect the output chunks"
                    );

                    for chunk in iter {
                        assert!(
                            chunk.is_empty(),
                            "chunks beyond the reported consumed length should be empty"
                        );
                    }
                }
            }
        }
    }
}


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

mod local_initiated;
mod remote_initiated;

use self::{local_initiated::LocalInitiated, remote_initiated::RemoteInitiated};
use crate::{
    connection,
    contexts::OnTransmitError,
    transmission,
    transmission::{interest::Provider, WriteContext},
};
use core::{
    task::{Context, Poll},
    time::Duration,
};
use futures_core::ready;
use s2n_quic_core::{
    ack, endpoint,
    frame::MaxStreams,
    stream::{self, iter::StreamIter, StreamId, StreamType},
    time::{timer, Timestamp},
    transport,
    transport::parameters::InitialFlowControlLimits,
    varint::VarInt,
};

pub use remote_initiated::MAX_STREAMS_SYNC_FRACTION;

/// This component manages stream concurrency limits.
///
/// It enforces both the local initiated stream limits and the peer initiated
/// stream limits.
///
/// It will also signal an increased max streams once streams have been consumed.
#[derive(Debug)]
pub struct Controller {
    local_endpoint_type: endpoint::Type,
    local_bidi_controller: LocalInitiated<
        stream::limits::LocalBidirectional,
        local_initiated::OpenNotifyBidirectional,
    >,
    remote_bidi_controller: RemoteInitiated,
    local_uni_controller: LocalInitiated<
        stream::limits::LocalUnidirectional,
        local_initiated::OpenNotifyUnidirectional,
    >,
    remote_uni_controller: RemoteInitiated,
}

impl Controller {
    /// Creates a new `stream::Controller`
    ///
    /// The peer will be allowed to open streams up to the given `initial_local_limits`.
    ///
    /// For local_initiated unidirectional streams, the local application will be allowed to open
    /// up to the minimum of the given local limits (`stream_limits`) and `initial_peer_limits`.
    ///
    /// For bidirectional streams, the local application will be allowed to open
    /// up to the minimum of the given `initial_local_limits` and `initial_peer_limits`.
    ///
    /// The peer may give additional credit to open more streams by delivering `MAX_STREAMS` frames.
    pub fn new(
        local_endpoint_type: endpoint::Type,
        initial_peer_limits: InitialFlowControlLimits,
        initial_local_limits: InitialFlowControlLimits,
        stream_limits: stream::Limits,
    ) -> Self {
        Self {
            local_endpoint_type,
            local_bidi_controller: LocalInitiated::new(
                initial_peer_limits.max_open_remote_bidirectional_streams,
                stream_limits.max_open_local_bidirectional_streams,
            ),
            remote_bidi_controller: RemoteInitiated::new(
                initial_local_limits.max_open_remote_bidirectional_streams,
            ),
            local_uni_controller: LocalInitiated::new(
                initial_peer_limits.max_open_remote_unidirectional_streams,
                stream_limits.max_open_local_unidirectional_streams,
            ),
            remote_uni_controller: RemoteInitiated::new(
                initial_local_limits.max_open_remote_unidirectional_streams,
            ),
        }
    }

    /// This method is called when a `MAX_STREAMS` frame is received,
    /// which signals an increase in the available streams budget.
    pub fn on_max_streams(&mut self, frame: &MaxStreams) {
        match frame.stream_type {
            StreamType::Bidirectional => self.local_bidi_controller.on_max_streams(frame),
            StreamType::Unidirectional => self.local_uni_controller.on_max_streams(frame),
        }
    }

    /// This method is called when the local application wishes to open the next stream
    /// of a type (Bidirectional/Unidirectional).
    ///
    /// `Poll::Pending` is returned when there isn't available capacity to open a stream,
    /// either because of local initiated concurrency limits or the peer's stream limits.
    /// If `Poll::Pending` is returned, the waker in the given `context` will be woken
    /// when additional stream capacity becomes available.
    pub fn poll_open_local_stream(
        &mut self,
        stream_type: StreamType,
        open_tokens: &mut connection::OpenToken,
        context: &Context,
    ) -> Poll<()> {
        let poll_open = match stream_type {
            StreamType::Bidirectional => self
                .local_bidi_controller
                .poll_open_stream(&mut open_tokens.bidirectional, context),
            StreamType::Unidirectional => self
                .local_uni_controller
                .poll_open_stream(&mut open_tokens.unidirectional, context),
        };

        // returns Pending if there is no capacity available
        ready!(poll_open);

        // only open streams if there is sufficient capacity based on limits
        let direction = self.direction(StreamId::initial(self.local_endpoint_type, stream_type));
        self.on_open_stream(direction);
        Poll::Ready(())
    }

    /// This method is called when the remote peer wishes to open a new stream.
    ///
    /// Opening a Stream also opens all lower Streams of the same type. Therefore
    /// this function validates if there is enough capacity to open all streams.
    ///
    /// A `STREAM_LIMIT_ERROR` will be returned if the peer has exceeded the
    /// stream limits that were communicated by transport parameters or
    /// MAX_STREAMS frames.
    pub fn on_open_remote_stream(
        &mut self,
        stream_iter: StreamIter,
    ) -> Result<(), transport::Error> {
        debug_assert!(
            self.direction(stream_iter.max_stream_id()).is_remote(),
            "should only be called for remote initiated streams"
        );

        // return early if there is not sufficient capacity based on limits
        match stream_iter.max_stream_id().stream_type() {
            StreamType::Bidirectional => self
                .remote_bidi_controller
                .on_remote_open_stream(stream_iter.max_stream_id())?,
            StreamType::Unidirectional => self
                .remote_uni_controller
                .on_remote_open_stream(stream_iter.max_stream_id())?,
        }

        let direction = self.direction(stream_iter.max_stream_id());
        // checked above that there is enough capacity to open all streams
        for _stream_id in stream_iter {
            self.on_open_stream(direction);
        }
        Ok(())
    }

    /// This method is called whenever a stream is opened, regardless of
    /// which side initiated.
    ///
    /// The caller is responsible for performing stream capacity checks
    /// prior to calling this function.
    fn on_open_stream(&mut self, direction: StreamDirection) {
        match direction {
            StreamDirection::LocalInitiatedBidirectional => {
                self.local_bidi_controller.on_open_stream()
            }
            StreamDirection::RemoteInitiatedBidirectional => {
                self.remote_bidi_controller.on_open_stream()
            }
            StreamDirection::LocalInitiatedUnidirectional => {
                self.local_uni_controller.on_open_stream()
            }
            StreamDirection::RemoteInitiatedUnidirectional => {
                self.remote_uni_controller.on_open_stream()
            }
        }
    }

    /// This method is called whenever a stream is closed.
    pub fn on_close_stream(&mut self, stream_id: StreamId) {
        match self.direction(stream_id) {
            StreamDirection::LocalInitiatedBidirectional => {
                self.local_bidi_controller.on_close_stream()
            }
            StreamDirection::RemoteInitiatedBidirectional => {
                self.remote_bidi_controller.on_close_stream()
            }
            StreamDirection::LocalInitiatedUnidirectional => {
                self.local_uni_controller.on_close_stream()
            }
            StreamDirection::RemoteInitiatedUnidirectional => {
                self.remote_uni_controller.on_close_stream()
            }
        }
    }

    /// This method is called when the stream manager is closed. All wakers will be woken
    /// to unblock waiting tasks.
    pub fn close(&mut self) {
        self.local_bidi_controller.close();
        self.remote_bidi_controller.close();
        self.local_uni_controller.close();
        self.remote_uni_controller.close();
    }

    /// This method is called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        self.local_bidi_controller.on_packet_ack(ack_set);
        self.remote_bidi_controller.on_packet_ack(ack_set);
        self.local_uni_controller.on_packet_ack(ack_set);
        self.remote_uni_controller.on_packet_ack(ack_set);
    }

    /// This method is called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        self.local_bidi_controller.on_packet_loss(ack_set);
        self.remote_bidi_controller.on_packet_loss(ack_set);
        self.local_uni_controller.on_packet_loss(ack_set);
        self.remote_uni_controller.on_packet_loss(ack_set);
    }

    /// Updates the period at which `STREAMS_BLOCKED` frames are sent to the peer
    /// if the application is blocked by peer limits.
    pub fn update_blocked_sync_period(&mut self, blocked_sync_period: Duration) {
        self.local_bidi_controller
            .update_sync_period(blocked_sync_period);
        self.local_uni_controller
            .update_sync_period(blocked_sync_period);
    }

    /// Queries the component for any local_initiated frames that need to get sent
    #[inline]
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) -> Result<(), OnTransmitError> {
        if !self.has_transmission_interest() {
            return Ok(());
        }

        let peer_endpoint_type = self.local_endpoint_type.peer_type();

        macro_rules! on_transmit {
            ($controller:ident, $endpoint:expr, $ty:expr) => {
                if let Some(nth) = self
                    .$controller
                    .total_open_stream_count()
                    .checked_sub(VarInt::from_u32(1))
                {
                    if let Some(stream_id) = StreamId::nth($endpoint, $ty, nth.as_u64()) {
                        self.$controller.on_transmit(stream_id, context)?;
                    }
                }
            };
        }

        on_transmit!(
            local_bidi_controller,
            self.local_endpoint_type,
            StreamType::Bidirectional
        );
        on_transmit!(
            remote_bidi_controller,
            peer_endpoint_type,
            StreamType::Bidirectional
        );

        on_transmit!(
            local_uni_controller,
            self.local_endpoint_type,
            StreamType::Unidirectional
        );
        on_transmit!(
            remote_uni_controller,
            peer_endpoint_type,
            StreamType::Unidirectional
        );

        Ok(())
    }

    /// Called when the connection timer expires
    pub fn on_timeout(&mut self, now: Timestamp) {
        self.local_bidi_controller.on_timeout(now);
        self.local_uni_controller.on_timeout(now);
    }

    #[inline]
    fn direction(&self, stream_id: StreamId) -> StreamDirection {
        let is_local_initiated = self.local_endpoint_type == stream_id.initiator();
        match (is_local_initiated, stream_id.stream_type()) {
            (true, StreamType::Bidirectional) => StreamDirection::LocalInitiatedBidirectional,
            (true, StreamType::Unidirectional) => StreamDirection::LocalInitiatedUnidirectional,
            (false, StreamType::Bidirectional) => StreamDirection::RemoteInitiatedBidirectional,
            (false, StreamType::Unidirectional) => StreamDirection::RemoteInitiatedUnidirectional,
        }
    }
}

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

/// Queries the component for interest in transmitting frames
impl transmission::interest::Provider for Controller {
    #[inline]
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.local_bidi_controller.transmission_interest(query)?;
        self.remote_bidi_controller.transmission_interest(query)?;
        self.local_uni_controller.transmission_interest(query)?;
        self.remote_uni_controller.transmission_interest(query)?;
        Ok(())
    }
}

#[derive(Debug, Copy, Clone)]
enum StreamDirection {
    // A bidirectional stream opened by the local application to send
    // and receive data
    LocalInitiatedBidirectional,

    // A bidirectional stream opened by the peer to send and receive
    // data
    RemoteInitiatedBidirectional,

    // A unidirectional stream opened by the local application to send
    // data
    LocalInitiatedUnidirectional,

    // A unidirectional stream opened by the peer to send data
    RemoteInitiatedUnidirectional,
}

impl StreamDirection {
    fn is_remote(&self) -> bool {
        match self {
            StreamDirection::LocalInitiatedBidirectional => false,
            StreamDirection::RemoteInitiatedBidirectional => true,
            StreamDirection::LocalInitiatedUnidirectional => false,
            StreamDirection::RemoteInitiatedUnidirectional => true,
        }
    }
}

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

    impl Controller {
        pub fn available_local_initiated_stream_capacity(&self, stream_type: StreamType) -> VarInt {
            match stream_type {
                StreamType::Bidirectional => self.local_bidi_controller.available_stream_capacity(),
                StreamType::Unidirectional => self.local_uni_controller.available_stream_capacity(),
            }
        }

        pub fn remote_initiated_max_streams_latest_value(&self, stream_type: StreamType) -> VarInt {
            match stream_type {
                StreamType::Bidirectional => self.remote_bidi_controller.latest_limit(),
                StreamType::Unidirectional => self.remote_uni_controller.latest_limit(),
            }
        }

        pub fn available_remote_initiated_stream_capacity(
            &self,
            stream_type: StreamType,
        ) -> VarInt {
            match stream_type {
                StreamType::Bidirectional => {
                    self.remote_initiated_max_streams_latest_value(stream_type)
                        - self.remote_bidi_controller.open_stream_count()
                }
                StreamType::Unidirectional => {
                    self.remote_initiated_max_streams_latest_value(stream_type)
                        - self.remote_uni_controller.open_stream_count()
                }
            }
        }
    }
}

#[cfg(test)]
mod fuzz_target;


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

//! Manages the per-connection flow-control window

use crate::{
    contexts::{OnTransmitError, WriteContext},
    sync::{IncrementalValueSync, ValueToFrameWriter},
    transmission,
};
use alloc::rc::Rc;
use core::cell::RefCell;
use s2n_quic_core::{
    ack, frame::max_data::MaxData, packet::number::PacketNumber, stream::StreamId, transport,
    varint::VarInt,
};

/// Writes `MAX_DATA` frames based on the connections flow control window.
#[derive(Default, Debug)]
pub(super) struct MaxDataToFrameWriter {}

impl ValueToFrameWriter<VarInt> for MaxDataToFrameWriter {
    fn write_value_as_frame<W: WriteContext>(
        &self,
        value: VarInt,
        _stream_id: StreamId,
        context: &mut W,
    ) -> Option<PacketNumber> {
        context.write_frame(&MaxData {
            maximum_data: value,
        })
    }
}

/// The actual implementation/state of the per Connection flow controller for
/// incoming data
#[derive(Debug)]
struct IncomingConnectionFlowControllerImpl {
    /// Synchronizes the read window to the remote peer
    pub(super) read_window_sync: IncrementalValueSync<VarInt, MaxDataToFrameWriter>,
    /// The relative flow control window we want to maintain
    pub(super) desired_flow_control_window: u32,
    /// The amount of flow control credits which already have been acquired by
    /// Streams.
    pub(super) acquired_window: VarInt,
    /// The amount of flow control credits which had been acquired and where the
    /// data had already been consumed by the application
    pub(super) consumed_window: VarInt,
}

impl IncomingConnectionFlowControllerImpl {
    pub fn new(initial_window_size: VarInt, desired_flow_control_window: u32) -> Self {
        Self {
            read_window_sync: IncrementalValueSync::new(
                VarInt::from_u32(desired_flow_control_window),
                initial_window_size,
                VarInt::from_u32(desired_flow_control_window / 10),
            ),
            desired_flow_control_window,
            acquired_window: VarInt::from_u32(0),
            consumed_window: VarInt::from_u32(0),
        }
    }

    pub fn remaining_window(&self) -> VarInt {
        self.read_window_sync.latest_value() - self.acquired_window
    }

    #[cfg(test)]
    pub(super) fn current_receive_window(&self) -> VarInt {
        self.read_window_sync.latest_value()
    }

    pub fn release_window(&mut self, amount: VarInt) {
        self.consumed_window += amount;
        debug_assert!(
            self.consumed_window <= self.acquired_window,
            "Can not consume more window than previously acquired"
        );

        self.read_window_sync.update_latest_value(
            self.consumed_window
                .saturating_add(VarInt::from_u32(self.desired_flow_control_window)),
        );
    }

    pub fn acquire_window(&mut self, desired: VarInt) -> Result<(), transport::Error> {
        if self.remaining_window() < desired {
            //= https://www.rfc-editor.org/rfc/rfc9000#section-4.1
            //# A receiver MUST close the connection with an error of type
            //# FLOW_CONTROL_ERROR if the sender violates the advertised connection
            //# or stream data limits; see Section 11 for details on error handling.
            return Err(transport::Error::FLOW_CONTROL_ERROR);
        }

        self.acquired_window += desired;
        Ok(())
    }

    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        self.read_window_sync.on_packet_ack(ack_set)
    }

    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        self.read_window_sync.on_packet_loss(ack_set)
    }

    #[inline]
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) -> Result<(), OnTransmitError> {
        // Stream ID does not matter here, since it does not get transmitted
        self.read_window_sync
            .on_transmit(StreamId::from_varint(VarInt::from_u32(0)), context)
    }
}

/// This component manages the flow control on the reception side.
///
/// It allows to check whether the peer sent more data on a per-connection base
/// then what was allowed through the flow control window.
///
/// It will also signal an increased window once data had been consumed.
#[derive(Clone, Debug)]
pub struct IncomingConnectionFlowController {
    inner: Rc<RefCell<IncomingConnectionFlowControllerImpl>>,
}

impl IncomingConnectionFlowController {
    /// Creates a new `IncomingConnectionFlowController`
    ///
    /// The connection flow controller will allow the peer to send up to
    /// `initial_window_size` bytes initially.
    ///
    /// The flow controller will try to maintain a window of
    /// `desired_flow_control_window`. This means if the window which is indicated
    /// to the peer is lower than this value the new value will be communicated
    /// to the peer.
    pub fn new(initial_window_size: VarInt, desired_flow_control_window: u32) -> Self {
        Self {
            inner: Rc::new(RefCell::new(IncomingConnectionFlowControllerImpl::new(
                initial_window_size,
                desired_flow_control_window,
            ))),
        }
    }

    /// Acquires a part of the window from the `IncomingConnectionFlowController` in
    /// in order to be able to use it for receiving data. `desired` is the window
    /// size that is intended to be borrowed.
    ///
    /// If the requested window size is not available the method will return
    /// an error in form of the `Err` variant.
    pub fn acquire_window(&mut self, desired: VarInt) -> Result<(), transport::Error> {
        self.inner.borrow_mut().acquire_window(desired)
    }

    pub fn release_window(&mut self, amount: VarInt) {
        self.inner.borrow_mut().release_window(amount)
    }

    /// This method gets called when a packet delivery got acknowledged
    pub fn on_packet_ack<A: ack::Set>(&mut self, ack_set: &A) {
        self.inner.borrow_mut().on_packet_ack(ack_set)
    }

    /// This method gets called when a packet loss is reported
    pub fn on_packet_loss<A: ack::Set>(&mut self, ack_set: &A) {
        self.inner.borrow_mut().on_packet_loss(ack_set)
    }

    /// Queries the component for any outgoing frames that need to get sent
    #[inline]
    pub fn on_transmit<W: WriteContext>(&mut self, context: &mut W) -> Result<(), OnTransmitError> {
        self.inner.borrow_mut().on_transmit(context)
    }

    pub fn acquired_window(&self) -> VarInt {
        self.inner.borrow().acquired_window
    }

    #[cfg(test)]
    pub fn remaining_window(&self) -> VarInt {
        self.inner.borrow_mut().remaining_window()
    }

    /// Returns the MAX_DATA window that is currently synchronized
    /// towards the peer.
    #[cfg(test)]
    pub(super) fn current_receive_window(&self) -> VarInt {
        self.inner.borrow().current_receive_window()
    }

    #[cfg(test)]
    pub fn desired_flow_control_window(&self) -> u32 {
        self.inner.borrow().desired_flow_control_window
    }

    #[cfg(test)]
    pub fn is_inflight(&self) -> bool {
        self.inner.borrow().read_window_sync.is_inflight()
    }
}

impl transmission::interest::Provider for IncomingConnectionFlowController {
    fn transmission_interest<Q: transmission::interest::Query>(
        &self,
        query: &mut Q,
    ) -> transmission::interest::Result {
        self.inner
            .borrow()
            .read_window_sync
            .transmission_interest(query)
    }
}