//! Declare DataStream, a type that wraps RawCellStream so as to be useful

//! for byte-oriented communication.

use crate::{Error, Result};

use static_assertions::assert_impl_all;

use tor_cell::relaycell::msg::EndReason;

use tor_cell::relaycell::{RelayCellFormat, RelayCmd};

use futures::io::{AsyncRead, AsyncWrite};

use futures::stream::StreamExt;

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

use futures::{Future, Stream};

use pin_project::pin_project;

use postage::watch;

#[cfg(feature = "tokio")]

use tokio_crate::io::ReadBuf;

#[cfg(feature = "tokio")]

use tokio_crate::io::{AsyncRead as TokioAsyncRead, AsyncWrite as TokioAsyncWrite};

#[cfg(feature = "tokio")]

use tokio_util::compat::{FuturesAsyncReadCompatExt, FuturesAsyncWriteCompatExt};

use tor_cell::restricted_msg;

use std::fmt::Debug;

use std::io::Result as IoResult;

use std::num::NonZero;

use std::pin::Pin;

#[cfg(any(feature = "stream-ctrl", feature = "experimental-api"))]

use std::sync::Arc;

#[cfg(feature = "stream-ctrl")]

use std::sync::{Mutex, Weak};

use educe::Educe;

use crate::client::ClientTunnel;

use crate::client::stream::StreamReceiver;

use crate::memquota::StreamAccount;

use crate::stream::StreamTarget;

use crate::stream::cmdcheck::{AnyCmdChecker, CmdChecker, StreamStatus};

use crate::stream::flow_ctrl::state::StreamRateLimit;

use crate::stream::flow_ctrl::xon_xoff::reader::{BufferIsEmpty, XonXoffReader, XonXoffReaderCtrl};

use crate::util::token_bucket::dynamic_writer::DynamicRateLimitedWriter;

use crate::util::token_bucket::writer::{RateLimitedWriter, RateLimitedWriterConfig};

use tor_basic_utils::skip_fmt;

use tor_cell::relaycell::msg::Data;

use tor_error::internal;

use tor_rtcompat::{CoarseTimeProvider, DynTimeProvider, SleepProvider};

/// A stream of [`RateLimitedWriterConfig`] used to update a [`DynamicRateLimitedWriter`].

///

/// Unfortunately we need to store the result of a [`StreamExt::map`] and [`StreamExt::fuse`] in

/// [`DataWriter`], which leaves us with this ugly type.

/// We use a type alias to make `DataWriter` a little nicer.

type RateConfigStream = futures::stream::Map<

    futures::stream::Fuse<watch::Receiver<StreamRateLimit>>,

    fn(StreamRateLimit) -> RateLimitedWriterConfig,

>;

/// An anonymized stream over the Tor network.

///

/// For most purposes, you can think of this type as an anonymized

/// TCP stream: it can read and write data, and get closed when it's done.

///

/// [`DataStream`] implements [`futures::io::AsyncRead`] and

/// [`futures::io::AsyncWrite`], so you can use it anywhere that those

/// traits are expected.

///

/// # Examples

///

/// Connecting to an HTTP server and sending a request, using

/// [`AsyncWriteExt::write_all`](futures::io::AsyncWriteExt::write_all):

///

/// ```ignore

/// let mut stream = tor_client.connect(("icanhazip.com", 80), None).await?;

///

/// use futures::io::AsyncWriteExt;

///

/// stream

///     .write_all(b"GET / HTTP/1.1\r\nHost: icanhazip.com\r\nConnection: close\r\n\r\n")

///     .await?;

///

/// // Flushing the stream is important; see below!

/// stream.flush().await?;

/// ```

///

/// Reading the result, using [`AsyncReadExt::read_to_end`](futures::io::AsyncReadExt::read_to_end):

///

/// ```ignore

/// use futures::io::AsyncReadExt;

///

/// let mut buf = Vec::new();

/// stream.read_to_end(&mut buf).await?;

///

/// println!("{}", String::from_utf8_lossy(&buf));

/// ```

///

/// # Usage with Tokio

///

/// If the `tokio` crate feature is enabled, this type also implements

/// [`tokio::io::AsyncRead`](tokio_crate::io::AsyncRead) and

/// [`tokio::io::AsyncWrite`](tokio_crate::io::AsyncWrite) for easier integration

/// with code that expects those traits.

///

/// # Remember to call `flush`!

///

/// DataStream buffers data internally, in order to write as few cells

/// as possible onto the network.  In order to make sure that your

/// data has actually been sent, you need to make sure that

/// [`AsyncWrite::poll_flush`] runs to completion: probably via

/// [`AsyncWriteExt::flush`](futures::io::AsyncWriteExt::flush).

///

/// # Splitting the type

///

/// This type is internally composed of a [`DataReader`] and a [`DataWriter`]; the

/// `DataStream::split` method can be used to split it into those two parts, for more

/// convenient usage with e.g. stream combinators.

///

/// # How long does a stream live?

///

/// A `DataStream` will live until all references to it are dropped,

/// or until it is closed explicitly.

///

/// If you split the stream into a `DataReader` and a `DataWriter`, it

/// will survive until _both_ are dropped, or until it is closed

/// explicitly.

///

/// A stream can also close because of a network error,

/// or because the other side of the stream decided to close it.

///

// # Semver note

//

// Note that this type is re-exported as a part of the public API of

// the `arti-client` crate.  Any changes to its API here in

// `tor-proto` need to be reflected above.

#[derive(Debug)]

pub struct DataStream {

    /// Underlying writer for this stream

    w: DataWriter,

    /// Underlying reader for this stream

    r: DataReader,

    /// A control object that can be used to monitor and control this stream

    /// without needing to own it.

    #[cfg(feature = "stream-ctrl")]

    ctrl: Arc<ClientDataStreamCtrl>,

}

assert_impl_all! { DataStream: Send, Sync }

/// An object used to control and monitor a data stream.

///

/// # Notes

///

/// This is a separate type from [`DataStream`] because it's useful to have

/// multiple references to this object, whereas a [`DataReader`] and [`DataWriter`]

/// need to have a single owner for the `AsyncRead` and `AsyncWrite` APIs to

/// work correctly.

#[cfg(feature = "stream-ctrl")]

#[cfg_attr(

    feature = "rpc",

    derive(derive_deftly::Deftly),

    derive_deftly(tor_rpcbase::templates::Object)

)]

#[derive(Debug)]

pub struct ClientDataStreamCtrl {

    /// The circuit to which this stream is attached.

    ///

    /// Note that the stream's reader and writer halves each contain a `StreamTarget`,

    /// which in turn has a strong reference to the `ClientCirc`.  So as long as any

    /// one of those is alive, this reference will be present.

    ///

    /// We make this a Weak reference so that once the stream itself is closed,

    /// we can't leak circuits.

    tunnel: Weak<ClientTunnel>,

    /// Shared user-visible information about the state of this stream.

    ///

    /// TODO RPC: This will probably want to be a `postage::Watch` or something

    /// similar, if and when it stops moving around.

    #[cfg(feature = "stream-ctrl")]

    status: Arc<Mutex<DataStreamStatus>>,

    /// The memory quota account that should be used for this stream's data

    ///

    /// Exists to keep the account alive

    _memquota: StreamAccount,

}

/// The inner writer for [`DataWriter`].

///

/// This type is responsible for taking bytes and packaging them into cells.

/// Rate limiting is implemented in [`DataWriter`] to avoid making this type more complex.

#[derive(Debug)]

struct DataWriterInner {

    /// Internal state for this writer

    ///

    /// This is stored in an Option so that we can mutate it in the

    /// AsyncWrite functions.  It might be possible to do better here,

    /// and we should refactor if so.

    state: Option<DataWriterState>,

    /// The memory quota account that should be used for this stream's data

    ///

    /// Exists to keep the account alive

    // If we liked, we could make this conditional; see DataReaderInner.memquota

    _memquota: StreamAccount,

    /// A control object that can be used to monitor and control this stream

    /// without needing to own it.

    #[cfg(feature = "stream-ctrl")]

    ctrl: Arc<ClientDataStreamCtrl>,

}

/// The write half of a [`DataStream`], implementing [`futures::io::AsyncWrite`].

///

/// See the [`DataStream`] docs for more information. In particular, note

/// that this writer requires `poll_flush` to complete in order to guarantee that

/// all data has been written.

///

/// # Usage with Tokio

///

/// If the `tokio` crate feature is enabled, this type also implements

/// [`tokio::io::AsyncWrite`](tokio_crate::io::AsyncWrite) for easier integration

/// with code that expects that trait.

///

/// # Drop and close

///

/// Note that dropping a `DataWriter` has no special effect on its own:

/// if the `DataWriter` is dropped, the underlying stream will still remain open

/// until the `DataReader` is also dropped.

///

/// If you want the stream to close earlier, use [`close`](futures::io::AsyncWriteExt::close)

/// (or [`shutdown`](tokio_crate::io::AsyncWriteExt::shutdown) with `tokio`).

///

/// Remember that Tor does not support half-open streams:

/// If you `close` or `shutdown` a stream,

/// the other side will not see the stream as half-open,

/// and so will (probably) not finish sending you any in-progress data.

/// Do not use `close`/`shutdown` to communicate anything besides

/// "I am done using this stream."

///

// # Semver note

//

// Note that this type is re-exported as a part of the public API of

// the `arti-client` crate.  Any changes to its API here in

// `tor-proto` need to be reflected above.

#[derive(Debug)]

pub struct DataWriter {

    /// A wrapper around [`DataWriterInner`] that adds rate limiting.

    writer: DynamicRateLimitedWriter<DataWriterInner, RateConfigStream, DynTimeProvider>,

}

impl DataWriter {

    /// Create a new rate-limited [`DataWriter`] from a [`DataWriterInner`].

        /// Converts a `rate` into a `RateLimitedWriterConfig`.

        // get the current rate from the `watch::Receiver`, which we'll use as the initial rate

        // map the rate update stream to the type required by `DynamicRateLimitedWriter`

        // build the rate limiter

    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and

    /// interact with this stream without holding the stream itself.

    #[cfg(feature = "stream-ctrl")]

}

impl AsyncWrite for DataWriter {

}

#[cfg(feature = "tokio")]

impl TokioAsyncWrite for DataWriter {

}

/// The read half of a [`DataStream`], implementing [`futures::io::AsyncRead`].

///

/// See the [`DataStream`] docs for more information.

///

/// # Usage with Tokio

///

/// If the `tokio` crate feature is enabled, this type also implements

/// [`tokio::io::AsyncRead`](tokio_crate::io::AsyncRead) for easier integration

/// with code that expects that trait.

//

// # Semver note

//

// Note that this type is re-exported as a part of the public API of

// the `arti-client` crate.  Any changes to its API here in

// `tor-proto` need to be reflected above.

#[derive(Debug)]

pub struct DataReader {

    /// The [`DataReaderInner`] with a wrapper to support XON/XOFF flow control.

    reader: XonXoffReader<DataReaderInner>,

}

impl DataReader {

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

    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and

    /// interact with this stream without holding the stream itself.

    #[cfg(feature = "stream-ctrl")]

}

impl AsyncRead for DataReader {

}

#[cfg(feature = "tokio")]

impl TokioAsyncRead for DataReader {

}

/// The inner reader for [`DataReader`].

///

/// This type is responsible for taking stream messages and extracting the stream data from them.

/// Flow control logic is implemented in [`DataReader`] to avoid making this type more complex.

#[derive(Debug)]

pub(crate) struct DataReaderInner {

    /// Internal state for this reader.

    ///

    /// This is stored in an Option so that we can mutate it in

    /// poll_read().  It might be possible to do better here, and we

    /// should refactor if so.

    state: Option<DataReaderState>,

    /// The memory quota account that should be used for this stream's data

    ///

    /// Exists to keep the account alive

    // If we liked, we could make this conditional on not(cfg(feature = "stream-ctrl"))

    // since, ClientDataStreamCtrl contains a StreamAccount clone too.  But that seems fragile.

    _memquota: StreamAccount,

    /// A control object that can be used to monitor and control this stream

    /// without needing to own it.

    #[cfg(feature = "stream-ctrl")]

    ctrl: Arc<ClientDataStreamCtrl>,

}

impl BufferIsEmpty for DataReaderInner {

    /// The result will become stale,

    /// so is most accurate immediately after a [`poll_read`](AsyncRead::poll_read).

        {

                // check if the partial cell in `pending` is empty,

                // and if the message stream is empty

            }

            // closed, so any data should have been discarded

        }

}

/// Shared status flags for tracking the status of as `DataStream`.

///

/// We expect to refactor this a bit, so it's not exposed at all.

//

// TODO RPC: Possibly instead of manipulating the fields of DataStreamStatus

// from various points in this module, we should instead construct

// DataStreamStatus as needed from information available elsewhere.  In any

// case, we should really  eliminate as much duplicate state here as we can.

// (See discussions at !1198 for some challenges with this.)

#[cfg(feature = "stream-ctrl")]

#[derive(Clone, Debug, Default)]

struct DataStreamStatus {

    /// True if we've received a CONNECTED message.

    //

    // TODO: This is redundant with `connected` in DataReaderImpl.

    received_connected: bool,

    /// True if we have decided to send an END message.

    //

    // TODO RPC: There is not an easy way to set this from this module!  Really,

    // the decision to send an "end" is made when the StreamTarget object is

    // dropped, but we don't currently have any way to see when that happens.

    // Perhaps we need a different shared StreamStatus object that the

    // StreamTarget holds?

    sent_end: bool,

    /// True if we have received an END message telling us to close the stream.

    received_end: bool,

    /// True if we have received an error.

    ///

    /// (This is not a subset or superset of received_end; some errors are END

    /// messages but some aren't; some END messages are errors but some aren't.)

    received_err: bool,

}

#[cfg(feature = "stream-ctrl")]

impl DataStreamStatus {

    /// Remember that we've received a connected message.

    /// Remember that we've received an error of some kind.

        // TODO: Probably we should remember the actual error in a box or

        // something.  But that means making a redundant copy of the error

        // even if nobody will want it.  Do we care?

        }

}

restricted_msg! {

    /// An allowable incoming message on a client data stream.

    enum ClientDataStreamMsg:RelayMsg {

        // SENDME is handled by the reactor.

        Data, End, Connected,

    }

}

// TODO RPC: Should we also implement this trait for everything that holds a

// ClientDataStreamCtrl?

#[cfg(feature = "stream-ctrl")]

impl super::ctrl::ClientStreamCtrl for ClientDataStreamCtrl {

}

#[cfg(feature = "stream-ctrl")]

impl ClientDataStreamCtrl {

    /// Return true if the underlying stream is connected. (That is, if it has

    /// received a `CONNECTED` message, and has not been closed.)

    // TODO RPC: Add more functions once we have the desired API more nailed

    // down.

}

impl DataStream {

    /// Wrap raw stream receiver and target parts as a DataStream.

    ///

    /// For non-optimistic stream, function `wait_for_connection`

    /// must be called after to make sure CONNECTED is received.

            false,

        )

    /// Wrap raw stream receiver and target parts as a connected DataStream.

    ///

    /// Unlike [`DataStream::new`], this creates a `DataStream` that does not expect to receive a

    /// CONNECTED cell.

    ///

    /// This is used by hidden services, exit relays, and directory servers to accept streams.

    #[cfg(any(feature = "hs-service", feature = "relay"))]

            true,

        )

    /// The shared implementation of the `new*()` functions.

        #[cfg(feature = "stream-ctrl")]

        };

        #[cfg(feature = "stream-ctrl")]

                #[cfg(feature = "relay")]

            };

        };

    /// Divide this DataStream into its constituent parts.

    /// Wait until a CONNECTED cell is received, or some other cell

    /// is received to indicate an error.

    ///

    /// Does nothing if this stream is already connected.

        // We must put state back before returning

            } else {

                // This succeeds if the cell is CONNECTED, and fails otherwise.

            };

            });

        } else {

        }

    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and

    /// interact with this stream without holding the stream itself.

    #[cfg(feature = "stream-ctrl")]

}

impl AsyncRead for DataStream {

}

#[cfg(feature = "tokio")]

impl TokioAsyncRead for DataStream {

}

impl AsyncWrite for DataStream {

}

#[cfg(feature = "tokio")]

impl TokioAsyncWrite for DataStream {

}

/// Helper type: Like BoxFuture, but also requires that the future be Sync.

type BoxSyncFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + Sync + 'a>>;

/// An enumeration for the state of a DataWriter.

///

/// We have to use an enum here because, for as long as we're waiting

/// for a flush operation to complete, the future returned by

/// `flush_cell()` owns the DataWriterImpl.

#[derive(Educe)]

#[educe(Debug)]

enum DataWriterState {

    /// The writer has closed or gotten an error: nothing more to do.

    Closed,

    /// The writer is not currently flushing; more data can get queued

    /// immediately.

    Ready(DataWriterImpl),

    /// The writer is flushing a cell.

    Flushing(

        #[educe(Debug(method = "skip_fmt"))] //

        BoxSyncFuture<'static, (DataWriterImpl, Result<()>)>,

    ),

}

/// Internal: the write part of a DataStream

#[derive(Educe)]

#[educe(Debug)]

struct DataWriterImpl {

    /// The underlying StreamTarget object.

    s: StreamTarget,

    /// Buffered data to send over the connection.

    ///

    /// This buffer is currently allocated using a number of bytes

    /// equal to the maximum that we can package at a time.

    //

    // TODO: this buffer is probably smaller than we want, but it's good

    // enough for now.  If we _do_ make it bigger, we'll have to change

    // our use of Data::split_from to handle the case where we can't fit

    // all the data.

    #[educe(Debug(method = "skip_fmt"))]

    buf: Box<[u8]>,

    /// Number of unflushed bytes in buf.

    n_pending: usize,

    /// Relay cell format in use

    relay_cell_format: RelayCellFormat,

    /// Shared user-visible information about the state of this stream.

    #[cfg(feature = "stream-ctrl")]

    status: Arc<Mutex<DataStreamStatus>>,

}

impl DataWriterInner {

    /// See [`DataWriter::client_stream_ctrl`].

    #[cfg(feature = "stream-ctrl")]

    /// Helper for poll_flush() and poll_close(): Performs a flush, then

    /// closes the stream if should_close is true.

        // TODO: this whole function is a bit copy-pasted.

                    // Nothing to flush!

                        // We need to actually continue with this function to do the closing.

                        // Thus, make a future that does nothing and is ready immediately.

                    } else {

                        // There's nothing more to do; we can return.

                    }

                } else {

                    // We need to flush the buffer's contents; Make a future for that.

                }

            }

            DataWriterState::Closed => {

            }

        };

            }

            }

            Poll::Pending => {

            }

        }

}

impl AsyncWrite for DataWriterInner {

                // we couldn't queue anything, so the current cell must be full.

            }

            DataWriterState::Closed => {

            }

        };

                #[cfg(feature = "stream-ctrl")]

            }

                // Great!  We're done flushing.  Queue as much as we can of this

                // cell.

            }

            Poll::Pending => {

            }

        }

}

#[cfg(feature = "tokio")]

impl TokioAsyncWrite for DataWriterInner {

}

impl DataWriterImpl {

    /// Try to flush the current buffer contents as a data cell.

        {

            // TODO: Eventually we may want a larger buffer; if we do,

            // this invariant will become false.

        } else {

        };

    /// Add as many bytes as possible from `b` to our internal buffer;

    /// return the number we were able to add.

            // that is, len == 0

}

impl DataReaderInner {

    /// Return a [`ClientDataStreamCtrl`] object that can be used to monitor and

    /// interact with this stream without holding the stream itself.

    #[cfg(feature = "stream-ctrl")]

}

/// An enumeration for the state of a [`DataReaderInner`].

// TODO: We don't need to implement the state in this way anymore now that we've removed the saved

// future. There are a few ways we could simplify this. See:

// https://gitlab.torproject.org/tpo/core/arti/-/merge_requests/3076#note_3218210

#[derive(Educe)]

#[educe(Debug)]

// We allow this since it's expected that streams will spend most of their time in the `Open` state,

// and will be cleaned up shortly after closing.

#[allow(clippy::large_enum_variant)]

enum DataReaderState {

    /// In this state we have received an end cell or an error.

    Closed,

    /// In this state the reader is open.

    Open(DataReaderImpl),

}

/// Wrapper for the read part of a [`DataStream`].

#[derive(Educe)]

#[educe(Debug)]

#[pin_project]

struct DataReaderImpl {

    /// The underlying StreamReceiver object.

    #[educe(Debug(method = "skip_fmt"))]

    #[pin]

    s: StreamReceiver,

    /// If present, data that we received on this stream but have not

    /// been able to send to the caller yet.

    // TODO: This data structure is probably not what we want, but

    // it's good enough for now.

    #[educe(Debug(method = "skip_fmt"))]

    pending: Vec<u8>,

    /// Index into pending to show what we've already read.

    offset: usize,

    /// If true, we have received a CONNECTED cell on this stream.

    connected: bool,

    /// Shared user-visible information about the state of this stream.

    #[cfg(feature = "stream-ctrl")]

    status: Arc<Mutex<DataStreamStatus>>,

}

impl AsyncRead for DataReaderInner {

        // We're pulling the state object out of the reader.  We MUST

        // put it back before this function returns.

        loop {

                    // There may be data to read already.

                        // We read data into the buffer, or the buffer was 0-len to begin with.

                        // Tell the caller.

                    // No data available!  We have to try reading.

                }

                DataReaderState::Closed => {

                    // TODO: Why are we returning an error rather than continuing to return EOF?

                }

            };

            // See if a cell is ready.

                    // There aren't any survivable errors in the current

                    // design.

                    #[cfg(feature = "stream-ctrl")]

                    } else {

                    };

                }

                Poll::Pending => {

                    // No cells ready, so tell the

                    // caller to get back to us later.

                }

            }

        }

}

#[cfg(feature = "tokio")]

impl TokioAsyncRead for DataReaderInner {

}

impl DataReaderImpl {

    /// Pull as many bytes as we can off of self.pending, and return that

    /// number of bytes.

    /// Return true iff there are no buffered bytes here to yield

    /// Load self.pending with the contents of a new data cell.

        use ClientDataStreamMsg::*;

                }

            },

            // TODO: This doesn't seem right to me, but seems to be the behaviour of the code before

            // the refactoring, so I've kept the same behaviour. I think if the cell stream is

            // terminated, we should be returning `None` here and not considering it as an error.

            // The `StreamReceiver` will have already returned an error if the cell stream was

            // terminated without an END message.

        };

                #[cfg(feature = "stream-ctrl")]

            }

            Connected(_) => {

            }

            }

            Data(_) => {

            }

        };

    /// Add the data from `d` to the end of our pending bytes.

}

/// A `CmdChecker` that enforces invariants for outbound data streams.

#[derive(Debug)]

pub(crate) struct OutboundDataCmdChecker {

    /// True if we are expecting to receive a CONNECTED message on this stream.

    expecting_connected: bool,

}

impl Default for OutboundDataCmdChecker {

}

impl CmdChecker for OutboundDataCmdChecker {

        use StreamStatus::*;

            RelayCmd::CONNECTED => {

                } else {

                }

            }

            RelayCmd::DATA => {

                } else {

                }

            }

        }

}

impl OutboundDataCmdChecker {

    /// Return a new boxed `DataCmdChecker` in a state suitable for a newly

    /// constructed connection.

}