Skip to main content

rustlavel_ws/
connection.rs

1//! A live WebSocket connection.
2//!
3//! [`WebSocket`] owns the read half and drives everything the protocol requires
4//! but an application should never have to think about: pongs are sent for
5//! pings, fragments are reassembled, text is validated as UTF-8, oversize
6//! messages are refused, and a close is answered with a close.
7//!
8//! Writing happens in a separate task fed by a bounded queue. That split is
9//! what makes broadcasting possible: a [`Sender`] can be cloned and handed to
10//! anything that wants to push to this client, while `recv` keeps exclusive
11//! ownership of the reader. It is also what makes backpressure expressible —
12//! the queue has a fixed depth, and a client that will not drain it is dropped
13//! rather than allowed to grow the server's memory.
14
15use 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/// Limits and timers for one socket.
25#[derive(Debug, Clone)]
26pub struct WebSocketConfig {
27    /// The largest reassembled message that will be accepted.
28    pub max_message_size: usize,
29    /// The largest single frame that will be accepted. Checked against the
30    /// declared length, so an absurd announcement costs no memory.
31    pub max_frame_size: usize,
32    /// How long a connection may be silent before it is pinged. A second
33    /// silent period with no pong closes it. `None` disables the check.
34    pub idle_timeout: Option<Duration>,
35    /// How long to wait for the peer's half of the close handshake.
36    pub close_timeout: Duration,
37    /// Depth of the outgoing queue. This is the backpressure budget: a client
38    /// this far behind is dropped.
39    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    /// Read the limits from the application's configuration, so an operator can
56    /// raise them in `.env` without recompiling.
57    ///
58    /// `websocket.idle_timeout` is in seconds; zero turns the check off.
59    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/// A cloneable handle for pushing messages at one client.
81///
82/// Handed to the broadcaster, to background tasks, to anything that produces
83/// events. It never blocks the reader, and it can outlive the socket: once the
84/// connection is gone every clone reports closed.
85#[derive(Clone)]
86pub struct Sender {
87    queue: mpsc::Sender<Message>,
88}
89
90impl Sender {
91    /// Queue a message, waiting for room. Use this from the connection's own
92    /// task, where waiting is what you want.
93    pub async fn send(&self, message: impl Into<Message>) -> WsResult<()> {
94        self.queue.send(message.into()).await.map_err(|_| WsError::Closed)
95    }
96
97    /// Queue a message only if there is room right now.
98    ///
99    /// This is the shape broadcasting needs: fan-out must never wait on the
100    /// slowest subscriber, so a full queue is a failure the caller acts on by
101    /// dropping the subscriber.
102    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    /// How many more messages fit before this client counts as too slow.
115    pub fn room(&self) -> usize {
116        self.queue.capacity()
117    }
118}
119
120/// A sender with no socket behind it, for tests and for anything that wants to
121/// collect what would have been written.
122pub 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    /// Bytes read but not yet framed. Seeded with whatever the HTTP server had
130    /// already buffered past the request — dropping those loses the first frame.
131    buffer: Vec<u8>,
132    outgoing: Sender,
133    config: WebSocketConfig,
134    /// The opcode and bytes of a message still arriving in fragments.
135    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    /// Take over a connection the HTTP server has finished with.
147    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    /// Label the socket for instrumentation with the route it was opened on.
168    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    /// A handle other tasks can push messages through.
179    pub fn sender(&self) -> Sender {
180        self.outgoing.clone()
181    }
182
183    pub fn path(&self) -> &str {
184        &self.path
185    }
186
187    /// The subprotocol agreed during the handshake, if any.
188    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    /// How the connection ended, once it has. Available after `recv` returns
197    /// `None`, which is where an application usually wants to log it.
198    pub fn close_frame(&self) -> Option<&CloseFrame> {
199        self.close_frame.as_ref()
200    }
201
202    /// How many data messages have arrived.
203    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    /// The next message from the peer, or `None` once the connection is over.
215    ///
216    /// Only `Text` and `Binary` come out of here. Pings, pongs and closes are
217    /// answered internally, because every application would have to write the
218    /// same three replies and one of them would get it wrong.
219    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                // The socket ended without a close frame: nothing to reply to.
228                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            // Any frame at all proves the peer is alive, not just a pong.
239            self.awaiting_pong = false;
240
241            match frame.opcode {
242                OpCode::Ping => {
243                    // The pong must carry the ping's application data back
244                    // unchanged; that is how the peer matches them up.
245                    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                    // The other half of the close handshake: echo the status
260                    // back so the peer knows the close was seen, not lost.
261                    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                        // Checked as fragments arrive, not at the end: otherwise
294                        // a peer could spend the whole limit one frame at a time.
295                        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    /// Close the connection politely: send a close frame, then wait — briefly —
314    /// for the peer's.
315    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        // Frames the peer already sent are drained rather than reset, which is
327        // what stops a clean close from looking like a crash on the other end.
328        // Bounded, because a peer that never answers must not pin this task.
329        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    /// Hand a reassembled payload to the application.
337    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                    // Size and kind only. A socket carrying chat, tokens or
343                    // order data must not spill its payloads into Telescope.
344                    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    /// Report a protocol failure to the peer and stop.
360    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    /// One frame, with the idle timer wrapped around it.
386    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            // Cancelling the read is safe: bytes already read live in
393            // `self.buffer`, and a read that had not completed consumed nothing.
394            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                    // Silent for a whole period: prod the peer once. A TCP
399                    // connection to a machine that has vanished stays "open"
400                    // for a long time, so silence has to be tested, not assumed.
401                    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    /// One frame, reading more bytes until the buffer holds a whole one.
411    /// `Ok(None)` means the peer closed the socket underneath us.
412    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
431/// Dispatched here so every socket reports exactly once, however it ended —
432/// including the ones an application drops without draining.
433impl 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
449/// Serialize everything queued for this client onto the socket.
450///
451/// One task per connection means writes are ordered and never interleave two
452/// half-written frames, which is exactly the bug an unsynchronised broadcaster
453/// would produce.
454async 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            // Nothing may follow a close frame.
467            break;
468        }
469    }
470
471    // Closing the queue makes every surviving `Sender` clone — the broadcaster
472    // holds some — report the client as gone instead of quietly filling up.
473    queue.close();
474    let _ = writer.shutdown().await;
475}
476
477#[cfg(test)]
478mod tests {
479    use super::*;
480    use tokio::io::DuplexStream;
481
482    /// A fixed mask, so the bytes a test writes are reproducible.
483    const MASK: [u8; 4] = [0xa1, 0xb2, 0xc3, 0xd4];
484
485    /// The other end of an in-memory socket: no ports, no sockets, no sharing
486    /// between tests.
487    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        /// A client frame: always masked.
498        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        /// The next frame the server wrote. Panics if the socket ends first,
507        /// which in a test is the failure we want to see.
508        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    /// The close frame the server sent, whatever else it sent first.
540    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        // The first frame arrived glued to the handshake request.
569        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        // The ping never surfaces; the message behind it does.
588        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        // A control frame between two fragments is legal, and is the classic
600        // place a hand-written implementation corrupts the message.
601        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        // A server frame sent by a client: unmasked, and therefore illegal.
639        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        // Each fragment is small; together they are over the limit.
668        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        // The socket remembers how it ended, for the application's log line.
692        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        // The peer's status wins: it is the more informative of the two.
718        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        // Nothing is sent, so the first quiet period produces a ping…
736        let ping = peer.recv().await;
737        assert_eq!(ping.opcode, OpCode::Ping);
738
739        // …and the second, with no pong in between, ends the connection.
740        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        // The third has nowhere to go, and waiting for room is exactly what
782        // broadcasting must not do.
783        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        // Zero seconds means "never ping", which is what a test harness wants.
810        assert_eq!(settings.idle_timeout, None);
811        // Anything unset keeps the built-in default.
812        assert_eq!(settings.max_frame_size, WebSocketConfig::default().max_frame_size);
813    }
814}