mercuryworkshop-mirror/epoxy-tls
1
0
Fork 0
mirror of https://github.com/MercuryWorkshop/epoxy-tls.git synced 2025-05-14 06:50:01 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

add wisp to client

Browse source
This commit is contained in:
Toshit Chawda 2024-01-30 21:15:17 -08:00
parent be7d92b4c5
commit c5cf95fcb1
No known key found for this signature in database
GPG key ID: 91480ED99E2B3D9D
12 changed files with 210 additions and 320 deletions
Download patch file Download diff file Expand all files Collapse all files

6
wisp/src/fastwebsockets.rs
View file

@ -53,10 +53,10 @@ impl From<WebSocketError> for crate::WispError {
}
}
impl<S: AsyncRead + Unpin> crate::ws::WebSocketRead for FragmentCollectorRead<S> {
impl<S: AsyncRead + Unpin + Send> crate::ws::WebSocketRead for FragmentCollectorRead<S> {
async fn wisp_read_frame(
&mut self,
tx: &mut crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite>,
tx: &crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite + Send>,
) -> Result<crate::ws::Frame, crate::WispError> {
Ok(self
.read_frame(&mut |frame| async { tx.write_frame(frame.into()).await })
@ -65,7 +65,7 @@ impl<S: AsyncRead + Unpin> crate::ws::WebSocketRead for FragmentCollectorRead<S>
}
}
impl<S: AsyncWrite + Unpin> crate::ws::WebSocketWrite for WebSocketWrite<S> {
impl<S: AsyncWrite + Unpin + Send> crate::ws::WebSocketWrite for WebSocketWrite<S> {
async fn wisp_write_frame(&mut self, frame: crate::ws::Frame) -> Result<(), crate::WispError> {
self.write_frame(frame.try_into()?).await.map_err(|e| e.into())
}

162
wisp/src/lib.rs
View file

@ -10,7 +10,7 @@ pub use crate::packet::*;
pub use crate::stream::*;
use dashmap::DashMap;
use futures::{channel::mpsc, StreamExt};
use futures::{channel::mpsc, Future, StreamExt};
use std::sync::{
atomic::{AtomicBool, AtomicU32, Ordering},
Arc,
@ -68,38 +68,66 @@ impl std::fmt::Display for WispError {
impl std::error::Error for WispError {}
pub struct ServerMux<R, W>
pub struct ServerMux<W>
where
R: ws::WebSocketRead,
W: ws::WebSocketWrite,
{
rx: R,
tx: ws::LockedWebSocketWrite<W>,
stream_map: Arc<DashMap<u32, mpsc::UnboundedSender<WsEvent>>>,
close_rx: mpsc::UnboundedReceiver<MuxEvent>,
close_tx: mpsc::UnboundedSender<MuxEvent>,
}
impl<R: ws::WebSocketRead, W: ws::WebSocketWrite> ServerMux<R, W> {
pub fn new(read: R, write: W) -> Self {
impl<W: ws::WebSocketWrite + Send + 'static> ServerMux<W> {
pub fn handle<'a, FR, R>(
read: R,
write: W,
handler_fn: &'a mut impl Fn(ConnectPacket, MuxStream<W>) -> FR,
) -> impl Future<Output = Result<(), WispError>> + 'a
where
FR: std::future::Future<Output = Result<(), WispError>> + 'a,
R: ws::WebSocketRead + 'a,
W: ws::WebSocketWrite + 'a,
{
let (tx, rx) = mpsc::unbounded::<MuxEvent>();
Self {
rx: read,
tx: ws::LockedWebSocketWrite::new(write),
stream_map: Arc::new(DashMap::new()),
close_rx: rx,
let write = ws::LockedWebSocketWrite::new(write);
let map = Arc::new(DashMap::new());
let inner = ServerMux {
stream_map: map.clone(),
tx: write.clone(),
close_tx: tx,
}
};
inner.into_future(read, rx, handler_fn)
}
pub async fn server_bg_loop(&mut self) {
while let Some(msg) = self.close_rx.next().await {
async fn into_future<R, FR>(
self,
rx: R,
close_rx: mpsc::UnboundedReceiver<MuxEvent>,
handler_fn: &mut impl Fn(ConnectPacket, MuxStream<W>) -> FR,
) -> Result<(), WispError>
where
R: ws::WebSocketRead,
FR: std::future::Future<Output = Result<(), WispError>>,
{
futures::try_join! {
self.server_close_loop(close_rx, self.stream_map.clone(), self.tx.clone()),
self.server_msg_loop(rx, handler_fn)
}
.map(|_| ())
}
async fn server_close_loop(
&self,
mut close_rx: mpsc::UnboundedReceiver<MuxEvent>,
stream_map: Arc<DashMap<u32, mpsc::UnboundedSender<WsEvent>>>,
tx: ws::LockedWebSocketWrite<W>,
) -> Result<(), WispError> {
while let Some(msg) = close_rx.next().await {
match msg {
MuxEvent::Close(stream_id, reason, channel) => {
if self.stream_map.clone().remove(&stream_id).is_some() {
if stream_map.clone().remove(&stream_id).is_some() {
let _ = channel.send(
self.tx
.write_frame(Packet::new_close(stream_id, reason).into())
tx.write_frame(Packet::new_close(stream_id, reason).into())
.await,
);
} else {
@ -108,20 +136,23 @@ impl<R: ws::WebSocketRead, W: ws::WebSocketWrite> ServerMux<R, W> {
}
}
}
Ok(())
}
pub async fn server_loop<FR>(
&mut self,
async fn server_msg_loop<R, FR>(
&self,
mut rx: R,
handler_fn: &mut impl Fn(ConnectPacket, MuxStream<W>) -> FR,
) -> Result<(), WispError>
where
FR: std::future::Future<Output = Result<(), crate::WispError>>,
R: ws::WebSocketRead,
FR: std::future::Future<Output = Result<(), WispError>>,
{
self.tx
.write_frame(Packet::new_continue(0, u32::MAX).into())
.await?;
while let Ok(frame) = self.rx.wisp_read_frame(&mut self.tx).await {
while let Ok(frame) = rx.wisp_read_frame(&self.tx).await {
if let Ok(packet) = Packet::try_from(frame) {
use PacketType::*;
match packet.packet {
@ -164,34 +195,31 @@ impl<R: ws::WebSocketRead, W: ws::WebSocketWrite> ServerMux<R, W> {
}
}
pub struct ClientMux<R, W>
pub struct ClientMuxInner<W>
where
R: ws::WebSocketRead,
W: ws::WebSocketWrite,
{
rx: R,
tx: ws::LockedWebSocketWrite<W>,
stream_map: Arc<DashMap<u32, mpsc::UnboundedSender<WsEvent>>>,
next_free_stream_id: AtomicU32,
close_rx: mpsc::UnboundedReceiver<MuxEvent>,
close_tx: mpsc::UnboundedSender<MuxEvent>,
}
impl<R: ws::WebSocketRead, W: ws::WebSocketWrite> ClientMux<R, W> {
pub fn new(read: R, write: W) -> Self {
let (tx, rx) = mpsc::unbounded::<MuxEvent>();
Self {
rx: read,
tx: ws::LockedWebSocketWrite::new(write),
stream_map: Arc::new(DashMap::new()),
next_free_stream_id: AtomicU32::new(1),
close_rx: rx,
close_tx: tx,
}
impl<W: ws::WebSocketWrite + Send> ClientMuxInner<W> {
pub async fn into_future<R>(
self,
rx: R,
close_rx: mpsc::UnboundedReceiver<MuxEvent>,
) -> Result<(), WispError>
where
R: ws::WebSocketRead,
{
futures::try_join!(self.client_bg_loop(close_rx), self.client_loop(rx)).map(|_| ())
}
pub async fn client_bg_loop(&mut self) {
while let Some(msg) = self.close_rx.next().await {
async fn client_bg_loop(
&self,
mut close_rx: mpsc::UnboundedReceiver<MuxEvent>,
) -> Result<(), WispError> {
while let Some(msg) = close_rx.next().await {
match msg {
MuxEvent::Close(stream_id, reason, channel) => {
if self.stream_map.clone().remove(&stream_id).is_some() {
@ -206,14 +234,14 @@ impl<R: ws::WebSocketRead, W: ws::WebSocketWrite> ClientMux<R, W> {
}
}
}
Ok(())
}
pub async fn client_loop(&mut self) -> Result<(), WispError> {
self.tx
.write_frame(Packet::new_continue(0, u32::MAX).into())
.await?;
while let Ok(frame) = self.rx.wisp_read_frame(&mut self.tx).await {
async fn client_loop<R>(&self, mut rx: R) -> Result<(), WispError>
where
R: ws::WebSocketRead,
{
while let Ok(frame) = rx.wisp_read_frame(&self.tx).await {
if let Ok(packet) = Packet::try_from(frame) {
use PacketType::*;
match packet.packet {
@ -235,12 +263,52 @@ impl<R: ws::WebSocketRead, W: ws::WebSocketWrite> ClientMux<R, W> {
}
Ok(())
}
}
pub struct ClientMux<W>
where
W: ws::WebSocketWrite,
{
tx: ws::LockedWebSocketWrite<W>,
stream_map: Arc<DashMap<u32, mpsc::UnboundedSender<WsEvent>>>,
next_free_stream_id: AtomicU32,
close_tx: mpsc::UnboundedSender<MuxEvent>,
}
impl<W: ws::WebSocketWrite + Send + 'static> ClientMux<W> {
pub fn new<R>(read: R, write: W) -> (Self, impl Future<Output = Result<(), WispError>>)
where
R: ws::WebSocketRead,
{
let (tx, rx) = mpsc::unbounded::<MuxEvent>();
let map = Arc::new(DashMap::new());
let write = ws::LockedWebSocketWrite::new(write);
(
Self {
tx: write.clone(),
stream_map: map.clone(),
next_free_stream_id: AtomicU32::new(1),
close_tx: tx,
},
ClientMuxInner {
tx: write.clone(),
stream_map: map.clone(),
}
.into_future(read, rx),
)
}
pub async fn client_new_stream(
&mut self,
stream_type: StreamType,
host: String,
port: u16,
) -> Result<MuxStream<impl ws::WebSocketWrite>, WispError> {
let (ch_tx, ch_rx) = mpsc::unbounded();
let stream_id = self.next_free_stream_id.load(Ordering::Acquire);
self.tx
.write_frame(Packet::new_connect(stream_id, stream_type, port, host).into())
.await?;
self.next_free_stream_id.store(
stream_id
.checked_add(1)

77
wisp/src/stream.rs
View file

@ -4,7 +4,7 @@ use futures::{
channel::{mpsc, oneshot},
sink, stream,
task::{Context, Poll},
AsyncRead, AsyncWrite, Sink, Stream, StreamExt,
Sink, Stream, StreamExt,
};
use pin_project_lite::pin_project;
use std::{
@ -44,7 +44,7 @@ impl MuxStreamRead {
}
}
pub(crate) fn into_stream(self) -> Pin<Box<dyn Stream<Item = Bytes>>> {
pub(crate) fn into_stream(self) -> Pin<Box<dyn Stream<Item = Bytes> + Send>> {
Box::pin(stream::unfold(self, |mut rx| async move {
let evt = rx.read().await?;
Some((
@ -68,7 +68,7 @@ where
is_closed: Arc<AtomicBool>,
}
impl<W: crate::ws::WebSocketWrite> MuxStreamWrite<W> {
impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStreamWrite<W> {
pub async fn write(&mut self, data: Bytes) -> Result<(), crate::WispError> {
if self.is_closed.load(Ordering::Acquire) {
return Err(crate::WispError::StreamAlreadyClosed);
@ -101,10 +101,7 @@ impl<W: crate::ws::WebSocketWrite> MuxStreamWrite<W> {
Ok(())
}
pub(crate) fn into_sink<'a>(self) -> Pin<Box<dyn Sink<Bytes, Error = crate::WispError> + 'a>>
where
W: 'a,
{
pub(crate) fn into_sink(self) -> Pin<Box<dyn Sink<Bytes, Error = crate::WispError> + Send>> {
Box::pin(sink::unfold(self, |mut tx, data| async move {
tx.write(data).await?;
Ok(tx)
@ -130,7 +127,7 @@ where
tx: MuxStreamWrite<W>,
}
impl<W: crate::ws::WebSocketWrite> MuxStream<W> {
impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStream<W> {
pub(crate) fn new(
stream_id: u32,
rx: mpsc::UnboundedReceiver<WsEvent>,
@ -174,10 +171,7 @@ impl<W: crate::ws::WebSocketWrite> MuxStream<W> {
(self.rx, self.tx)
}
pub fn into_io<'a>(self) -> MuxStreamIo<'a>
where
W: 'a,
{
pub fn into_io(self) -> MuxStreamIo {
MuxStreamIo {
rx: self.rx.into_stream(),
tx: self.tx.into_sink(),
@ -208,55 +202,54 @@ impl MuxStreamCloser {
}
pin_project! {
pub struct MuxStreamIo<'a> {
pub struct MuxStreamIo {
#[pin]
rx: Pin<Box<dyn Stream<Item = Bytes> + 'a>>,
rx: Pin<Box<dyn Stream<Item = Bytes> + Send>>,
#[pin]
tx: Pin<Box<dyn Sink<Bytes, Error = crate::WispError> + 'a>>,
tx: Pin<Box<dyn Sink<Bytes, Error = crate::WispError> + Send>>,
}
}
impl<'a> MuxStreamIo<'a> {
pub fn into_asyncrw(self) -> impl AsyncRead + AsyncWrite + 'a {
IoStream::new(self.map(|x| Ok::<Vec<u8>, std::io::Error>(x.to_vec())))
impl MuxStreamIo {
pub fn into_asyncrw(self) -> IoStream<MuxStreamIo, Vec<u8>> {
IoStream::new(self)
}
}
impl Stream for MuxStreamIo<'_> {
type Item = Bytes;
impl Stream for MuxStreamIo {
type Item = Result<Vec<u8>, std::io::Error>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.project().rx.poll_next(cx)
self.project()
.rx
.poll_next(cx)
.map(|x| x.map(|x| Ok(x.to_vec())))
}
}
impl Sink<Bytes> for MuxStreamIo<'_> {
type Error = crate::WispError;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project().tx.poll_ready(cx)
}
fn start_send(self: Pin<&mut Self>, item: Bytes) -> Result<(), Self::Error> {
self.project().tx.start_send(item)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project().tx.poll_flush(cx)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project().tx.poll_close(cx)
}
}
impl Sink<Vec<u8>> for MuxStreamIo<'_> {
impl Sink<Vec<u8>> for MuxStreamIo {
type Error = std::io::Error;
fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project().tx.poll_ready(cx).map_err(std::io::Error::other)
self.project()
.tx
.poll_ready(cx)
.map_err(std::io::Error::other)
}
fn start_send(self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
self.project().tx.start_send(item.into()).map_err(std::io::Error::other)
self.project()
.tx
.start_send(item.into())
.map_err(std::io::Error::other)
}
fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project().tx.poll_flush(cx).map_err(std::io::Error::other)
self.project()
.tx
.poll_flush(cx)
.map_err(std::io::Error::other)
}
fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.project().tx.poll_close(cx).map_err(std::io::Error::other)
self.project()
.tx
.poll_close(cx)
.map_err(std::io::Error::other)
}
}

8
wisp/src/ws.rs
View file

@ -46,20 +46,20 @@ impl Frame {
pub trait WebSocketRead {
fn wisp_read_frame(
&mut self,
tx: &mut crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite>,
) -> impl std::future::Future<Output = Result<Frame, crate::WispError>>;
tx: &crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite + Send>,
) -> impl std::future::Future<Output = Result<Frame, crate::WispError>> + Send;
}
pub trait WebSocketWrite {
fn wisp_write_frame(
&mut self,
frame: Frame,
) -> impl std::future::Future<Output = Result<(), crate::WispError>>;
) -> impl std::future::Future<Output = Result<(), crate::WispError>> + Send;
}
pub struct LockedWebSocketWrite<S>(Arc<Mutex<S>>);
impl<S: WebSocketWrite> LockedWebSocketWrite<S> {
impl<S: WebSocketWrite + Send> LockedWebSocketWrite<S> {
pub fn new(ws: S) -> Self {
Self(Arc::new(Mutex::new(ws)))
}

13
wisp/src/ws_stream_wasm.rs
View file

@ -1,4 +1,4 @@
use futures::{SinkExt, StreamExt};
use futures::{stream::{SplitStream, SplitSink}, SinkExt, StreamExt};
use ws_stream_wasm::{WsErr, WsMessage, WsStream};
impl From<WsMessage> for crate::ws::Frame {
@ -37,10 +37,10 @@ impl From<WsErr> for crate::WispError {
}
}
impl crate::ws::WebSocketRead for WsStream {
impl crate::ws::WebSocketRead for SplitStream<WsStream> {
async fn wisp_read_frame(
&mut self,
_: &mut crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite>,
_: &crate::ws::LockedWebSocketWrite<impl crate::ws::WebSocketWrite>,
) -> Result<crate::ws::Frame, crate::WispError> {
Ok(self
.next()
@ -50,8 +50,11 @@ impl crate::ws::WebSocketRead for WsStream {
}
}
impl crate::ws::WebSocketWrite for WsStream {
impl crate::ws::WebSocketWrite for SplitSink<WsStream, WsMessage> {
async fn wisp_write_frame(&mut self, frame: crate::ws::Frame) -> Result<(), crate::WispError> {
self.send(frame.try_into()?).await.map_err(|e| e.into())
self
.send(frame.try_into()?)
.await
.map_err(|e| e.into())
}
}