tokio_websockets/proto/
types.rs

1//! Types required for the WebSocket protocol implementation.
2use std::{fmt, mem::replace, num::NonZeroU16, ops::Deref};
3
4use bytes::{BufMut, Bytes, BytesMut};
5
6use super::error::ProtocolError;
7use crate::utf8;
8
9/// The opcode of a WebSocket frame. It denotes the type of the frame or an
10/// assembled message.
11///
12/// A fully assembled [`Message`] will never have a continuation opcode.
13#[derive(Debug, PartialEq, Eq, Clone, Copy)]
14pub(super) enum OpCode {
15    /// A continuation opcode. This will never be encountered in a full
16    /// [`Message`].
17    Continuation,
18    /// A text opcode.
19    Text,
20    /// A binary opcode.
21    Binary,
22    /// A close opcode.
23    Close,
24    /// A ping opcode.
25    Ping,
26    /// A pong opcode.
27    Pong,
28}
29
30impl OpCode {
31    /// Whether this is a control opcode (i.e. close, ping or pong).
32    pub(super) fn is_control(self) -> bool {
33        matches!(self, Self::Close | Self::Ping | Self::Pong)
34    }
35}
36
37impl TryFrom<u8> for OpCode {
38    type Error = ProtocolError;
39
40    fn try_from(value: u8) -> Result<Self, Self::Error> {
41        match value {
42            0 => Ok(Self::Continuation),
43            1 => Ok(Self::Text),
44            2 => Ok(Self::Binary),
45            8 => Ok(Self::Close),
46            9 => Ok(Self::Ping),
47            10 => Ok(Self::Pong),
48            _ => Err(ProtocolError::InvalidOpcode),
49        }
50    }
51}
52
53impl From<OpCode> for u8 {
54    fn from(value: OpCode) -> Self {
55        match value {
56            OpCode::Continuation => 0,
57            OpCode::Text => 1,
58            OpCode::Binary => 2,
59            OpCode::Close => 8,
60            OpCode::Ping => 9,
61            OpCode::Pong => 10,
62        }
63    }
64}
65
66/// Close status code.
67#[derive(Clone, Copy, Debug, Eq, PartialEq)]
68pub struct CloseCode(NonZeroU16);
69
70// rustfmt reorders these alphabetically
71#[rustfmt::skip]
72impl CloseCode {
73    /// Normal closure, meaning that the purpose for which the connection was
74    /// established has been fulfilled.
75    pub const NORMAL_CLOSURE: Self = Self::constant(1000);
76    /// Endpoint is "going away", such as a server going down or a browser
77    /// having navigated away from a page.
78    pub const GOING_AWAY: Self = Self::constant(1001);
79    /// Endpoint is terminating the connection due to a protocol error.
80    pub const PROTOCOL_ERROR: Self = Self::constant(1002);
81    /// Endpoint is terminating the connection because it has received a type of
82    /// data it cannot accept.
83    pub const UNSUPPORTED_DATA: Self = Self::constant(1003);
84    /// No status code was actually present.
85    pub const NO_STATUS_RECEIVED: Self = Self::constant(1005);
86    /// Endpoint is terminating the connection because it has received data
87    /// within a message that was not consistent with the type of the message.
88    pub const INVALID_FRAME_PAYLOAD_DATA: Self = Self::constant(1007);
89    /// Endpoint is terminating the connection because it has received a message
90    /// that violates its policy.
91    pub const POLICY_VIOLATION: Self = Self::constant(1008);
92    /// Endpoint is terminating the connection because it has received a message
93    /// that is too big for it to process.
94    pub const MESSAGE_TOO_BIG: Self = Self::constant(1009);
95    /// Client is terminating the connection because it has expected the server
96    /// to negotiate one or more extension, but the server didn't return them in
97    /// the response message of the WebSocket handshake.
98    pub const MANDATORY_EXTENSION: Self = Self::constant(1010);
99    /// Server is terminating the connection because it encountered an
100    /// unexpected condition that prevented it from fulfilling the request.
101    pub const INTERNAL_SERVER_ERROR: Self = Self::constant(1011);
102    /// Service is restarted. A client may reconnect, and if it choses to do,
103    /// should reconnect using a randomized delay of 5--30s.
104    pub const SERVICE_RESTART: Self = Self::constant(1012);
105    /// Service is experiencing overload. A client should only connect to a
106    /// different IP (when there are multiple for the target) or reconnect to
107    /// the same IP upon user action.
108    pub const SERVICE_OVERLOAD: Self = Self::constant(1013);
109    /// The server was acting as a gateway or proxy and received an invalid
110    /// response from the upstream server. This is similar to the HTTP 502
111    /// status code.
112    pub const BAD_GATEWAY: Self = Self::constant(1014);
113}
114
115impl CloseCode {
116    /// Try to construct [`CloseCode`] from `u16`
117    ///
118    /// Returns `None` if `code` is not a valid `CloseCode`
119    const fn try_from_u16(code: u16) -> Option<Self> {
120        match code {
121            1000..=1015 | 3000..=4999 => Some(Self(NonZeroU16::new(code).unwrap())),
122            0..=999 | 1016..=2999 | 5000..=u16::MAX => None,
123        }
124    }
125
126    /// Try to construct [`CloseCode`] from `u16`
127    ///
128    /// Panics if `code` is not a valid `CloseCode`
129    const fn constant(code: u16) -> Self {
130        Self::try_from_u16(code).unwrap()
131    }
132
133    /// Whether the close code is reserved and cannot be sent over the wire.
134    #[must_use]
135    pub fn is_reserved(self) -> bool {
136        match self.0.get() {
137            1004 | 1005 | 1006 | 1015 => true,
138            1000..=4999 => false,
139            // `TryFrom` is the only way to acquire self and it errors for these values
140            0..=999 | 5000..=u16::MAX => {
141                debug_assert!(false, "unexpected CloseCode");
142                false
143            }
144        }
145    }
146}
147
148impl From<CloseCode> for u16 {
149    fn from(value: CloseCode) -> Self {
150        value.0.get()
151    }
152}
153
154impl TryFrom<u16> for CloseCode {
155    type Error = ProtocolError;
156
157    fn try_from(value: u16) -> Result<Self, Self::Error> {
158        Self::try_from_u16(value).ok_or(ProtocolError::InvalidCloseCode)
159    }
160}
161
162/// The websocket message payload storage.
163///
164/// Payloads can be created by using the `From<T>` implementations.
165///
166/// Sending the payloads or calling [`Into<BytesMut>`] is zero-copy, except when
167/// sending a payload created from a static slice or when the payload buffer is
168/// not unique. All conversions to other types are zero-cost.
169///
170/// [`Into<BytesMut>`]: #impl-From<Payload>-for-BytesMut
171#[derive(Clone)]
172pub struct Payload {
173    /// The raw payload data.
174    data: Bytes,
175    /// Whether the payload data was validated to be valid UTF-8.
176    utf8_validated: bool,
177}
178
179impl Payload {
180    /// Creates a new shared `Payload` from a static slice.
181    const fn from_static(bytes: &'static [u8]) -> Self {
182        Self {
183            data: Bytes::from_static(bytes),
184            utf8_validated: false,
185        }
186    }
187
188    /// Marks whether the payload contents were validated to be valid UTF-8.
189    pub(super) fn set_utf8_validated(&mut self, value: bool) {
190        self.utf8_validated = value;
191    }
192
193    /// Shortens the buffer, keeping the first `len` bytes and dropping the
194    /// rest.
195    pub(super) fn truncate(&mut self, len: usize) {
196        self.data.truncate(len);
197    }
198
199    /// Splits the buffer into two at the given index.
200    fn split_to(&mut self, at: usize) -> Self {
201        // This is only used by the outgoing message frame iterator, so we do not care
202        // about the value of utf8_validated. For the sake of correctness (in case we
203        // split a utf8 codepoint), we set it to false.
204        self.utf8_validated = false;
205        Self {
206            data: self.data.split_to(at),
207            utf8_validated: false,
208        }
209    }
210}
211
212impl Deref for Payload {
213    type Target = [u8];
214
215    fn deref(&self) -> &Self::Target {
216        &self.data
217    }
218}
219
220impl fmt::Debug for Payload {
221    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
222        f.debug_tuple("Payload").field(&self.data).finish()
223    }
224}
225
226impl From<Bytes> for Payload {
227    fn from(value: Bytes) -> Self {
228        Self {
229            data: value,
230            utf8_validated: false,
231        }
232    }
233}
234
235impl From<BytesMut> for Payload {
236    fn from(value: BytesMut) -> Self {
237        Self {
238            data: value.freeze(),
239            utf8_validated: false,
240        }
241    }
242}
243
244impl From<Payload> for Bytes {
245    fn from(value: Payload) -> Self {
246        value.data
247    }
248}
249
250impl From<Payload> for BytesMut {
251    fn from(value: Payload) -> Self {
252        value.data.into()
253    }
254}
255
256impl From<Vec<u8>> for Payload {
257    fn from(value: Vec<u8>) -> Self {
258        // BytesMut::from_iter goes through a specialization in std if the iterator is a
259        // Vec, effectively allowing us to use BytesMut::from_vec which isn't
260        // exposed in bytes. See https://github.com/tokio-rs/bytes/issues/723 for details.
261        Self {
262            data: BytesMut::from_iter(value).freeze(),
263            utf8_validated: false,
264        }
265    }
266}
267
268impl From<String> for Payload {
269    fn from(value: String) -> Self {
270        // See From<Vec<u8>> impl for reasoning behind this.
271        Self {
272            data: BytesMut::from_iter(value.into_bytes()).freeze(),
273            utf8_validated: true,
274        }
275    }
276}
277
278impl From<&'static [u8]> for Payload {
279    fn from(value: &'static [u8]) -> Self {
280        Self {
281            data: Bytes::from_static(value),
282            utf8_validated: false,
283        }
284    }
285}
286
287impl From<&'static str> for Payload {
288    fn from(value: &'static str) -> Self {
289        Self {
290            data: Bytes::from_static(value.as_bytes()),
291            utf8_validated: true,
292        }
293    }
294}
295
296/// A WebSocket message. This is cheaply clonable and uses [`Payload`] as the
297/// payload storage underneath.
298///
299/// Received messages are always validated prior to dealing with them, so all
300/// the type casting methods are either almost or fully zero cost.
301#[derive(Debug, Clone)]
302pub struct Message {
303    /// The [`OpCode`] of the message.
304    pub(super) opcode: OpCode,
305    /// The payload of the message.
306    pub(super) payload: Payload,
307}
308
309impl Message {
310    /// Create a new text message. The payload contents must be valid UTF-8.
311    #[must_use]
312    pub fn text<P: Into<Payload>>(payload: P) -> Self {
313        Self {
314            opcode: OpCode::Text,
315            payload: payload.into(),
316        }
317    }
318
319    /// Create a new binary message.
320    #[must_use]
321    pub fn binary<P: Into<Payload>>(payload: P) -> Self {
322        Self {
323            opcode: OpCode::Binary,
324            payload: payload.into(),
325        }
326    }
327
328    /// Create a new close message. If an non-empty reason is specified, a
329    /// [`CloseCode`] must be specified for it to be included.
330    ///
331    /// # Panics
332    /// - If the `code` is reserved so it cannot be sent.
333    /// - If `code` is present and the `reason` exceeds 123 bytes, the
334    ///   protocol-imposed limit.
335    #[must_use]
336    #[track_caller]
337    pub fn close(code: Option<CloseCode>, reason: &str) -> Self {
338        let mut payload = BytesMut::with_capacity((2 + reason.len()) * usize::from(code.is_some()));
339
340        if let Some(code) = code {
341            assert!(!code.is_reserved());
342            payload.put_u16(code.into());
343
344            assert!(reason.len() <= 123);
345            payload.extend_from_slice(reason.as_bytes());
346        }
347
348        Self {
349            opcode: OpCode::Close,
350            payload: payload.into(),
351        }
352    }
353
354    /// Create a new ping message.
355    ///
356    /// # Panics
357    /// If the payload exceeds 125 bytes, the protocol-imposed limit.
358    #[must_use]
359    #[track_caller]
360    pub fn ping<P: Into<Payload>>(payload: P) -> Self {
361        let payload = payload.into();
362        assert!(payload.len() <= 125);
363        Self {
364            opcode: OpCode::Ping,
365            payload,
366        }
367    }
368
369    /// Create a new pong message.
370    ///
371    /// # Panics
372    /// If the payload exceeds 125 bytes, the protocol-imposed limit.
373    #[must_use]
374    #[track_caller]
375    pub fn pong<P: Into<Payload>>(payload: P) -> Self {
376        let payload = payload.into();
377        assert!(payload.len() <= 125);
378        Self {
379            opcode: OpCode::Pong,
380            payload,
381        }
382    }
383
384    /// Whether the message is a text message.
385    #[must_use]
386    pub fn is_text(&self) -> bool {
387        self.opcode == OpCode::Text
388    }
389
390    /// Whether the message is a binary message.
391    #[must_use]
392    pub fn is_binary(&self) -> bool {
393        self.opcode == OpCode::Binary
394    }
395
396    /// Whether the message is a close message.
397    #[must_use]
398    pub fn is_close(&self) -> bool {
399        self.opcode == OpCode::Close
400    }
401
402    /// Whether the message is a ping message.
403    #[must_use]
404    pub fn is_ping(&self) -> bool {
405        self.opcode == OpCode::Ping
406    }
407
408    /// Whether the message is a pong message.
409    #[must_use]
410    pub fn is_pong(&self) -> bool {
411        self.opcode == OpCode::Pong
412    }
413
414    /// Returns the message payload and consumes the message, regardless of
415    /// type.
416    #[must_use]
417    pub fn into_payload(self) -> Payload {
418        self.payload
419    }
420
421    /// Returns a reference to the message payload, regardless of message type.
422    pub fn as_payload(&self) -> &Payload {
423        &self.payload
424    }
425
426    /// Returns a reference to the message payload as a string if it is a text
427    /// message.
428    ///
429    /// # Panics
430    ///
431    /// This method will panic when the message was created via
432    /// [`Message::text`] with invalid UTF-8.
433    pub fn as_text(&self) -> Option<&str> {
434        // SAFETY: Received messages were validated to be valid UTF-8, otherwise
435        // we check if it is valid UTF-8.
436        (self.opcode == OpCode::Text).then(|| {
437            assert!(
438                self.payload.utf8_validated || utf8::parse_str(&self.payload).is_ok(),
439                "called as_text on message created from payload with invalid utf-8"
440            );
441            unsafe { std::str::from_utf8_unchecked(&self.payload) }
442        })
443    }
444
445    /// Returns the [`CloseCode`] and close reason if the message is a close
446    /// message.
447    pub fn as_close(&self) -> Option<(CloseCode, &str)> {
448        (self.opcode == OpCode::Close).then(|| {
449            let code = if self.payload.is_empty() {
450                CloseCode::NO_STATUS_RECEIVED
451            } else {
452                // SAFETY: Opcode is Close with a non-empty payload so it's at least 2 bytes
453                // long
454                unsafe {
455                    CloseCode::try_from(u16::from_be_bytes(
456                        self.payload
457                            .get_unchecked(0..2)
458                            .try_into()
459                            .unwrap_unchecked(),
460                    ))
461                    .unwrap_unchecked()
462                }
463            };
464
465            // SAFETY: Opcode is Close so the rest of the payload is valid UTF-8
466            let reason =
467                unsafe { std::str::from_utf8_unchecked(self.payload.get(2..).unwrap_or_default()) };
468
469            (code, reason)
470        })
471    }
472
473    /// Returns an iterator over frames of `frame_size` length to split this
474    /// message into.
475    pub(super) fn into_frames(self, frame_size: usize) -> MessageFrames {
476        MessageFrames {
477            frame_size,
478            payload: self.payload,
479            opcode: self.opcode,
480        }
481    }
482}
483
484/// Iterator over frames of a chunked message.
485pub(super) struct MessageFrames {
486    /// Iterator over payload chunks.
487    frame_size: usize,
488    /// The full message payload this iterates over.
489    payload: Payload,
490    /// Opcode for the next frame.
491    opcode: OpCode,
492}
493
494impl Iterator for MessageFrames {
495    type Item = Frame;
496
497    fn next(&mut self) -> Option<Self::Item> {
498        let is_empty = self.payload.is_empty() && self.opcode == OpCode::Continuation;
499
500        (!is_empty).then(|| {
501            let payload = self
502                .payload
503                .split_to(self.frame_size.min(self.payload.len()));
504
505            Frame {
506                opcode: replace(&mut self.opcode, OpCode::Continuation),
507                is_final: self.payload.is_empty(),
508                payload,
509            }
510        })
511    }
512}
513
514/// Configuration for limitations on reading of [`Message`]s from a
515/// [`WebSocketStream`] to prevent high memory usage caused by malicious actors.
516///
517/// [`WebSocketStream`]: super::WebSocketStream
518#[derive(Debug, Clone, Copy)]
519pub struct Limits {
520    /// The maximum allowed payload length. The default is 64 MiB.
521    pub(super) max_payload_len: usize,
522}
523
524impl Limits {
525    /// A limit configuration without any limits.
526    #[must_use]
527    pub fn unlimited() -> Self {
528        Self {
529            max_payload_len: usize::MAX,
530        }
531    }
532
533    /// Sets the maximum allowed payload length. `None` equals no limit.
534    ///
535    /// The default is 64 MiB.
536    #[must_use]
537    pub fn max_payload_len(mut self, size: Option<usize>) -> Self {
538        self.set_max_payload_len(size);
539
540        self
541    }
542
543    /// See [`max_payload_len`](Self::max_payload_len).
544    pub fn set_max_payload_len(&mut self, size: Option<usize>) {
545        self.max_payload_len = size.unwrap_or(usize::MAX);
546    }
547}
548
549impl Default for Limits {
550    fn default() -> Self {
551        Self {
552            max_payload_len: 64 * 1024 * 1024,
553        }
554    }
555}
556
557/// Low-level configuration for a [`WebSocketStream`] that allows configuring
558/// behavior for sending and receiving messages.
559///
560/// [`WebSocketStream`]: super::WebSocketStream
561#[derive(Debug, Clone, Copy)]
562pub struct Config {
563    /// Frame payload size to split outgoing messages into.
564    ///
565    /// Consider decreasing this if the remote imposes a limit on the frame
566    /// payload size. The default is 4MiB.
567    pub(super) frame_size: usize,
568    /// Threshold of queued up bytes after which the underlying I/O is flushed
569    /// before the sink is declared ready. The default is 8 KiB.
570    pub(super) flush_threshold: usize,
571}
572
573impl Config {
574    /// Set the frame payload size to split outgoing messages into.
575    ///
576    /// Consider decreasing this if the remote imposes a limit on the frame
577    /// payload size. The default is 4MiB.
578    ///
579    /// # Panics
580    ///
581    /// If `frame_size` is `0`.
582    #[must_use]
583    pub fn frame_size(mut self, frame_size: usize) -> Self {
584        assert_ne!(frame_size, 0, "frame_size must be non-zero");
585        self.frame_size = frame_size;
586
587        self
588    }
589
590    /// Sets the threshold of queued up bytes after which the underlying I/O is
591    /// flushed before the sink is declared ready. The default is 8 KiB.
592    #[must_use]
593    pub fn flush_threshold(mut self, threshold: usize) -> Self {
594        self.flush_threshold = threshold;
595
596        self
597    }
598}
599
600impl Default for Config {
601    fn default() -> Self {
602        Self {
603            frame_size: 4 * 1024 * 1024,
604            flush_threshold: 8 * 1024,
605        }
606    }
607}
608
609/// Role assumed by the [`WebSocketStream`] in a connection.
610#[derive(Debug, PartialEq, Eq, Clone, Copy)]
611pub(crate) enum Role {
612    /// The client end.
613    Client,
614    /// The server end.
615    Server,
616}
617
618/// The connection state of the stream.
619#[derive(Debug, PartialEq)]
620pub(super) enum StreamState {
621    /// The connection is fully active and no close has been initiated.
622    Active,
623    /// The connection has been closed by the peer, but not yet acknowledged by
624    /// us.
625    ClosedByPeer,
626    /// The connection has been closed by us, but not yet acknowledged.
627    ClosedByUs,
628    /// The close has been acknowledged by the end that did not initiate the
629    /// close.
630    CloseAcknowledged,
631}
632
633/// A frame of a WebSocket [`Message`].
634#[derive(Clone, Debug)]
635pub(super) struct Frame {
636    /// The [`OpCode`] of the frame.
637    pub opcode: OpCode,
638    /// Whether this is the last frame of a message.
639    pub is_final: bool,
640    /// The payload bytes of the frame.
641    pub payload: Payload,
642}
643
644impl Frame {
645    /// Default close frame.
646    #[allow(clippy::declare_interior_mutable_const)]
647    pub const DEFAULT_CLOSE: Self = Self {
648        opcode: OpCode::Close,
649        is_final: true,
650        payload: Payload::from_static(&CloseCode::NORMAL_CLOSURE.0.get().to_be_bytes()),
651    };
652
653    /// Encode the frame head into `out`, returning a subslice where the mask
654    /// should be written to.
655    pub fn encode<'a>(&self, out: &'a mut [u8; 14]) -> &'a mut [u8; 4] {
656        out[0] = (u8::from(self.is_final) << 7) | u8::from(self.opcode);
657        let mask_slice = if u16::try_from(self.payload.len()).is_err() {
658            out[1] = 127;
659            let len = u64::try_from(self.payload.len()).unwrap();
660            out[2..10].copy_from_slice(&len.to_be_bytes());
661            &mut out[10..14]
662        } else if self.payload.len() > 125 {
663            out[1] = 126;
664            let len = u16::try_from(self.payload.len()).expect("checked by previous branch");
665            out[2..4].copy_from_slice(&len.to_be_bytes());
666            &mut out[4..8]
667        } else {
668            out[1] = u8::try_from(self.payload.len()).expect("checked by previous branch");
669            &mut out[2..6]
670        };
671        mask_slice.try_into().unwrap()
672    }
673}
674
675impl From<Message> for Frame {
676    fn from(value: Message) -> Self {
677        Self {
678            opcode: value.opcode,
679            is_final: true,
680            payload: value.payload,
681        }
682    }
683}
684
685impl From<&ProtocolError> for Frame {
686    fn from(val: &ProtocolError) -> Self {
687        match val {
688            ProtocolError::InvalidUtf8 => {
689                Message::close(Some(CloseCode::INVALID_FRAME_PAYLOAD_DATA), "invalid utf8")
690            }
691            _ => Message::close(Some(CloseCode::PROTOCOL_ERROR), val.as_str()),
692        }
693        .into()
694    }
695}
tokio_websockets/proto/types.rs

tokio_websockets/proto/
types.rs
