http_kit/
ws.rs

1//! WebSocket message and configuration types.
2
3use alloc::{borrow::ToOwned, string::String, vec::Vec};
4use bytes::Bytes;
5use bytestr::ByteStr;
6
7/// Message transmitted over a websocket connection.
8#[derive(Clone, Debug, PartialEq, Eq)]
9pub enum WebSocketMessage {
10    /// UTF-8 text payload.
11    Text(ByteStr),
12    /// Binary payload.
13    Binary(Bytes),
14
15    /// Ping control frame.
16    Ping(Bytes),
17    /// Pong control frame.
18    Pong(Bytes),
19
20    /// Close control frame.
21    Close,
22}
23
24/// Configuration applied when establishing a websocket connection.
25#[derive(Clone, Debug)]
26#[non_exhaustive]
27pub struct WebSocketConfig {
28    /// Maximum incoming websocket message size in bytes.
29    /// `None` means no limit.
30    pub max_message_size: Option<usize>,
31
32    /// Maximum incoming websocket frame size in bytes.
33    /// `None` means no limit.
34    pub max_frame_size: Option<usize>,
35}
36
37const DEFAULT_MAX_MESSAGE_SIZE: Option<usize> = Some(64 << 20);
38const DEFAULT_MAX_FRAME_SIZE: Option<usize> = Some(16 << 20);
39
40impl Default for WebSocketConfig {
41    fn default() -> Self {
42        Self {
43            max_message_size: DEFAULT_MAX_MESSAGE_SIZE,
44            max_frame_size: DEFAULT_MAX_FRAME_SIZE,
45        }
46    }
47}
48
49impl WebSocketConfig {
50    /// Override the maximum incoming websocket message size in bytes.
51    ///
52    /// `None` means no limit.
53    ///
54    /// Defaults to 64 MiB.
55    #[must_use]
56    pub const fn with_max_message_size(mut self, max_message_size: Option<usize>) -> Self {
57        self.max_message_size = max_message_size;
58        self
59    }
60
61    /// Override the maximum incoming websocket frame size in bytes.
62    ///
63    /// `None` means no limit.
64    ///
65    /// Defaults to 16 MiB.
66    #[must_use]
67    pub const fn with_max_frame_size(mut self, max_frame_size: Option<usize>) -> Self {
68        self.max_frame_size = max_frame_size;
69        self
70    }
71}
72
73impl WebSocketMessage {
74    /// Construct a text message.
75    #[must_use]
76    pub fn text(value: impl Into<ByteStr>) -> Self {
77        Self::Text(value.into())
78    }
79
80    /// Construct a ping message.
81    #[must_use]
82    pub fn ping(value: impl Into<Bytes>) -> Self {
83        Self::Ping(value.into())
84    }
85
86    /// Construct a pong message.
87    #[must_use]
88    pub fn pong(value: impl Into<Bytes>) -> Self {
89        Self::Pong(value.into())
90    }
91
92    /// Construct a close message.
93    pub fn close() -> Self {
94        Self::Close
95    }
96
97    /// Construct a binary message.
98    #[must_use]
99    pub fn binary(value: impl Into<Bytes>) -> Self {
100        Self::Binary(value.into())
101    }
102
103    /// Construct a JSON message.
104    #[cfg(feature = "json")]
105    pub fn json<T: serde::Serialize>(value: &T) -> serde_json::Result<Self> {
106        let json_string = serde_json::to_string(value)?;
107        Ok(Self::Text(json_string.into()))
108    }
109
110    /// Returns the payload as text when possible.
111    #[must_use]
112    pub fn as_text(&self) -> Option<&str> {
113        if let Self::Text(text) = self {
114            Some(text)
115        } else {
116            None
117        }
118    }
119
120    /// Returns the payload as raw bytes when possible.
121    #[must_use]
122    pub fn as_bytes(&self) -> Option<&[u8]> {
123        if let Self::Binary(bytes) = self {
124            Some(bytes)
125        } else {
126            None
127        }
128    }
129
130    /// Converts the payload into owned text when possible.
131    #[must_use]
132    pub fn into_text(self) -> Option<ByteStr> {
133        if let Self::Text(text) = self {
134            Some(text)
135        } else {
136            None
137        }
138    }
139
140    /// Converts the payload into a JSON value when possible.
141    #[cfg(feature = "json")]
142    pub fn into_json<T>(self) -> Option<Result<T, serde_json::Error>>
143    where
144        T: serde::de::DeserializeOwned,
145    {
146        if let Self::Text(text) = self {
147            Some(serde_json::from_str(&text))
148        } else {
149            None
150        }
151    }
152
153    /// Converts the payload into owned bytes when possible.
154    #[must_use]
155    pub fn into_bytes(self) -> Option<Bytes> {
156        if let Self::Binary(bytes) = self {
157            Some(bytes)
158        } else {
159            None
160        }
161    }
162}
163
164impl From<String> for WebSocketMessage {
165    fn from(value: String) -> Self {
166        Self::Text(value.into())
167    }
168}
169
170impl From<ByteStr> for WebSocketMessage {
171    fn from(value: ByteStr) -> Self {
172        Self::Text(value)
173    }
174}
175
176impl From<&str> for WebSocketMessage {
177    fn from(value: &str) -> Self {
178        Self::Text(value.to_owned().into())
179    }
180}
181
182impl From<Bytes> for WebSocketMessage {
183    fn from(value: Bytes) -> Self {
184        Self::Binary(value)
185    }
186}
187
188impl From<Vec<u8>> for WebSocketMessage {
189    fn from(value: Vec<u8>) -> Self {
190        Self::Binary(value.into())
191    }
192}
193
194impl From<&[u8]> for WebSocketMessage {
195    fn from(value: &[u8]) -> Self {
196        Self::Binary(value.to_vec().into())
197    }
198}