diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 7ba1638..2703011 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -29,11 +29,11 @@ jobs: - name: "Clippy" cmd: clippy args: --workspace --all-features --all-targets -- -D warnings - rust: stable + rust: nightly - name: "Formatting" cmd: fmt args: --all -- --check - rust: stable + rust: nightly - name: "Tests" cmd: nextest args: run --workspace --all-features --retries 3 diff --git a/Cargo.toml b/Cargo.toml index 3d6a648..f5aefd4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -23,6 +23,7 @@ full = [ "geoip", "metrics", "rate_limit", + "websocket" ] alloc = ["dep:alloc"] analytics = ["dep:analytics"] @@ -35,6 +36,7 @@ metrics_exporter_prometheus = ["metrics/exporter_prometheus"] alloc_metrics = ["alloc/metrics"] profiler = ["alloc/profiler"] rate_limit = ["dep:rate_limit"] +websocket = ["dep:websocket"] [workspace.dependencies] aws-sdk-s3 = "1.21.0" @@ -47,6 +49,7 @@ future = { path = "./crates/future", optional = true } geoip = { path = "./crates/geoip", optional = true } metrics = { package = "wc_metrics", path = "./crates/metrics", optional = true } rate_limit = { path = "./crates/rate_limit", optional = true } +websocket = { path = "./crates/websocket", optional = true } [dev-dependencies] anyhow = "1" diff --git a/clippy.toml b/clippy.toml index 7fb6452..48125e0 100644 --- a/clippy.toml +++ b/clippy.toml @@ -1,4 +1,5 @@ await-holding-invalid-types = [ - "tracing::trace::Entered", - "tracing::trace::EnteredSpan", + "tracing::span::Entered", + "tracing::span::EnteredSpan", ] +absolute-paths-max-segments = 4 diff --git a/crates/websocket/Cargo.toml b/crates/websocket/Cargo.toml new file mode 100644 index 0000000..0a112da --- /dev/null +++ b/crates/websocket/Cargo.toml @@ -0,0 +1,38 @@ +[package] +name = "websocket" +version = "0.1.0" +edition = "2024" + +[features] +default = [] +full = ["axum", "tungstenite", "warp", "json"] +axum = ["dep:axum"] +tungstenite = ["dep:tungstenite", "dep:tokio-tungstenite"] +warp = ["dep:warp"] +json = ["dep:serde", "dep:serde_json"] + +[dependencies] +axum = { optional = true, version = "0.8", features = ["ws"] } +tokio-tungstenite = { optional = true, version = "0.28" } +tungstenite = { optional = true, version = "0.28" } +warp = { optional = true, version = "0.4", features = ["websocket"] } +serde = { optional = true, version = "1.0" } +serde_json = { optional = true, version = "1.0" } +futures-concurrency = "7.6" +futures-timer = "3.0" +futures-util = "0.3" +pin-project = "1.1" +tap = "1.0" +thiserror = "2.0" +tokio = { version = "1.48", default-features = false, features = ["rt", "time"] } +tokio-stream = "0.1" +tokio-util = "0.7" +enum-as-inner = "0.6" +bytes = "1.11" +derive_more = { version = "2.1", features = ["into", "from"] } + +[dev-dependencies] +tokio = { version = "1.48", features = ["macros"] } +serde = { version = "1.0", features = ["derive"] } +tokio-tungstenite = "0.28" +websocket = { path = ".", features = ["tungstenite", "json"] } diff --git a/crates/websocket/src/backend.rs b/crates/websocket/src/backend.rs new file mode 100644 index 0000000..51ffe47 --- /dev/null +++ b/crates/websocket/src/backend.rs @@ -0,0 +1,6 @@ +#[cfg(feature = "axum")] +mod axum; +#[cfg(feature = "tungstenite")] +mod tungstenite; +#[cfg(feature = "warp")] +mod warp; diff --git a/crates/websocket/src/backend/axum.rs b/crates/websocket/src/backend/axum.rs new file mode 100644 index 0000000..18aeff6 --- /dev/null +++ b/crates/websocket/src/backend/axum.rs @@ -0,0 +1,61 @@ +use { + crate::{Backend, CloseFrame, Message}, + axum::{ + Error as NativeError, + body::Bytes, + extract::ws::{ + CloseFrame as NativeCloseFrame, + Message as NativeMessage, + Utf8Bytes, + WebSocket as NativeWebSocket, + }, + }, +}; + +impl Backend for NativeWebSocket { + type Error = NativeError; + type Message = NativeMessage; + type Transport = Self; + + fn into_transport(self) -> Self::Transport { + self + } + + fn encode_message(msg: Message) -> Self::Message { + match msg { + Message::Binary(data) => NativeMessage::binary(data), + Message::Text(data) => NativeMessage::text(data), + Message::Ping(data) => NativeMessage::Ping(data), + Message::Pong(data) => NativeMessage::Pong(data), + Message::Close(msg) => { + let frame = msg.map(|frame| NativeCloseFrame { + code: frame.code, + reason: frame.reason.into(), + }); + + NativeMessage::Close(frame) + } + } + } + + fn decode_message(msg: Self::Message) -> Message { + match msg { + NativeMessage::Binary(data) => Message::Binary(data), + NativeMessage::Text(data) => Message::Text(bytes_to_string(data)), + NativeMessage::Ping(data) => Message::Ping(data), + NativeMessage::Pong(data) => Message::Pong(data), + NativeMessage::Close(frame) => { + let frame = frame.map(|frame| CloseFrame { + code: frame.code, + reason: bytes_to_string(frame.reason), + }); + + Message::Close(frame) + } + } + } +} + +fn bytes_to_string(data: Utf8Bytes) -> String { + String::from_utf8(Bytes::from(data).into()).unwrap_or_default() +} diff --git a/crates/websocket/src/backend/tungstenite.rs b/crates/websocket/src/backend/tungstenite.rs new file mode 100644 index 0000000..1496778 --- /dev/null +++ b/crates/websocket/src/backend/tungstenite.rs @@ -0,0 +1,64 @@ +use { + crate::{Backend, CloseFrame, Message}, + bytes::Bytes, + tokio::io::{AsyncRead, AsyncWrite}, + tokio_tungstenite::WebSocketStream, + tungstenite::{ + Error as NativeError, + Message as NativeMessage, + Utf8Bytes, + protocol::CloseFrame as NativeCloseFrame, + }, +}; + +impl Backend for WebSocketStream +where + T: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + type Error = NativeError; + type Message = NativeMessage; + type Transport = Self; + + fn into_transport(self) -> Self::Transport { + self + } + + fn encode_message(msg: Message) -> Self::Message { + match msg { + Message::Binary(data) => NativeMessage::binary(data), + Message::Text(data) => NativeMessage::text(data), + Message::Ping(data) => NativeMessage::Ping(data), + Message::Pong(data) => NativeMessage::Pong(data), + Message::Close(msg) => { + let frame = msg.map(|frame| NativeCloseFrame { + code: frame.code.into(), + reason: frame.reason.into(), + }); + + NativeMessage::Close(frame) + } + } + } + + fn decode_message(msg: Self::Message) -> Message { + match msg { + NativeMessage::Binary(data) => Message::Binary(data), + NativeMessage::Text(data) => Message::Text(bytes_to_string(data)), + NativeMessage::Ping(data) => Message::Ping(data), + NativeMessage::Pong(data) => Message::Pong(data), + NativeMessage::Close(frame) => { + let frame = frame.map(|frame| CloseFrame { + code: frame.code.into(), + reason: bytes_to_string(frame.reason), + }); + + Message::Close(frame) + } + NativeMessage::Frame(_) => Message::Close(None), + } + } +} + +fn bytes_to_string(data: Utf8Bytes) -> String { + String::from_utf8(Bytes::from(data).into()).unwrap_or_default() +} diff --git a/crates/websocket/src/backend/warp.rs b/crates/websocket/src/backend/warp.rs new file mode 100644 index 0000000..6af7d0b --- /dev/null +++ b/crates/websocket/src/backend/warp.rs @@ -0,0 +1,49 @@ +use { + crate::{Backend, CloseFrame, Message}, + warp::{ + Error as NativeError, + ws::{Message as NativeMessage, WebSocket as NativeWebSocket}, + }, +}; + +impl Backend for NativeWebSocket { + type Error = NativeError; + type Message = NativeMessage; + type Transport = Self; + + fn into_transport(self) -> Self::Transport { + self + } + + fn encode_message(msg: Message) -> Self::Message { + match msg { + Message::Binary(data) => NativeMessage::binary(data), + Message::Text(data) => NativeMessage::text(data), + Message::Ping(data) => NativeMessage::ping(data), + Message::Pong(data) => NativeMessage::pong(data), + Message::Close(msg) => match msg { + Some(msg) => NativeMessage::close_with(msg.code, msg.reason), + None => NativeMessage::close(), + }, + } + } + + fn decode_message(msg: Self::Message) -> Message { + if msg.is_binary() { + Message::Binary(msg.into_bytes()) + } else if msg.is_text() { + Message::Text(String::from_utf8(msg.into_bytes().into()).unwrap_or_default()) + } else if msg.is_ping() { + Message::Ping(msg.into_bytes()) + } else if msg.is_pong() { + Message::Pong(msg.into_bytes()) + } else { + let msg = msg.close_frame().map(|(code, msg)| CloseFrame { + code, + reason: msg.to_owned(), + }); + + Message::Close(msg) + } + } +} diff --git a/crates/websocket/src/lib.rs b/crates/websocket/src/lib.rs new file mode 100644 index 0000000..e2b510e --- /dev/null +++ b/crates/websocket/src/lib.rs @@ -0,0 +1,384 @@ +#[cfg(feature = "json")] +pub use json::Json; +#[cfg(feature = "tungstenite")] +pub use tokio_tungstenite; +use { + crate::wrapper::Config, + derive_more::{From, Into}, + enum_as_inner::EnumAsInner, + futures_util::{Sink, Stream}, + std::{error::Error as StdError, time::Duration}, +}; +pub use {bytes::Bytes, wrapper::WebSocket}; + +mod backend; +mod transport; +mod wrapper; + +type BoxError = Box; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Encoding failed: {0}")] + Encoding(BoxError), + + #[error("Decoding failed: {0}")] + Decoding(BoxError), + + #[error("Invalid payload: {0}")] + InvalidPayload(BoxError), + + #[error("Transport is closed")] + Closed, + + #[error("Transport error: {0}")] + Transport(BoxError), + + #[error("Internal error: {0}")] + Internal(BoxError), +} + +impl Error { + pub fn encoding(err: T) -> Self { + Self::Encoding(Box::new(err)) + } + + pub fn decoding(err: T) -> Self { + Self::Decoding(Box::new(err)) + } + + pub fn transport(err: T) -> Self { + Self::Transport(Box::new(err)) + } + + pub fn internal(err: T) -> Self { + Self::Internal(Box::new(err)) + } + + pub fn invalid_payload(err: T) -> Self { + Self::InvalidPayload(Box::new(err)) + } +} + +#[derive(Debug, EnumAsInner)] +pub enum Message { + Text(String), + Binary(Bytes), + Ping(Bytes), + Pong(Bytes), + Close(Option), +} + +impl Message { + pub fn as_bytes(&self) -> &[u8] { + match self { + Message::Binary(data) => data, + Message::Text(data) => data.as_bytes(), + Message::Ping(data) => data, + Message::Pong(data) => data, + Message::Close(_) => &[], + } + } +} + +#[derive(Debug)] +pub struct CloseFrame { + pub code: u16, + pub reason: String, +} + +/// Embeddable observer for monitoring WebSocket messages. +pub trait Observer: Send + Sync + 'static { + /// Called when an inbound message is received. + fn inbound_message(&self, _msg: &Message) {} + + /// Called when an outbound message is sent. + fn outbound_message(&self, _msg: &Message) {} + + /// Called when round-trip latency is measured. + /// + /// Latency measurement is based on ping-pong messages, which are triggered + /// on the heartbeat interval, so it should roughly correspond to that + /// interval. + fn latency(&self, _rtt: Duration) {} +} + +impl Observer for () {} + +/// Codec for encoding and decoding the payload sent and received over the +/// WebSocket. +pub trait DataCodec { + /// Associated message codec used to wrap and unwrap the encoded payload + /// into the [`Message`] transmitted. + type Message: Into + TryFrom; + + /// Payload type that can be sent and received. Assumes a symmetrical + /// payload format for both directions. + type Payload: Send + 'static; + + /// Encode the given payload into [`Message`] for transmission. + fn encode(&self, data: Self::Payload) -> Result; + + /// Decode the given [`Message`] into the payload. + fn decode(&self, data: Self::Message) -> Result; +} + +/// Backend for integrating different WebSocket transport implementations. +pub trait Backend: Send + 'static { + type Error: StdError + Send; + type Message: Send; + type Transport: Sink + + Stream> + + Send; + + /// Convert the backend into the underlying transport. + fn into_transport(self) -> Self::Transport; + + /// Encode the given [`Message`] into the transport-specific message type. + fn encode_message(msg: Message) -> Self::Message; + + /// Decode the given transport-specific message type into [`Message`]. + fn decode_message(msg: Self::Message) -> Message; +} + +#[cfg(feature = "json")] +mod json { + use { + super::*, + serde::{Serialize, de::DeserializeOwned}, + std::marker::PhantomData, + }; + + /// Generic JSON data codec using [`serde_json`] for all payloads that + /// implement [`serde`]'s [`Serialize`] and [`DeserializeOwned`]. + #[derive(Debug)] + pub struct Json(PhantomData); + + impl Default for Json { + fn default() -> Self { + Self(PhantomData) + } + } + + impl DataCodec for Json + where + T: Serialize + DeserializeOwned + Send + 'static, + { + type Message = TextMessage; + type Payload = T; + + fn encode(&self, data: Self::Payload) -> Result { + serde_json::to_string(&data) + .map(TextMessage) + .map_err(Error::encoding) + } + + fn decode(&self, data: Self::Message) -> Result { + serde_json::from_slice(data.as_bytes()).map_err(Error::decoding) + } + } +} + +/// Generic binary data codec that transmits raw bytes as-is using WebSocket +/// binary messages. +#[derive(Debug, Default)] +pub struct Binary; + +impl DataCodec for Binary { + type Message = BinaryMessage; + type Payload = Bytes; + + fn encode(&self, data: Self::Payload) -> Result { + Ok(data.into()) + } + + fn decode(&self, data: Self::Message) -> Result { + Ok(data.into()) + } +} + +/// Generic plaintext data codec that transmits UTF-8 strings as-is using +/// WebSocket text messages. +#[derive(Debug, Default)] +pub struct Plaintext; + +impl DataCodec for Plaintext { + type Message = TextMessage; + type Payload = String; + + fn encode(&self, data: Self::Payload) -> Result { + Ok(data.into()) + } + + fn decode(&self, data: Self::Message) -> Result { + Ok(data.into()) + } +} + +/// WebSocket message codec for binary data. Encodes the payload into +/// [`Message::Binary`]. +#[derive(Into, From)] +pub struct BinaryMessage(Bytes); + +impl BinaryMessage { + pub fn as_bytes(&self) -> &[u8] { + self.0.as_ref() + } +} + +impl From for Message { + fn from(msg: BinaryMessage) -> Self { + Message::Binary(msg.into()) + } +} + +#[derive(Debug, thiserror::Error)] +#[error("Message is not binary")] +struct InvalidBinaryError; + +impl TryFrom for BinaryMessage { + type Error = Error; + + fn try_from(data: Message) -> Result { + data.into_binary() + .map(Self) + .map_err(|_| Error::decoding(InvalidBinaryError)) + } +} + +/// WebSocket message codec for text data. Encodes the payload into +/// [`Message::Text`]. +#[derive(Into, From)] +pub struct TextMessage(String); + +impl TextMessage { + pub fn as_bytes(&self) -> &[u8] { + self.0.as_bytes() + } +} + +impl From for Message { + fn from(msg: TextMessage) -> Self { + Message::Text(msg.into()) + } +} + +#[derive(Debug, thiserror::Error)] +#[error("Message is not UTF-8")] +struct InvalidUtf8Error; + +impl TryFrom for TextMessage { + type Error = Error; + + fn try_from(data: Message) -> Result { + data.into_text() + .map(Self) + .map_err(|_| Error::decoding(InvalidUtf8Error)) + } +} + +/// Builder for configuring and constructing a [`WebSocket`] instance. +pub struct Builder { + backend: B, + codec: C, + observer: O, + config: Config, +} + +impl Builder<(), (), ()> { + /// Create a new [`WebSocket`] builder instance. + pub fn new() -> Self { + Self { + backend: (), + codec: (), + observer: (), + config: Default::default(), + } + } +} + +impl Default for Builder<(), (), ()> { + fn default() -> Self { + Self::new() + } +} + +impl Builder { + /// Set the [`Backend`] for the WebSocket. + pub fn backend(self, backend: T) -> Builder + where + T: Backend, + { + Builder { + backend, + codec: self.codec, + observer: self.observer, + config: self.config, + } + } + + /// Set the [`DataCodec`] for the WebSocket. + pub fn codec(self, codec: T) -> Builder + where + T: DataCodec, + { + Builder { + backend: self.backend, + codec, + observer: self.observer, + config: self.config, + } + } + + /// Set the [`Observer`] for the WebSocket. + pub fn observer(self, observer: T) -> Builder + where + T: Observer, + { + Builder { + backend: self.backend, + codec: self.codec, + observer, + config: self.config, + } + } + + /// Set the internal channel capacity for the WebSocket. The channel is used + /// to buffer messages sent and received. + /// + /// Default value: `64`. + pub fn channel_capacity(mut self, capacity: usize) -> Self { + self.config.channel_capacity = capacity; + self + } + + /// Set the heartbeat interval for the WebSocket. Heartbeat messages are + /// sent as [`Message::Ping`] and act as a keep-alive mechanism as well + /// as to measure the round-trip time latency (see [`Observer::latency`]). + /// + /// Default value: `5s`. + pub fn heartbeat_interval(mut self, interval: Duration) -> Self { + self.config.heartbeat_interval = interval; + self + } + + /// Set the idle timeout for the WebSocket. If no messages are received + /// within the timeout duration, the WebSocket connection is closed. This + /// should always be higher than the heartbeat interval. + /// + /// Default value: `15s`. + pub fn idle_timeout(mut self, timeout: Duration) -> Self { + self.config.idle_timeout = timeout; + self + } + + /// Build the configured [`WebSocket`] instance. + pub fn build(self) -> WebSocket + where + B: Backend, + C: DataCodec, + O: Observer, + { + WebSocket::new_internal(self.backend, self.codec, self.observer, self.config) + } +} diff --git a/crates/websocket/src/transport.rs b/crates/websocket/src/transport.rs new file mode 100644 index 0000000..d37f5b1 --- /dev/null +++ b/crates/websocket/src/transport.rs @@ -0,0 +1,226 @@ +use { + crate::{Backend, Error, Message, Observer}, + bytes::Bytes, + futures_concurrency::future::{Join as _, Race}, + futures_timer::Delay, + futures_util::{FutureExt as _, Sink, SinkExt as _, Stream, StreamExt as _, TryStreamExt}, + pin_project::pin_project, + std::{ + pin::Pin, + sync::Arc, + task::{self, Context, Poll}, + time::{Duration, SystemTime, UNIX_EPOCH}, + }, + tap::Pipe as _, + tokio::sync::{Notify, mpsc}, + tokio_stream::wrappers::{IntervalStream, ReceiverStream}, + tokio_util::sync::PollSender, +}; + +pub struct DropGuard(Arc); + +impl Drop for DropGuard { + fn drop(&mut self) { + self.0.notify_waiters(); + } +} + +/// Spawn the transport task, which handles forwarding messages between the +/// native transport and [`Core`] via [`tokio`] channels. +pub fn spawn( + transport: B::Transport, + observer: O, + capacity: usize, + heartbeat_interval: Duration, +) -> (mpsc::Sender, mpsc::Receiver, DropGuard) +where + B: Backend, + O: Observer, +{ + let (trans_tx, trans_rx) = transport.split(); + let (in_tx, in_rx) = mpsc::channel(capacity); + let (out_tx, out_rx) = mpsc::channel(capacity); + + // External shutdown is triggered when the `WebSocket` is dropped. + let shutdown = Arc::new(Notify::new()); + + tokio::spawn({ + let external_shutdown = shutdown.clone(); + + async move { + // Internal shutdown is triggered when the receiving stream from the underlying + // transport has ended. + let internal_shutdown = Notify::new(); + + let in_rx = ReceiverStream::new(in_rx); + + // Since we're merging multiple streams below, we need to end the heartbeat + // stream with both internal and external triggers. Otherwise the heartbeat + // stream will keep the channels alive indefinitely. + let heartbeat = heartbeat_stream(heartbeat_interval) + .take_until((external_shutdown.notified(), internal_shutdown.notified()).race()); + + // Forward messages from the `WebSocket` instance into the native transport. + let fwd_in = tokio_stream::StreamExt::merge(in_rx, heartbeat) + .inspect(|msg| { + observer.outbound_message(msg); + }) + .map(B::encode_message) + .map(Ok) + .forward(trans_tx); + + // Forward messages from the native transport to the `WebSocket` instance. + let fwd_out = trans_rx + .take_until(external_shutdown.notified()) + .map_ok(B::decode_message) + .map_err(Error::transport) + .inspect_ok(|msg| { + observer.inbound_message(msg); + + // We've received a `Pong` message with the timestamp from the heartbeat stream. + // Decode the timestamp and calculate the round-trip time. + if let Message::Pong(data) = msg { + let rtt = timestamp() + .saturating_sub(decode_timestamp(data)) + .pipe(Duration::from_millis); + + observer.latency(rtt); + } + }) + .forward(PollSender::new(out_tx).sink_map_err(Error::internal)) + .map(|_| internal_shutdown.notify_one()); + + (fwd_in, fwd_out).join().map(drop).await; + } + }); + + (in_tx, out_rx, DropGuard(shutdown)) +} + +/// Core transport that handles sending and receiving [`Message`]s with +/// heartbeat and idle timeout support. +#[pin_project] +pub struct Core { + #[pin] + tx: PollSender, + rx: ReceiverStream, + timeout_fut: Option, + timeout: Duration, +} + +impl Core { + pub fn new(tx: mpsc::Sender, rx: mpsc::Receiver, timeout: Duration) -> Self { + Self { + tx: PollSender::new(tx), + rx: ReceiverStream::new(rx), + timeout_fut: Some(Delay::new(timeout)), + timeout, + } + } + + #[inline] + fn poll_timeout(&mut self, cx: &mut Context<'_>) -> Poll<()> { + if let Some(delay) = &mut self.timeout_fut { + if delay.poll_unpin(cx).is_ready() { + self.timeout_fut = None; + + Poll::Ready(()) + } else { + Poll::Pending + } + } else { + Poll::Ready(()) + } + } + + #[inline] + fn reset_timeout(&mut self) { + let timeout = self.timeout; + + if let Some(delay) = &mut self.timeout_fut { + delay.reset(timeout); + } + } +} + +impl Sink for Core { + type Error = Error; + + fn start_send(self: Pin<&mut Self>, msg: Message) -> Result<(), Self::Error> { + self.project().tx.start_send(msg).map_err(|_| Error::Closed) + } + + fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.poll_timeout(cx).is_ready() { + Poll::Ready(Err(Error::Closed)) + } else { + self.project().tx.poll_ready(cx).map_err(|_| Error::Closed) + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().tx.poll_flush(cx).map_err(|_| Error::Closed) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().tx.poll_close(cx).map_err(|_| Error::Closed) + } +} + +impl Stream for Core { + type Item = Message; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.poll_timeout(cx).is_ready() { + return Poll::Ready(None); + } + + loop { + let result = task::ready!(self.rx.poll_next_unpin(cx)); + + self.reset_timeout(); + + let result = match result { + Some(msg) => match msg { + Message::Binary(_) | Message::Text(_) => Some(msg), + Message::Close(_) => None, + Message::Ping(_) | Message::Pong(_) => continue, + }, + + None => None, + }; + + return Poll::Ready(result); + } + } + + fn size_hint(&self) -> (usize, Option) { + self.rx.size_hint() + } +} + +/// Creates a stream that yields heartbeat [`Message::Ping`] messages at the +/// specified period. +fn heartbeat_stream(period: Duration) -> impl Stream { + let mut interval = tokio::time::interval(period); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + interval.reset(); + + IntervalStream::new(interval).map(|_| encode_timestamp(timestamp()).pipe(Message::Ping)) +} + +fn encode_timestamp(timestamp: u64) -> Bytes { + timestamp.to_be_bytes().to_vec().into() +} + +fn decode_timestamp(data: &[u8]) -> u64 { + data.try_into().map(u64::from_be_bytes).unwrap_or_default() +} + +fn timestamp() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .pipe(|timestamp| timestamp as u64) +} diff --git a/crates/websocket/src/wrapper.rs b/crates/websocket/src/wrapper.rs new file mode 100644 index 0000000..4ddfa2f --- /dev/null +++ b/crates/websocket/src/wrapper.rs @@ -0,0 +1,143 @@ +use { + crate::{ + Backend, + Builder, + DataCodec, + Error, + Observer, + transport::{self, Core, DropGuard}, + }, + futures_util::{Sink, Stream}, + pin_project::pin_project, + std::{ + pin::Pin, + task::{self, Context, Poll}, + time::Duration, + }, +}; + +/// Configuration options for the WebSocket transport. +/// +/// This should not be used directly. Instead, use the [`Builder`] to configure +/// and create a [`WebSocket`] instance. +pub struct Config { + pub channel_capacity: usize, + pub heartbeat_interval: Duration, + pub idle_timeout: Duration, +} + +impl Default for Config { + fn default() -> Self { + Self { + channel_capacity: 64, + heartbeat_interval: Duration::from_secs(5), + idle_timeout: Duration::from_secs(15), + } + } +} + +/// WebSocket wrapper that supports sending and receiving messages using the +/// specified data codec. +/// +/// This is a high-level wrapper around the provided [`Backend`] that adds +/// serialization and implements [`Sink`] and [`Stream`]. +/// +/// The underlying transport is closed when the [`WebSocket`] is dropped. +#[pin_project] +pub struct WebSocket { + #[pin] + inner: Core, + codec: C, + _guard: DropGuard, +} + +impl WebSocket<()> { + /// Creates a new [`Builder`] for configuring and constructing a + /// [`WebSocket`] instance. + pub fn builder() -> Builder<(), (), ()> { + Builder::new() + } +} + +impl WebSocket +where + C: DataCodec, +{ + /// Creates a new [`WebSocket`] instance with the specified backend and + /// codec using default configuration. + pub fn new(backend: B, codec: C) -> Self + where + B: Backend, + { + Self::new_internal(backend, codec, (), Default::default()) + } + + pub(crate) fn new_internal(backend: B, codec: C, observer: O, config: Config) -> Self + where + B: Backend, + O: Observer, + { + let (tx, rx, _guard) = transport::spawn::( + backend.into_transport(), + observer, + config.channel_capacity, + config.heartbeat_interval, + ); + + Self { + inner: transport::Core::new(tx, rx, config.idle_timeout), + codec, + _guard, + } + } +} + +impl Sink for WebSocket +where + C: DataCodec, +{ + type Error = Error; + + fn start_send(self: Pin<&mut Self>, item: C::Payload) -> Result<(), Self::Error> { + let item = self.codec.encode(item)?.into(); + + self.project().inner.start_send(item) + } + + fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_ready(cx) + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_flush(cx) + } + + fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.project().inner.poll_close(cx) + } +} + +impl Stream for WebSocket +where + C: DataCodec, +{ + type Item = C::Payload; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.project(); + + let data = task::ready!(this.inner.poll_next(cx)) + .map(|msg| { + C::Message::try_from(msg) + .and_then(|data| this.codec.decode(data)) + .ok() + }) + .flatten(); + + Poll::Ready(data) + } + + fn size_hint(&self) -> (usize, Option) { + self.inner.size_hint() + } +} diff --git a/crates/websocket/tests/integration.rs b/crates/websocket/tests/integration.rs new file mode 100644 index 0000000..d800e03 --- /dev/null +++ b/crates/websocket/tests/integration.rs @@ -0,0 +1,273 @@ +use { + bytes::Bytes, + futures_concurrency::future::Race, + futures_util::{FutureExt as _, SinkExt as _, Stream, StreamExt as _}, + serde::{Deserialize, Serialize}, + std::{ + marker::PhantomData, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, + }, + tokio::{ + net::{TcpListener, TcpStream}, + sync::Notify, + }, + tokio_tungstenite::{MaybeTlsStream, WebSocketStream}, + websocket::{Binary, DataCodec, Json, Message, Observer, Plaintext, WebSocket}, +}; + +struct EchoServer { + address: String, + notify: Arc, + num_connections: Arc, + _marker: PhantomData, +} + +impl EchoServer +where + C: DataCodec + Default + Send + Sync + 'static, +{ + async fn new() -> Self { + Self::with_builder(|socket| WebSocket::new(socket, C::default())).await + } + + async fn with_builder(builder: B) -> Self + where + B: Fn(WebSocketStream) -> WebSocket + Send + Sync + 'static, + { + let builder = Arc::new(builder); + let socket = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = socket.local_addr().unwrap(); + let address = format!("ws://{address}"); + let notify = Arc::new(Notify::new()); + let num_connections = Arc::new(AtomicUsize::new(0)); + + tokio::spawn({ + let notify = notify.clone(); + let num_connections = num_connections.clone(); + + async move { + while let Ok((stream, _)) = socket.accept().await { + let builder = builder.clone(); + let notify = notify.clone(); + let num_connections = num_connections.clone(); + + tokio::spawn(async move { + num_connections.fetch_add(1, Ordering::SeqCst); + + let socket = tokio_tungstenite::accept_async(stream).await.unwrap(); + let socket = builder(socket); + let (tx, rx) = socket.split(); + let _ = rx.take_until(notify.notified()).map(Ok).forward(tx).await; + + num_connections.fetch_sub(1, Ordering::SeqCst); + }); + } + } + }); + + Self { + address, + notify, + num_connections, + _marker: PhantomData, + } + } + + fn disconnect_clients(&self) { + self.notify.notify_waiters(); + } + + fn num_connections(&self) -> usize { + self.num_connections.load(Ordering::SeqCst) + } + + async fn connect(&self) -> WebSocketStream> { + let (socket, _) = tokio_tungstenite::connect_async(&self.address) + .await + .unwrap(); + + socket + } +} + +#[tokio::test] +async fn basic() { + let server = EchoServer::::new().await; + let mut socket = WebSocket::new(server.connect().await, Plaintext); + + assert_eq!(server.num_connections(), 1); + + let sent_payload = "hello world".to_owned(); + socket.send(sent_payload.clone()).await.unwrap(); + + assert_eq!(next(&mut socket).await, Some(sent_payload)); + server.disconnect_clients(); + assert_eq!(next(&mut socket).await, None); + assert_eq!(server.num_connections(), 0); +} + +#[derive(Default, Clone)] +struct HeartbeatObserver { + inbound_pings: Arc<AtomicUsize>, + inbound_pongs: Arc<AtomicUsize>, + outbound_pings: Arc<AtomicUsize>, + latency_reports: Arc<AtomicUsize>, +} + +impl HeartbeatObserver { + fn inbound_pings(&self) -> usize { + fetch(&self.inbound_pings) + } + + fn inbound_pongs(&self) -> usize { + fetch(&self.inbound_pongs) + } + + fn outbound_pings(&self) -> usize { + fetch(&self.outbound_pings) + } + + fn latency_reports(&self) -> usize { + fetch(&self.latency_reports) + } +} + +impl Observer for HeartbeatObserver { + fn inbound_message(&self, msg: &Message) { + match msg { + Message::Ping(_) => inc(&self.inbound_pings), + Message::Pong(_) => inc(&self.inbound_pongs), + _ => {} + } + } + + fn outbound_message(&self, msg: &Message) { + if let Message::Ping(_) = msg { + inc(&self.outbound_pings); + } + } + + fn latency(&self, _: Duration) { + inc(&self.latency_reports); + } +} + +#[tokio::test] +async fn heartbeat() { + let server = EchoServer::with_builder(|socket| { + WebSocket::builder() + .backend(socket) + .codec(Plaintext) + .heartbeat_interval(Duration::from_millis(300)) + .build() + }) + .await; + + let observer = HeartbeatObserver::default(); + + let mut socket = WebSocket::builder() + .backend(server.connect().await) + .observer(observer.clone()) + .codec(Plaintext) + .heartbeat_interval(Duration::from_millis(500)) + .build(); + + (socket.next().map(drop), sleep(1100)).race().await; + + assert_eq!(observer.inbound_pings(), 3); + assert_eq!(observer.outbound_pings(), 2); + assert_eq!(observer.inbound_pongs(), 2); + assert_eq!(observer.latency_reports(), 2); +} + +#[tokio::test] +async fn timeout_server() { + let server = EchoServer::with_builder(|socket| { + WebSocket::builder() + .backend(socket) + .codec(Plaintext) + .heartbeat_interval(Duration::from_secs(5)) + .idle_timeout(Duration::from_secs(1)) + .build() + }) + .await; + + let mut socket = WebSocket::new(server.connect().await, Plaintext); + + assert_eq!(server.num_connections(), 1); + sleep(1100).await; + assert_eq!(server.num_connections(), 0); + assert_eq!(next(&mut socket).await, None); +} + +#[tokio::test] +async fn timeout_client() { + let server = EchoServer::with_builder(|socket| { + WebSocket::builder() + .backend(socket) + .codec(Plaintext) + .heartbeat_interval(Duration::from_secs(5)) + .idle_timeout(Duration::from_secs(15)) + .build() + }) + .await; + + let mut socket = WebSocket::builder() + .backend(server.connect().await) + .codec(Plaintext) + .heartbeat_interval(Duration::from_secs(5)) + .idle_timeout(Duration::from_secs(1)) + .build(); + + sleep(1100).await; + assert!(socket.send("hello world".to_owned()).await.is_err()); + assert_eq!(next(&mut socket).await, None); +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +struct Payload(u32); + +#[tokio::test] +async fn json() { + let server = EchoServer::<Json<Payload>>::new().await; + let mut socket = WebSocket::new(server.connect().await, Json::<Payload>::default()); + + let sent_payload = Payload(42); + socket.send(sent_payload.clone()).await.unwrap(); + assert_eq!(next(&mut socket).await, Some(sent_payload)); +} + +#[tokio::test] +async fn binary() { + let server = EchoServer::<Binary>::new().await; + let mut socket = WebSocket::new(server.connect().await, Binary); + + let sent_payload = Bytes::from("hello world"); + socket.send(sent_payload.clone()).await.unwrap(); + assert_eq!(next(&mut socket).await, Some(sent_payload)); +} + +async fn sleep(millis: u64) { + tokio::time::sleep(Duration::from_millis(millis)).await +} + +async fn next<T, S>(stream: &mut S) -> Option<T> +where + S: Stream<Item = T> + Unpin, +{ + tokio::time::timeout(Duration::from_secs(1), stream.next()) + .await + .unwrap() +} + +fn inc(counter: &Arc<AtomicUsize>) { + counter.fetch_add(1, Ordering::SeqCst); +} + +fn fetch(counter: &Arc<AtomicUsize>) -> usize { + counter.load(Ordering::SeqCst) +} diff --git a/justfile b/justfile index 81e8a83..4c096dd 100644 --- a/justfile +++ b/justfile @@ -61,7 +61,7 @@ clippy: if command -v cargo-clippy >/dev/null; then echo '==> Running clippy' - cargo clippy --workspace --all-features --all-targets -- -D warnings + cargo +nightly clippy --workspace --all-features --all-targets -- -D warnings else echo '==> clippy not found in PATH, skipping' echo ' ^^^^^^ To install `rustup component add clippy`, see https://github.com/rust-lang/rust-clippy for details' diff --git a/rustfmt.toml b/rustfmt.toml index c39c39a..0087f01 100644 --- a/rustfmt.toml +++ b/rustfmt.toml @@ -1,4 +1,4 @@ -edition = "2021" +edition = "2024" fn_single_line = false format_code_in_doc_comments = true format_strings = true @@ -7,7 +7,6 @@ imports_granularity = "One" normalize_comments = true normalize_doc_attributes = true reorder_imports = true -reorder_impl_items = true group_imports = "StdExternalCrate" use_try_shorthand = true wrap_comments = true diff --git a/src/lib.rs b/src/lib.rs index e823040..a9ba382 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -13,3 +13,5 @@ pub use geoip; pub use metrics; #[cfg(feature = "rate_limit")] pub use rate_limit; +#[cfg(feature = "websocket")] +pub use websocket;