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
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
//! WebSocket connection handling
use super::frame::{Frame, Message, OpCode};
use super::pubsub::ChannelManager;
use super::WebSocketConfig;
use bytes::{Bytes, BytesMut};
use hyper::upgrade::Upgraded;
use hyper_util::rt::TokioIo;
use serde::Serialize;
use std::io::{self, ErrorKind};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::sync::mpsc;
use tokio::time;
/// WebSocket connection with typed context data
pub struct WebSocket<T = ()> {
data: T,
sender: mpsc::Sender<Message>,
channel_manager: Arc<ChannelManager>,
connection_id: uuid::Uuid,
remote_addr: Option<SocketAddr>,
config: Arc<WebSocketConfig>,
}
impl<T> WebSocket<T> {
/// Create a new WebSocket connection
pub(crate) fn new(
data: T,
sender: mpsc::Sender<Message>,
channel_manager: Arc<ChannelManager>,
connection_id: uuid::Uuid,
remote_addr: Option<SocketAddr>,
config: Arc<WebSocketConfig>,
) -> Self {
Self {
data,
sender,
channel_manager,
connection_id,
remote_addr,
config,
}
}
/// Get reference to typed context data
pub fn data(&self) -> &T {
&self.data
}
/// Get mutable reference to typed context data
pub fn data_mut(&mut self) -> &mut T {
&mut self.data
}
/// Get reference to WebSocket configuration
pub fn config(&self) -> &WebSocketConfig {
&self.config
}
/// Send text message
///
/// Returns `Err` if the connection is closed or the write buffer is full.
/// When the buffer is full, consider waiting for `on_drain` callback before retrying.
pub async fn send(&self, text: impl Into<String>) -> Result<(), std::io::Error> {
self.sender
.try_send(Message::Text(text.into()))
.map_err(|e| match e {
mpsc::error::TrySendError::Full(_) => {
std::io::Error::new(std::io::ErrorKind::WouldBlock, "write buffer full")
}
mpsc::error::TrySendError::Closed(_) => {
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "connection closed")
}
})
}
/// Send binary message
///
/// Returns `Err` if the connection is closed or the write buffer is full.
/// When the buffer is full, consider waiting for `on_drain` callback before retrying.
pub async fn send_binary(&self, data: impl Into<Bytes>) -> Result<(), std::io::Error> {
self.sender
.try_send(Message::Binary(data.into()))
.map_err(|e| match e {
mpsc::error::TrySendError::Full(_) => {
std::io::Error::new(std::io::ErrorKind::WouldBlock, "write buffer full")
}
mpsc::error::TrySendError::Closed(_) => {
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "connection closed")
}
})
}
/// Send JSON message
pub async fn send_json<S: Serialize>(&self, data: &S) -> Result<(), std::io::Error> {
let json = serde_json::to_string(data)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
self.send(json).await
}
/// Subscribe to a topic/channel
pub async fn subscribe(&self, topic: impl AsRef<str>) -> Result<(), std::io::Error> {
self.channel_manager
.subscribe(self.connection_id, topic.as_ref(), self.sender.clone())
.await
}
/// Unsubscribe from a topic/channel
pub async fn unsubscribe(&self, topic: impl AsRef<str>) -> Result<(), std::io::Error> {
self.channel_manager
.unsubscribe(self.connection_id, topic.as_ref())
.await
}
/// Publish message to all subscribers of a topic
pub async fn publish<S: Serialize>(
&self,
topic: impl AsRef<str>,
data: &S,
) -> Result<usize, std::io::Error> {
let json = serde_json::to_string(data)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
self.channel_manager
.publish(topic.as_ref(), Message::Text(json))
.await
}
/// Close the WebSocket connection
pub async fn close(
&self,
code: Option<u16>,
reason: Option<&str>,
) -> Result<(), std::io::Error> {
let close_frame = if let Some(c) = code {
let r = reason.unwrap_or("");
let mut payload = BytesMut::new();
payload.extend_from_slice(&c.to_be_bytes());
payload.extend_from_slice(r.as_bytes());
Message::Close(Some(super::frame::CloseFrame {
code: c,
reason: r.to_string(),
}))
} else {
Message::Close(None)
};
self.sender.try_send(close_frame).map_err(|e| match e {
mpsc::error::TrySendError::Full(_) => {
std::io::Error::new(std::io::ErrorKind::WouldBlock, "write buffer full")
}
mpsc::error::TrySendError::Closed(_) => {
std::io::Error::new(std::io::ErrorKind::BrokenPipe, "connection closed")
}
})
}
/// Get remote address
pub fn remote_addr(&self) -> Option<SocketAddr> {
self.remote_addr
}
/// Check if connection is writable (for backpressure)
pub fn is_writable(&self) -> bool {
!self.sender.is_closed()
}
/// Get maximum capacity of the write queue
pub fn max_capacity(&self) -> usize {
self.sender.max_capacity()
}
/// Get current number of messages in the write queue
pub fn capacity(&self) -> usize {
self.sender.capacity()
}
/// Check if the write buffer has available capacity
pub fn has_capacity(&self) -> bool {
self.sender.capacity() > 0
}
}
/// WebSocket connection handler that manages the connection lifecycle
pub(crate) struct ConnectionHandler {
upgraded: Upgraded,
receiver: mpsc::Receiver<Message>,
incoming_tx: mpsc::UnboundedSender<Message>,
drain_tx: mpsc::UnboundedSender<()>,
channel_manager: Arc<ChannelManager>,
connection_id: uuid::Uuid,
config: Arc<WebSocketConfig>,
}
/// Fragment accumulator for reassembling fragmented messages
struct FragmentAccumulator {
opcode: Option<OpCode>,
fragments: BytesMut,
total_size: usize,
}
impl ConnectionHandler {
pub fn new(
upgraded: Upgraded,
channel_manager: Arc<ChannelManager>,
config: Arc<WebSocketConfig>,
) -> (
Self,
mpsc::Sender<Message>,
mpsc::UnboundedReceiver<Message>,
mpsc::UnboundedReceiver<()>,
) {
// Use bounded channel for outgoing messages (backpressure)
let (tx, rx) = mpsc::channel(config.max_write_queue_size);
// Use unbounded channel for incoming messages (from client)
let (incoming_tx, incoming_rx) = mpsc::unbounded_channel();
// Use unbounded channel for drain notifications
let (drain_tx, drain_rx) = mpsc::unbounded_channel();
let connection_id = uuid::Uuid::new_v4();
let handler = Self {
upgraded,
receiver: rx,
incoming_tx,
drain_tx,
channel_manager,
connection_id,
config,
};
(handler, tx, incoming_rx, drain_rx)
}
pub async fn handle(self) -> Result<(), std::io::Error> {
tracing::info!("ConnectionHandler::handle() started");
let mut read_buf = BytesMut::with_capacity(8192);
let io = TokioIo::new(self.upgraded);
let (mut reader, mut writer) = tokio::io::split(io);
let mut receiver = self.receiver;
let channel_manager = self.channel_manager;
let connection_id = self.connection_id;
let incoming_tx = self.incoming_tx;
let drain_tx = self.drain_tx;
let config = self.config;
let mut fragment_accumulator: Option<FragmentAccumulator> = None;
// Track write buffer state for backpressure
let mut was_full = false;
// Setup ping interval if configured
let mut ping_interval = config
.ping_interval
.map(|interval_secs| time::interval(Duration::from_secs(interval_secs)));
let mut last_pong_received = Instant::now();
let ping_timeout = Duration::from_secs(config.ping_timeout);
tracing::info!("Entering main WebSocket loop");
loop {
// Check for ping timeout
if config.ping_interval.is_some() && last_pong_received.elapsed() > ping_timeout {
tracing::warn!("Ping timeout - closing connection");
let close_frame = Frame::close(Some(1001), Some("Ping timeout"));
let _ = writer.write_all(&close_frame.encode()).await;
break;
}
tokio::select! {
// Ping interval
_ = async {
match &mut ping_interval {
Some(interval) => interval.tick().await,
None => std::future::pending().await,
}
} => {
// Send ping frame
let ping = Frame::ping(Bytes::new());
if let Err(e) = writer.write_all(&ping.encode()).await {
tracing::error!("Error sending ping: {}", e);
break;
}
tracing::trace!("Sent ping frame");
}
// Read frames from client
result = reader.read_buf(&mut read_buf) => {
match result {
Ok(0) => break, // Connection closed
Ok(_) => {
// Try to parse frames with size limits
while let Some(frame) = Frame::parse_with_limits(&mut read_buf, Some(config.max_frame_size))? {
match frame.opcode {
OpCode::Text | OpCode::Binary => {
if frame.fin {
// Single unfragmented message
if let Ok(message) = Message::from_frame_with_limit(frame, Some(config.max_message_size)) {
let _ = incoming_tx.send(message);
}
} else {
// Start of fragmented message
if fragment_accumulator.is_some() {
return Err(io::Error::new(
ErrorKind::InvalidData,
"received new fragment before previous completed",
));
}
fragment_accumulator = Some(FragmentAccumulator {
opcode: Some(frame.opcode),
fragments: BytesMut::from(frame.payload.as_ref()),
total_size: frame.payload.len(),
});
}
}
OpCode::Continue => {
// Continuation frame
let should_clear = if let Some(ref mut accumulator) = fragment_accumulator {
accumulator.total_size += frame.payload.len();
// Check message size limit
if accumulator.total_size > config.max_message_size {
return Err(io::Error::new(
ErrorKind::InvalidData,
format!("fragmented message size {} exceeds maximum {}",
accumulator.total_size, config.max_message_size),
));
}
accumulator.fragments.extend_from_slice(&frame.payload);
frame.fin // Clear accumulator if this is the final fragment
} else {
return Err(io::Error::new(
ErrorKind::InvalidData,
"received continuation frame without initial fragment",
));
};
if should_clear {
// Take ownership and reassemble
if let Some(accumulator) = fragment_accumulator.take() {
let reassembled_frame = Frame {
fin: true,
opcode: accumulator.opcode.unwrap(),
mask: None,
payload: accumulator.fragments.freeze(),
};
if let Ok(message) = Message::from_frame(reassembled_frame) {
let _ = incoming_tx.send(message);
}
}
}
}
OpCode::Close => {
// Send close frame to handler
if let Ok(message) = Message::from_frame(frame) {
let _ = incoming_tx.send(message);
}
// Echo close frame back
let close_frame = Frame::close(Some(1000), Some("Normal closure"));
let _ = writer.write_all(&close_frame.encode()).await;
break;
}
OpCode::Ping => {
// Respond with pong
let pong = Frame::pong(frame.payload);
let _ = writer.write_all(&pong.encode()).await;
}
OpCode::Pong => {
// Update last pong received time
last_pong_received = Instant::now();
tracing::trace!("Received pong frame");
}
}
}
}
Err(e) => {
tracing::error!("Error reading from socket: {}", e);
break;
}
}
}
// Send frames to client
Some(message) = receiver.recv() => {
// Check if buffer was full before this recv
let is_now_draining = was_full;
// Use fragmentation if message exceeds max frame size
let frames = message.to_fragmented_frames(config.max_frame_size);
for frame in frames {
let encoded = frame.encode();
if let Err(e) = writer.write_all(&encoded).await {
tracing::error!("Error writing to socket: {}", e);
break;
}
}
// Update full state - buffer is full if no capacity remains
let has_capacity = !receiver.is_empty() || receiver.is_closed();
was_full = !has_capacity;
// Notify on_drain if buffer was full and now has drained
if is_now_draining && has_capacity {
let _ = drain_tx.send(());
tracing::trace!("Write buffer drained, notified handler");
}
}
}
}
// Cleanup on disconnect
channel_manager.disconnect(connection_id).await;
Ok(())
}
}