use futures::{future, prelude::*, ready, stream::{BoxStream, LocalBoxStream}};
use tetsy_libp2p_core::muxing::{StreamMuxer, StreamMuxerEvent};
use tetsy_libp2p_core::upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
use parking_lot::Mutex;
use std::{fmt, io, iter, pin::Pin, task::{Context, Poll}};
use thiserror::Error;
pub struct Remux<S>(Mutex<Inner<S>>);
impl<S> fmt::Debug for Remux<S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("Remux")
}
}
struct Inner<S> {
incoming: S,
control: remux::Control,
}
#[derive(Debug)]
pub struct OpenSubstreamToken(());
impl<C> Remux<Incoming<C>>
where
C: AsyncRead + AsyncWrite + Send + Unpin + 'static
{
fn new(io: C, cfg: remux::Config, mode: remux::Mode) -> Self {
let conn = remux::Connection::new(io, cfg, mode);
let ctrl = conn.control();
let inner = Inner {
incoming: Incoming {
stream: remux::into_stream(conn).err_into().boxed(),
_marker: std::marker::PhantomData
},
control: ctrl,
};
Remux(Mutex::new(inner))
}
}
impl<C> Remux<LocalIncoming<C>>
where
C: AsyncRead + AsyncWrite + Unpin + 'static
{
fn local(io: C, cfg: remux::Config, mode: remux::Mode) -> Self {
let conn = remux::Connection::new(io, cfg, mode);
let ctrl = conn.control();
let inner = Inner {
incoming: LocalIncoming {
stream: remux::into_stream(conn).err_into().boxed_local(),
_marker: std::marker::PhantomData
},
control: ctrl,
};
Remux(Mutex::new(inner))
}
}
pub type RemuxResult<T> = Result<T, RemuxError>;
impl<S> StreamMuxer for Remux<S>
where
S: Stream<Item = Result<remux::Stream, RemuxError>> + Unpin
{
type Substream = remux::Stream;
type OutboundSubstream = OpenSubstreamToken;
type Error = RemuxError;
fn poll_event(&self, c: &mut Context<'_>)
-> Poll<RemuxResult<StreamMuxerEvent<Self::Substream>>>
{
let mut inner = self.0.lock();
match ready!(inner.incoming.poll_next_unpin(c)) {
Some(Ok(s)) => Poll::Ready(Ok(StreamMuxerEvent::InboundSubstream(s))),
Some(Err(e)) => Poll::Ready(Err(e)),
None => Poll::Ready(Err(remux::ConnectionError::Closed.into()))
}
}
fn open_outbound(&self) -> Self::OutboundSubstream {
OpenSubstreamToken(())
}
fn poll_outbound(&self, c: &mut Context<'_>, _: &mut OpenSubstreamToken)
-> Poll<RemuxResult<Self::Substream>>
{
let mut inner = self.0.lock();
Pin::new(&mut inner.control).poll_open_stream(c).map_err(RemuxError)
}
fn destroy_outbound(&self, _: Self::OutboundSubstream) {
self.0.lock().control.abort_open_stream()
}
fn read_substream(&self, c: &mut Context<'_>, s: &mut Self::Substream, b: &mut [u8])
-> Poll<RemuxResult<usize>>
{
Pin::new(s).poll_read(c, b).map_err(|e| RemuxError(e.into()))
}
fn write_substream(&self, c: &mut Context<'_>, s: &mut Self::Substream, b: &[u8])
-> Poll<RemuxResult<usize>>
{
Pin::new(s).poll_write(c, b).map_err(|e| RemuxError(e.into()))
}
fn flush_substream(&self, c: &mut Context<'_>, s: &mut Self::Substream)
-> Poll<RemuxResult<()>>
{
Pin::new(s).poll_flush(c).map_err(|e| RemuxError(e.into()))
}
fn shutdown_substream(&self, c: &mut Context<'_>, s: &mut Self::Substream)
-> Poll<RemuxResult<()>>
{
Pin::new(s).poll_close(c).map_err(|e| RemuxError(e.into()))
}
fn destroy_substream(&self, _: Self::Substream) { }
fn close(&self, c: &mut Context<'_>) -> Poll<RemuxResult<()>> {
let mut inner = self.0.lock();
if let std::task::Poll::Ready(x) = Pin::new(&mut inner.control).poll_close(c) {
return Poll::Ready(x.map_err(RemuxError))
}
while let std::task::Poll::Ready(x) = inner.incoming.poll_next_unpin(c) {
match x {
Some(Ok(_)) => {}
Some(Err(e)) => return Poll::Ready(Err(e)),
None => return Poll::Ready(Ok(()))
}
}
Poll::Pending
}
fn flush_all(&self, _: &mut Context<'_>) -> Poll<RemuxResult<()>> {
Poll::Ready(Ok(()))
}
}
#[derive(Clone)]
pub struct RemuxConfig {
inner: remux::Config,
mode: Option<remux::Mode>
}
pub struct WindowUpdateMode(remux::WindowUpdateMode);
impl WindowUpdateMode {
pub fn on_receive() -> Self {
WindowUpdateMode(remux::WindowUpdateMode::OnReceive)
}
pub fn on_read() -> Self {
WindowUpdateMode(remux::WindowUpdateMode::OnRead)
}
}
#[derive(Clone)]
pub struct RemuxLocalConfig(RemuxConfig);
impl RemuxConfig {
pub fn client() -> Self {
let mut cfg = Self::default();
cfg.mode = Some(remux::Mode::Client);
cfg
}
pub fn server() -> Self {
let mut cfg = Self::default();
cfg.mode = Some(remux::Mode::Server);
cfg
}
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) -> RemuxLocalConfig {
RemuxLocalConfig(self)
}
}
impl Default for RemuxConfig {
fn default() -> Self {
let mut inner = remux::Config::default();
inner.set_read_after_close(false);
RemuxConfig { inner, mode: None }
}
}
impl UpgradeInfo for RemuxConfig {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(b"/remux/1.0.0")
}
}
impl UpgradeInfo for RemuxLocalConfig {
type Info = &'static [u8];
type InfoIter = iter::Once<Self::Info>;
fn protocol_info(&self) -> Self::InfoIter {
iter::once(b"/remux/1.0.0")
}
}
impl<C> InboundUpgrade<C> for RemuxConfig
where
C: AsyncRead + AsyncWrite + Send + Unpin + 'static
{
type Output = Remux<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(remux::Mode::Server);
future::ready(Ok(Remux::new(io, self.inner, mode)))
}
}
impl<C> InboundUpgrade<C> for RemuxLocalConfig
where
C: AsyncRead + AsyncWrite + Unpin + 'static
{
type Output = Remux<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(remux::Mode::Server);
future::ready(Ok(Remux::local(io, cfg.inner, mode)))
}
}
impl<C> OutboundUpgrade<C> for RemuxConfig
where
C: AsyncRead + AsyncWrite + Send + Unpin + 'static
{
type Output = Remux<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(remux::Mode::Client);
future::ready(Ok(Remux::new(io, self.inner, mode)))
}
}
impl<C> OutboundUpgrade<C> for RemuxLocalConfig
where
C: AsyncRead + AsyncWrite + Unpin + 'static
{
type Output = Remux<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(remux::Mode::Client);
future::ready(Ok(Remux::local(io, cfg.inner, mode)))
}
}
#[derive(Debug, Error)]
#[error("remux error: {0}")]
pub struct RemuxError(#[from] remux::ConnectionError);
impl Into<io::Error> for RemuxError {
fn into(self: RemuxError) -> io::Error {
match self.0 {
remux::ConnectionError::Io(e) => e,
e => io::Error::new(io::ErrorKind::Other, e)
}
}
}
pub struct Incoming<T> {
stream: BoxStream<'static, Result<remux::Stream, RemuxError>>,
_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<remux::Stream, RemuxError>>,
_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<remux::Stream, RemuxError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> std::task::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<remux::Stream, RemuxError>;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> std::task::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> {
}