1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
use super::ConnectionHandler;
use crate::error::Result;
use crate::protocol::constants::PONG_TIMEOUT;
use crate::protocol::messages::PusherMessage;
use crate::websocket::SocketId;
use bytes::Bytes;
use sockudo_ws::Message;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::sleep;
use tracing::{debug, warn};
impl ConnectionHandler {
pub async fn setup_initial_timeouts(
&self,
socket_id: &SocketId,
app_config: &crate::app::config::App,
) -> Result<()> {
// Set activity timeout
self.set_activity_timeout(&app_config.id, socket_id).await?;
// Set user authentication timeout if required
if app_config.enable_user_authentication.unwrap_or(false) {
let auth_timeout = self.server_options.user_authentication_timeout;
self.set_user_authentication_timeout(&app_config.id, socket_id, auth_timeout)
.await?;
}
Ok(())
}
pub async fn set_activity_timeout(&self, app_id: &str, socket_id: &SocketId) -> Result<()> {
let socket_id_clone = *socket_id;
let app_id_clone = app_id.to_string();
let connection_manager = self.connection_manager.clone();
let activity_timeout = self.server_options.activity_timeout;
// Clear any existing timeout
self.clear_activity_timeout(app_id, socket_id).await?;
let timeout_handle = tokio::spawn(async move {
// Initial sleep before first check
sleep(Duration::from_secs(activity_timeout)).await;
loop {
// Check if connection still exists and get actual inactivity time
let conn_manager = Arc::clone(&connection_manager);
let conn = match conn_manager
.get_connection(&socket_id_clone, &app_id_clone)
.await
{
Some(c) => c,
None => {
// Connection already cleaned up, nothing to do
return;
}
};
// Check actual time since last activity
let time_since_activity = {
let ws = conn.inner.lock().await;
ws.state.time_since_last_ping()
};
// If less than activity timeout seconds have passed since last activity, wait more
if time_since_activity < Duration::from_secs(activity_timeout) {
let remaining = Duration::from_secs(activity_timeout) - time_since_activity;
debug!(
"Socket {} still active ({}s ago), waiting {} more seconds",
socket_id_clone,
time_since_activity.as_secs(),
remaining.as_secs()
);
drop(conn_manager);
sleep(remaining).await;
// Continue to check again without additional delay
continue;
}
// Truly inactive for activity timeout duration, send ping
let ping_result = {
let mut ws = conn.inner.lock().await;
// Update connection status to indicate ping sent
ws.state.status =
crate::websocket::ConnectionStatus::PingSent(std::time::Instant::now());
let ping_message = PusherMessage::ping();
ws.send_message(&ping_message)
};
match ping_result {
Ok(_) => {
debug!(
"Sent ping to socket {} due to activity timeout",
socket_id_clone
);
// Release locks before waiting for pong
drop(conn_manager);
// Wait for pong response
sleep(Duration::from_secs(PONG_TIMEOUT)).await;
// Re-acquire lock to check pong status
let conn_manager = Arc::clone(&connection_manager);
if let Some(conn) = conn_manager
.get_connection(&socket_id_clone, &app_id_clone)
.await
{
let mut ws = conn.inner.lock().await;
// Check if we're still in PingSent state (no pong received)
if let crate::websocket::ConnectionStatus::PingSent(ping_time) =
ws.state.status
&& ping_time.elapsed() > Duration::from_secs(PONG_TIMEOUT)
{
// No pong received, close connection gracefully
warn!(
"No pong received from socket {} after ping, closing connection",
socket_id_clone
);
let _ = ws
.close(4201, "Pong reply not received in time".to_string())
.await;
}
}
// After handling ping/pong, wait full activity timeout before next check
drop(conn_manager);
sleep(Duration::from_secs(activity_timeout)).await;
}
Err(e) => {
// Connection is broken (e.g., broken pipe)
// This is expected when client disconnects abruptly
debug!(
"Failed to send ping to socket {} (connection likely closed by client): {}",
socket_id_clone, e
);
// Clean up the connection since it's broken
// Note: cleanup_connection expects the connection to still exist
conn_manager.cleanup_connection(&app_id_clone, conn).await;
break; // Exit the loop after cleanup
}
}
}
});
// Store the timeout handle
let conn_manager = &self.connection_manager;
if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
let mut ws = conn.inner.lock().await;
ws.state.timeouts.activity_timeout_handle = Some(timeout_handle);
}
Ok(())
}
pub async fn clear_activity_timeout(&self, app_id: &str, socket_id: &SocketId) -> Result<()> {
let conn_manager = &self.connection_manager;
if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
let mut ws = conn.inner.lock().await;
ws.state.timeouts.clear_activity_timeout();
}
Ok(())
}
pub async fn update_activity_timeout(&self, app_id: &str, socket_id: &SocketId) -> Result<()> {
// Update last activity time
let conn_manager = &self.connection_manager;
if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
let mut ws = conn.inner.lock().await;
ws.update_activity();
}
Ok(())
}
pub async fn set_user_authentication_timeout(
&self,
app_id: &str,
socket_id: &SocketId,
timeout_seconds: u64,
) -> Result<()> {
let socket_id_clone = *socket_id;
let app_id_clone = app_id.to_string();
let connection_manager = self.connection_manager.clone();
// Clear any existing auth timeout
self.clear_user_authentication_timeout(app_id, socket_id)
.await?;
let timeout_handle = tokio::spawn(async move {
sleep(Duration::from_secs(timeout_seconds)).await;
let conn_manager = connection_manager;
if let Some(conn) = conn_manager
.get_connection(&socket_id_clone, &app_id_clone)
.await
{
let mut ws = conn.inner.lock().await;
// Check if user is still not authenticated
if !ws.state.is_authenticated() {
let _ = ws
.close(
4009,
"Connection not authorized within timeout.".to_string(),
)
.await;
}
}
});
// Store the timeout handle
let conn_manager = &self.connection_manager;
if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
let mut ws = conn.inner.lock().await;
ws.state.timeouts.auth_timeout_handle = Some(timeout_handle);
}
Ok(())
}
pub async fn clear_user_authentication_timeout(
&self,
app_id: &str,
socket_id: &SocketId,
) -> Result<()> {
let conn_manager = &self.connection_manager;
if let Some(conn) = conn_manager.get_connection(socket_id, app_id).await {
let mut ws = conn.inner.lock().await;
ws.state.timeouts.clear_auth_timeout();
}
Ok(())
}
pub async fn handle_ping_frame(
&self,
socket_id: &SocketId,
app_config: &crate::app::config::App,
payload: Bytes,
) -> Result<()> {
// Update activity and send pong
self.update_activity_timeout(&app_config.id, socket_id)
.await?;
let conn_manager = &self.connection_manager;
if let Some(conn) = conn_manager.get_connection(socket_id, &app_config.id).await {
let mut ws = conn.inner.lock().await;
// Reset connection status to Active when we receive a ping (client is alive)
ws.state.status = crate::websocket::ConnectionStatus::Active;
// Send low-level Pong frame in response because the auto response is disabled to allow custom handling
ws.send_frame(Message::pong(payload))?;
}
Ok(())
}
}