fix udp

wisp/Cargo.toml

@@ -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"

wisp/src/lib.rs

@@ -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(),
             },

wisp/src/packet.rs

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

wisp/src/stream.rs

@@ -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,
+    pub stream_type: StreamType,
-    role: crate::Role,
+    role: Role,
-    tx: crate::ws::LockedWebSocketWrite<W>,
+    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(
+                        .write_frame(Packet::new_continue(self.stream_id, old_val + 1).into())
-                            crate::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,
+    /// Type of the stream.
-    tx: crate::ws::LockedWebSocketWrite<W>,
+    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
+                self.flow_control.load(Ordering::Acquire).saturating_sub(1),
-                    .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)))?;
+            .map_err(|x| WispError::Other(Box::new(x)))?;
-        rx.await
+        rx.await.map_err(|x| WispError::Other(Box::new(x)))??;
-            .map_err(|x| crate::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 (tx, _) = oneshot::channel::<Result<(), WispError>>();
-        let _ = self
+        let _ = self.close_channel.unbounded_send(WsEvent::Close(
-            .close_channel
+            self.stream_id,
-            .unbounded_send(WsEvent::Close(self.stream_id, crate::CloseReason::Unknown, tx));
+            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,
+        role: Role,
-        stream_type: crate::StreamType,
+        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)))?;
+            .map_err(|x| WispError::Other(Box::new(x)))?;
-        rx.await
+        rx.await.map_err(|x| WispError::Other(Box::new(x)))??;
-            .map_err(|x| crate::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>>,
     }
 }