Add a ByteReader wrapper that implements AsyncRead for a Stream of messages

This is useful for programs that want to treat a WebSocket as a stream
of bytes.

Author: joshtriplett

examples/client-bytes.rs

@@ -15,10 +15,9 @@ use std::env;
 use futures::StreamExt;
 
 use async_std::io;
-use async_std::prelude::*;
 use async_std::task;
 use async_tungstenite::async_std::connect_async;
-use async_tungstenite::ByteWriter;
+use async_tungstenite::{ByteReader, ByteWriter};
 
 async fn run() {
     let connect_addr = env::args()
@@ -32,15 +31,11 @@ async fn run() {
 
     let (write, read) = ws_stream.split();
     let byte_writer = ByteWriter::new(write);
-    let stdin_to_ws = task::spawn(async {
-        io::copy(io::stdin(), byte_writer).await.unwrap();
-    });
-    let ws_to_stdout = task::spawn(read.for_each(|message| async {
-        let data = message.unwrap().into_data();
-        async_std::io::stdout().write_all(&data).await.unwrap();
-    }));
-    stdin_to_ws.await;
-    ws_to_stdout.await;
+    let byte_reader = ByteReader::new(read);
+    let stdin_to_ws = task::spawn(io::copy(io::stdin(), byte_writer));
+    let ws_to_stdout = task::spawn(io::copy(byte_reader, io::stdout()));
+    stdin_to_ws.await.unwrap();
+    ws_to_stdout.await.unwrap();
 }
 
 fn main() {

src/bytes.rs

@@ -1,22 +1,24 @@
-//! Provides an abstraction to use `AsyncWrite` to write bytes to a `WebSocketStream`.
+//! Provides abstractions to use `AsyncRead` and `AsyncWrite` with a `WebSocketStream`.
 
 use std::{
     io,
     pin::Pin,
     task::{Context, Poll},
 };
 
-use futures_util::Sink;
+use futures_core::stream::Stream;
 
-use crate::{Message, WsError};
+use crate::{tungstenite::Bytes, Message, WsError};
 
 /// Treat a `WebSocketStream` as an `AsyncWrite` implementation.
 ///
 /// Every write sends a binary message. If you want to group writes together, consider wrapping
 /// this with a `BufWriter`.
+#[cfg(feature = "futures-03-sink")]
 #[derive(Debug)]
 pub struct ByteWriter<S>(S);
 
+#[cfg(feature = "futures-03-sink")]
 impl<S> ByteWriter<S> {
     /// Create a new `ByteWriter` from a `Sink` that accepts a WebSocket `Message`
     #[inline(always)]
@@ -31,13 +33,14 @@ impl<S> ByteWriter<S> {
     }
 }
 
+#[cfg(feature = "futures-03-sink")]
 fn poll_write_helper<S>(
     mut s: Pin<&mut ByteWriter<S>>,
     cx: &mut Context<'_>,
     buf: &[u8],
 ) -> Poll<io::Result<usize>>
 where
-    S: Sink<Message, Error = WsError> + Unpin,
+    S: futures_util::Sink<Message, Error = WsError> + Unpin,
 {
     match Pin::new(&mut s.0).poll_ready(cx).map_err(convert_err) {
         Poll::Ready(Ok(())) => {}
@@ -54,9 +57,10 @@ where
     )
 }
 
+#[cfg(feature = "futures-03-sink")]
 impl<S> futures_io::AsyncWrite for ByteWriter<S>
 where
-    S: Sink<Message, Error = WsError> + Unpin,
+    S: futures_util::Sink<Message, Error = WsError> + Unpin,
 {
     fn poll_write(
         self: Pin<&mut Self>,
@@ -75,10 +79,11 @@ where
     }
 }
 
+#[cfg(feature = "futures-03-sink")]
 #[cfg(feature = "tokio-runtime")]
 impl<S> tokio::io::AsyncWrite for ByteWriter<S>
 where
-    S: Sink<Message, Error = WsError> + Unpin,
+    S: futures_util::Sink<Message, Error = WsError> + Unpin,
 {
     fn poll_write(
         self: Pin<&mut Self>,
@@ -97,6 +102,65 @@ where
     }
 }
 
+/// Treat a `WebSocketStream` as an `AsyncRead` implementation.
+///
+/// This also works with any other `Stream` of `Message`, such as a `SplitStream`.
+///
+/// Each read will only return data from one message. If you want to combine data from multiple
+/// messages into one read, consider wrapping this in a `BufReader`.
+#[derive(Debug)]
+pub struct ByteReader<S> {
+    stream: S,
+    bytes: Option<Bytes>,
+}
+
+impl<S> ByteReader<S> {
+    /// Create a new `ByteReader` from a `Stream` that returns a WebSocket `Message`
+    #[inline(always)]
+    pub fn new(stream: S) -> Self {
+        Self {
+            stream,
+            bytes: None,
+        }
+    }
+}
+
+impl<S> futures_io::AsyncRead for ByteReader<S>
+where
+    S: Stream<Item = Result<Message, WsError>> + Unpin,
+{
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut [u8],
+    ) -> Poll<io::Result<usize>> {
+        let buf_len = buf.len();
+        let bytes_to_read = match self.bytes {
+            None => match Pin::new(&mut self.stream).poll_next(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(None) => return Poll::Ready(Ok(0)),
+                Poll::Ready(Some(Err(e))) => return Poll::Ready(Err(convert_err(e))),
+                Poll::Ready(Some(Ok(msg))) => {
+                    let bytes = msg.into_data();
+                    if bytes.len() > buf_len {
+                        self.bytes.insert(bytes).split_to(buf_len)
+                    } else {
+                        bytes
+                    }
+                }
+            },
+            Some(ref mut bytes) if bytes.len() > buf_len => bytes.split_to(buf_len),
+            Some(ref mut bytes) => {
+                let bytes = bytes.clone();
+                self.bytes = None;
+                bytes
+            }
+        };
+        buf.copy_from_slice(&bytes_to_read);
+        Poll::Ready(Ok(bytes_to_read.len()))
+    }
+}
+
 fn convert_err(e: WsError) -> io::Error {
     match e {
         WsError::Io(io) => io,

src/lib.rs

@@ -92,8 +92,8 @@ pub mod gio;
 #[cfg(feature = "tokio-runtime")]
 pub mod tokio;
 
-#[cfg(feature = "futures-03-sink")]
 pub mod bytes;
+pub use bytes::ByteReader;
 #[cfg(feature = "futures-03-sink")]
 pub use bytes::ByteWriter;