use bytes::Bytes;
use futures::{Future, FutureExt, Sink, Stream, StreamExt, TryStreamExt, future::BoxFuture};
use std::{
convert::Infallible,
error::Error,
fmt, io,
pin::Pin,
task::{Context, Poll},
};
use tokio::io::{AsyncRead, AsyncWrite};
use tokio_util::codec::{FramedRead, FramedWrite};
use crate::{
RemoteSend,
chmux::{ChMux, ChMuxError},
codec,
rch::base,
};
pub mod ext;
mod io_transport;
#[cfg_attr(docsrs, doc(cfg(feature = "rch")))]
#[derive(Debug, Clone)]
pub enum ConnectError<TransportSinkError, TransportStreamError> {
ChMux(ChMuxError<TransportSinkError, TransportStreamError>),
RemoteConnect(base::ConnectError),
}
impl<TransportSinkError, TransportStreamError> fmt::Display
for ConnectError<TransportSinkError, TransportStreamError>
where
TransportSinkError: fmt::Display,
TransportStreamError: fmt::Display,
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self {
Self::ChMux(err) => write!(f, "chmux error: {err}"),
Self::RemoteConnect(err) => write!(f, "channel connect failed: {err}"),
}
}
}
impl<TransportSinkError, TransportStreamError> Error for ConnectError<TransportSinkError, TransportStreamError>
where
TransportSinkError: Error,
TransportStreamError: Error,
{
}
impl<TransportSinkError, TransportStreamError> From<ChMuxError<TransportSinkError, TransportStreamError>>
for ConnectError<TransportSinkError, TransportStreamError>
{
fn from(err: ChMuxError<TransportSinkError, TransportStreamError>) -> Self {
Self::ChMux(err)
}
}
impl<TransportSinkError, TransportStreamError> From<base::ConnectError>
for ConnectError<TransportSinkError, TransportStreamError>
{
fn from(err: base::ConnectError) -> Self {
Self::RemoteConnect(err)
}
}
#[cfg_attr(docsrs, doc(cfg(feature = "rch")))]
#[must_use = "You must poll or spawn the Connect future for the connection to work."]
pub struct Connect<'transport, TransportSinkError, TransportStreamError>(
BoxFuture<'transport, Result<(), ChMuxError<TransportSinkError, TransportStreamError>>>,
);
impl<'transport, TransportSinkError, TransportStreamError>
Connect<'transport, TransportSinkError, TransportStreamError>
{
pub async fn framed<TransportSink, TransportStream, Tx, Rx, Codec>(
cfg: crate::Cfg, transport_sink: TransportSink, transport_stream: TransportStream,
) -> Result<
(
Connect<'transport, TransportSinkError, TransportStreamError>,
base::Sender<Tx, Codec>,
base::Receiver<Rx, Codec>,
),
ConnectError<TransportSinkError, TransportStreamError>,
>
where
TransportSink: Sink<Bytes, Error = TransportSinkError> + Send + Sync + Unpin + 'transport,
TransportSinkError: Error + Send + Sync + 'static,
TransportStream: Stream<Item = Result<Bytes, TransportStreamError>> + Send + Sync + Unpin + 'transport,
TransportStreamError: Error + Send + Sync + 'static,
Tx: RemoteSend,
Rx: RemoteSend,
Codec: codec::Codec,
{
let (mux, client, mut listener) = ChMux::new(cfg, transport_sink, transport_stream).await?;
let mut connection = Self(mux.run().boxed());
tokio::select! {
biased;
Err(err) = &mut connection => Err(err.into()),
result = base::connect(&client, &mut listener) => {
match result {
Ok((tx, rx)) => Ok((connection, tx, rx)),
Err(err) => Err(err.into()),
}
}
}
}
}
impl<'transport> Connect<'transport, io::Error, io::Error> {
pub async fn io<Read, Write, Tx, Rx, Codec>(
cfg: crate::Cfg, input: Read, output: Write,
) -> Result<
(Connect<'transport, io::Error, io::Error>, base::Sender<Tx, Codec>, base::Receiver<Rx, Codec>),
ConnectError<io::Error, io::Error>,
>
where
Read: AsyncRead + Send + Sync + Unpin + 'transport,
Write: AsyncWrite + Send + Sync + Unpin + 'transport,
Tx: RemoteSend,
Rx: RemoteSend,
Codec: codec::Codec,
{
let encoder = io_transport::LengthCodec::new(u32::MAX);
let transport_sink = io_transport::FilterFlushOuter(FramedWrite::with_capacity(
io_transport::FilterFlushInner::new(output),
encoder,
cfg.io_buffer_size,
));
let decoder = io_transport::LengthCodec::new(cfg.max_frame_length());
let transport_stream =
FramedRead::with_capacity(input, decoder, cfg.io_buffer_size).map_ok(|item| item.freeze());
Self::framed(cfg, transport_sink, transport_stream).await
}
}
impl<TransportSinkError, TransportStreamError> Future for Connect<'_, TransportSinkError, TransportStreamError> {
type Output = Result<(), ChMuxError<TransportSinkError, TransportStreamError>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
Pin::into_inner(self).0.poll_unpin(cx)
}
}
type LoopbackSendError = futures::channel::mpsc::SendError;
type LoopbackRecvError = Infallible;
pub type LoopbackConnect = Connect<'static, LoopbackSendError, LoopbackRecvError>;
impl LoopbackConnect {
pub async fn loopback<Tx, Rx, Codec>(
cfg: crate::Cfg,
) -> (LoopbackConnect, base::Sender<Tx, Codec>, base::Receiver<Rx, Codec>)
where
Tx: RemoteSend,
Rx: RemoteSend,
Codec: codec::Codec,
{
let (a_transport_tx, a_transport_rx) = futures::channel::mpsc::channel(cfg.transport_send_queue);
let (b_transport_tx, b_transport_rx) = futures::channel::mpsc::channel(cfg.transport_send_queue);
let a_transport_rx = a_transport_rx.map(Ok);
let b_transport_rx = b_transport_rx.map(Ok);
let ((a_connect, a_base_tx, _a_base_rx), (b_connect, _b_base_tx, b_base_rx)) = tokio::try_join!(
Self::framed::<_, _, _, (), _>(cfg.clone(), a_transport_tx, b_transport_rx),
Self::framed::<_, _, (), _, _>(cfg.clone(), b_transport_tx, a_transport_rx),
)
.unwrap();
let connection = Self(
async move {
tokio::try_join!(a_connect, b_connect)?;
Ok(())
}
.boxed(),
);
(connection, a_base_tx, b_base_rx)
}
}