// Copyright (c) 2018-2019 Parity Technologies (UK) Ltd.
//
// Licensed under the Apache License, Version 2.0 or MIT license, at your option.
//
// A copy of the Apache License, Version 2.0 is included in the software as
// LICENSE-APACHE and a copy of the MIT license is included in the software
// as LICENSE-MIT. You may also obtain a copy of the Apache License, Version 2.0
// at https://www.apache.org/licenses/LICENSE-2.0 and a copy of the MIT license
// at https://opensource.org/licenses/MIT.

// This module contains the `Connection` type and associated helpers.
// A `Connection` wraps an underlying (async) I/O resource and multiplexes
// `Stream`s over it.
//
// The overall idea is as follows: The `Connection` makes progress via calls
// to its `next_stream` method which polls several futures, one that decodes
// `Frame`s from the I/O resource, one that consumes `ControlCommand`s
// from an MPSC channel and another one that consumes `StreamCommand`s from
// yet another MPSC channel. The latter channel is shared with every `Stream`
// created and whenever a `Stream` wishes to send a `Frame` to the remote end,
// it enqueues it into this channel (waiting if the channel is full). The
// former is shared with every `Control` clone and used to open new outbound
// streams or to trigger a connection close.
//
// The `Connection` updates the `Stream` state based on incoming frames, e.g.
// it pushes incoming data to the `Stream`'s buffer or increases the sending
// credit if the remote has sent us a corresponding `Frame::<WindowUpdate>`.
// Updating a `Stream`'s state acquires a `Mutex`, which every `Stream` has
// around its `Shared` state. While blocking, we make sure the lock is only
// held for brief moments and *never* while doing I/O. The only contention is
// between the `Connection` and a single `Stream`, which should resolve
// quickly. Ideally, we could use `futures::lock::Mutex` but it does not offer
// a poll-based API as of futures-preview 0.3.0-alpha.19, which makes it
// difficult to use in a `Stream`'s `AsyncRead` and `AsyncWrite` trait
// implementations.
//
// Closing a `Connection`
// ----------------------
//
// Every `Control` may send a `ControlCommand::Close` at any time and then
// waits on a `oneshot::Receiver` for confirmation that the connection is
// closed. The closing proceeds as follows:
//
// 1. As soon as we receive the close command we close the MPSC receiver
//    of `StreamCommand`s. We want to process any stream commands which are
//    already enqueued at this point but no more.
// 2. We change the internal shutdown state to `Shutdown::InProgress` which
//    contains the `oneshot::Sender` of the `Control` which triggered the
//    closure and which we need to notify eventually.
// 3. Crucially -- while closing -- we no longer process further control
//    commands, because opening new streams should no longer be allowed
//    and further close commands would mean we need to save those
//    `oneshot::Sender`s for later. On the other hand we also do not simply
//    close the control channel as this would signal to `Control`s that
//    try to send close commands, that the connection is already closed,
//    which it is not. So we just pause processing control commands which
//    means such `Control`s will wait.
// 4. We keep processing I/O and stream commands until the remaining stream
//    commands have all been consumed, at which point we transition the
//    shutdown state to `Shutdown::Complete`, which entails sending the
//    final termination frame to the remote, informing the `Control` and
//    now also closing the control channel.
// 5. Now that we are closed we go through all pending control commands
//    and tell the `Control`s that we are closed and we are finally done.
//
// While all of this may look complicated, it ensures that `Control`s are
// only informed about a closed connection when it really is closed.
//
// Potential improvements
// ----------------------
//
// There is always more work that can be done to make this a better crate,
// for example:
//
// - Instead of `futures::mpsc` a more efficient channel implementation
//   could be used, e.g. `tokio-sync`. Unfortunately `tokio-sync` is about
//   to be merged into `tokio` and depending on this large crate is not
//   attractive, especially given the dire situation around cargo's flag
//   resolution.
// - Flushing could be optimised. This would also require adding a
//   `StreamCommand::Flush` so that `Stream`s can trigger a flush, which
//   they would have to when they run out of credit, or else a series of
//   send operations might never finish.
// - If Rust gets async destructors, the `garbage_collect()` method can be
//   removed. Instead a `Stream` would send a `StreamCommand::Dropped(..)`
//   or something similar and the removal logic could happen within regular
//   command processing instead of having to scan the whole collection of
//   `Stream`s on each loop iteration, which is not great.

mod control;
mod stream;

use crate::{
    Config,
    DEFAULT_CREDIT,
    WindowUpdateMode,
    error::ConnectionError,
    frame::{self, Frame},
    frame::header::{self, CONNECTION_ID, Data, GoAway, Header, Ping, StreamId, Tag, WindowUpdate},
    pause::Pausable
};
use futures::{
    channel::{mpsc, oneshot},
    future::{self, Either},
    prelude::*,
    stream::{Fuse, FusedStream}
};
use nohash_hasher::IntMap;
use std::{fmt, sync::Arc, task::{Context, Poll}};

pub use control::Control;
pub use stream::{Packet, State, Stream};

/// Arbitrary limit of our internal command channels.
///
/// Since each `mpsc::Sender` gets a guaranteed slot in a channel the
/// actual upper bound is this value + number of clones.
const MAX_COMMAND_BACKLOG: usize = 32;

type Result<T> = std::result::Result<T, ConnectionError>;

/// How the connection is used.
#[derive(Copy, Clone, Debug, Hash, PartialEq, Eq)]
pub enum Mode {
    /// Client to server connection.
    Client,
    /// Server to client connection.
    Server
}

/// The connection identifier.
///
/// Randomly generated, this is mainly intended to improve log output.
#[derive(Clone, Copy)]
pub(crate) struct Id(u32);

impl fmt::Debug for Id {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:08x}", self.0)
    }
}

impl fmt::Display for Id {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:08x}", self.0)
    }
}

/// A Yamux connection object.
///
/// Wraps the underlying I/O resource and makes progress via its
/// [`Connection::next_stream`] method which must be called repeatedly
/// until `Ok(None)` signals EOF or an error is encountered.
pub struct Connection<T> {
    id: Id,
    mode: Mode,
    config: Arc<Config>,
    socket: Fuse<frame::Io<T>>,
    next_id: u32,
    streams: IntMap<StreamId, Stream>,
    control_sender: mpsc::Sender<ControlCommand>,
    control_receiver: Pausable<mpsc::Receiver<ControlCommand>>,
    stream_sender: mpsc::Sender<StreamCommand>,
    stream_receiver: mpsc::Receiver<StreamCommand>,
    garbage: Vec<StreamId>, // see `Connection::garbage_collect()`
    shutdown: Shutdown
}

/// `Control` to `Connection` commands.
#[derive(Debug)]
pub(crate) enum ControlCommand {
    /// Open a new stream to the remote end.
    OpenStream(oneshot::Sender<Result<Stream>>),
    /// Close the whole connection.
    CloseConnection(oneshot::Sender<()>)
}

/// `Stream` to `Connection` commands.
#[derive(Debug)]
pub(crate) enum StreamCommand {
    /// A new frame should be sent to the remote.
    SendFrame(Frame<Either<Data, WindowUpdate>>),
    /// Close a stream.
    CloseStream { id: StreamId, ack: bool }
}

/// Possible actions as a result of incoming frame handling.
#[derive(Debug)]
enum Action {
    /// Nothing to be done.
    None,
    /// A new stream has been opened by the remote.
    New(Stream),
    /// A window update should be sent to the remote.
    Update(Frame<WindowUpdate>),
    /// A ping should be answered.
    Ping(Frame<Ping>),
    /// A stream should be reset.
    Reset(Frame<Data>),
    /// The connection should be terminated.
    Terminate(Frame<GoAway>)
}

/// This enum captures the various stages of shutting down the connection.
#[derive(Debug)]
enum Shutdown {
    /// We are open for business.
    NotStarted,
    /// We have received a `ControlCommand::Close` and are shutting
    /// down operations. The `Sender` will be informed once we are done.
    InProgress(oneshot::Sender<()>),
    /// The shutdown is complete and we are closed for good.
    Complete
}

impl Shutdown {
    fn has_not_started(&self) -> bool {
        if let Shutdown::NotStarted = self {
            true
        } else {
            false
        }
    }

    fn is_in_progress(&self) -> bool {
        if let Shutdown::InProgress(_) = self {
            true
        } else {
            false
        }
    }

    fn is_complete(&self) -> bool {
        if let Shutdown::Complete = self {
            true
        } else {
            false
        }
    }
}

impl<T> fmt::Debug for Connection<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        f.debug_struct("Connection")
            .field("id", &self.id)
            .field("mode", &self.mode)
            .field("streams", &self.streams.len())
            .field("next_id", &self.next_id)
            .finish()
    }
}

impl<T> fmt::Display for Connection<T> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "(Connection {} {:?} (streams {}))", self.id, self.mode, self.streams.len())
    }
}

impl<T: AsyncRead + AsyncWrite + Unpin> Connection<T> {
    /// Create a new `Connection` from the given I/O resource.
    pub fn new(socket: T, cfg: Config, mode: Mode) -> Self {
        let id = Id(rand::random());
        log::debug!("new connection: {} ({:?})", id, mode);
        let (stream_sender, stream_receiver) = mpsc::channel(MAX_COMMAND_BACKLOG);
        let (control_sender, control_receiver) = mpsc::channel(MAX_COMMAND_BACKLOG);
        let socket = frame::Io::new(socket, cfg.max_buffer_size).fuse();
        Connection {
            id,
            mode,
            config: Arc::new(cfg),
            socket,
            streams: IntMap::default(),
            control_sender,
            control_receiver: Pausable::new(control_receiver),
            stream_sender,
            stream_receiver,
            next_id: match mode {
                Mode::Client => 1,
                Mode::Server => 2
            },
            garbage: Vec::new(),
            shutdown: Shutdown::NotStarted
        }
    }

    /// Get a controller for this connection.
    pub fn control(&self) -> Control {
        Control::new(self.control_sender.clone())
    }

    /// Get the next incoming stream, opened by the remote.
    ///
    /// This must be called repeatedly in order to make progress.
    /// Once `Ok(None)` or `Err(_)` is returned the connection is
    /// considered closed and no further invocation of this method
    /// must be attempted.
    ///
    /// # Cancellation
    ///
    /// Please note that if you poll the returned [`Future`] it *must
    /// not be cancelled* but polled until [`Poll::Ready`] is returned.
    pub async fn next_stream(&mut self) -> Result<Option<Stream>> {
        let result = self.next().await;

        if let Ok(Some(_)) = result {
            return result
        }

        // At this point we are either at EOF or encountered an error.
        // We close all streams and wake up the associated tasks before
        // closing the socket. The connection is then considered closed.

        self.drop_all_streams();

        if let Err(ConnectionError::Closed) = result {
            return Ok(None)
        }

        result
    }

    /// Get the next inbound `Stream` and make progress along the way.
    ///
    /// This is called from `Connection::next_stream` instead of being a
    /// public method itself in order to guarantee proper closing in
    /// case of an error or at EOF.
    async fn next(&mut self) -> Result<Option<Stream>> {
        loop {
            self.garbage_collect().await?;

            // For each channel and the socket we create a future that gets
            // the next item. We will poll each future and if any one of them
            // yields an item, we return the tuple of poll results which are
            // then all processed.
            //
            // For terminated sources we create non-finishing futures.
            // This guarantees that if the remaining futures are pending
            // we properly wait until woken up because we actually can make
            // progress.
            //
            // While it should never happen that all futures are terminated
            // we nevertheless count them and return early if they are.

            let mut num_terminated = 0;

            let mut next_inbound_frame =
                if self.socket.is_terminated() {
                    num_terminated += 1;
                    Either::Left(future::pending())
                } else {
                    Either::Right(self.socket.try_next().err_into())
                };

            let mut next_stream_command =
                if self.stream_receiver.is_terminated() {
                    num_terminated += 1;
                    Either::Left(future::pending())
                } else {
                    Either::Right(self.stream_receiver.next())
                };

            let mut next_control_command =
                if self.control_receiver.is_terminated() {
                    num_terminated += 1;
                    Either::Left(future::pending())
                } else {
                    Either::Right(self.control_receiver.next())
                };

            if num_terminated == 3 {
                log::error!("{}: all futures are terminated", self.id);
                return Err(ConnectionError::Closed)
            }

            let next_item =
                future::poll_fn(move |cx: &mut Context| {
                    let a = next_stream_command.poll_unpin(cx);
                    let b = next_control_command.poll_unpin(cx);
                    let c = next_inbound_frame.poll_unpin(cx);
                    if a.is_pending() && b.is_pending() && c.is_pending() {
                        return Poll::Pending
                    }
                    Poll::Ready((a, b, c))
                });

            let (stream_command, control_command, inbound_frame) = next_item.await;

            if let Poll::Ready(cmd) = control_command {
                self.on_control_command(cmd).await?
            }

            if let Poll::Ready(cmd) = stream_command {
                self.on_stream_command(cmd).await?
            }

            if let Poll::Ready(frame) = inbound_frame {
                if let Some(stream) = self.on_frame(frame).await? {
                    self.socket.get_mut().flush().await.or(Err(ConnectionError::Closed))?;
                    return Ok(Some(stream))
                }
            }

            self.socket.get_mut().flush().await.or(Err(ConnectionError::Closed))?
        }
    }

    /// Process a command from a `Control`.
    ///
    /// We only process control commands if we are not in the process of closing
    /// the connection. Only once we finished closing will we drain the remaining
    /// commands and reply back that we are closed.
    async fn on_control_command(&mut self, cmd: Option<ControlCommand>) -> Result<()> {
        match cmd {
            Some(ControlCommand::OpenStream(reply)) => {
                if self.shutdown.is_complete() {
                    // We are already closed so just inform the control.
                    let _ = reply.send(Err(ConnectionError::Closed));
                    return Ok(())
                }
                if self.streams.len() >= self.config.max_num_streams {
                    log::error!("{}: maximum number of streams reached", self.id);
                    let _ = reply.send(Err(ConnectionError::TooManyStreams));
                    return Ok(())
                }
                log::trace!("{}: creating new outbound stream", self.id);
                let id = self.next_stream_id()?;
                if !self.config.lazy_open {
                    let mut frame = Frame::window_update(id, self.config.receive_window);
                    frame.header_mut().syn();
                    log::trace!("{}: sending initial {}", self.id, frame.header());
                    self.socket.get_mut().send(&frame).await.or(Err(ConnectionError::Closed))?
                }
                let stream = {
                    let config = self.config.clone();
                    let sender = self.stream_sender.clone();
                    let window = self.config.receive_window;
                    let mut stream = Stream::new(id, self.id, config, window, DEFAULT_CREDIT, sender);
                    if self.config.lazy_open {
                        stream.set_flag(stream::Flag::Syn)
                    }
                    stream
                };
                if reply.send(Ok(stream.clone())).is_ok() {
                    log::debug!("{}: new outbound {} of {}", self.id, stream, self);
                    self.streams.insert(id, stream);
                } else {
                    log::debug!("{}: open stream {} has been cancelled", self.id, id);
                    if !self.config.lazy_open {
                        let mut header = Header::data(id, 0);
                        header.rst();
                        let frame = Frame::new(header);
                        self.socket.get_mut().send(&frame).await.or(Err(ConnectionError::Closed))?
                    }
                }
            }
            Some(ControlCommand::CloseConnection(reply)) => {
                if self.shutdown.is_complete() {
                    // We are already closed so just inform the control.
                    let _ = reply.send(());
                    return Ok(())
                }
                // Handle initial close command.
                debug_assert!(self.shutdown.has_not_started());
                self.shutdown = Shutdown::InProgress(reply);
                log::trace!("{}: shutting down connection", self.id);
                self.control_receiver.pause();
                self.stream_receiver.close()
            }
            None => {
                // We only get here after the whole connection shutdown is complete.
                // No further processing of commands of any kind or incoming frames
                // will happen.
                debug_assert!(self.shutdown.is_complete());
                self.socket.get_mut().close().await.or(Err(ConnectionError::Closed))?;
                return Err(ConnectionError::Closed)
            }
        }
        Ok(())
    }

    /// Process a command from one of our `Stream`s.
    async fn on_stream_command(&mut self, cmd: Option<StreamCommand>) -> Result<()> {
        match cmd {
            Some(StreamCommand::SendFrame(frame)) => {
                log::trace!("{}: sending: {}", self.id, frame.header());
                self.socket.get_mut().send(&frame).await.or(Err(ConnectionError::Closed))?
            }
            Some(StreamCommand::CloseStream { id, ack }) => {
                log::trace!("{}: closing stream {} of {}", self.id, id, self);
                let mut header = Header::data(id, 0);
                header.fin();
                if ack { header.ack() }
                let frame = Frame::new(header);
                self.socket.get_mut().send(&frame).await.or(Err(ConnectionError::Closed))?
            }
            None => {
                // We only get to this point when `self.stream_receiver`
                // was closed which only happens in response to a close control
                // command. Now that we are at the end of the stream command queue,
                // we send the final term frame to the remote and complete the
                // closure.
                debug_assert!(self.shutdown.is_in_progress());
                log::debug!("{}: closing {}", self.id, self);
                let frame = Frame::term();
                self.socket.get_mut().send(&frame).await.or(Err(ConnectionError::Closed))?;
                let shutdown = std::mem::replace(&mut self.shutdown, Shutdown::Complete);
                if let Shutdown::InProgress(tx) = shutdown {
                    // Inform the `Control` that initiated the shutdown.
                    let _ = tx.send(());
                }
                debug_assert!(self.control_receiver.is_paused());
                self.control_receiver.unpause();
                self.control_receiver.stream().close()
            }
        }
        Ok(())
    }

    /// Process the result of reading from the socket.
    ///
    /// Unless `frame` is `Ok(Some(_))` we will assume the connection got closed
    /// and return a corresponding error, which terminates the connection.
    /// Otherwise we process the frame and potentially return a new `Stream`
    /// if one was opened by the remote.
    async fn on_frame(&mut self, frame: Result<Option<Frame<()>>>) -> Result<Option<Stream>> {
        match frame {
            Ok(Some(frame)) => {
                log::trace!("{}: received: {}", self.id, frame.header());
                let action = match frame.header().tag() {
                    Tag::Data => self.on_data(frame.into_data()),
                    Tag::WindowUpdate => self.on_window_update(&frame.into_window_update()),
                    Tag::Ping => self.on_ping(&frame.into_ping()),
                    Tag::GoAway => return Err(ConnectionError::Closed)
                };
                match action {
                    Action::None => {}
                    Action::New(stream) => {
                        log::trace!("{}: new inbound {} of {}", self.id, stream, self);
                        return Ok(Some(stream))
                    }
                    Action::Update(f) => {
                        log::trace!("{}/{}: sending update", self.id, f.header().stream_id());
                        self.socket.get_mut().send(&f).await.or(Err(ConnectionError::Closed))?
                    }
                    Action::Ping(f) => {
                        log::trace!("{}/{}: pong", self.id, f.header().stream_id());
                        self.socket.get_mut().send(&f).await.or(Err(ConnectionError::Closed))?
                    }
                    Action::Reset(f) => {
                        log::trace!("{}/{}: sending reset", self.id, f.header().stream_id());
                        self.socket.get_mut().send(&f).await.or(Err(ConnectionError::Closed))?
                    }
                    Action::Terminate(f) => {
                        log::trace!("{}: sending term", self.id);
                        self.socket.get_mut().send(&f).await.or(Err(ConnectionError::Closed))?
                    }
                }
                Ok(None)
            }
            Ok(None) => {
                log::debug!("{}: socket eof", self.id);
                Err(ConnectionError::Closed)
            }
            Err(e) if e.io_kind() == Some(std::io::ErrorKind::ConnectionReset) => {
                log::debug!("{}: connection reset", self.id);
                Err(ConnectionError::Closed)
            }
            Err(e) => {
                log::error!("{}: socket error: {}", self.id, e);
                Err(e)
            }
        }
    }

    fn on_data(&mut self, frame: Frame<Data>) -> Action {
        let stream_id = frame.header().stream_id();

        if frame.header().flags().contains(header::RST) { // stream reset
            if let Some(s) = self.streams.get_mut(&stream_id) {
                let mut shared = s.shared();
                shared.update_state(self.id, stream_id, State::Closed);
                if let Some(w) = shared.reader.take() {
                    w.wake()
                }
                if let Some(w) = shared.writer.take() {
                    w.wake()
                }
            }
            return Action::None
        }

        let is_finish = frame.header().flags().contains(header::FIN); // half-close

        if frame.header().flags().contains(header::SYN) { // new stream
            if !self.is_valid_remote_id(stream_id, Tag::Data) {
                log::error!("{}: invalid stream id {}", self.id, stream_id);
                return Action::Terminate(Frame::protocol_error())
            }
            if frame.body().len() > crate::u32_as_usize(DEFAULT_CREDIT) {
                log::error!("{}/{}: 1st body of stream exceeds default credit", self.id, stream_id);
                return Action::Terminate(Frame::protocol_error())
            }
            if self.streams.contains_key(&stream_id) {
                log::error!("{}/{}: stream already exists", self.id, stream_id);
                return Action::Terminate(Frame::protocol_error())
            }
            if self.streams.len() == self.config.max_num_streams {
                log::error!("{}: maximum number of streams reached", self.id);
                return Action::Terminate(Frame::internal_error())
            }
            let stream = {
                let config = self.config.clone();
                let credit = DEFAULT_CREDIT;
                let sender = self.stream_sender.clone();
                let mut stream = Stream::new(stream_id, self.id, config, credit, credit, sender);
                stream.set_flag(stream::Flag::Ack);
                stream
            };
            {
                let mut shared = stream.shared();
                if is_finish {
                    shared.update_state(self.id, stream_id, State::RecvClosed);
                }
                shared.window = shared.window.saturating_sub(frame.body_len());
                shared.buffer.push(frame.into_body());
            }
            self.streams.insert(stream_id, stream.clone());
            return Action::New(stream)
        }

        if let Some(stream) = self.streams.get_mut(&stream_id) {
            let mut shared = stream.shared();
            if frame.body().len() > crate::u32_as_usize(shared.window) {
                log::error!("{}/{}: frame body larger than window of stream", self.id, stream_id);
                return Action::Terminate(Frame::protocol_error())
            }
            if is_finish {
                shared.update_state(self.id, stream_id, State::RecvClosed);
            }
            let max_buffer_size = self.config.max_buffer_size;
            if shared.buffer.len().map(move |n| n >= max_buffer_size).unwrap_or(true) {
                log::error!("{}/{}: buffer of stream grows beyond limit", self.id, stream_id);
                let mut header = Header::data(stream_id, 0);
                header.rst();
                return Action::Reset(Frame::new(header))
            }
            shared.window = shared.window.saturating_sub(frame.body_len());
            shared.buffer.push(frame.into_body());
            if let Some(w) = shared.reader.take() {
                w.wake()
            }
            if !is_finish
                && shared.window == 0
                && self.config.window_update_mode == WindowUpdateMode::OnReceive
            {
                shared.window = self.config.receive_window;
                let frame = Frame::window_update(stream_id, self.config.receive_window);
                return Action::Update(frame)
            }
        } else if !is_finish {
            log::debug!("{}/{}: data for unknown stream", self.id, stream_id);
            let mut header = Header::data(stream_id, 0);
            header.rst();
            return Action::Reset(Frame::new(header))
        }

        Action::None
    }

    fn on_window_update(&mut self, frame: &Frame<WindowUpdate>) -> Action {
        let stream_id = frame.header().stream_id();

        if frame.header().flags().contains(header::RST) { // stream reset
            if let Some(s) = self.streams.get_mut(&stream_id) {
                let mut shared = s.shared();
                shared.update_state(self.id, stream_id, State::Closed);
                if let Some(w) = shared.reader.take() {
                    w.wake()
                }
                if let Some(w) = shared.writer.take() {
                    w.wake()
                }
            }
            return Action::None
        }

        let is_finish = frame.header().flags().contains(header::FIN); // half-close

        if frame.header().flags().contains(header::SYN) { // new stream
            if !self.is_valid_remote_id(stream_id, Tag::WindowUpdate) {
                log::error!("{}: invalid stream id {}", self.id, stream_id);
                return Action::Terminate(Frame::protocol_error())
            }
            if self.streams.contains_key(&stream_id) {
                log::error!("{}/{}: stream already exists", self.id, stream_id);
                return Action::Terminate(Frame::protocol_error())
            }
            if self.streams.len() == self.config.max_num_streams {
                log::error!("{}: maximum number of streams reached", self.id);
                return Action::Terminate(Frame::protocol_error())
            }
            let stream = {
                let credit = frame.header().credit();
                let config = self.config.clone();
                let sender = self.stream_sender.clone();
                let mut stream = Stream::new(stream_id, self.id, config, DEFAULT_CREDIT, credit, sender);
                stream.set_flag(stream::Flag::Ack);
                stream
            };
            if is_finish {
                stream.shared().update_state(self.id, stream_id, State::RecvClosed);
            }
            self.streams.insert(stream_id, stream.clone());
            return Action::New(stream)
        }

        if let Some(stream) = self.streams.get_mut(&stream_id) {
            let mut shared = stream.shared();
            shared.credit += frame.header().credit();
            if is_finish {
                shared.update_state(self.id, stream_id, State::RecvClosed);
            }
            if let Some(w) = shared.writer.take() {
                w.wake()
            }
        } else if !is_finish {
            log::debug!("{}/{}: window update for unknown stream", self.id, stream_id);
            let mut header = Header::data(stream_id, 0);
            header.rst();
            return Action::Reset(Frame::new(header))
        }

        Action::None
    }

    fn on_ping(&mut self, frame: &Frame<Ping>) -> Action {
        let stream_id = frame.header().stream_id();
        if frame.header().flags().contains(header::ACK) { // pong
            return Action::None
        }
        if stream_id == CONNECTION_ID || self.streams.contains_key(&stream_id) {
            let mut hdr = Header::ping(frame.header().nonce());
            hdr.ack();
            return Action::Ping(Frame::new(hdr))
        }
        log::debug!("{}/{}: ping for unknown stream", self.id, stream_id);
        let mut header = Header::data(stream_id, 0);
        header.rst();
        Action::Reset(Frame::new(header))
    }

    fn next_stream_id(&mut self) -> Result<StreamId> {
        let proposed = StreamId::new(self.next_id);
        self.next_id = self.next_id.checked_add(2).ok_or(ConnectionError::NoMoreStreamIds)?;
        match self.mode {
            Mode::Client => assert!(proposed.is_client()),
            Mode::Server => assert!(proposed.is_server())
        }
        Ok(proposed)
    }

    // Check if the given stream ID is valid w.r.t. the provided tag and our connection mode.
    fn is_valid_remote_id(&self, id: StreamId, tag: Tag) -> bool {
        if tag == Tag::Ping || tag == Tag::GoAway {
            return id.is_session()
        }
        match self.mode {
            Mode::Client => id.is_server(),
            Mode::Server => id.is_client()
        }
    }

    /// Remove stale streams and send necessary messages to the remote.
    ///
    /// If we ever get async destructors we can replace this with streams
    /// sending a proper command when dropped.
    async fn garbage_collect(&mut self) -> Result<()> {
        let conn_id = self.id;
        let win_update_mode = self.config.window_update_mode;
        for stream in self.streams.values_mut() {
            if stream.strong_count() > 1 {
                continue
            }
            log::trace!("{}: removing dropped {}", conn_id, stream);
            let stream_id = stream.id();
            let frame = {
                let mut shared = stream.shared();
                let frame = match shared.update_state(conn_id, stream_id, State::Closed) {
                    // The stream was dropped without calling `poll_close`.
                    // We reset the stream to inform the remote of the closure.
                    State::Open => {
                        let mut header = Header::data(stream_id, 0);
                        header.rst();
                        Some(Frame::new(header))
                    }
                    // The stream was dropped without calling `poll_close`.
                    // We have already received a FIN from remote and send one
                    // back which closes the stream for good.
                    State::RecvClosed => {
                        let mut header = Header::data(stream_id, 0);
                        header.fin();
                        Some(Frame::new(header))
                    }
                    // The stream was properly closed. We either already have
                    // or will at some later point send our FIN frame.
                    // The remote may be out of credit though and blocked on
                    // writing more data. We may need to reset the stream.
                    State::SendClosed =>
                        if win_update_mode == WindowUpdateMode::OnRead && !shared.buffer.is_empty() {
                            // The stream has unconsumed data left when closed.
                            // The remote may be waiting for a window update
                            // which we will never send, so reset the stream now.
                            let mut header = Header::data(stream_id, 0);
                            header.rst();
                            Some(Frame::new(header))
                        } else {
                            // The remote is not blocked as we send window updates
                            // for as long as we know the stream. For unknown streams
                            // we send a RST in `Connection::on_data`.
                            // For `OnRead` and empty stream buffers we have or will
                            // send another window update too.
                            None
                        }
                    // The stream was properly closed. We either already have
                    // or will at some later point send our FIN frame. The
                    // remote end has already done so in the past.
                    State::Closed => None
                };
                if let Some(w) = shared.reader.take() {
                    w.wake()
                }
                if let Some(w) = shared.writer.take() {
                    w.wake()
                }
                frame
            };
            if let Some(f) = frame {
                self.socket.get_mut().send(&f).await.or(Err(ConnectionError::Closed))?
            }
            self.garbage.push(stream_id)
        }
        for id in self.garbage.drain(..) {
            self.streams.remove(&id);
        }
        Ok(())
    }
}

impl<T> Connection<T> {
    /// Close and drop all `Stream`s and wake any pending `Waker`s.
    fn drop_all_streams(&mut self) {
        for (id, s) in self.streams.drain() {
            let mut shared = s.shared();
            shared.update_state(self.id, id, State::Closed);
            if let Some(w) = shared.reader.take() {
                w.wake()
            }
            if let Some(w) = shared.writer.take() {
                w.wake()
            }
        }
    }
}

impl<T> Drop for Connection<T> {
    fn drop(&mut self) {
        self.drop_all_streams()
    }
}

/// Turn a Yamux [`Connection`] into a [`futures::Stream`].
pub fn into_stream<T>(c: Connection<T>) -> impl futures::stream::Stream<Item = Result<Stream>>
where
    T: AsyncRead + AsyncWrite + Unpin
{
    futures::stream::unfold(c, |mut c| async {
        match c.next_stream().await {
            Ok(None) => None,
            Ok(Some(stream)) => Some((Ok(stream), c)),
            Err(e) => Some((Err(e), c))
        }
    })
}