#![allow(deprecated)]
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tokio::sync::mpsc;
#[derive(Debug, Clone)]
pub struct PingPongConfig {
ping_interval: Duration,
pong_timeout: Duration,
}
impl Default for PingPongConfig {
fn default() -> Self {
Self {
ping_interval: Duration::from_secs(30),
pong_timeout: Duration::from_secs(10),
}
}
}
impl PingPongConfig {
pub fn new(ping_interval: Duration, pong_timeout: Duration) -> Self {
Self {
ping_interval,
pong_timeout,
}
}
pub fn ping_interval(&self) -> Duration {
self.ping_interval
}
pub fn pong_timeout(&self) -> Duration {
self.pong_timeout
}
pub fn with_ping_interval(mut self, interval: Duration) -> Self {
self.ping_interval = interval;
self
}
pub fn with_pong_timeout(mut self, timeout: Duration) -> Self {
self.pong_timeout = timeout;
self
}
}
#[deprecated(
since = "0.2.0",
note = "Use `ConnectionSettings` with the `#[settings]` macro instead."
)]
#[derive(Debug, Clone)]
pub struct ConnectionConfig {
idle_timeout: Duration,
handshake_timeout: Duration,
cleanup_interval: Duration,
max_connections: Option<usize>,
ping_config: PingPongConfig,
}
impl Default for ConnectionConfig {
fn default() -> Self {
Self {
idle_timeout: Duration::from_secs(300), handshake_timeout: Duration::from_secs(10), cleanup_interval: Duration::from_secs(30), max_connections: None, ping_config: PingPongConfig::default(),
}
}
}
impl ConnectionConfig {
pub fn new() -> Self {
Self::default()
}
pub fn with_idle_timeout(mut self, timeout: Duration) -> Self {
self.idle_timeout = timeout;
self
}
pub fn with_handshake_timeout(mut self, timeout: Duration) -> Self {
self.handshake_timeout = timeout;
self
}
pub fn with_cleanup_interval(mut self, interval: Duration) -> Self {
self.cleanup_interval = interval;
self
}
pub fn idle_timeout(&self) -> Duration {
self.idle_timeout
}
pub fn handshake_timeout(&self) -> Duration {
self.handshake_timeout
}
pub fn cleanup_interval(&self) -> Duration {
self.cleanup_interval
}
pub fn with_max_connections(mut self, max: Option<usize>) -> Self {
self.max_connections = max;
self
}
pub fn max_connections(&self) -> Option<usize> {
self.max_connections
}
pub fn with_ping_config(mut self, config: PingPongConfig) -> Self {
self.ping_config = config;
self
}
pub fn ping_config(&self) -> &PingPongConfig {
&self.ping_config
}
pub fn no_timeout() -> Self {
Self {
idle_timeout: Duration::MAX,
handshake_timeout: Duration::MAX,
cleanup_interval: Duration::from_secs(30),
max_connections: None,
ping_config: PingPongConfig::default(),
}
}
pub fn strict() -> Self {
Self {
idle_timeout: Duration::from_secs(30),
handshake_timeout: Duration::from_secs(5),
cleanup_interval: Duration::from_secs(10),
max_connections: None,
ping_config: PingPongConfig::new(Duration::from_secs(10), Duration::from_secs(5)),
}
}
pub fn permissive() -> Self {
Self {
idle_timeout: Duration::from_secs(3600),
handshake_timeout: Duration::from_secs(30),
cleanup_interval: Duration::from_secs(60),
max_connections: None,
ping_config: PingPongConfig::new(Duration::from_secs(60), Duration::from_secs(30)),
}
}
}
#[derive(Debug, thiserror::Error)]
pub enum WebSocketError {
#[error("Connection error")]
Connection(String),
#[error("Send failed")]
Send(String),
#[error("Receive failed")]
Receive(String),
#[error("Protocol error")]
Protocol(String),
#[error("Internal error")]
Internal(String),
#[error("Connection timed out")]
Timeout(Duration),
#[error("Reconnection failed")]
ReconnectFailed(u32),
#[error("Invalid binary payload: {0}")]
BinaryPayload(String),
#[error("Heartbeat timeout: no pong received within {0:?}")]
HeartbeatTimeout(Duration),
#[error("Slow consumer: send timed out after {0:?}")]
SlowConsumer(Duration),
}
impl WebSocketError {
pub fn client_message(&self) -> &'static str {
match self {
Self::Connection(_) => "Connection error",
Self::Send(_) => "Failed to send message",
Self::Receive(_) => "Failed to receive message",
Self::Protocol(_) => "Protocol error",
Self::Internal(_) => "Internal server error",
Self::Timeout(_) => "Connection timed out",
Self::ReconnectFailed(_) => "Reconnection failed",
Self::BinaryPayload(_) => "Invalid message format",
Self::HeartbeatTimeout(_) => "Connection timed out",
Self::SlowConsumer(_) => "Server overloaded",
}
}
pub fn internal_detail(&self) -> String {
match self {
Self::Connection(msg) => format!("Connection error: {}", msg),
Self::Send(msg) => format!("Send error: {}", msg),
Self::Receive(msg) => format!("Receive error: {}", msg),
Self::Protocol(msg) => format!("Protocol error: {}", msg),
Self::Internal(msg) => format!("Internal error: {}", msg),
Self::Timeout(d) => format!("Connection timeout: idle for {:?}", d),
Self::ReconnectFailed(n) => format!("Reconnection failed after {} attempts", n),
Self::BinaryPayload(msg) => format!("Invalid binary payload: {}", msg),
Self::HeartbeatTimeout(d) => format!("Heartbeat timeout: no pong within {:?}", d),
Self::SlowConsumer(d) => format!("Slow consumer: send timed out after {:?}", d),
}
}
}
pub type WebSocketResult<T> = Result<T, WebSocketError>;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(tag = "type")]
pub enum Message {
Text {
data: String,
},
Binary {
data: Vec<u8>,
},
Ping,
Pong,
Close {
code: u16,
reason: String,
},
}
impl Message {
pub fn text(data: String) -> Self {
Self::Text { data }
}
pub fn binary(data: Vec<u8>) -> Self {
Self::Binary { data }
}
pub fn json<T: serde::Serialize>(data: &T) -> WebSocketResult<Self> {
let json =
serde_json::to_string(data).map_err(|e| WebSocketError::Protocol(e.to_string()))?;
Ok(Self::text(json))
}
pub fn parse_json<T: serde::de::DeserializeOwned>(&self) -> WebSocketResult<T> {
match self {
Message::Text { data } => {
serde_json::from_str(data).map_err(|e| WebSocketError::Protocol(e.to_string()))
}
_ => Err(WebSocketError::Protocol("Not a text message".to_string())),
}
}
}
pub struct WebSocketConnection {
id: String,
tx: mpsc::UnboundedSender<Message>,
closed: Arc<RwLock<bool>>,
subprotocol: Option<String>,
last_activity: Arc<RwLock<Instant>>,
config: ConnectionConfig,
}
impl WebSocketConnection {
pub fn new(id: String, tx: mpsc::UnboundedSender<Message>) -> Self {
Self {
id,
tx,
closed: Arc::new(RwLock::new(false)),
subprotocol: None,
last_activity: Arc::new(RwLock::new(Instant::now())),
config: ConnectionConfig::default(),
}
}
pub fn with_config(
id: String,
tx: mpsc::UnboundedSender<Message>,
config: ConnectionConfig,
) -> Self {
Self {
id,
tx,
closed: Arc::new(RwLock::new(false)),
subprotocol: None,
last_activity: Arc::new(RwLock::new(Instant::now())),
config,
}
}
pub fn with_subprotocol(
id: String,
tx: mpsc::UnboundedSender<Message>,
subprotocol: Option<String>,
) -> Self {
Self {
id,
tx,
closed: Arc::new(RwLock::new(false)),
subprotocol,
last_activity: Arc::new(RwLock::new(Instant::now())),
config: ConnectionConfig::default(),
}
}
pub fn subprotocol(&self) -> Option<&str> {
self.subprotocol.as_deref()
}
pub fn id(&self) -> &str {
&self.id
}
pub fn config(&self) -> &ConnectionConfig {
&self.config
}
pub async fn record_activity(&self) {
*self.last_activity.write().await = Instant::now();
}
pub async fn idle_duration(&self) -> Duration {
self.last_activity.read().await.elapsed()
}
pub async fn is_idle(&self) -> bool {
self.idle_duration().await > self.config.idle_timeout
}
pub async fn send(&self, message: Message) -> WebSocketResult<()> {
if *self.closed.read().await {
return Err(WebSocketError::Send("Connection closed".to_string()));
}
let result = self
.tx
.send(message)
.map_err(|e| WebSocketError::Send(e.to_string()));
if result.is_ok() {
self.record_activity().await;
}
result
}
pub async fn send_text(&self, text: String) -> WebSocketResult<()> {
self.send(Message::text(text)).await
}
pub async fn send_binary(&self, data: Vec<u8>) -> WebSocketResult<()> {
self.send(Message::binary(data)).await
}
pub async fn send_json<T: serde::Serialize>(&self, data: &T) -> WebSocketResult<()> {
let message = Message::json(data)?;
self.send(message).await
}
pub async fn close(&self) -> WebSocketResult<()> {
*self.closed.write().await = true;
self.tx
.send(Message::Close {
code: 1000,
reason: "Normal closure".to_string(),
})
.map_err(|e| WebSocketError::Send(e.to_string()))
}
pub async fn close_with_reason(&self, code: u16, reason: String) -> WebSocketResult<()> {
*self.closed.write().await = true;
self.tx
.send(Message::Close { code, reason })
.map_err(|e| WebSocketError::Send(e.to_string()))
}
pub async fn force_close(&self) {
*self.closed.write().await = true;
}
pub async fn is_closed(&self) -> bool {
*self.closed.read().await
}
}
pub struct ConnectionTimeoutMonitor {
connections: Arc<RwLock<HashMap<String, Arc<WebSocketConnection>>>>,
config: ConnectionConfig,
}
impl ConnectionTimeoutMonitor {
pub fn new(config: ConnectionConfig) -> Self {
Self {
connections: Arc::new(RwLock::new(HashMap::new())),
config,
}
}
pub async fn register(
&self,
connection: Arc<WebSocketConnection>,
) -> Result<(), WebSocketError> {
let mut connections = self.connections.write().await;
if let Some(max) = self.config.max_connections
&& connections.len() >= max
{
return Err(WebSocketError::Connection(format!(
"maximum connection limit reached ({})",
max
)));
}
connections.insert(connection.id().to_string(), connection);
Ok(())
}
pub async fn unregister(&self, connection_id: &str) {
self.connections.write().await.remove(connection_id);
}
pub async fn connection_count(&self) -> usize {
self.connections.read().await.len()
}
pub async fn check_idle_connections(&self) -> Vec<String> {
let connections = self.connections.read().await;
let mut timed_out = Vec::new();
for (id, conn) in connections.iter() {
if conn.is_closed().await {
timed_out.push(id.clone());
continue;
}
let idle_duration = conn.idle_duration().await;
if idle_duration > self.config.idle_timeout {
let reason = format!(
"Idle timeout: connection idle for {}s (limit: {}s)",
idle_duration.as_secs(),
self.config.idle_timeout.as_secs()
);
let _ = conn.close_with_reason(1001, reason).await;
timed_out.push(id.clone());
}
}
drop(connections);
if !timed_out.is_empty() {
let mut connections = self.connections.write().await;
for id in &timed_out {
connections.remove(id);
}
}
timed_out
}
pub async fn shutdown_all(&self) -> Vec<String> {
let mut connections = self.connections.write().await;
let mut shut_down = Vec::with_capacity(connections.len());
for (id, conn) in connections.drain() {
if !conn.is_closed().await {
let _ = conn
.close_with_reason(1001, "Server shutting down".to_string())
.await;
}
shut_down.push(id);
}
shut_down
}
pub fn start(self: &Arc<Self>) -> tokio::task::JoinHandle<()> {
let monitor = Arc::clone(self);
tokio::spawn(async move {
let mut interval = tokio::time::interval(monitor.config.cleanup_interval);
loop {
interval.tick().await;
monitor.check_idle_connections().await;
}
})
}
}
#[derive(Debug, Clone)]
pub struct HeartbeatConfig {
ping_interval: Duration,
pong_timeout: Duration,
}
impl HeartbeatConfig {
pub fn new(ping_interval: Duration, pong_timeout: Duration) -> Self {
Self {
ping_interval,
pong_timeout,
}
}
pub fn ping_interval(&self) -> Duration {
self.ping_interval
}
pub fn pong_timeout(&self) -> Duration {
self.pong_timeout
}
}
impl Default for HeartbeatConfig {
fn default() -> Self {
Self {
ping_interval: Duration::from_secs(30),
pong_timeout: Duration::from_secs(10),
}
}
}
pub struct HeartbeatMonitor {
connection: Arc<WebSocketConnection>,
config: HeartbeatConfig,
last_pong: Arc<RwLock<Instant>>,
timed_out: Arc<RwLock<bool>>,
pong_notify: Arc<tokio::sync::Notify>,
}
impl HeartbeatMonitor {
pub fn new(connection: Arc<WebSocketConnection>, config: HeartbeatConfig) -> Self {
Self {
connection,
config,
last_pong: Arc::new(RwLock::new(Instant::now())),
timed_out: Arc::new(RwLock::new(false)),
pong_notify: Arc::new(tokio::sync::Notify::new()),
}
}
pub async fn record_pong(&self) {
*self.last_pong.write().await = Instant::now();
self.pong_notify.notify_one();
}
pub async fn time_since_last_pong(&self) -> Duration {
self.last_pong.read().await.elapsed()
}
pub async fn is_timed_out(&self) -> bool {
*self.timed_out.read().await
}
pub async fn check_heartbeat(&self) -> bool {
let since_pong = self.time_since_last_pong().await;
if since_pong > self.config.pong_timeout {
self.connection.force_close().await;
*self.timed_out.write().await = true;
return true;
}
false
}
pub async fn send_ping(&self) -> WebSocketResult<()> {
self.connection.send(Message::Ping).await
}
pub fn config(&self) -> &HeartbeatConfig {
&self.config
}
pub fn connection(&self) -> &Arc<WebSocketConnection> {
&self.connection
}
pub fn start(self: &Arc<Self>) -> tokio::task::JoinHandle<()> {
let monitor = Arc::clone(self);
tokio::spawn(async move {
let mut interval = tokio::time::interval(monitor.config.ping_interval);
loop {
interval.tick().await;
if monitor.connection.is_closed().await {
break;
}
let _ = monitor.send_ping().await;
tokio::select! {
() = tokio::time::sleep(monitor.config.pong_timeout) => {
if monitor.check_heartbeat().await {
break;
}
}
() = monitor.pong_notify.notified() => {
}
}
}
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
#[rstest]
fn test_message_text() {
let text = "Hello".to_string();
let msg = Message::text(text);
match msg {
Message::Text { data } => assert_eq!(data, "Hello"),
_ => panic!("Expected text message"),
}
}
#[rstest]
fn test_message_json() {
#[derive(serde::Serialize)]
struct TestData {
value: i32,
}
let data = TestData { value: 42 };
let msg = Message::json(&data).unwrap();
match msg {
Message::Text { data } => assert!(data.contains("42")),
_ => panic!("Expected text message"),
}
}
#[rstest]
#[tokio::test]
async fn test_connection_send() {
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::new("test".to_string(), tx);
conn.send_text("Hello".to_string()).await.unwrap();
let received = rx.recv().await.unwrap();
match received {
Message::Text { data } => assert_eq!(data, "Hello"),
_ => panic!("Expected text message"),
}
}
#[rstest]
fn test_connection_config_default() {
let config = ConnectionConfig::new();
assert_eq!(config.idle_timeout(), Duration::from_secs(300));
assert_eq!(config.handshake_timeout(), Duration::from_secs(10));
assert_eq!(config.cleanup_interval(), Duration::from_secs(30));
}
#[rstest]
fn test_connection_config_strict() {
let config = ConnectionConfig::strict();
assert_eq!(config.idle_timeout(), Duration::from_secs(30));
assert_eq!(config.handshake_timeout(), Duration::from_secs(5));
assert_eq!(config.cleanup_interval(), Duration::from_secs(10));
}
#[rstest]
fn test_connection_config_permissive() {
let config = ConnectionConfig::permissive();
assert_eq!(config.idle_timeout(), Duration::from_secs(3600));
assert_eq!(config.handshake_timeout(), Duration::from_secs(30));
assert_eq!(config.cleanup_interval(), Duration::from_secs(60));
}
#[rstest]
fn test_connection_config_no_timeout() {
let config = ConnectionConfig::no_timeout();
assert_eq!(config.idle_timeout(), Duration::MAX);
assert_eq!(config.handshake_timeout(), Duration::MAX);
}
#[rstest]
fn test_connection_config_builder() {
let config = ConnectionConfig::new()
.with_idle_timeout(Duration::from_secs(120))
.with_handshake_timeout(Duration::from_secs(15))
.with_cleanup_interval(Duration::from_secs(20));
assert_eq!(config.idle_timeout(), Duration::from_secs(120));
assert_eq!(config.handshake_timeout(), Duration::from_secs(15));
assert_eq!(config.cleanup_interval(), Duration::from_secs(20));
}
#[rstest]
#[tokio::test]
async fn test_connection_with_config() {
let config = ConnectionConfig::new().with_idle_timeout(Duration::from_secs(60));
let (tx, _rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
assert_eq!(conn.config().idle_timeout(), Duration::from_secs(60));
assert!(!conn.is_idle().await);
}
#[rstest]
#[tokio::test]
async fn test_connection_record_activity_resets_idle() {
let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50));
let (tx, _rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
tokio::time::sleep(Duration::from_millis(60)).await;
assert!(conn.is_idle().await);
conn.record_activity().await;
assert!(!conn.is_idle().await);
}
#[rstest]
#[tokio::test]
async fn test_connection_becomes_idle_after_timeout() {
let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50));
let (tx, _rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
tokio::time::sleep(Duration::from_millis(60)).await;
assert!(conn.is_idle().await);
assert!(conn.idle_duration().await >= Duration::from_millis(50));
}
#[rstest]
#[tokio::test]
async fn test_send_resets_activity() {
let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(100));
let (tx, mut _rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::with_config("test".to_string(), tx, config);
tokio::time::sleep(Duration::from_millis(50)).await;
conn.send_text("ping".to_string()).await.unwrap();
assert!(conn.idle_duration().await < Duration::from_millis(30));
assert!(!conn.is_idle().await);
}
#[rstest]
#[tokio::test]
async fn test_close_with_reason() {
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::new("test".to_string(), tx);
conn.close_with_reason(1001, "Idle timeout".to_string())
.await
.unwrap();
assert!(conn.is_closed().await);
let msg = rx.recv().await.unwrap();
match msg {
Message::Close { code, reason } => {
assert_eq!(code, 1001);
assert_eq!(reason, "Idle timeout");
}
_ => panic!("Expected close message"),
}
}
#[rstest]
#[tokio::test]
async fn test_timeout_monitor_register_and_count() {
let config = ConnectionConfig::new();
let monitor = ConnectionTimeoutMonitor::new(config);
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx));
monitor.register(conn).await.unwrap();
assert_eq!(monitor.connection_count().await, 1);
}
#[rstest]
#[tokio::test]
async fn test_timeout_monitor_unregister() {
let config = ConnectionConfig::new();
let monitor = ConnectionTimeoutMonitor::new(config);
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx));
monitor.register(conn).await.unwrap();
monitor.unregister("conn_1").await;
assert_eq!(monitor.connection_count().await, 0);
}
#[rstest]
#[tokio::test]
async fn test_timeout_monitor_closes_idle_connections() {
let config = ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50));
let monitor = ConnectionTimeoutMonitor::new(config);
let (tx1, mut rx1) = mpsc::unbounded_channel();
let conn1 = Arc::new(WebSocketConnection::with_config(
"idle_conn".to_string(),
tx1,
ConnectionConfig::new().with_idle_timeout(Duration::from_millis(50)),
));
let (tx2, _rx2) = mpsc::unbounded_channel();
let conn2 = Arc::new(WebSocketConnection::with_config(
"active_conn".to_string(),
tx2,
ConnectionConfig::new().with_idle_timeout(Duration::from_secs(300)),
));
monitor.register(conn1).await.unwrap();
monitor.register(conn2.clone()).await.unwrap();
tokio::time::sleep(Duration::from_millis(60)).await;
conn2.record_activity().await;
let timed_out = monitor.check_idle_connections().await;
assert_eq!(timed_out.len(), 1);
assert_eq!(timed_out[0], "idle_conn");
assert_eq!(monitor.connection_count().await, 1);
let msg = rx1.recv().await.unwrap();
match msg {
Message::Close { code, reason } => {
assert_eq!(code, 1001);
assert!(reason.contains("Idle timeout"));
}
_ => panic!("Expected close message for idle connection"),
}
}
#[rstest]
#[tokio::test]
async fn test_timeout_monitor_removes_already_closed_connections() {
let config = ConnectionConfig::new();
let monitor = ConnectionTimeoutMonitor::new(config);
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx));
conn.close().await.unwrap();
monitor.register(conn).await.unwrap();
let timed_out = monitor.check_idle_connections().await;
assert_eq!(timed_out.len(), 1);
assert_eq!(timed_out[0], "conn_1");
assert_eq!(monitor.connection_count().await, 0);
}
#[rstest]
#[tokio::test]
async fn test_timeout_monitor_background_task() {
let config = ConnectionConfig::new()
.with_idle_timeout(Duration::from_millis(30))
.with_cleanup_interval(Duration::from_millis(20));
let monitor = Arc::new(ConnectionTimeoutMonitor::new(config));
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::with_config(
"bg_conn".to_string(),
tx,
ConnectionConfig::new().with_idle_timeout(Duration::from_millis(30)),
));
monitor.register(conn).await.unwrap();
let handle = monitor.start();
tokio::time::sleep(Duration::from_millis(120)).await;
assert_eq!(monitor.connection_count().await, 0);
let msg = rx.recv().await.unwrap();
assert!(matches!(msg, Message::Close { .. }));
handle.abort();
}
#[rstest]
fn test_ping_pong_config_default() {
let config = PingPongConfig::default();
assert_eq!(config.ping_interval(), Duration::from_secs(30));
assert_eq!(config.pong_timeout(), Duration::from_secs(10));
}
#[rstest]
fn test_ping_pong_config_custom() {
let config = PingPongConfig::new(Duration::from_secs(15), Duration::from_secs(5));
assert_eq!(config.ping_interval(), Duration::from_secs(15));
assert_eq!(config.pong_timeout(), Duration::from_secs(5));
}
#[rstest]
fn test_ping_pong_config_builder() {
let config = PingPongConfig::default()
.with_ping_interval(Duration::from_secs(60))
.with_pong_timeout(Duration::from_secs(20));
assert_eq!(config.ping_interval(), Duration::from_secs(60));
assert_eq!(config.pong_timeout(), Duration::from_secs(20));
}
#[rstest]
fn test_connection_config_has_default_ping_config() {
let config = ConnectionConfig::new();
assert_eq!(
config.ping_config().ping_interval(),
Duration::from_secs(30)
);
assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(10));
}
#[rstest]
fn test_connection_config_with_custom_ping_config() {
let ping_config = PingPongConfig::new(Duration::from_secs(15), Duration::from_secs(5));
let config = ConnectionConfig::new().with_ping_config(ping_config);
assert_eq!(
config.ping_config().ping_interval(),
Duration::from_secs(15)
);
assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(5));
}
#[rstest]
fn test_strict_config_has_aggressive_ping() {
let config = ConnectionConfig::strict();
assert_eq!(
config.ping_config().ping_interval(),
Duration::from_secs(10)
);
assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(5));
}
#[rstest]
fn test_permissive_config_has_relaxed_ping() {
let config = ConnectionConfig::permissive();
assert_eq!(
config.ping_config().ping_interval(),
Duration::from_secs(60)
);
assert_eq!(config.ping_config().pong_timeout(), Duration::from_secs(30));
}
#[rstest]
#[tokio::test]
async fn test_timeout_monitor_rejects_when_max_connections_reached() {
let config = ConnectionConfig::new().with_max_connections(Some(1));
let monitor = ConnectionTimeoutMonitor::new(config);
let (tx1, _rx1) = mpsc::unbounded_channel();
let conn1 = Arc::new(WebSocketConnection::new("conn_1".to_string(), tx1));
let (tx2, _rx2) = mpsc::unbounded_channel();
let conn2 = Arc::new(WebSocketConnection::new("conn_2".to_string(), tx2));
monitor.register(conn1).await.unwrap();
let result = monitor.register(conn2).await;
assert!(result.is_err());
assert_eq!(monitor.connection_count().await, 1);
}
#[rstest]
#[tokio::test]
async fn test_force_close_marks_connection_closed() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::new("test".to_string(), tx);
conn.force_close().await;
assert!(conn.is_closed().await);
}
#[rstest]
#[tokio::test]
async fn test_close_marks_closed_even_when_channel_dropped() {
let (tx, rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::new("test".to_string(), tx);
drop(rx);
let result = conn.close().await;
assert!(result.is_err()); assert!(conn.is_closed().await); }
#[rstest]
#[tokio::test]
async fn test_close_with_reason_marks_closed_even_when_channel_dropped() {
let (tx, rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::new("test".to_string(), tx);
drop(rx);
let result = conn
.close_with_reason(1006, "Abnormal close".to_string())
.await;
assert!(result.is_err());
assert!(conn.is_closed().await);
}
#[rstest]
#[tokio::test]
async fn test_send_after_force_close_returns_error() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = WebSocketConnection::new("test".to_string(), tx);
conn.force_close().await;
let result = conn.send_text("should fail".to_string()).await;
assert!(result.is_err());
assert!(matches!(result.unwrap_err(), WebSocketError::Send(_)));
}
#[rstest]
fn test_heartbeat_config_default() {
let config = HeartbeatConfig::default();
assert_eq!(config.ping_interval(), Duration::from_secs(30));
assert_eq!(config.pong_timeout(), Duration::from_secs(10));
}
#[rstest]
fn test_heartbeat_config_custom() {
let config = HeartbeatConfig::new(Duration::from_secs(15), Duration::from_secs(5));
assert_eq!(config.ping_interval(), Duration::from_secs(15));
assert_eq!(config.pong_timeout(), Duration::from_secs(5));
}
#[rstest]
#[tokio::test]
async fn test_heartbeat_monitor_initial_state() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("hb_test".to_string(), tx));
let config = HeartbeatConfig::default();
let monitor = HeartbeatMonitor::new(conn, config);
assert!(!monitor.is_timed_out().await);
assert!(monitor.time_since_last_pong().await < Duration::from_secs(1));
}
#[rstest]
#[tokio::test]
async fn test_heartbeat_monitor_record_pong_resets_timer() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("hb_pong".to_string(), tx));
let config = HeartbeatConfig::new(Duration::from_millis(50), Duration::from_millis(30));
let monitor = HeartbeatMonitor::new(conn, config);
tokio::time::sleep(Duration::from_millis(20)).await;
monitor.record_pong().await;
assert!(monitor.time_since_last_pong().await < Duration::from_millis(10));
}
#[rstest]
#[tokio::test]
async fn test_heartbeat_monitor_timeout_closes_connection() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("hb_timeout".to_string(), tx));
let config = HeartbeatConfig::new(Duration::from_millis(50), Duration::from_millis(30));
let monitor = HeartbeatMonitor::new(conn.clone(), config);
tokio::time::sleep(Duration::from_millis(40)).await;
let timed_out = monitor.check_heartbeat().await;
assert!(timed_out);
assert!(monitor.is_timed_out().await);
assert!(conn.is_closed().await);
}
#[rstest]
#[tokio::test]
async fn test_heartbeat_monitor_no_timeout_when_pong_received() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("hb_ok".to_string(), tx));
let config = HeartbeatConfig::new(Duration::from_millis(100), Duration::from_millis(50));
let monitor = HeartbeatMonitor::new(conn.clone(), config);
tokio::time::sleep(Duration::from_millis(20)).await;
monitor.record_pong().await;
let timed_out = monitor.check_heartbeat().await;
assert!(!timed_out);
assert!(!monitor.is_timed_out().await);
assert!(!conn.is_closed().await);
}
#[rstest]
#[tokio::test]
async fn test_heartbeat_monitor_send_ping() {
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("hb_ping".to_string(), tx));
let config = HeartbeatConfig::default();
let monitor = HeartbeatMonitor::new(conn, config);
monitor.send_ping().await.unwrap();
let msg = rx.recv().await.unwrap();
assert!(matches!(msg, Message::Ping));
}
#[rstest]
#[tokio::test]
async fn test_heartbeat_monitor_early_pong_skips_full_sleep() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("hb_early".to_string(), tx));
let config = HeartbeatConfig {
ping_interval: Duration::from_secs(60),
pong_timeout: Duration::from_secs(10),
};
let monitor = Arc::new(HeartbeatMonitor::new(conn, config));
let monitor_clone = Arc::clone(&monitor);
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(50)).await;
monitor_clone.record_pong().await;
});
let _ = monitor.send_ping().await;
let start = Instant::now();
tokio::select! {
() = tokio::time::sleep(monitor.config.pong_timeout) => {
panic!("Should not reach full timeout");
}
() = monitor.pong_notify.notified() => {
}
}
let elapsed = start.elapsed();
assert!(
elapsed < Duration::from_secs(2),
"Expected early wakeup but elapsed {:?}",
elapsed
);
}
#[rstest]
fn test_websocket_error_binary_payload_variant() {
let err = WebSocketError::BinaryPayload("invalid data".to_string());
assert_eq!(err.to_string(), "Invalid binary payload: invalid data");
}
#[rstest]
fn test_websocket_error_heartbeat_timeout_variant() {
let err = WebSocketError::HeartbeatTimeout(Duration::from_secs(10));
assert_eq!(
err.to_string(),
"Heartbeat timeout: no pong received within 10s"
);
}
#[rstest]
fn test_websocket_error_slow_consumer_variant() {
let err = WebSocketError::SlowConsumer(Duration::from_secs(5));
assert_eq!(err.to_string(), "Slow consumer: send timed out after 5s");
}
}