Skip to main content

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    #[must_use]
143    pub fn into_json<T>(self) -> Option<Result<T, serde_json::Error>>
144    where
145        T: serde::de::DeserializeOwned,
146    {
147        if let Self::Text(text) = self {
148            Some(serde_json::from_str(&text))
149        } else {
150            None
151        }
152    }
153
154    /// Converts the payload into owned bytes when possible.
155    #[must_use]
156    pub fn into_bytes(self) -> Option<Bytes> {
157        if let Self::Binary(bytes) = self {
158            Some(bytes)
159        } else {
160            None
161        }
162    }
163}
164
165impl From<String> for WebSocketMessage {
166    fn from(value: String) -> Self {
167        Self::Text(value.into())
168    }
169}
170
171impl From<ByteStr> for WebSocketMessage {
172    fn from(value: ByteStr) -> Self {
173        Self::Text(value)
174    }
175}
176
177impl From<&str> for WebSocketMessage {
178    fn from(value: &str) -> Self {
179        Self::Text(value.to_owned().into())
180    }
181}
182
183impl From<Bytes> for WebSocketMessage {
184    fn from(value: Bytes) -> Self {
185        Self::Binary(value)
186    }
187}
188
189impl From<Vec<u8>> for WebSocketMessage {
190    fn from(value: Vec<u8>) -> Self {
191        Self::Binary(value.into())
192    }
193}
194
195impl From<&[u8]> for WebSocketMessage {
196    fn from(value: &[u8]) -> Self {
197        Self::Binary(value.to_vec().into())
198    }
199}