feat: websocket crate (#28)

Author: heilhead

.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

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"

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

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"] }

crates/websocket/src/backend.rs

@@ -0,0 +1,6 @@
+#[cfg(feature = "axum")]
+mod axum;
+#[cfg(feature = "tungstenite")]
+mod tungstenite;
+#[cfg(feature = "warp")]
+mod warp;

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()
+}

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<T> Backend for WebSocketStream<T>
+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()
+}

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)
+        }
+    }
+}