Detailed changes
@@ -286,6 +286,19 @@ dependencies = [
"syn",
]
+[[package]]
+name = "async-tungstenite"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8645e929ec7964448a901db9da30cd2ae8c7fecf4d6176af427837531dbbb63b"
+dependencies = [
+ "futures-io",
+ "futures-util",
+ "log",
+ "pin-project-lite",
+ "tungstenite",
+]
+
[[package]]
name = "atomic"
version = "0.5.0"
@@ -1713,6 +1726,12 @@ dependencies = [
"url",
]
+[[package]]
+name = "httparse"
+version = "1.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f3a87b616e37e93c22fb19bcd386f02f3af5ea98a25670ad0fce773de23c5e68"
+
[[package]]
name = "humantime"
version = "2.1.0"
@@ -1806,6 +1825,15 @@ dependencies = [
"adler32",
]
+[[package]]
+name = "input_buffer"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f97967975f448f1a7ddb12b0bc41069d09ed6a1c161a92687e057325db35d413"
+dependencies = [
+ "bytes 1.0.1",
+]
+
[[package]]
name = "instant"
version = "0.1.9"
@@ -3304,6 +3332,19 @@ dependencies = [
"pkg-config",
]
+[[package]]
+name = "sha-1"
+version = "0.9.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8c4cfa741c5832d0ef7fab46cabed29c2aae926db0b11bb2069edd8db5e64e16"
+dependencies = [
+ "block-buffer",
+ "cfg-if 1.0.0",
+ "cpufeatures",
+ "digest",
+ "opaque-debug",
+]
+
[[package]]
name = "sha1"
version = "0.2.0"
@@ -3920,6 +3961,26 @@ version = "0.12.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "85e00391c1f3d171490a3f8bd79999b0002ae38d3da0d6a3a306c754b053d71b"
+[[package]]
+name = "tungstenite"
+version = "0.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5fe8dada8c1a3aeca77d6b51a4f1314e0f4b8e438b7b1b71e3ddaca8080e4093"
+dependencies = [
+ "base64 0.13.0",
+ "byteorder",
+ "bytes 1.0.1",
+ "http",
+ "httparse",
+ "input_buffer",
+ "log",
+ "rand 0.8.3",
+ "sha-1",
+ "thiserror",
+ "url",
+ "utf-8",
+]
+
[[package]]
name = "typenum"
version = "1.13.0"
@@ -4057,6 +4118,12 @@ dependencies = [
"xmlwriter",
]
+[[package]]
+name = "utf-8"
+version = "0.7.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9"
+
[[package]]
name = "uuid"
version = "0.5.1"
@@ -4356,6 +4423,7 @@ version = "0.1.0"
dependencies = [
"anyhow",
"async-lock",
+ "async-tungstenite",
"base64 0.13.0",
"futures",
"log",
@@ -7,6 +7,7 @@ version = "0.1.0"
[dependencies]
anyhow = "1.0"
async-lock = "2.4"
+async-tungstenite = "0.14"
base64 = "0.13"
futures = "0.3"
log = "0.4"
@@ -2,5 +2,7 @@ pub mod auth;
mod peer;
pub mod proto;
pub mod rest;
+#[cfg(test)]
+mod test;
pub use peer::*;
@@ -1,7 +1,12 @@
use crate::proto::{self, EnvelopedMessage, MessageStream, RequestMessage};
use anyhow::{anyhow, Context, Result};
use async_lock::{Mutex, RwLock};
-use futures::{future::BoxFuture, AsyncRead, AsyncWrite, FutureExt};
+use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
+use futures::{
+ future::BoxFuture,
+ stream::{SplitSink, SplitStream},
+ FutureExt, StreamExt,
+};
use postage::{
mpsc,
prelude::{Sink, Stream},
@@ -72,13 +77,13 @@ struct Connection {
response_channels: ResponseChannels,
}
-pub struct ConnectionHandler<Conn> {
+pub struct ConnectionHandler<W, R> {
peer: Arc<Peer>,
connection_id: ConnectionId,
response_channels: ResponseChannels,
outgoing_rx: mpsc::Receiver<proto::Envelope>,
- reader: MessageStream<Conn>,
- writer: MessageStream<Conn>,
+ writer: MessageStream<W>,
+ reader: MessageStream<R>,
}
type ResponseChannels = Arc<Mutex<HashMap<u32, mpsc::Sender<proto::Envelope>>>>;
@@ -131,10 +136,16 @@ impl Peer {
pub async fn add_connection<Conn>(
self: &Arc<Self>,
conn: Conn,
- ) -> (ConnectionId, ConnectionHandler<Conn>)
+ ) -> (
+ ConnectionId,
+ ConnectionHandler<SplitSink<Conn, WebSocketMessage>, SplitStream<Conn>>,
+ )
where
- Conn: Clone + AsyncRead + AsyncWrite + Unpin + Send + 'static,
+ Conn: futures::Sink<WebSocketMessage, Error = WebSocketError>
+ + futures::Stream<Item = Result<WebSocketMessage, WebSocketError>>
+ + Unpin,
{
+ let (tx, rx) = conn.split();
let connection_id = ConnectionId(
self.next_connection_id
.fetch_add(1, atomic::Ordering::SeqCst),
@@ -150,8 +161,8 @@ impl Peer {
connection_id,
response_channels: connection.response_channels.clone(),
outgoing_rx,
- reader: MessageStream::new(conn.clone()),
- writer: MessageStream::new(conn),
+ writer: MessageStream::new(tx),
+ reader: MessageStream::new(rx),
};
self.connections
.write()
@@ -291,9 +302,10 @@ impl Peer {
}
}
-impl<Conn> ConnectionHandler<Conn>
+impl<W, R> ConnectionHandler<W, R>
where
- Conn: Clone + AsyncRead + AsyncWrite + Unpin + Send + 'static,
+ W: futures::Sink<WebSocketMessage, Error = WebSocketError> + Unpin,
+ R: futures::Stream<Item = Result<WebSocketMessage, WebSocketError>> + Unpin,
{
pub async fn run(mut self) -> Result<()> {
loop {
@@ -402,38 +414,25 @@ impl fmt::Display for PeerId {
#[cfg(test)]
mod tests {
use super::*;
+ use crate::test;
use postage::oneshot;
- use smol::{
- io::AsyncWriteExt,
- net::unix::{UnixListener, UnixStream},
- };
- use std::io;
- use tempdir::TempDir;
#[test]
fn test_request_response() {
smol::block_on(async move {
- // create socket
- let socket_dir_path = TempDir::new("test-request-response").unwrap();
- let socket_path = socket_dir_path.path().join("test.sock");
- let listener = UnixListener::bind(&socket_path).unwrap();
-
// create 2 clients connected to 1 server
let server = Peer::new();
let client1 = Peer::new();
let client2 = Peer::new();
- let (client1_conn_id, task1) = client1
- .add_connection(UnixStream::connect(&socket_path).await.unwrap())
- .await;
- let (client2_conn_id, task2) = client2
- .add_connection(UnixStream::connect(&socket_path).await.unwrap())
- .await;
- let (_, task3) = server
- .add_connection(listener.accept().await.unwrap().0)
- .await;
- let (_, task4) = server
- .add_connection(listener.accept().await.unwrap().0)
- .await;
+
+ let (client1_to_server_conn, server_to_client_1_conn) = test::Channel::bidirectional();
+ let (client1_conn_id, task1) = client1.add_connection(client1_to_server_conn).await;
+ let (_, task2) = server.add_connection(server_to_client_1_conn).await;
+
+ let (client2_to_server_conn, server_to_client_2_conn) = test::Channel::bidirectional();
+ let (client2_conn_id, task3) = client2.add_connection(client2_to_server_conn).await;
+ let (_, task4) = server.add_connection(server_to_client_2_conn).await;
+
smol::spawn(task1.run()).detach();
smol::spawn(task2.run()).detach();
smol::spawn(task3.run()).detach();
@@ -553,11 +552,7 @@ mod tests {
#[test]
fn test_disconnect() {
smol::block_on(async move {
- let socket_dir_path = TempDir::new("drop-client").unwrap();
- let socket_path = socket_dir_path.path().join(".sock");
- let listener = UnixListener::bind(&socket_path).unwrap();
- let client_conn = UnixStream::connect(&socket_path).await.unwrap();
- let (mut server_conn, _) = listener.accept().await.unwrap();
+ let (client_conn, mut server_conn) = test::Channel::bidirectional();
let client = Peer::new();
let (connection_id, handler) = client.add_connection(client_conn).await;
@@ -571,20 +566,19 @@ mod tests {
client.disconnect(connection_id).await;
incoming_messages_ended_rx.recv().await;
-
- let err = server_conn.write(&[]).await.unwrap_err();
- assert_eq!(err.kind(), io::ErrorKind::BrokenPipe);
+ assert!(
+ futures::SinkExt::send(&mut server_conn, WebSocketMessage::Binary(vec![]))
+ .await
+ .is_err()
+ );
});
}
#[test]
fn test_io_error() {
smol::block_on(async move {
- let socket_dir_path = TempDir::new("io-error").unwrap();
- let socket_path = socket_dir_path.path().join(".sock");
- let _listener = UnixListener::bind(&socket_path).unwrap();
- let mut client_conn = UnixStream::connect(&socket_path).await.unwrap();
- client_conn.close().await.unwrap();
+ let (client_conn, server_conn) = test::Channel::bidirectional();
+ drop(server_conn);
let client = Peer::new();
let (connection_id, handler) = client.add_connection(client_conn).await;
@@ -1,7 +1,7 @@
-use futures::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt as _};
+use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
+use futures::{SinkExt as _, StreamExt as _};
use prost::Message;
use std::{
- convert::TryInto,
io,
time::{Duration, SystemTime, UNIX_EPOCH},
};
@@ -81,66 +81,52 @@ message!(AddPeer);
message!(RemovePeer);
/// A stream of protobuf messages.
-pub struct MessageStream<T> {
- byte_stream: T,
- buffer: Vec<u8>,
- upcoming_message_len: Option<usize>,
+pub struct MessageStream<S> {
+ stream: S,
}
-impl<T> MessageStream<T> {
- pub fn new(byte_stream: T) -> Self {
- Self {
- byte_stream,
- buffer: Default::default(),
- upcoming_message_len: None,
- }
+impl<S> MessageStream<S> {
+ pub fn new(stream: S) -> Self {
+ Self { stream }
}
- pub fn inner_mut(&mut self) -> &mut T {
- &mut self.byte_stream
+ pub fn inner_mut(&mut self) -> &mut S {
+ &mut self.stream
}
}
-impl<T> MessageStream<T>
+impl<S> MessageStream<S>
where
- T: AsyncWrite + Unpin,
+ S: futures::Sink<WebSocketMessage, Error = WebSocketError> + Unpin,
{
/// Write a given protobuf message to the stream.
- pub async fn write_message(&mut self, message: &Envelope) -> io::Result<()> {
- let message_len: u32 = message
- .encoded_len()
- .try_into()
- .map_err(|_| io::Error::new(io::ErrorKind::InvalidData, "message is too large"))?;
- self.buffer.clear();
- self.buffer.extend_from_slice(&message_len.to_be_bytes());
- message.encode(&mut self.buffer)?;
- self.byte_stream.write_all(&self.buffer).await
+ pub async fn write_message(&mut self, message: &Envelope) -> Result<(), WebSocketError> {
+ let mut buffer = Vec::with_capacity(message.encoded_len());
+ message
+ .encode(&mut buffer)
+ .map_err(|err| io::Error::from(err))?;
+ self.stream.send(WebSocketMessage::Binary(buffer)).await?;
+ Ok(())
}
}
-impl<T> MessageStream<T>
+impl<S> MessageStream<S>
where
- T: AsyncRead + Unpin,
+ S: futures::Stream<Item = Result<WebSocketMessage, WebSocketError>> + Unpin,
{
/// Read a protobuf message of the given type from the stream.
- pub async fn read_message(&mut self) -> io::Result<Envelope> {
- loop {
- if let Some(upcoming_message_len) = self.upcoming_message_len {
- self.buffer.resize(upcoming_message_len, 0);
- self.byte_stream.read_exact(&mut self.buffer).await?;
- self.upcoming_message_len = None;
- return Ok(Envelope::decode(self.buffer.as_slice())?);
- } else {
- self.buffer.resize(4, 0);
- self.byte_stream.read_exact(&mut self.buffer).await?;
- self.upcoming_message_len = Some(u32::from_be_bytes([
- self.buffer[0],
- self.buffer[1],
- self.buffer[2],
- self.buffer[3],
- ]) as usize);
+ pub async fn read_message(&mut self) -> Result<Envelope, WebSocketError> {
+ while let Some(bytes) = self.stream.next().await {
+ match bytes? {
+ WebSocketMessage::Binary(bytes) => {
+ let envelope = Envelope::decode(bytes.as_slice()).map_err(io::Error::from)?;
+ return Ok(envelope);
+ }
+ WebSocketMessage::Close(_) => break,
+ _ => {}
}
}
+ Err(WebSocketError::ConnectionClosed)
}
}
@@ -165,20 +151,12 @@ impl From<SystemTime> for Timestamp {
#[cfg(test)]
mod tests {
use super::*;
- use std::{
- pin::Pin,
- task::{Context, Poll},
- };
+ use crate::test;
#[test]
fn test_round_trip_message() {
smol::block_on(async {
- let byte_stream = ChunkedStream {
- bytes: Vec::new(),
- read_offset: 0,
- chunk_size: 3,
- };
-
+ let stream = test::Channel::new();
let message1 = Auth {
user_id: 5,
access_token: "the-access-token".into(),
@@ -191,7 +169,7 @@ mod tests {
}
.into_envelope(5, None, None);
- let mut message_stream = MessageStream::new(byte_stream);
+ let mut message_stream = MessageStream::new(stream);
message_stream.write_message(&message1).await.unwrap();
message_stream.write_message(&message2).await.unwrap();
let decoded_message1 = message_stream.read_message().await.unwrap();
@@ -200,47 +178,4 @@ mod tests {
assert_eq!(decoded_message2, message2);
});
}
-
- struct ChunkedStream {
- bytes: Vec<u8>,
- read_offset: usize,
- chunk_size: usize,
- }
-
- impl AsyncWrite for ChunkedStream {
- fn poll_write(
- mut self: Pin<&mut Self>,
- _: &mut Context<'_>,
- buf: &[u8],
- ) -> Poll<io::Result<usize>> {
- let bytes_written = buf.len().min(self.chunk_size);
- self.bytes.extend_from_slice(&buf[0..bytes_written]);
- Poll::Ready(Ok(bytes_written))
- }
-
- fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
- Poll::Ready(Ok(()))
- }
-
- fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<io::Result<()>> {
- Poll::Ready(Ok(()))
- }
- }
-
- impl AsyncRead for ChunkedStream {
- fn poll_read(
- mut self: Pin<&mut Self>,
- _: &mut Context<'_>,
- buf: &mut [u8],
- ) -> Poll<io::Result<usize>> {
- let bytes_read = buf
- .len()
- .min(self.chunk_size)
- .min(self.bytes.len() - self.read_offset);
- let end_offset = self.read_offset + bytes_read;
- buf[0..bytes_read].copy_from_slice(&self.bytes[self.read_offset..end_offset]);
- self.read_offset = end_offset;
- Poll::Ready(Ok(bytes_read))
- }
- }
}
@@ -0,0 +1,64 @@
+use async_tungstenite::tungstenite::{Error as WebSocketError, Message as WebSocketMessage};
+use std::{
+ io,
+ pin::Pin,
+ task::{Context, Poll},
+};
+
+pub struct Channel {
+ tx: futures::channel::mpsc::UnboundedSender<WebSocketMessage>,
+ rx: futures::channel::mpsc::UnboundedReceiver<WebSocketMessage>,
+}
+
+impl Channel {
+ pub fn new() -> Self {
+ let (tx, rx) = futures::channel::mpsc::unbounded();
+ Self { tx, rx }
+ }
+
+ pub fn bidirectional() -> (Self, Self) {
+ let (a_tx, a_rx) = futures::channel::mpsc::unbounded();
+ let (b_tx, b_rx) = futures::channel::mpsc::unbounded();
+ let a = Self { tx: a_tx, rx: b_rx };
+ let b = Self { tx: b_tx, rx: a_rx };
+ (a, b)
+ }
+}
+
+impl futures::Sink<WebSocketMessage> for Channel {
+ type Error = WebSocketError;
+
+ fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+ Pin::new(&mut self.tx)
+ .poll_ready(cx)
+ .map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
+ }
+
+ fn start_send(mut self: Pin<&mut Self>, item: WebSocketMessage) -> Result<(), Self::Error> {
+ Pin::new(&mut self.tx)
+ .start_send(item)
+ .map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
+ }
+
+ fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+ Pin::new(&mut self.tx)
+ .poll_flush(cx)
+ .map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
+ }
+
+ fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+ Pin::new(&mut self.tx)
+ .poll_close(cx)
+ .map_err(|err| io::Error::new(io::ErrorKind::Other, err).into())
+ }
+}
+
+impl futures::Stream for Channel {
+ type Item = Result<WebSocketMessage, WebSocketError>;
+
+ fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+ Pin::new(&mut self.rx)
+ .poll_next(cx)
+ .map(|i| i.map(|i| Ok(i)))
+ }
+}