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

fix udp

Browse source
This commit is contained in:
Toshit Chawda 2024-03-10 07:57:28 -07:00
parent 19d9544e83
commit 53a399856f
4 changed files with 86 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

2
wisp/Cargo.toml
View file

@ -1,6 +1,6 @@
[package]
name = "wisp-mux"
version = "1.2.1"
version = "1.2.2"
license = "LGPL-3.0-only"
description = "A library for easily creating Wisp servers and clients."
homepage = "https://github.com/MercuryWorkshop/epoxy-tls/tree/multiplexed/wisp"

8
wisp/src/lib.rs
View file

@ -1,5 +1,4 @@
#![deny(missing_docs)]
#![feature(impl_trait_in_assoc_type)]
#![cfg_attr(docsrs, feature(doc_cfg))]
//! A library for easily creating [Wisp] clients and servers.
//!
@ -109,6 +108,7 @@ impl std::error::Error for WispError {}
struct MuxMapValue {
stream: mpsc::UnboundedSender<MuxEvent>,
stream_type: StreamType,
flow_control: Arc<AtomicU32>,
flow_control_event: Arc<Event>,
}
@ -200,6 +200,7 @@ impl<W: ws::WebSocketWrite + Send + 'static> ServerMuxInner<W> {
packet.stream_id,
MuxMapValue {
stream: ch_tx,
stream_type,
flow_control: flow_control.clone(),
flow_control_event: flow_control_event.clone(),
},
@ -224,6 +225,7 @@ impl<W: ws::WebSocketWrite + Send + 'static> ServerMuxInner<W> {
Data(data) => {
if let Some(stream) = self.stream_map.lock().await.get(&packet.stream_id) {
let _ = stream.stream.unbounded_send(MuxEvent::Send(data));
if stream.stream_type == StreamType::Tcp {
stream.flow_control.store(
stream
.flow_control
@ -233,6 +235,7 @@ impl<W: ws::WebSocketWrite + Send + 'static> ServerMuxInner<W> {
);
}
}
}
Continue(_) => unreachable!(),
Close(inner_packet) => {
if let Some(stream) = self.stream_map.lock().await.get(&packet.stream_id) {
@ -388,12 +391,14 @@ impl<W: ws::WebSocketWrite + Send> ClientMuxInner<W> {
}
Continue(inner_packet) => {
if let Some(stream) = self.stream_map.lock().await.get(&packet.stream_id) {
if stream.stream_type == StreamType::Tcp {
stream
.flow_control
.store(inner_packet.buffer_remaining, Ordering::Release);
let _ = stream.flow_control_event.notify(u32::MAX);
}
}
}
Close(inner_packet) => {
if let Some(stream) = self.stream_map.lock().await.get(&packet.stream_id) {
let _ = stream.stream.unbounded_send(MuxEvent::Close(inner_packet));
@ -490,6 +495,7 @@ impl<W: ws::WebSocketWrite + Send + 'static> ClientMux<W> {
stream_id,
MuxMapValue {
stream: ch_tx,
stream_type,
flow_control: flow_control.clone(),
flow_control_event: evt.clone(),
},

3
wisp/src/packet.rs
View file

@ -1,5 +1,4 @@
use crate::ws;
use crate::WispError;
use crate::{ws, WispError};
use bytes::{Buf, BufMut, Bytes};
/// Wisp stream type.

122
wisp/src/stream.rs
View file

@ -1,3 +1,5 @@
use crate::{sink_unfold, ws, ClosePacket, CloseReason, Packet, Role, StreamType, WispError};
use async_io_stream::IoStream;
use bytes::Bytes;
use event_listener::Event;
@ -21,31 +23,31 @@ pub enum MuxEvent {
/// The other side has sent data.
Send(Bytes),
/// The other side has closed.
Close(crate::ClosePacket),
Close(ClosePacket),
}
pub(crate) enum WsEvent {
Close(u32, crate::CloseReason, oneshot::Sender<Result<(), crate::WispError>>),
Close(u32, CloseReason, oneshot::Sender<Result<(), WispError>>),
EndFut,
}
/// Read side of a multiplexor stream.
pub struct MuxStreamRead<W>
where
W: crate::ws::WebSocketWrite,
W: ws::WebSocketWrite,
{
/// ID of the stream.
pub stream_id: u32,
/// Type of the stream.
pub stream_type: crate::StreamType,
role: crate::Role,
tx: crate::ws::LockedWebSocketWrite<W>,
pub stream_type: StreamType,
role: Role,
tx: ws::LockedWebSocketWrite<W>,
rx: mpsc::UnboundedReceiver<MuxEvent>,
is_closed: Arc<AtomicBool>,
flow_control: Arc<AtomicU32>,
}
impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStreamRead<W> {
impl<W: ws::WebSocketWrite + Send + 'static> MuxStreamRead<W> {
/// Read an event from the stream.
pub async fn read(&mut self) -> Option<MuxEvent> {
if self.is_closed.load(Ordering::Acquire) {
@ -53,12 +55,10 @@ impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStreamRead<W> {
}
match self.rx.next().await? {
MuxEvent::Send(bytes) => {
if self.role == crate::Role::Server && self.stream_type == crate::StreamType::Tcp {
if self.role == Role::Server && self.stream_type == StreamType::Tcp {
let old_val = self.flow_control.fetch_add(1, Ordering::AcqRel);
self.tx
.write_frame(
crate::Packet::new_continue(self.stream_id, old_val + 1).into(),
)
.write_frame(Packet::new_continue(self.stream_id, old_val + 1).into())
.await
.ok()?;
}
@ -88,35 +88,38 @@ impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStreamRead<W> {
/// Write side of a multiplexor stream.
pub struct MuxStreamWrite<W>
where
W: crate::ws::WebSocketWrite,
W: ws::WebSocketWrite,
{
/// ID of the stream.
pub stream_id: u32,
role: crate::Role,
tx: crate::ws::LockedWebSocketWrite<W>,
/// Type of the stream.
pub stream_type: StreamType,
role: Role,
tx: ws::LockedWebSocketWrite<W>,
close_channel: mpsc::UnboundedSender<WsEvent>,
is_closed: Arc<AtomicBool>,
continue_recieved: Arc<Event>,
flow_control: Arc<AtomicU32>,
}
impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStreamWrite<W> {
impl<W: ws::WebSocketWrite + Send + 'static> MuxStreamWrite<W> {
/// Write data to the stream.
pub async fn write(&self, data: Bytes) -> Result<(), crate::WispError> {
pub async fn write(&self, data: Bytes) -> Result<(), WispError> {
if self.is_closed.load(Ordering::Acquire) {
return Err(crate::WispError::StreamAlreadyClosed);
return Err(WispError::StreamAlreadyClosed);
}
if self.role == crate::Role::Client && self.flow_control.load(Ordering::Acquire) == 0 {
if self.role == Role::Client
&& self.stream_type == StreamType::Tcp
&& self.flow_control.load(Ordering::Acquire) == 0
{
self.continue_recieved.listen().await;
}
self.tx
.write_frame(crate::Packet::new_data(self.stream_id, data).into())
.write_frame(Packet::new_data(self.stream_id, data).into())
.await?;
if self.role == crate::Role::Client {
if self.role == Role::Client && self.stream_type == StreamType::Tcp {
self.flow_control.store(
self.flow_control
.load(Ordering::Acquire)
.saturating_sub(1),
self.flow_control.load(Ordering::Acquire).saturating_sub(1),
Ordering::Release,
);
}
@ -143,45 +146,48 @@ impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStreamWrite<W> {
}
/// Close the stream. You will no longer be able to write or read after this has been called.
pub async fn close(&self, reason: crate::CloseReason) -> Result<(), crate::WispError> {
pub async fn close(&self, reason: CloseReason) -> Result<(), WispError> {
if self.is_closed.load(Ordering::Acquire) {
return Err(crate::WispError::StreamAlreadyClosed);
return Err(WispError::StreamAlreadyClosed);
}
let (tx, rx) = oneshot::channel::<Result<(), crate::WispError>>();
let (tx, rx) = oneshot::channel::<Result<(), WispError>>();
self.close_channel
.unbounded_send(WsEvent::Close(self.stream_id, reason, tx))
.map_err(|x| crate::WispError::Other(Box::new(x)))?;
rx.await
.map_err(|x| crate::WispError::Other(Box::new(x)))??;
.map_err(|x| WispError::Other(Box::new(x)))?;
rx.await.map_err(|x| WispError::Other(Box::new(x)))??;
self.is_closed.store(true, Ordering::Release);
Ok(())
}
pub(crate) fn into_sink(self) -> Pin<Box<dyn Sink<Bytes, Error = crate::WispError> + Send>> {
pub(crate) fn into_sink(self) -> Pin<Box<dyn Sink<Bytes, Error = WispError> + Send>> {
let handle = self.get_close_handle();
Box::pin(crate::sink_unfold::unfold(self, |tx, data| async move {
Box::pin(sink_unfold::unfold(
self,
|tx, data| async move {
tx.write(data).await?;
Ok(tx)
}, move || {
handle.close_sync(crate::CloseReason::Unknown)
}))
},
move || handle.close_sync(CloseReason::Unknown),
))
}
}
impl<W: crate::ws::WebSocketWrite> Drop for MuxStreamWrite<W> {
impl<W: ws::WebSocketWrite> Drop for MuxStreamWrite<W> {
fn drop(&mut self) {
let (tx, _) = oneshot::channel::<Result<(), crate::WispError>>();
let _ = self
.close_channel
.unbounded_send(WsEvent::Close(self.stream_id, crate::CloseReason::Unknown, tx));
let (tx, _) = oneshot::channel::<Result<(), WispError>>();
let _ = self.close_channel.unbounded_send(WsEvent::Close(
self.stream_id,
CloseReason::Unknown,
tx,
));
}
}
/// Multiplexor stream.
pub struct MuxStream<W>
where
W: crate::ws::WebSocketWrite,
W: ws::WebSocketWrite,
{
/// ID of the stream.
pub stream_id: u32,
@ -189,18 +195,18 @@ where
tx: MuxStreamWrite<W>,
}
impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStream<W> {
impl<W: ws::WebSocketWrite + Send + 'static> MuxStream<W> {
#[allow(clippy::too_many_arguments)]
pub(crate) fn new(
stream_id: u32,
role: crate::Role,
stream_type: crate::StreamType,
role: Role,
stream_type: StreamType,
rx: mpsc::UnboundedReceiver<MuxEvent>,
tx: crate::ws::LockedWebSocketWrite<W>,
tx: ws::LockedWebSocketWrite<W>,
close_channel: mpsc::UnboundedSender<WsEvent>,
is_closed: Arc<AtomicBool>,
flow_control: Arc<AtomicU32>,
continue_recieved: Arc<Event>
continue_recieved: Arc<Event>,
) -> Self {
Self {
stream_id,
@ -215,6 +221,7 @@ impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStream<W> {
},
tx: MuxStreamWrite {
stream_id,
stream_type,
role,
tx,
close_channel,
@ -231,7 +238,7 @@ impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStream<W> {
}
/// Write data to the stream.
pub async fn write(&self, data: Bytes) -> Result<(), crate::WispError> {
pub async fn write(&self, data: Bytes) -> Result<(), WispError> {
self.tx.write(data).await
}
@ -251,7 +258,7 @@ impl<W: crate::ws::WebSocketWrite + Send + 'static> MuxStream<W> {
}
/// Close the stream. You will no longer be able to write or read after this has been called.
pub async fn close(&self, reason: crate::CloseReason) -> Result<(), crate::WispError> {
pub async fn close(&self, reason: CloseReason) -> Result<(), WispError> {
self.tx.close(reason).await
}
@ -280,29 +287,28 @@ pub struct MuxStreamCloser {
impl MuxStreamCloser {
/// Close the stream. You will no longer be able to write or read after this has been called.
pub async fn close(&self, reason: crate::CloseReason) -> Result<(), crate::WispError> {
pub async fn close(&self, reason: CloseReason) -> Result<(), WispError> {
if self.is_closed.load(Ordering::Acquire) {
return Err(crate::WispError::StreamAlreadyClosed);
return Err(WispError::StreamAlreadyClosed);
}
let (tx, rx) = oneshot::channel::<Result<(), crate::WispError>>();
let (tx, rx) = oneshot::channel::<Result<(), WispError>>();
self.close_channel
.unbounded_send(WsEvent::Close(self.stream_id, reason, tx))
.map_err(|x| crate::WispError::Other(Box::new(x)))?;
rx.await
.map_err(|x| crate::WispError::Other(Box::new(x)))??;
.map_err(|x| WispError::Other(Box::new(x)))?;
rx.await.map_err(|x| WispError::Other(Box::new(x)))??;
self.is_closed.store(true, Ordering::Release);
Ok(())
}
/// Close the stream. This function does not check if it was actually closed.
pub(crate) fn close_sync(&self, reason: crate::CloseReason) -> Result<(), crate::WispError> {
pub(crate) fn close_sync(&self, reason: CloseReason) -> Result<(), WispError> {
if self.is_closed.load(Ordering::Acquire) {
return Err(crate::WispError::StreamAlreadyClosed);
return Err(WispError::StreamAlreadyClosed);
}
let (tx, _) = oneshot::channel::<Result<(), crate::WispError>>();
let (tx, _) = oneshot::channel::<Result<(), WispError>>();
self.close_channel
.unbounded_send(WsEvent::Close(self.stream_id, reason, tx))
.map_err(|x| crate::WispError::Other(Box::new(x)))?;
.map_err(|x| WispError::Other(Box::new(x)))?;
self.is_closed.store(true, Ordering::Release);
Ok(())
}
@ -314,7 +320,7 @@ pin_project! {
#[pin]
rx: Pin<Box<dyn Stream<Item = Bytes> + Send>>,
#[pin]
tx: Pin<Box<dyn Sink<Bytes, Error = crate::WispError> + Send>>,
tx: Pin<Box<dyn Sink<Bytes, Error = WispError> + Send>>,
}
}