#![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
use futures::{
future,
prelude::*,
ready,
stream::{BoxStream, LocalBoxStream},
};
use libp2p_core::muxing::{StreamMuxer, StreamMuxerEvent};
use libp2p_core::upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use std::collections::VecDeque;
use std::task::Waker;
use std::{
fmt, io, iter, mem,
pin::Pin,
task::{Context, Poll},
};
use thiserror::Error;
use yamux::ConnectionError;
#[deprecated(note = "Import the `yamux` module instead and refer to this type as `yamux::Muxer`.")]
pub type Yamux<S> = Muxer<S>;
pub struct Muxer<S> {
incoming: S,
control: yamux::Control,
inbound_stream_buffer: VecDeque<yamux::Stream>,
inbound_stream_waker: Option<Waker>,
}
const MAX_BUFFERED_INBOUND_STREAMS: usize = 25;
impl<S> fmt::Debug for Muxer<S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Yamux")
}
}
impl<C> Muxer<Incoming<C>>
where
C: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
fn new(io: C, cfg: yamux::Config, mode: yamux::Mode) -> Self {
let conn = yamux::Connection::new(io, cfg, mode);
let ctrl = conn.control();
Self {
incoming: Incoming {
stream: yamux::into_stream(conn).err_into().boxed(),
_marker: std::marker::PhantomData,
},
control: ctrl,
inbound_stream_buffer: VecDeque::default(),
inbound_stream_waker: None,
}
}
}
impl<C> Muxer<LocalIncoming<C>>
where
C: AsyncRead + AsyncWrite + Unpin + 'static,
{
fn local(io: C, cfg: yamux::Config, mode: yamux::Mode) -> Self {
let conn = yamux::Connection::new(io, cfg, mode);
let ctrl = conn.control();
Self {
incoming: LocalIncoming {
stream: yamux::into_stream(conn).err_into().boxed_local(),
_marker: std::marker::PhantomData,
},
control: ctrl,
inbound_stream_buffer: VecDeque::default(),
inbound_stream_waker: None,
}
}
}
#[deprecated(note = "Use `Result<T, yamux::Error>` instead.")]
pub type YamuxResult<T> = Result<T, Error>;
impl<S> StreamMuxer for Muxer<S>
where
S: Stream<Item = Result<yamux::Stream, Error>> + Unpin,
{
type Substream = yamux::Stream;
type Error = Error;
fn poll_inbound(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<Self::Substream, Self::Error>> {
if let Some(stream) = self.inbound_stream_buffer.pop_front() {
return Poll::Ready(Ok(stream));
}
self.inbound_stream_waker = Some(cx.waker().clone());
self.poll_inner(cx)
}
fn poll_outbound(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<Self::Substream, Self::Error>> {
Pin::new(&mut self.control)
.poll_open_stream(cx)
.map_err(Error)
}
fn poll(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
let this = self.get_mut();
let inbound_stream = ready!(this.poll_inner(cx))?;
if this.inbound_stream_buffer.len() >= MAX_BUFFERED_INBOUND_STREAMS {
log::warn!("dropping {inbound_stream} because buffer is full");
drop(inbound_stream);
} else {
this.inbound_stream_buffer.push_back(inbound_stream);
if let Some(waker) = this.inbound_stream_waker.take() {
waker.wake()
}
}
cx.waker().wake_by_ref();
Poll::Pending
}
fn poll_close(mut self: Pin<&mut Self>, c: &mut Context<'_>) -> Poll<Result<(), Error>> {
if let Poll::Ready(()) = Pin::new(&mut self.control).poll_close(c).map_err(Error)? {
return Poll::Ready(Ok(()));
}
while let Poll::Ready(maybe_inbound_stream) = self.incoming.poll_next_unpin(c)? {
match maybe_inbound_stream {
Some(inbound_stream) => mem::drop(inbound_stream),
None => return Poll::Ready(Ok(())),
}
}
Poll::Pending
}
}
impl<S> Muxer<S>
where
S: Stream<Item = Result<yamux::Stream, Error>> + Unpin,
{
fn poll_inner(&mut self, cx: &mut Context<'_>) -> Poll<Result<yamux::Stream, Error>> {
self.incoming.poll_next_unpin(cx).map(|maybe_stream| {
let stream = maybe_stream
.transpose()?
.ok_or(Error(ConnectionError::Closed))?;
Ok(stream)
})
}
}
#[deprecated(note = "Import the `yamux` module and refer to this type as `yamux::Config` instead.")]
pub type YamuxConfig = Config;
#[derive(Debug, Clone)]
pub struct Config {
inner: yamux::Config,
mode: Option<yamux::Mode>,
}
pub struct WindowUpdateMode(yamux::WindowUpdateMode);
impl WindowUpdateMode {
pub fn on_receive() -> Self {
WindowUpdateMode(yamux::WindowUpdateMode::OnReceive)
}
pub fn on_read() -> Self {
WindowUpdateMode(yamux::WindowUpdateMode::OnRead)
}
}
#[deprecated(
note = "Import the `yamux` module and refer to this type as `yamux::LocalConfig` instead."
)]
pub type YamuxLocalConfig = LocalConfig;
#[derive(Clone)]
pub struct LocalConfig(Config);
impl Config {
pub fn client() -> Self {
Self {
mode: Some(yamux::Mode::Client),
..Default::default()
}
}
pub fn server() -> Self {
Self {
mode: Some(yamux::Mode::Server),
..Default::default()
}
}
pub fn set_receive_window_size(&mut self, num_bytes: u32) -> &mut Self {
self.inner.set_receive_window(num_bytes);
self
}
pub fn set_max_buffer_size(&mut self, num_bytes: usize) -> &mut Self {
self.inner.set_max_buffer_size(num_bytes);
self
}
pub fn set_max_num_streams(&mut self, num_streams: usize) -> &mut Self {
self.inner.set_max_num_streams(num_streams);
self
}
pub fn set_window_update_mode(&mut self, mode: WindowUpdateMode) -> &mut Self {
self.inner.set_window_update_mode(mode.0);
self
}
pub fn into_local(self) -> LocalConfig {
LocalConfig(self)
}
}
impl Default for Config {
fn default() -> Self {
let mut inner = yamux::Config::default();
inner.set_read_after_close(false);
Config { inner, mode: None }
}
}
impl UpgradeInfo for Config {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(b"/yamux/1.0.0")
}
}
impl UpgradeInfo for LocalConfig {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(b"/yamux/1.0.0")
}
}
impl<C> InboundUpgrade<C> for Config
where
C: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
type Output = Muxer<Incoming<C>>;
type Error = io::Error;
type Future = future::Ready<Result<Self::Output, Self::Error>>;
fn upgrade_inbound(self, io: C, _: Self::Info) -> Self::Future {
let mode = self.mode.unwrap_or(yamux::Mode::Server);
future::ready(Ok(Muxer::new(io, self.inner, mode)))
}
}
impl<C> InboundUpgrade<C> for LocalConfig
where
C: AsyncRead + AsyncWrite + Unpin + 'static,
{
type Output = Muxer<LocalIncoming<C>>;
type Error = io::Error;
type Future = future::Ready<Result<Self::Output, Self::Error>>;
fn upgrade_inbound(self, io: C, _: Self::Info) -> Self::Future {
let cfg = self.0;
let mode = cfg.mode.unwrap_or(yamux::Mode::Server);
future::ready(Ok(Muxer::local(io, cfg.inner, mode)))
}
}
impl<C> OutboundUpgrade<C> for Config
where
C: AsyncRead + AsyncWrite + Send + Unpin + 'static,
{
type Output = Muxer<Incoming<C>>;
type Error = io::Error;
type Future = future::Ready<Result<Self::Output, Self::Error>>;
fn upgrade_outbound(self, io: C, _: Self::Info) -> Self::Future {
let mode = self.mode.unwrap_or(yamux::Mode::Client);
future::ready(Ok(Muxer::new(io, self.inner, mode)))
}
}
impl<C> OutboundUpgrade<C> for LocalConfig
where
C: AsyncRead + AsyncWrite + Unpin + 'static,
{
type Output = Muxer<LocalIncoming<C>>;
type Error = io::Error;
type Future = future::Ready<Result<Self::Output, Self::Error>>;
fn upgrade_outbound(self, io: C, _: Self::Info) -> Self::Future {
let cfg = self.0;
let mode = cfg.mode.unwrap_or(yamux::Mode::Client);
future::ready(Ok(Muxer::local(io, cfg.inner, mode)))
}
}
#[deprecated(note = "Import the `yamux` module and refer to this type as `yamux::Error` instead.")]
pub type YamuxError = Error;
#[derive(Debug, Error)]
#[error("yamux error: {0}")]
pub struct Error(#[from] yamux::ConnectionError);
impl From<Error> for io::Error {
fn from(err: Error) -> Self {
match err.0 {
yamux::ConnectionError::Io(e) => e,
e => io::Error::new(io::ErrorKind::Other, e),
}
}
}
pub struct Incoming<T> {
stream: BoxStream<'static, Result<yamux::Stream, Error>>,
_marker: std::marker::PhantomData<T>,
}
impl<T> fmt::Debug for Incoming<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Incoming")
}
}
pub struct LocalIncoming<T> {
stream: LocalBoxStream<'static, Result<yamux::Stream, Error>>,
_marker: std::marker::PhantomData<T>,
}
impl<T> fmt::Debug for LocalIncoming<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("LocalIncoming")
}
}
impl<T> Stream for Incoming<T> {
type Item = Result<yamux::Stream, Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.stream.as_mut().poll_next_unpin(cx)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.stream.size_hint()
}
}
impl<T> Unpin for Incoming<T> {}
impl<T> Stream for LocalIncoming<T> {
type Item = Result<yamux::Stream, Error>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.stream.as_mut().poll_next_unpin(cx)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.stream.size_hint()
}
}
impl<T> Unpin for LocalIncoming<T> {}