use std::hash::{DefaultHasher, Hash, Hasher};
use futures::stream::{BoxStream, SplitSink, SplitStream};
use futures::{SinkExt as _, StreamExt as _, stream};
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::protocol::CloseFrame;
use tokio_tungstenite::{MaybeTlsStream, WebSocketStream, connect_async};
use super::{SubscriptionId, SubscriptionSource};
#[derive(Debug, Clone)]
pub enum WebSocketCommand {
SendText(String),
SendBinary(Vec<u8>),
Close(Option<CloseFrame>),
}
#[derive(Debug, Clone)]
pub enum WebSocketMessage {
Connected {
sender: mpsc::UnboundedSender<WebSocketCommand>,
},
Disconnected,
Received(Message),
Error { error: String },
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct WebSocket {
url: String,
}
impl WebSocket {
#[must_use]
pub fn new(url: impl Into<String>) -> Self {
Self { url: url.into() }
}
}
type WsSink = SplitSink<WebSocketStream<MaybeTlsStream<TcpStream>>, Message>;
type WsRead = SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>;
enum WsStreamState {
Connecting {
url: String,
cmd_tx: mpsc::UnboundedSender<WebSocketCommand>,
cmd_rx: mpsc::UnboundedReceiver<WebSocketCommand>,
},
Running {
write: WsSink,
read: WsRead,
cmd_rx: mpsc::UnboundedReceiver<WebSocketCommand>,
},
Done,
}
impl SubscriptionSource for WebSocket {
type Output = WebSocketMessage;
#[allow(clippy::too_many_lines)]
fn stream(&self) -> BoxStream<'static, WebSocketMessage> {
let (cmd_tx, cmd_rx) = mpsc::unbounded_channel();
stream::unfold(
WsStreamState::Connecting {
url: self.url.clone(),
cmd_tx,
cmd_rx,
},
|state| async move {
match state {
WsStreamState::Connecting { url, cmd_tx, cmd_rx } => {
let trace_url = trace_url(&url);
tracing::debug!(
target: "tears::subscription::websocket",
url = %trace_url,
"websocket connecting"
);
match connect_async(&url).await {
Ok((ws, _)) => {
tracing::debug!(
target: "tears::subscription::websocket",
url = %trace_url,
"websocket connected"
);
let (write, read) = ws.split();
Some((
WebSocketMessage::Connected { sender: cmd_tx },
WsStreamState::Running { write, read, cmd_rx },
))
}
Err(e) => {
tracing::debug!(
target: "tears::subscription::websocket",
url = %trace_url,
error = %e,
"websocket connection failed"
);
Some((
WebSocketMessage::Error {
error: format!("Connection failed: {e}"),
},
WsStreamState::Done,
))
}
}
}
WsStreamState::Running {
mut write,
mut read,
mut cmd_rx,
} => loop {
tokio::select! {
msg = read.next() => {
match msg {
Some(Ok(Message::Close(_))) => {
tracing::debug!(target: "tears::subscription::websocket", reason = "server_close", "websocket disconnected");
let _ = write.close().await;
break Some((WebSocketMessage::Disconnected, WsStreamState::Done));
}
Some(Ok(message)) => {
tracing::trace!(
target: "tears::subscription::websocket",
message_type = trace_message_type(&message),
message_size = trace_message_size(&message),
"websocket message received"
);
break Some((
WebSocketMessage::Received(message),
WsStreamState::Running { write, read, cmd_rx },
));
}
Some(Err(e)) => {
tracing::debug!(
target: "tears::subscription::websocket",
error = %e,
"websocket read failed"
);
let _ = write.close().await;
break Some((
WebSocketMessage::Error { error: e.to_string() },
WsStreamState::Done,
));
}
None => {
tracing::debug!(target: "tears::subscription::websocket", reason = "read_stream_ended", "websocket disconnected");
break Some((WebSocketMessage::Disconnected, WsStreamState::Done));
}
}
}
cmd = cmd_rx.recv() => {
match cmd {
Some(WebSocketCommand::Close(frame)) => {
tracing::debug!(target: "tears::subscription::websocket", reason = "close_command", "websocket disconnecting");
let _ = write.send(Message::Close(frame)).await;
let _ = write.close().await;
break Some((WebSocketMessage::Disconnected, WsStreamState::Done));
}
Some(WebSocketCommand::SendText(text)) => {
tracing::trace!(
target: "tears::subscription::websocket",
message_type = "text",
message_size = text.len(),
"websocket message sending"
);
if let Err(e) = write.send(Message::Text(text.into())).await {
tracing::debug!(
target: "tears::subscription::websocket",
error = %e,
message_type = "text",
"websocket write failed"
);
break Some((
WebSocketMessage::Error { error: e.to_string() },
WsStreamState::Running { write, read, cmd_rx },
));
}
}
Some(WebSocketCommand::SendBinary(data)) => {
tracing::trace!(
target: "tears::subscription::websocket",
message_type = "binary",
message_size = data.len(),
"websocket message sending"
);
if let Err(e) = write.send(Message::Binary(data.into())).await {
tracing::debug!(
target: "tears::subscription::websocket",
error = %e,
message_type = "binary",
"websocket write failed"
);
break Some((
WebSocketMessage::Error { error: e.to_string() },
WsStreamState::Running { write, read, cmd_rx },
));
}
}
None => {
tracing::debug!(target: "tears::subscription::websocket", reason = "command_sender_dropped", "websocket disconnected");
let _ = write.close().await;
break Some((WebSocketMessage::Disconnected, WsStreamState::Done));
}
}
}
}
},
WsStreamState::Done => None,
}
},
)
.boxed()
}
fn id(&self) -> SubscriptionId {
let mut hasher = DefaultHasher::new();
self.hash(&mut hasher);
SubscriptionId::of::<Self>(hasher.finish())
}
}
const fn trace_message_type(message: &Message) -> &'static str {
match message {
Message::Text(_) => "text",
Message::Binary(_) => "binary",
Message::Ping(_) => "ping",
Message::Pong(_) => "pong",
Message::Close(_) => "close",
Message::Frame(_) => "frame",
}
}
fn trace_message_size(message: &Message) -> usize {
match message {
Message::Text(text) => text.len(),
Message::Binary(data) | Message::Ping(data) | Message::Pong(data) => data.len(),
Message::Close(_) => 0,
Message::Frame(frame) => frame.payload().len(),
}
}
fn trace_url(url: &str) -> String {
let base = url.split_once(['?', '#']).map_or(url, |(base, _)| base);
let Some((scheme, rest)) = base.split_once("://") else {
return base.to_string();
};
let (authority, path) = rest
.split_once('/')
.map_or((rest, ""), |(authority, path)| (authority, path));
let Some((_, host)) = authority.split_once('@') else {
return base.to_string();
};
if path.is_empty() {
format!("{scheme}://<redacted>@{host}")
} else {
format!("{scheme}://<redacted>@{host}/{path}")
}
}
#[cfg(test)]
mod tests {
use super::*;
use futures::StreamExt;
#[test]
fn test_ws_new() {
let ws = WebSocket::new("wss://example.com");
assert_eq!(ws.url, "wss://example.com");
}
#[test]
fn test_ws_id_consistency() {
let ws1 = WebSocket::new("wss://example.com");
let ws2 = WebSocket::new("wss://example.com");
assert_eq!(ws1.id(), ws2.id());
}
#[test]
fn test_ws_id_different_urls() {
let ws1 = WebSocket::new("wss://example.com");
let ws2 = WebSocket::new("wss://different.com");
assert_ne!(ws1.id(), ws2.id());
}
#[test]
fn test_trace_url_redacts_userinfo_only_in_authority() {
assert_eq!(
trace_url("wss://user:password@example.com/socket?token=secret#fragment"),
"wss://<redacted>@example.com/socket"
);
assert_eq!(
trace_url("wss://example.com/@user/feed?token=secret"),
"wss://example.com/@user/feed"
);
}
#[tokio::test]
async fn test_stream_emits_error_on_connection_failure() {
let ws = WebSocket::new("ws://localhost:1");
let mut stream = ws.stream();
assert!(matches!(
stream.next().await,
Some(WebSocketMessage::Error { .. }),
));
}
#[test]
fn test_message_variants() {
let (tx, _rx) = mpsc::unbounded_channel();
matches!(
WebSocketMessage::Connected { sender: tx },
WebSocketMessage::Connected { .. }
);
matches!(
WebSocketMessage::Disconnected,
WebSocketMessage::Disconnected
);
matches!(
WebSocketMessage::Received(Message::Text("test".into())),
WebSocketMessage::Received(_)
);
matches!(
WebSocketMessage::Error {
error: "test".to_string()
},
WebSocketMessage::Error { .. }
);
}
}