//! Types and code for mapping StreamIDs to streams on a circuit.

mod halfstream;

use crate::congestion::sendme;

use crate::stream::RECV_WINDOW_INIT;

use crate::stream::StreamMpscReceiver;

use crate::stream::cmdcheck::AnyCmdChecker;

use crate::stream::flow_ctrl::state::{FlowCtrlHooks, StreamFlowCtrl};

use crate::stream::queue::StreamQueueSender;

use crate::util::stream_poll_set::{KeyAlreadyInsertedError, StreamPollSet};

use crate::{Error, Result};

use pin_project::pin_project;

use tor_async_utils::peekable_stream::{PeekableStream, UnobtrusivePeekableStream};

use tor_async_utils::stream_peek::StreamUnobtrusivePeeker;

use tor_cell::relaycell::flow_ctrl::{Xoff, Xon, XonKbpsEwma};

use tor_cell::relaycell::{RelayMsg, UnparsedRelayMsg};

use tor_cell::relaycell::{StreamId, msg::AnyRelayMsg};

use std::collections::HashMap;

use std::collections::hash_map;

use std::num::NonZeroU16;

use std::pin::Pin;

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

use std::time::Instant;

use tor_error::{bad_api_usage, internal};

use rand::Rng;

use tracing::debug;

use halfstream::HalfStream;

/// Entry for an open stream

///

/// (For the purposes of this module, an open stream is one where we have not

/// sent or received any message indicating that the stream is ended.)

#[derive(Debug)]

#[pin_project]

pub(super) struct OpenStreamEnt {

    /// Sink to send relay cells tagged for this stream into.

    pub(super) sink: StreamQueueSender,

    /// Number of cells dropped due to the stream disappearing before we can

    /// transform this into an `EndSent`.

    pub(super) dropped: u16,

    /// A `CmdChecker` used to tell whether cells on this stream are valid.

    pub(super) cmd_checker: AnyCmdChecker,

    /// Flow control for this stream.

    // Non-pub because we need to proxy `put_for_incoming_sendme` to ensure

    // `flow_ctrl_waker` is woken.

    flow_ctrl: StreamFlowCtrl,

    /// Stream for cells that should be sent down this stream.

    // Not directly exposed. This should only be polled via

    // `OpenStreamEntStream`s implementation of `Stream`, which in turn should

    // only be used through `StreamPollSet`.

    #[pin]

    rx: StreamUnobtrusivePeeker<StreamMpscReceiver<AnyRelayMsg>>,

    /// Waker to be woken when more sending capacity becomes available (e.g.

    /// receiving a SENDME).

    flow_ctrl_waker: Option<Waker>,

}

impl OpenStreamEnt {

    /// Whether this stream is ready to send `msg`.

    /// Handle an incoming sendme.

    ///

    /// On failure, return an error: the caller should close the stream or

    /// circuit with a protocol error.

        // Wake the stream if it was blocked on flow control.

    /// The approximate number of stream inbound data bytes buffered.

        // NOTE: Here we want to know the total number of buffered incoming stream data bytes. We

        // have access to the `StreamQueueSender` and can get how many bytes are buffered in that

        // queue.

        // But this isn't always the total number of buffered bytes since some bytes might be

        // buffered outside of this queue.

        // For example `DataReaderImpl` stores some stream bytes in its `pending` buffer, and we

        // have no way to access that from here in the reactor. So it's impossible to know the total

        // number of incoming stream data bytes that are buffered.

        //

        // This isn't really an issue in practice since *most* of the bytes will be queued in the

        // `StreamQueueSender`, the XOFF threshold is very large, and we don't need to be exact.

    /// Check if we should send an XON message.

    ///

    /// If we should, then returns the XON message that should be sent.

    /// Returns an error if XON/XOFF messages aren't supported for this type of flow control.

    /// Check if we should send an XOFF message.

    ///

    /// If we should, then returns the XOFF message that should be sent.

    /// Returns an error if XON/XOFF messages aren't supported for this type of flow control.

    /// Handle an incoming XON message.

    ///

    /// On failure, return an error: the caller should close the stream or

    /// circuit with a protocol error.

    /// Handle an incoming XOFF message.

    ///

    /// On failure, return an error: the caller should close the stream or

    /// circuit with a protocol error.

    /// Inform the flow control code that we're about to send `msg`.

    /// Should be called at the point we've fully committed to sending the message.

    /// Returns an error if we can't send `msg` and should close the circuit.

    //

    // TODO: Consider not exposing this, and instead calling in

    // `StreamMap::take_ready_msg`.

}

/// Private wrapper over `OpenStreamEnt`. We implement `futures::Stream` for

/// this wrapper, and not directly for `OpenStreamEnt`, so that client code

/// can't directly access the stream.

#[derive(Debug)]

#[pin_project]

struct OpenStreamEntStream {

    /// Inner value.

    #[pin]

    inner: OpenStreamEnt,

}

impl futures::Stream for OpenStreamEntStream {

    type Item = AnyRelayMsg;

        // TODO: consider calling `inner.flow_ctrl.about_to_send` here;

        // particularly if we change it to return a wrapper type that proves

        // we've taken the capacity. Otherwise it'd make it tricky in the reactor

        // to be sure we've correctly taken the capacity, since messages can originate

        // in other parts of the code (currently none of those should be of types that

        // count towards flow control, but that may change).

}

impl PeekableStream for OpenStreamEntStream {

        };

}

impl UnobtrusivePeekableStream for OpenStreamEntStream {

        } else {

        }

}

/// Entry for a stream where we have sent an END, or other message

/// indicating that the stream is terminated.

#[derive(Debug)]

pub(super) struct EndSentStreamEnt {

    /// A "half-stream" that we use to check the validity of incoming

    /// messages on this stream.

    pub(super) half_stream: HalfStream,

    /// True if the sender on this stream has been explicitly dropped;

    /// false if we got an explicit close from `close_pending`

    explicitly_dropped: bool,

    /// When this entry should be removed from the stream map.

    ///

    /// This is the amount of time we are willing to wait for

    /// an END ack before removing the half-stream from the map.

    pub(super) expiry: Instant,

}

/// The entry for a stream.

#[derive(Debug)]

enum ClosedStreamEnt {

    /// A stream for which we have received an END cell, but not yet

    /// had the stream object get dropped.

    EndReceived,

    /// A stream for which we have sent an END cell but not yet received an END

    /// cell.

    ///

    /// TODO(arti#264) Can we ever throw this out? Do we really get END cells for

    /// these?

    EndSent(EndSentStreamEnt),

}

/// Mutable reference to a stream entry.

pub(super) enum StreamEntMut<'a> {

    /// An open stream.

    Open(&'a mut OpenStreamEnt),

    /// A stream for which we have received an END cell, but not yet

    /// had the stream object get dropped.

    EndReceived,

    /// A stream for which we have sent an END cell but not yet received an END

    /// cell.

    EndSent(&'a mut EndSentStreamEnt),

}

impl<'a> From<&'a mut ClosedStreamEnt> for StreamEntMut<'a> {

        }

}

impl<'a> From<&'a mut OpenStreamEntStream> for StreamEntMut<'a> {

}

/// Return value to indicate whether or not we send an END cell upon

/// terminating a given stream.

#[derive(Debug, Copy, Clone, Eq, PartialEq)]

pub(super) enum ShouldSendEnd {

    /// An END cell should be sent.

    Send,

    /// An END cell should not be sent.

    DontSend,

}

/// A priority for use with [`StreamPollSet`].

#[derive(Debug, Copy, Clone, Eq, PartialEq, PartialOrd, Ord)]

struct Priority(u64);

/// A map from stream IDs to stream entries. Each circuit has one for each

/// hop.

pub(crate) struct StreamMap {

    /// Open streams.

    // Invariants:

    // * Keys are disjoint with `closed_streams`.

    open_streams: StreamPollSet<StreamId, Priority, OpenStreamEntStream>,

    /// Closed streams.

    // Invariants:

    // * Keys are disjoint with `open_streams`.

    closed_streams: HashMap<StreamId, ClosedStreamEnt>,

    /// The next StreamId that we should use for a newly allocated

    /// circuit.

    next_stream_id: StreamId,

    /// Next priority to use in `open_streams`. We implement round-robin scheduling of

    /// handling outgoing messages from streams by assigning a stream the next

    /// priority whenever an outgoing message is processed from that stream,

    /// putting it last in line.

    next_priority: Priority,

}

impl StreamMap {

    /// Make a new empty StreamMap.

    /// Return the number of open streams in this map.

    /// Return a [`TunnelActivity`](crate::util::tunnel_activity::TunnelActivity) for this hop.

    /// Return the next available priority.

    /// Add an entry to this map; return the newly allocated StreamId.

        // This "65536" seems too aggressive, but it's what tor does.

        //

        // Also, going around in a loop here is (sadly) needed in order

        // to look like Tor clients.

                Err(KeyAlreadyInsertedError {

                    key: _,

                    priority: _,

            };

        }

    /// Add an entry to this map using the specified StreamId.

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

    /// Return the entry for `id` in this map, if any.

    /// Note that we received an END message (or other message indicating the end of

    /// the stream) on the stream with `id`.

    ///

    /// Returns true if there was really a stream there.

        };

        // Progress the stream's state machine accordingly

            ClosedStreamEnt::EndSent { .. } => {

                // We got an END, and we already sent an END. Great!

                // we can forget about this stream.

            }

        }

    /// Handle a termination of the stream with `id` from this side of

    /// the circuit. Return true if the stream was open and an END

    /// ought to be sent.

        use TerminateReason as TR;

            let OpenStreamEntStream {

                inner:

                    OpenStreamEnt {

                        // notably absent: the channels for sink and stream, which will get dropped and

                        // closed (meaning reads/writes from/to this stream will now fail)

                        ..

                    },

            // FIXME(eta): we don't copy the receive window, instead just creating a new one,

            //             so a malicious peer can send us slightly more data than they should

            //             be able to; see arti#230.

            // TODO: would be nice to avoid new_ref.

            );

        // Progress the stream's state machine accordingly

        {

            ClosedStreamEnt::EndSent(EndSentStreamEnt {

                ..

                (false, TR::StreamTargetClosed) => {

                }

                (true, TR::StreamTargetClosed) => {

                }

            },

        }

    /// Get an up-to-date iterator of streams with ready items. `Option<AnyRelayMsg>::None`

    /// indicates that the local sender has been dropped.

    ///

    /// Conceptually all streams are in a queue; new streams are added to the

    /// back of the queue, and a stream is sent to the back of the queue

    /// whenever a ready message is taken from it (via

    /// [`Self::take_ready_msg`]). The returned iterator is an ordered view of

    /// this queue, showing the subset of streams that have a message ready to

    /// send, or whose sender has been dropped.

    /// If the stream `sid` has a message ready, take it, and reprioritize `sid`

    /// to the "back of the line" with respect to

    /// [`Self::poll_ready_streams_iter`].

    /// Remove all halfstreams that are expired at `now`.

}

/// A reason for terminating a stream.

///

/// We use this type in order to ensure that we obey the API restrictions of [`StreamMap::terminate`]

#[derive(Copy, Clone, Debug, PartialEq, Eq)]

pub(super) enum TerminateReason {

    /// Closing a stream because the receiver got `Ok(None)`, indicating that the

    /// corresponding senders were all dropped.

    StreamTargetClosed,

    /// Closing a stream because we were explicitly told to end it via

    /// [`StreamTarget::close_pending`](crate::stream::StreamTarget::close_pending).

    ExplicitEnd,

}

/// Convenience function for doing a wrapping increment of a `StreamId`.

#[cfg(test)]

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)]

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

    use super::*;

    use crate::client::circuit::test::fake_mpsc;

    use crate::stream::queue::fake_stream_queue;

    use crate::{client::stream::OutboundDataCmdChecker, congestion::sendme::StreamSendWindow};

    #[test]

    fn test_wrapping_next_stream_id() {

        let one = StreamId::new(1).unwrap();

        let two = StreamId::new(2).unwrap();

        let max = StreamId::new(0xffff).unwrap();

        assert_eq!(wrapping_next_stream_id(one), two);

        assert_eq!(wrapping_next_stream_id(max), one);

    }

    #[test]

    #[allow(clippy::cognitive_complexity)]

    fn streammap_basics() -> Result<()> {

        let mut map = StreamMap::new();

        let mut next_id = map.next_stream_id;

        let mut ids = Vec::new();

        assert_eq!(map.n_open_streams(), 0);

        // Try add_ent

        for n in 1..=128 {

            let (sink, _) = fake_stream_queue(

                #[cfg(not(feature = "flowctl-cc"))]

                128,

            );

            let (_, rx) = fake_mpsc(2);

            let id = map.add_ent(

                sink,

                rx,

                StreamFlowCtrl::new_window(StreamSendWindow::new(500)),

                OutboundDataCmdChecker::new_any(),

            )?;

            let expect_id: StreamId = next_id;

            assert_eq!(expect_id, id);

            next_id = wrapping_next_stream_id(next_id);

            ids.push(id);

            assert_eq!(map.n_open_streams(), n);

        }

        // Test get_mut.

        let nonesuch_id = next_id;

        assert!(matches!(

            map.get_mut(ids[0]),

            Some(StreamEntMut::Open { .. })

        ));

        assert!(map.get_mut(nonesuch_id).is_none());

        // Test end_received

        assert!(map.ending_msg_received(nonesuch_id).is_err());

        assert_eq!(map.n_open_streams(), 128);

        assert!(map.ending_msg_received(ids[1]).is_ok());

        assert_eq!(map.n_open_streams(), 127);

        assert!(matches!(

            map.get_mut(ids[1]),

            Some(StreamEntMut::EndReceived)

        ));

        assert!(map.ending_msg_received(ids[1]).is_err());

        // Test terminate

        use TerminateReason as TR;

        let expiry = Instant::now(); // dummy value, unused outside of the reactor

        assert!(map.terminate(nonesuch_id, TR::ExplicitEnd, expiry).is_err());

        assert_eq!(map.n_open_streams(), 127);

        assert_eq!(

            map.terminate(ids[2], TR::ExplicitEnd, expiry).unwrap(),

            ShouldSendEnd::Send

        );

        assert_eq!(map.n_open_streams(), 126);

        assert!(matches!(

            map.get_mut(ids[2]),

            Some(StreamEntMut::EndSent { .. })

        ));

        assert_eq!(

            map.terminate(ids[1], TR::ExplicitEnd, expiry).unwrap(),

            ShouldSendEnd::DontSend

        );

        // This stream was already closed when we called `ending_msg_received`

        // above.

        assert_eq!(map.n_open_streams(), 126);

        assert!(map.get_mut(ids[1]).is_none());

        // Try receiving an end after a terminate.

        assert!(map.ending_msg_received(ids[2]).is_ok());

        assert!(map.get_mut(ids[2]).is_none());

        assert_eq!(map.n_open_streams(), 126);

        Ok(())

    }

}