1use crate::error::{WsError, WsResult};
16use crate::frame::{CloseCode, CloseFrame, Frame, OpCode, Role};
17use crate::message::Message;
18use rustlavel_core::{Config, Event, events};
19use rustlavel_http::Upgraded;
20use std::time::{Duration, Instant};
21use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
22use tokio::sync::mpsc;
23
24#[derive(Debug, Clone)]
26pub struct WebSocketConfig {
27 pub max_message_size: usize,
29 pub max_frame_size: usize,
32 pub idle_timeout: Option<Duration>,
35 pub close_timeout: Duration,
37 pub send_queue: usize,
40}
41
42impl Default for WebSocketConfig {
43 fn default() -> Self {
44 WebSocketConfig {
45 max_message_size: 1 << 20,
46 max_frame_size: 1 << 20,
47 idle_timeout: Some(Duration::from_secs(60)),
48 close_timeout: Duration::from_secs(5),
49 send_queue: 64,
50 }
51 }
52}
53
54impl WebSocketConfig {
55 pub fn from_config(config: &Config) -> Self {
60 let defaults = WebSocketConfig::default();
61 let idle = config.int("websocket.idle_timeout", 60).max(0) as u64;
62
63 WebSocketConfig {
64 max_message_size: config
65 .int("websocket.max_message_size", defaults.max_message_size as i64)
66 .max(1) as usize,
67 max_frame_size: config
68 .int("websocket.max_frame_size", defaults.max_frame_size as i64)
69 .max(1) as usize,
70 idle_timeout: (idle > 0).then(|| Duration::from_secs(idle)),
71 close_timeout: Duration::from_secs(
72 config.int("websocket.close_timeout", 5).max(1) as u64
73 ),
74 send_queue: config.int("websocket.send_queue", defaults.send_queue as i64).max(1)
75 as usize,
76 }
77 }
78}
79
80#[derive(Clone)]
86pub struct Sender {
87 queue: mpsc::Sender<Message>,
88}
89
90impl Sender {
91 pub async fn send(&self, message: impl Into<Message>) -> WsResult<()> {
94 self.queue.send(message.into()).await.map_err(|_| WsError::Closed)
95 }
96
97 pub fn try_send(&self, message: impl Into<Message>) -> WsResult<()> {
103 match self.queue.try_send(message.into()) {
104 Ok(()) => Ok(()),
105 Err(mpsc::error::TrySendError::Full(_)) => Err(WsError::Full),
106 Err(mpsc::error::TrySendError::Closed(_)) => Err(WsError::Closed),
107 }
108 }
109
110 pub fn is_closed(&self) -> bool {
111 self.queue.is_closed()
112 }
113
114 pub fn room(&self) -> usize {
116 self.queue.capacity()
117 }
118}
119
120pub fn channel(capacity: usize) -> (Sender, mpsc::Receiver<Message>) {
123 let (queue, receiver) = mpsc::channel(capacity.max(1));
124 (Sender { queue }, receiver)
125}
126
127pub struct WebSocket {
128 reader: Box<dyn AsyncRead + Send + Unpin>,
129 buffer: Vec<u8>,
132 outgoing: Sender,
133 config: WebSocketConfig,
134 partial: Option<(OpCode, Vec<u8>)>,
136 awaiting_pong: bool,
137 closed: bool,
138 close_frame: Option<CloseFrame>,
139 path: String,
140 protocol: Option<String>,
141 opened_at: Instant,
142 received: u64,
143}
144
145impl WebSocket {
146 pub fn new(upgraded: Upgraded, config: WebSocketConfig) -> WebSocket {
148 let (queue, receiver) = mpsc::channel(config.send_queue.max(1));
149 tokio::spawn(write_loop(upgraded.writer, receiver));
150
151 WebSocket {
152 reader: upgraded.reader,
153 buffer: upgraded.buffered,
154 outgoing: Sender { queue },
155 config,
156 partial: None,
157 awaiting_pong: false,
158 closed: false,
159 close_frame: None,
160 path: String::new(),
161 protocol: None,
162 opened_at: Instant::now(),
163 received: 0,
164 }
165 }
166
167 pub fn on_path(mut self, path: impl Into<String>) -> Self {
169 self.path = path.into();
170 self
171 }
172
173 pub fn with_protocol(mut self, protocol: Option<String>) -> Self {
174 self.protocol = protocol;
175 self
176 }
177
178 pub fn sender(&self) -> Sender {
180 self.outgoing.clone()
181 }
182
183 pub fn path(&self) -> &str {
184 &self.path
185 }
186
187 pub fn protocol(&self) -> Option<&str> {
189 self.protocol.as_deref()
190 }
191
192 pub fn is_closed(&self) -> bool {
193 self.closed
194 }
195
196 pub fn close_frame(&self) -> Option<&CloseFrame> {
199 self.close_frame.as_ref()
200 }
201
202 pub fn received(&self) -> u64 {
204 self.received
205 }
206
207 pub async fn send(&mut self, message: impl Into<Message>) -> WsResult<()> {
208 if self.closed {
209 return Err(WsError::Closed);
210 }
211 self.outgoing.send(message).await
212 }
213
214 pub async fn recv(&mut self) -> Option<Message> {
220 loop {
221 if self.closed {
222 return None;
223 }
224
225 let frame = match self.next_frame().await {
226 Ok(Some(frame)) => frame,
227 Ok(None) => {
229 self.mark_closed(Some(CloseFrame::new(CloseCode::ABNORMAL, "")));
230 return None;
231 }
232 Err(error) => {
233 self.fail(error).await;
234 return None;
235 }
236 };
237
238 self.awaiting_pong = false;
240
241 match frame.opcode {
242 OpCode::Ping => {
243 if self.outgoing.send(Message::Pong(frame.payload)).await.is_err() {
246 self.mark_closed(Some(CloseFrame::new(CloseCode::ABNORMAL, "")));
247 return None;
248 }
249 }
250 OpCode::Pong => {}
251 OpCode::Close => {
252 let peer = match CloseFrame::parse(&frame.payload) {
253 Ok(close) => close,
254 Err(error) => {
255 self.fail(error).await;
256 return None;
257 }
258 };
259 let echo =
262 peer.clone().unwrap_or_else(|| CloseFrame::new(CloseCode::NORMAL, ""));
263 let _ = self.outgoing.send(Message::Close(Some(echo))).await;
264 self.mark_closed(
265 peer.or_else(|| Some(CloseFrame::new(CloseCode::NO_STATUS, ""))),
266 );
267 return None;
268 }
269 OpCode::Text | OpCode::Binary => {
270 if self.partial.is_some() {
271 self.fail(WsError::protocol(
272 "a new data frame arrived while a fragmented message was still open",
273 ))
274 .await;
275 return None;
276 }
277 if frame.fin {
278 return self.deliver(frame.opcode, frame.payload).await;
279 }
280 self.partial = Some((frame.opcode, frame.payload));
281 }
282 OpCode::Continuation => {
283 let Some((opcode, mut assembled)) = self.partial.take() else {
284 self.fail(WsError::protocol(
285 "a continuation frame arrived with no message to continue",
286 ))
287 .await;
288 return None;
289 };
290
291 let total = assembled.len() + frame.payload.len();
292 if total > self.config.max_message_size {
293 self.fail(WsError::TooLarge {
296 size: total,
297 limit: self.config.max_message_size,
298 })
299 .await;
300 return None;
301 }
302
303 assembled.extend_from_slice(&frame.payload);
304 if frame.fin {
305 return self.deliver(opcode, assembled).await;
306 }
307 self.partial = Some((opcode, assembled));
308 }
309 }
310 }
311 }
312
313 pub async fn close(&mut self, code: CloseCode, reason: impl Into<String>) {
316 if self.closed {
317 return;
318 }
319 let close = CloseFrame::new(code, reason);
320
321 if self.outgoing.send(Message::Close(Some(close.clone()))).await.is_err() {
322 self.mark_closed(Some(close));
323 return;
324 }
325
326 let peer = tokio::time::timeout(self.config.close_timeout, self.drain_until_close())
330 .await
331 .ok()
332 .flatten();
333 self.mark_closed(peer.or(Some(close)));
334 }
335
336 async fn deliver(&mut self, opcode: OpCode, payload: Vec<u8>) -> Option<Message> {
338 match Message::from_data(opcode, payload) {
339 Ok(message) => {
340 self.received += 1;
341 if events::has_subscribers() {
342 Event::new("ws.message")
345 .with("path", self.path.as_str())
346 .with("kind", message.kind())
347 .with("bytes", message.len())
348 .dispatch();
349 }
350 Some(message)
351 }
352 Err(error) => {
353 self.fail(error).await;
354 None
355 }
356 }
357 }
358
359 async fn fail(&mut self, error: WsError) {
361 let close = CloseFrame::new(error.close_code(), error.to_string());
362 let _ = self.outgoing.send(Message::Close(Some(close.clone()))).await;
363 self.mark_closed(Some(close));
364 }
365
366 fn mark_closed(&mut self, close: Option<CloseFrame>) {
367 self.closed = true;
368 if self.close_frame.is_none() {
369 self.close_frame = close;
370 }
371 }
372
373 async fn drain_until_close(&mut self) -> Option<CloseFrame> {
374 loop {
375 match self.read_frame().await {
376 Ok(Some(frame)) if frame.opcode == OpCode::Close => {
377 return CloseFrame::parse(&frame.payload).ok().flatten();
378 }
379 Ok(Some(_)) => continue,
380 Ok(None) | Err(_) => return None,
381 }
382 }
383 }
384
385 async fn next_frame(&mut self) -> WsResult<Option<Frame>> {
387 let Some(idle) = self.config.idle_timeout else {
388 return self.read_frame().await;
389 };
390
391 loop {
392 match tokio::time::timeout(idle, self.read_frame()).await {
395 Ok(result) => return result,
396 Err(_) if self.awaiting_pong => return Err(WsError::Idle),
397 Err(_) => {
398 self.awaiting_pong = true;
402 if self.outgoing.send(Message::Ping(Vec::new())).await.is_err() {
403 return Ok(None);
404 }
405 }
406 }
407 }
408 }
409
410 async fn read_frame(&mut self) -> WsResult<Option<Frame>> {
413 loop {
414 if let Some((frame, used)) =
415 Frame::decode(&self.buffer, Role::Server, self.config.max_frame_size)?
416 {
417 self.buffer.drain(..used);
418 return Ok(Some(frame));
419 }
420
421 let mut chunk = [0u8; 8192];
422 let read = self.reader.read(&mut chunk).await?;
423 if read == 0 {
424 return Ok(None);
425 }
426 self.buffer.extend_from_slice(&chunk[..read]);
427 }
428 }
429}
430
431impl Drop for WebSocket {
434 fn drop(&mut self) {
435 if !events::has_subscribers() {
436 return;
437 }
438 let code = self.close_frame.as_ref().map_or(CloseCode::ABNORMAL, |close| close.code);
439
440 Event::new("ws.disconnected")
441 .with("path", self.path.as_str())
442 .with("messages", self.received)
443 .with("code", code.code())
444 .took(self.opened_at.elapsed())
445 .dispatch();
446 }
447}
448
449async fn write_loop(
455 mut writer: Box<dyn AsyncWrite + Send + Unpin>,
456 mut queue: mpsc::Receiver<Message>,
457) {
458 while let Some(message) = queue.recv().await {
459 let closing = matches!(message, Message::Close(_));
460 let bytes = message.into_frame().encode();
461
462 if writer.write_all(&bytes).await.is_err() || writer.flush().await.is_err() {
463 break;
464 }
465 if closing {
466 break;
468 }
469 }
470
471 queue.close();
474 let _ = writer.shutdown().await;
475}
476
477#[cfg(test)]
478mod tests {
479 use super::*;
480 use tokio::io::DuplexStream;
481
482 const MASK: [u8; 4] = [0xa1, 0xb2, 0xc3, 0xd4];
484
485 struct Peer {
488 stream: DuplexStream,
489 buffer: Vec<u8>,
490 }
491
492 impl Peer {
493 async fn send(&mut self, frame: Frame) {
494 self.stream.write_all(&frame.encode()).await.unwrap();
495 }
496
497 async fn send_masked(&mut self, opcode: OpCode, payload: &[u8]) {
499 self.send(Frame::client(opcode, payload.to_vec(), MASK)).await;
500 }
501
502 async fn send_fragment(&mut self, opcode: OpCode, payload: &[u8], fin: bool) {
503 self.send(Frame { fin, opcode, mask: Some(MASK), payload: payload.to_vec() }).await;
504 }
505
506 async fn recv(&mut self) -> Frame {
509 loop {
510 if let Some((frame, used)) =
511 Frame::decode(&self.buffer, Role::Client, 1 << 20).unwrap()
512 {
513 self.buffer.drain(..used);
514 return frame;
515 }
516 let mut chunk = [0u8; 4096];
517 let read = self.stream.read(&mut chunk).await.unwrap();
518 assert!(read > 0, "the server closed without sending an expected frame");
519 self.buffer.extend_from_slice(&chunk[..read]);
520 }
521 }
522
523 async fn hang_up(&mut self) {
524 self.stream.shutdown().await.unwrap();
525 }
526 }
527
528 fn pair(config: WebSocketConfig) -> (WebSocket, Peer) {
529 let (client, server) = tokio::io::duplex(64 * 1024);
530 let (reader, writer) = tokio::io::split(server);
531 let upgraded = Upgraded {
532 reader: Box::new(reader),
533 writer: Box::new(writer),
534 buffered: Vec::new(),
535 };
536 (WebSocket::new(upgraded, config), Peer { stream: client, buffer: Vec::new() })
537 }
538
539 async fn close_from_server(peer: &mut Peer) -> CloseFrame {
541 loop {
542 let frame = peer.recv().await;
543 if frame.opcode == OpCode::Close {
544 return CloseFrame::parse(&frame.payload).unwrap().unwrap();
545 }
546 }
547 }
548
549 #[tokio::test]
550 async fn echoes_a_message_the_client_sent() {
551 let (mut socket, mut peer) = pair(WebSocketConfig::default());
552
553 peer.send_masked(OpCode::Text, b"hello").await;
554 let message = socket.recv().await.unwrap();
555 assert_eq!(message, Message::text("hello"));
556
557 socket.send(message).await.unwrap();
558 let echoed = peer.recv().await;
559 assert_eq!(echoed.opcode, OpCode::Text);
560 assert_eq!(echoed.mask, None, "a server must never mask");
561 assert_eq!(echoed.payload, b"hello");
562 }
563
564 #[tokio::test]
565 async fn bytes_already_buffered_by_the_http_server_are_not_lost() {
566 let (client, server) = tokio::io::duplex(1024);
567 let (reader, writer) = tokio::io::split(server);
568 let early = Frame::client(OpCode::Text, b"first".to_vec(), MASK).encode();
570
571 let mut socket = WebSocket::new(
572 Upgraded { reader: Box::new(reader), writer: Box::new(writer), buffered: early },
573 WebSocketConfig::default(),
574 );
575
576 assert_eq!(socket.recv().await.unwrap(), Message::text("first"));
577 drop(client);
578 }
579
580 #[tokio::test]
581 async fn answers_a_ping_with_a_pong_carrying_the_same_payload() {
582 let (mut socket, mut peer) = pair(WebSocketConfig::default());
583
584 peer.send_masked(OpCode::Ping, b"keepalive").await;
585 peer.send_masked(OpCode::Text, b"after").await;
586
587 assert_eq!(socket.recv().await.unwrap(), Message::text("after"));
589
590 let pong = peer.recv().await;
591 assert_eq!(pong.opcode, OpCode::Pong);
592 assert_eq!(pong.payload, b"keepalive");
593 }
594
595 #[tokio::test]
596 async fn reassembles_a_fragmented_message_around_an_interleaved_ping() {
597 let (mut socket, mut peer) = pair(WebSocketConfig::default());
598
599 peer.send_fragment(OpCode::Text, b"Hel", false).await;
602 peer.send_masked(OpCode::Ping, b"mid").await;
603 peer.send_fragment(OpCode::Continuation, b"lo ", false).await;
604 peer.send_fragment(OpCode::Continuation, b"there", true).await;
605
606 assert_eq!(socket.recv().await.unwrap(), Message::text("Hello there"));
607
608 let pong = peer.recv().await;
609 assert_eq!(pong.opcode, OpCode::Pong);
610 assert_eq!(pong.payload, b"mid");
611 }
612
613 #[tokio::test]
614 async fn a_continuation_with_nothing_to_continue_closes_with_1002() {
615 let (mut socket, mut peer) = pair(WebSocketConfig::default());
616
617 peer.send_fragment(OpCode::Continuation, b"orphan", true).await;
618
619 assert!(socket.recv().await.is_none());
620 assert_eq!(close_from_server(&mut peer).await.code, CloseCode::PROTOCOL_ERROR);
621 }
622
623 #[tokio::test]
624 async fn a_text_frame_that_is_not_utf8_closes_with_1007() {
625 let (mut socket, mut peer) = pair(WebSocketConfig::default());
626
627 peer.send_masked(OpCode::Text, &[0xf0, 0x28, 0x8c, 0x28]).await;
628
629 assert!(socket.recv().await.is_none());
630 assert_eq!(close_from_server(&mut peer).await.code, CloseCode::INVALID_PAYLOAD);
631 assert_eq!(socket.close_frame().unwrap().code, CloseCode::INVALID_PAYLOAD);
632 }
633
634 #[tokio::test]
635 async fn an_unmasked_client_frame_closes_with_1002() {
636 let (mut socket, mut peer) = pair(WebSocketConfig::default());
637
638 peer.send(Frame::server(OpCode::Text, b"naked".to_vec())).await;
640
641 assert!(socket.recv().await.is_none());
642 let close = close_from_server(&mut peer).await;
643 assert_eq!(close.code, CloseCode::PROTOCOL_ERROR);
644 assert!(close.reason.contains("must be masked"), "{}", close.reason);
645 }
646
647 #[tokio::test]
648 async fn an_oversize_frame_closes_with_1009() {
649 let config = WebSocketConfig { max_frame_size: 64, ..WebSocketConfig::default() };
650 let (mut socket, mut peer) = pair(config);
651
652 peer.send_masked(OpCode::Binary, &[7u8; 200]).await;
653
654 assert!(socket.recv().await.is_none());
655 assert_eq!(close_from_server(&mut peer).await.code, CloseCode::TOO_LARGE);
656 }
657
658 #[tokio::test]
659 async fn fragments_that_add_up_to_more_than_the_limit_close_with_1009() {
660 let config = WebSocketConfig {
661 max_message_size: 8,
662 max_frame_size: 1024,
663 ..WebSocketConfig::default()
664 };
665 let (mut socket, mut peer) = pair(config);
666
667 peer.send_fragment(OpCode::Binary, &[0u8; 5], false).await;
669 peer.send_fragment(OpCode::Continuation, &[0u8; 5], true).await;
670
671 assert!(socket.recv().await.is_none());
672 assert_eq!(close_from_server(&mut peer).await.code, CloseCode::TOO_LARGE);
673 }
674
675 #[tokio::test]
676 async fn a_close_from_the_peer_is_echoed_back() {
677 let (mut socket, mut peer) = pair(WebSocketConfig::default());
678
679 peer.send(Frame::client(
680 OpCode::Close,
681 CloseFrame::new(CloseCode::NORMAL, "bye").to_payload(),
682 MASK,
683 ))
684 .await;
685
686 assert!(socket.recv().await.is_none());
687
688 let echo = close_from_server(&mut peer).await;
689 assert_eq!(echo.code, CloseCode::NORMAL);
690 assert_eq!(echo.reason, "bye");
691 assert_eq!(socket.close_frame().unwrap().code, CloseCode::NORMAL);
693 }
694
695 #[tokio::test]
696 async fn closing_waits_for_the_peers_half_of_the_handshake() {
697 let (mut socket, mut peer) = pair(WebSocketConfig::default());
698
699 let reply = async {
700 let frame = peer.recv().await;
701 assert_eq!(frame.opcode, OpCode::Close);
702 assert_eq!(
703 CloseFrame::parse(&frame.payload).unwrap().unwrap().code,
704 CloseCode::GOING_AWAY
705 );
706 peer.send(Frame::client(
707 OpCode::Close,
708 CloseFrame::new(CloseCode(4001), "acknowledged").to_payload(),
709 MASK,
710 ))
711 .await;
712 };
713
714 tokio::join!(socket.close(CloseCode::GOING_AWAY, "restarting"), reply);
715
716 assert!(socket.is_closed());
717 assert_eq!(socket.close_frame().unwrap().code, CloseCode(4001));
719 assert!(socket.send(Message::text("late")).await.is_err());
720 }
721
722 #[tokio::test(start_paused = true)]
723 async fn an_idle_socket_is_pinged_and_then_closed() {
724 let config = WebSocketConfig {
725 idle_timeout: Some(Duration::from_secs(30)),
726 ..WebSocketConfig::default()
727 };
728 let (mut socket, mut peer) = pair(config);
729
730 let reading = tokio::spawn(async move {
731 let message = socket.recv().await;
732 (message, socket.close_frame().cloned())
733 });
734
735 let ping = peer.recv().await;
737 assert_eq!(ping.opcode, OpCode::Ping);
738
739 assert_eq!(close_from_server(&mut peer).await.code, CloseCode::GOING_AWAY);
741
742 let (message, close) = reading.await.unwrap();
743 assert!(message.is_none());
744 assert_eq!(close.unwrap().code, CloseCode::GOING_AWAY);
745 }
746
747 #[tokio::test(start_paused = true)]
748 async fn a_pong_keeps_an_otherwise_idle_socket_open() {
749 let config = WebSocketConfig {
750 idle_timeout: Some(Duration::from_secs(30)),
751 ..WebSocketConfig::default()
752 };
753 let (mut socket, mut peer) = pair(config);
754
755 let reading = tokio::spawn(async move { socket.recv().await });
756
757 let ping = peer.recv().await;
758 assert_eq!(ping.opcode, OpCode::Ping);
759 peer.send_masked(OpCode::Pong, &ping.payload).await;
760 peer.send_masked(OpCode::Text, b"still here").await;
761
762 assert_eq!(reading.await.unwrap(), Some(Message::text("still here")));
763 }
764
765 #[tokio::test]
766 async fn a_socket_that_just_ends_reports_an_abnormal_close() {
767 let (mut socket, mut peer) = pair(WebSocketConfig::default());
768
769 peer.hang_up().await;
770
771 assert!(socket.recv().await.is_none());
772 assert_eq!(socket.close_frame().unwrap().code, CloseCode::ABNORMAL);
773 }
774
775 #[tokio::test]
776 async fn a_full_send_queue_is_reported_rather_than_waited_on() {
777 let (sender, mut receiver) = channel(2);
778
779 assert!(sender.try_send(Message::text("one")).is_ok());
780 assert!(sender.try_send(Message::text("two")).is_ok());
781 assert!(matches!(sender.try_send(Message::text("three")), Err(WsError::Full)));
784
785 assert_eq!(receiver.recv().await, Some(Message::text("one")));
786 assert!(sender.try_send(Message::text("three")).is_ok());
787 }
788
789 #[tokio::test]
790 async fn senders_report_closed_once_the_socket_is_gone() {
791 let (sender, receiver) = channel(2);
792 drop(receiver);
793
794 assert!(sender.is_closed());
795 assert!(matches!(sender.try_send(Message::text("x")), Err(WsError::Closed)));
796 }
797
798 #[test]
799 fn limits_come_from_the_application_configuration() {
800 let config = Config::new();
801 config.set("websocket.max_message_size", 4096);
802 config.set("websocket.idle_timeout", 0);
803 config.set("websocket.send_queue", 8);
804
805 let settings = WebSocketConfig::from_config(&config);
806
807 assert_eq!(settings.max_message_size, 4096);
808 assert_eq!(settings.send_queue, 8);
809 assert_eq!(settings.idle_timeout, None);
811 assert_eq!(settings.max_frame_size, WebSocketConfig::default().max_frame_size);
813 }
814}