//! A wrapper for an [`AsyncRead`] to support XON/XOFF flow control.

//!

//! This allows any `AsyncRead` that implements [`BufferIsEmpty`] to be used with XON/XOFF flow

//! control.

use std::io::Error;

use std::pin::Pin;

use std::task::{Context, Poll};

use futures::{AsyncRead, Stream};

use pin_project::pin_project;

use tor_basic_utils::assert_val_impl_trait;

use tor_cell::relaycell::flow_ctrl::XonKBpsEwma;

use crate::stream::StreamTarget;

use crate::util::notify::NotifyReceiver;

/// A wrapper for an [`AsyncRead`] to support XON/XOFF flow control.

///

/// This reader will take care of communicating with the circuit reactor to handle XON/XOFF-related

/// events.

#[derive(Debug)]

#[pin_project]

pub(crate) struct XonXoffReader<R, T: DrainRateNotifier = StreamTarget> {

    /// How we communicate with the circuit reactor.

    #[pin]

    ctrl: XonXoffReaderCtrl<T>,

    /// The inner reader.

    #[pin]

    reader: R,

    /// Have we received a drain rate request notification from the reactor,

    /// but haven't yet sent a drain rate update back to the reactor?

    pending_drain_rate_update: bool,

}

impl<R, T: DrainRateNotifier> XonXoffReader<R, T> {

    /// Create a new [`XonXoffReader`].

    ///

    /// The reader must implement [`BufferIsEmpty`], which allows the `XonXoffReader` to check if

    /// the incoming stream buffer is empty or not.

    /// Get a reference to the inner [`AsyncRead`].

    ///

    /// NOTE: This will bypass the [`XonXoffReader`] and may cause incorrect behaviour depending on

    /// how you use the returned reader (for example if it uses interior mutability).

    /// Get a mutable reference to the inner [`AsyncRead`].

    ///

    /// NOTE: This will bypass the [`XonXoffReader`] and may cause incorrect behaviour depending on

    /// how you use the returned reader (for example if you read bytes directly).

}

impl<R: AsyncRead + BufferIsEmpty, T: DrainRateNotifier> AsyncRead for XonXoffReader<R, T> {

        // ensure that `drain_rate_request_stream` is a `FusedStream`,

        // which means that we don't need to worry about calling `poll_next()` repeatedly

            futures::stream::FusedStream,

        );

        // check if the circuit reactor has requested a drain rate update

        // try reading from the inner reader

        // if we need to send a drain rate update and the stream buffer is empty, inform the reactor

            // TODO(arti#534): in the future we want to do rate estimation, but for now we'll just

            // send an "unlimited" drain rate

}

/// Something that sends drain rate updates to the flow control logic (the `XonXoffFlowCtrl`).

pub(crate) trait DrainRateNotifier {

    /// Send the drain rate update.

    fn notify(&mut self, rate: XonKBpsEwma) -> Result<(), Error>;

}

impl DrainRateNotifier for StreamTarget {

}

/// The control structure for a stream that partakes in XON/XOFF flow control.

///

/// Used to construct an [`XonXoffReader`].

///

/// This contains a mechanism for us to be asked for our drain rate,

/// and a mechanism of sending the drain rate in response.

///

/// The `DrainRateNotifier` is typically a `StreamTarget`,

/// which sends the drain rate to the circuit reactor so that it can be sent in an XON message.

/// We make this a trait to make unit testing possible.

#[derive(Debug)]

#[pin_project]

pub(crate) struct XonXoffReaderCtrl<T: DrainRateNotifier = StreamTarget> {

    /// Receive notifications when the reactor requests a new drain rate.

    /// When we do, we should begin waiting for the receive buffer to clear.

    /// Then when the buffer clears, we should send a new drain rate update to the reactor.

    #[pin]

    drain_rate_request_stream: NotifyReceiver<DrainRateRequest>,

    /// An abstract handle to the reactor for this stream.

    /// This allows us to send drain rate updates to the circuit reactor.

    drain_rate_notifier: T,

}

impl<T: DrainRateNotifier> XonXoffReaderCtrl<T> {

    /// Create a new [`XonXoffReaderCtrl`].

    ///

    /// The `drain_rate_request_stream` informs us when we need to send our drain rate,

    /// and `drain_rate_notifier` allows us to send that drain rate.

}

/// Used by the [`XonXoffReader`] to decide when to send a drain rate update

/// (typically resulting in an XON message).

pub(crate) trait BufferIsEmpty {

    /// Returns `true` if there are no incoming bytes buffered on this stream.

    ///

    /// This takes a `&mut` so that implementers can

    /// [`unobtrusive_peek()`](tor_async_utils::peekable_stream::UnobtrusivePeekableStream::unobtrusive_peek)

    /// a stream if necessary.

    fn is_empty(self: Pin<&mut Self>) -> bool;

}

/// A marker type for a [`NotifySender`](crate::util::notify::NotifySender)

/// indicating that notifications are for new drain rate requests.

#[derive(Debug)]

pub(crate) struct DrainRateRequest;

#[cfg(test)]

// We use some tokio-specific types here to make the test easier to write.

#[cfg(feature = "tokio")]

mod test {

    // @@ begin test lint list maintained by maint/add_warning @@

    #![allow(clippy::bool_assert_comparison)]

    #![allow(clippy::clone_on_copy)]

    #![allow(clippy::dbg_macro)]

    #![allow(clippy::mixed_attributes_style)]

    #![allow(clippy::print_stderr)]

    #![allow(clippy::print_stdout)]

    #![allow(clippy::single_char_pattern)]

    #![allow(clippy::unwrap_used)]

    #![allow(clippy::unchecked_time_subtraction)]

    #![allow(clippy::useless_vec)]

    #![allow(clippy::needless_pass_by_value)]

    #![allow(clippy::string_slice)] // See arti#2571

    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->

    use super::*;

    use std::sync::Arc;

    use std::sync::atomic::{AtomicU64, Ordering};

    use crate::stream::flow_ctrl::params::FlowCtrlParameters;

    use crate::stream::flow_ctrl::state::{

        FlowCtrlHooks, StreamRateLimit, WithSidechannelMitigations,

    };

    use crate::stream::flow_ctrl::xon_xoff::state::XonXoffFlowCtrl;

    use crate::util::notify::NotifySender;

    use futures::channel::mpsc::{self, TryRecvError};

    use futures::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};

    use tokio_crate::io::{DuplexStream, duplex};

    use tokio_util::compat::{Compat, TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};

    /// The type that will be stored by the [`XonXoffReader`] and used to send drain rate updates.

    ///

    /// This essentially mocks what the [`StreamTarget`] would do.

    struct TestingDrainRateUpdates(mpsc::UnboundedSender<XonKBpsEwma>);

    impl TestingDrainRateUpdates {

        pub(crate) fn new(sender: mpsc::UnboundedSender<XonKBpsEwma>) -> Self {

            Self(sender)

        }

    }

    impl DrainRateNotifier for TestingDrainRateUpdates {

        fn notify(&mut self, rate: XonKBpsEwma) -> Result<(), Error> {

            self.0.unbounded_send(rate).unwrap();

            Ok(())

        }

    }

    /// The writer for a data stream that tracks the length.

    #[pin_project::pin_project]

    struct WriterWithLength<W> {

        #[pin]

        writer: W,

        length: Arc<AtomicU64>,

    }

    /// The reader for a data stream that tracks the length.

    #[pin_project::pin_project]

    struct ReaderWithLength<R> {

        #[pin]

        reader: R,

        length: Arc<AtomicU64>,

    }

    /// Wraps a writer and reader to track the queue length.

    fn with_length<W, R>(writer: W, reader: R) -> (WriterWithLength<W>, ReaderWithLength<R>) {

        let length = Arc::new(AtomicU64::new(0));

        let writer = WriterWithLength {

            writer,

            length: Arc::clone(&length),

        };

        let reader = ReaderWithLength { reader, length };

        (writer, reader)

    }

    impl<W> WriterWithLength<W> {

        /// Amount of bytes queued.

        pub(crate) fn len(&self) -> u64 {

            self.length.load(Ordering::Acquire)

        }

    }

    impl<R> BufferIsEmpty for ReaderWithLength<R> {

        fn is_empty(self: Pin<&mut Self>) -> bool {

            self.length.load(Ordering::Acquire) == 0

        }

    }

    impl<W: AsyncWrite> AsyncWrite for WriterWithLength<W> {

        fn poll_write(

            self: Pin<&mut Self>,

            cx: &mut Context<'_>,

            buf: &[u8],

        ) -> Poll<std::io::Result<usize>> {

            let self_ = self.project();

            let rv = self_.writer.poll_write(cx, buf);

            // NOTE: There's a race condition here since we don't write to the writer and update the

            // length as one atomic operation.

            // But this is good enough for our test where the mock runtime is deterministic and

            // single-threaded.

            //

            // We ignore the possibility of overflowing the 64-bit integer here.

            if let Poll::Ready(Ok(len)) = rv {

                let len: u64 = len.try_into().expect("usize should fit into u64");

                // The effect of `poll_write()` above will be visible after another thread checks

                // the length with `load(Acquire)`.

                self_.length.fetch_add(len, Ordering::Release);

            }

            rv

        }

        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {

            self.project().writer.poll_flush(cx)

        }

        fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {

            self.project().writer.poll_close(cx)

        }

    }

    impl<R: AsyncRead> AsyncRead for ReaderWithLength<R> {

        fn poll_read(

            self: Pin<&mut Self>,

            cx: &mut Context<'_>,

            buf: &mut [u8],

        ) -> Poll<std::io::Result<usize>> {

            let self_ = self.project();

            let rv = self_.reader.poll_read(cx, buf);

            // NOTE: There's a race condition here since we don't read from the reader and update

            // the length as one atomic operation.

            // But this is good enough for our test where the mock runtime is deterministic and

            // single-threaded.

            //

            // We ignore the possibility of underflowing the integer here.

            if let Poll::Ready(Ok(len)) = rv {

                let len: u64 = len.try_into().expect("usize should fit into u64");

                // The effect of `poll_read()` above will be visible after another thread checks

                // the length with `load(Acquire)`.

                self_.length.fetch_sub(len, Ordering::Release);

            }

            rv

        }

    }

    /// Set up all of the flow control stuff needed to test the [`XonXoffReader`].

    ///

    /// Returns:

    ///

    /// 1. The stream writer (as would be held by the circuit/stream reactor).

    /// 2. The stream reader (as would be held in a user-facing `DataStream`).

    /// 3. An MPSC receiver of drain rate updates.

    /// 4. The flow control logic.

    #[allow(clippy::type_complexity)]

    fn init_flow_ctrl(

        with_sidechannel_mitigations: WithSidechannelMitigations,

    ) -> (

        WriterWithLength<Compat<DuplexStream>>,

        XonXoffReader<ReaderWithLength<Compat<DuplexStream>>, TestingDrainRateUpdates>,

        mpsc::UnboundedReceiver<XonKBpsEwma>,

        XonXoffFlowCtrl,

    ) {

        let params = FlowCtrlParameters::defaults_for_tests();

        // For the flow control logic to send rate limit changes to the stream writer.

        // We don't use this in this test, but the `XonXoffFlowCtrl` needs the tx side.

        let (rate_limit_tx, _rate_limit_rx) = postage::watch::channel_with(StreamRateLimit::MAX);

        // For the flow control logic to request a new drain rate update from the stream reader.

        let mut drain_rate_request_tx = NotifySender::new_typed();

        let drain_rate_request_rx = drain_rate_request_tx.subscribe();

        // The flow control logic.

        let flow_ctrl = XonXoffFlowCtrl::new(

            Arc::new(params),

            with_sidechannel_mitigations,

            rate_limit_tx,

            drain_rate_request_tx,

        );

        // For the `XonXoffReader` to send a drain rate update.

        let (drain_rate_sender, drain_rate_receiver) = mpsc::unbounded();

        let drain_rate_updates = TestingDrainRateUpdates::new(drain_rate_sender);

        // All of the information needed to build a `XonXoffReader`.

        let reader_ctrl = XonXoffReaderCtrl::new(drain_rate_request_rx, drain_rate_updates);

        // This is the stream queue for incoming data.

        // So the `reader` is the stream reader and the `writer` would be within the reactor.

        //

        // In arti this stream should be unbounded, so here we use a max size of `usize::MAX`.

        let (writer, reader) = duplex(/* max_buf_size= */ usize::MAX);

        let writer = writer.compat_write();

        let reader = reader.compat();

        // Make the reader+writer pair track the length of the buffer so that it can support

        // `BufferIsEmpty`.

        let (writer, reader) = with_length(writer, reader);

        // The reader for incoming stream data, with XON/XOFF support.

        let reader = XonXoffReader::new(reader_ctrl, reader);

        (writer, reader, drain_rate_receiver, flow_ctrl)

    }

    /// Buffer `num_bytes` as if the bytes arrived on the stream.

    ///

    /// Returns whether the flow control logic wanted to send an XOFF.

    async fn buffer_incoming_data(

        writer: &mut WriterWithLength<impl AsyncWrite + Unpin>,

        mut num_bytes: usize,

        flow_ctrl: &mut XonXoffFlowCtrl,

    ) -> bool {

        let mut wants_to_send_xoff = false;

        // Write the requested number of bytes.

        while num_bytes > 0 {

            // Write 100_000 bytes at a time.

            let buf_size = num_bytes.min(100_000);

            writer.write_all(&vec![0; buf_size]).await.unwrap();

            num_bytes -= buf_size;

            // Inform the flow control logic.

            let xoff = flow_ctrl.maybe_send_xoff(writer.len() as usize).unwrap();

            wants_to_send_xoff |= xoff.is_some();

        }

        wants_to_send_xoff

    }

    /// Read `num_bytes` from the stream.

    async fn read_incoming_data(mut reader: impl AsyncRead + Unpin, mut num_bytes: usize) {

        // Read the requested number of bytes.

        while num_bytes > 0 {

            // Read 100_000 bytes at a time.

            let buf_size = num_bytes.min(100_000);

            reader.read_exact(&mut vec![0; buf_size]).await.unwrap();

            num_bytes -= buf_size;

        }

    }

    /// This test is meant to test the drain rate update.

    /// It adds a lot of data to the stream queue so that it triggers sending an XOFF

    /// and sends a drain rate request to the [`XonXoffReader`],

    /// then it reads from the stream until it's empty

    /// and the `XonXoffReader` sends a drain rate update.

    /// The flow control logic receives the drain rate update and sends an XON.

    #[test]

    fn drain_rate_update() {

        tor_rtmock::MockRuntime::test_with_various(|_rt| async move {

            // This is the stream queue for incoming data.

            // So the `reader` is the stream reader and the `writer` would be within the reactor.

            let (mut writer, mut reader, mut drain_rate_receiver, mut flow_ctrl) =

                init_flow_ctrl(WithSidechannelMitigations::Enabled);

            // Data has arrived on the stream.

            // We always consider sending an XOFF when a stream has received data.

            // The amount of incoming data wasn't very large,

            // so we don't expect that it would actually want to send an XOFF.

            let wants_to_send_xoff =

                buffer_incoming_data(&mut writer, 10_000, &mut flow_ctrl).await;

            assert!(!wants_to_send_xoff);

            // We didn't want to send an XOFF,

            // so the stream reader will never have been asked for a drain rate update.

            assert!(!reader.pending_drain_rate_update);

            // The stream reader reads all of the incoming data.

            read_incoming_data(&mut reader, 10_000).await;

            // Check `pending_drain_rate_update` again,

            // and also ensure that we didn't send a drain rate update.

            assert!(!reader.pending_drain_rate_update);

            assert_eq!(drain_rate_receiver.try_recv(), Err(TryRecvError::Empty));

            // Data has arrived on the stream.

            // We always consider sending an XOFF when a stream has received data.

            // The amount of incoming data was large,

            // so we expect that it would want to send an XOFF.

            let wants_to_send_xoff =

                buffer_incoming_data(&mut writer, 800_000, &mut flow_ctrl).await;

            assert!(wants_to_send_xoff);

            // The above code should have sent an XOFF and asked the reader for a drain rate update,

            // but the reader hasn't realized this yet.

            assert!(!reader.pending_drain_rate_update);

            assert_eq!(drain_rate_receiver.try_recv(), Err(TryRecvError::Empty));

            // The reader won't realize it was asked for a drain rate update until after it's tried

            // reading once.

            let _ = reader.read(&mut [0; 0]).await.unwrap();

            assert!(reader.pending_drain_rate_update);

            // The drain rate update is only sent once we've drained the buffer,

            // so an update should not have been sent yet.

            assert_eq!(drain_rate_receiver.try_recv(), Err(TryRecvError::Empty));

            // Read most (but not all) of the data on the stream.

            read_incoming_data(&mut reader, 700_000).await;

            // We haven't read *all* of the data,

            // so should still not have sent a drain rate update.

            assert!(!Pin::new(reader.inner_mut()).is_empty());

            assert!(reader.pending_drain_rate_update);

            assert_eq!(drain_rate_receiver.try_recv(), Err(TryRecvError::Empty));

            // Read the last of the data on the stream.

            read_incoming_data(&mut reader, 100_000).await;

            // Now that the buffer is empty,

            // we should have sent a drain rate update.

            assert!(Pin::new(reader.inner_mut()).is_empty());

            assert!(!reader.pending_drain_rate_update);

            let xon_rate = drain_rate_receiver.try_recv().unwrap();

            assert_eq!(xon_rate, XonKBpsEwma::Unlimited);

            // The buffer is still empty,

            // so the flow control logic should want to send an XON.

            let xon = flow_ctrl

                .maybe_send_xon(xon_rate, writer.len() as usize)

                .unwrap()

                .unwrap();

            assert_eq!(xon.kbytes_per_sec_ewma(), xon_rate);

        });

    }

    /// Like the `drain_rate_update()` test,

    /// this test causes the `XonXoffReader` to send a drain rate update.

    /// But in this case the buffer refills again past the high-water mark

    /// before the drain rate update can be processed by the flow control logic,

    /// so it *does not* send an XON.

    /// Instead it re-requests a drain rate from the `XonXoffReader`.

    #[test]

    fn drain_rate_update_then_buffer_refill() {

        tor_rtmock::MockRuntime::test_with_various(|_rt| async move {

            // This is the stream queue for incoming data.

            // So the `reader` is the stream reader and the `writer` would be within the reactor.

            let (mut writer, mut reader, mut drain_rate_receiver, mut flow_ctrl) =

                init_flow_ctrl(WithSidechannelMitigations::Enabled);

            // Data has arrived on the stream.

            // We always consider sending an XOFF when a stream has received data.

            // The amount of incoming data was large,

            // so we expect that it would want to send an XOFF.

            let wants_to_send_xoff =

                buffer_incoming_data(&mut writer, 800_000, &mut flow_ctrl).await;

            assert!(wants_to_send_xoff);

            // Read all of the data on the stream.

            read_incoming_data(&mut reader, 700_000).await;

            assert!(reader.pending_drain_rate_update);

            read_incoming_data(&mut reader, 100_000).await;

            // Now that the buffer is empty,

            // we should have sent a drain rate update.

            assert!(Pin::new(reader.inner_mut()).is_empty());

            assert!(!reader.pending_drain_rate_update);

            // Before this drain rate update can make it to the

            // flow control logic with `maybe_send_xon()`,

            // the buffer fills again past the high-water mark.

            let wants_to_send_xoff =

                buffer_incoming_data(&mut writer, 800_000, &mut flow_ctrl).await;

            assert!(!wants_to_send_xoff);

            // Now the drain rate update makes it to the flow control logic.

            // Since the buffer is past the high-water mark,

            // we won't want to send an XON.

            let xon_rate = drain_rate_receiver.try_recv().unwrap();

            assert_eq!(xon_rate, XonKBpsEwma::Unlimited);

            let xon = flow_ctrl

                .maybe_send_xon(xon_rate, writer.len() as usize)

                .unwrap();

            assert!(xon.is_none());

            // Instead the reader will have been asked for a drain rate update again,

            // which restarts the entire process.

            assert!(!reader.pending_drain_rate_update);

            let _ = reader.read(&mut [0; 0]).await.unwrap();

            assert!(reader.pending_drain_rate_update);

        });

    }

}