actix_http/ws/
codec.rs

1use bitflags::bitflags;
2use bytes::{Bytes, BytesMut};
3use bytestring::ByteString;
4use tokio_util::codec::{Decoder, Encoder};
5
6use super::{
7    frame::Parser,
8    proto::{CloseReason, OpCode},
9    ProtocolError,
10};
11
12/// A WebSocket message.
13#[derive(Debug, PartialEq, Eq)]
14pub enum Message {
15    /// Text message.
16    Text(ByteString),
17
18    /// Binary message.
19    Binary(Bytes),
20
21    /// Continuation.
22    Continuation(Item),
23
24    /// Ping message.
25    Ping(Bytes),
26
27    /// Pong message.
28    Pong(Bytes),
29
30    /// Close message with optional reason.
31    Close(Option<CloseReason>),
32
33    /// No-op. Useful for low-level services.
34    Nop,
35}
36
37/// A WebSocket frame.
38#[derive(Debug, PartialEq, Eq)]
39pub enum Frame {
40    /// Text frame. Note that the codec does not validate UTF-8 encoding.
41    Text(Bytes),
42
43    /// Binary frame.
44    Binary(Bytes),
45
46    /// Continuation.
47    Continuation(Item),
48
49    /// Ping message.
50    Ping(Bytes),
51
52    /// Pong message.
53    Pong(Bytes),
54
55    /// Close message with optional reason.
56    Close(Option<CloseReason>),
57}
58
59/// A WebSocket continuation item.
60#[derive(Debug, PartialEq, Eq)]
61pub enum Item {
62    FirstText(Bytes),
63    FirstBinary(Bytes),
64    Continue(Bytes),
65    Last(Bytes),
66}
67
68/// WebSocket protocol codec.
69#[derive(Debug, Clone)]
70pub struct Codec {
71    flags: Flags,
72    max_size: usize,
73}
74
75bitflags! {
76    #[derive(Debug, Clone, Copy)]
77    struct Flags: u8 {
78        const SERVER         = 0b0000_0001;
79        const CONTINUATION   = 0b0000_0010;
80        const W_CONTINUATION = 0b0000_0100;
81    }
82}
83
84/// Default maximum frame size (64 KiB).
85pub const DEFAULT_MAX_FRAME_SIZE: usize = 64 * 1024;
86
87impl Codec {
88    /// Create new WebSocket frames decoder.
89    pub const fn new() -> Codec {
90        Codec {
91            max_size: DEFAULT_MAX_FRAME_SIZE,
92            flags: Flags::SERVER,
93        }
94    }
95
96    /// Set max frame size.
97    ///
98    /// By default max size is set to 64KiB.
99    #[must_use = "This returns the a new Codec, without modifying the original."]
100    pub fn max_size(mut self, size: usize) -> Self {
101        self.max_size = size;
102        self
103    }
104
105    /// Set decoder to client mode.
106    ///
107    /// By default decoder works in server mode.
108    #[must_use = "This returns the a new Codec, without modifying the original."]
109    pub fn client_mode(mut self) -> Self {
110        self.flags.remove(Flags::SERVER);
111        self
112    }
113}
114
115impl Default for Codec {
116    fn default() -> Self {
117        Self::new()
118    }
119}
120
121impl Encoder<Message> for Codec {
122    type Error = ProtocolError;
123
124    fn encode(&mut self, item: Message, dst: &mut BytesMut) -> Result<(), Self::Error> {
125        match item {
126            Message::Text(txt) => Parser::write_message(
127                dst,
128                txt,
129                OpCode::Text,
130                true,
131                !self.flags.contains(Flags::SERVER),
132            ),
133            Message::Binary(bin) => Parser::write_message(
134                dst,
135                bin,
136                OpCode::Binary,
137                true,
138                !self.flags.contains(Flags::SERVER),
139            ),
140            Message::Ping(txt) => Parser::write_message(
141                dst,
142                txt,
143                OpCode::Ping,
144                true,
145                !self.flags.contains(Flags::SERVER),
146            ),
147            Message::Pong(txt) => Parser::write_message(
148                dst,
149                txt,
150                OpCode::Pong,
151                true,
152                !self.flags.contains(Flags::SERVER),
153            ),
154            Message::Close(reason) => {
155                Parser::write_close(dst, reason, !self.flags.contains(Flags::SERVER))
156            }
157            Message::Continuation(cont) => match cont {
158                Item::FirstText(data) => {
159                    if self.flags.contains(Flags::W_CONTINUATION) {
160                        return Err(ProtocolError::ContinuationStarted);
161                    } else {
162                        self.flags.insert(Flags::W_CONTINUATION);
163                        Parser::write_message(
164                            dst,
165                            &data[..],
166                            OpCode::Text,
167                            false,
168                            !self.flags.contains(Flags::SERVER),
169                        )
170                    }
171                }
172                Item::FirstBinary(data) => {
173                    if self.flags.contains(Flags::W_CONTINUATION) {
174                        return Err(ProtocolError::ContinuationStarted);
175                    } else {
176                        self.flags.insert(Flags::W_CONTINUATION);
177                        Parser::write_message(
178                            dst,
179                            &data[..],
180                            OpCode::Binary,
181                            false,
182                            !self.flags.contains(Flags::SERVER),
183                        )
184                    }
185                }
186                Item::Continue(data) => {
187                    if self.flags.contains(Flags::W_CONTINUATION) {
188                        Parser::write_message(
189                            dst,
190                            &data[..],
191                            OpCode::Continue,
192                            false,
193                            !self.flags.contains(Flags::SERVER),
194                        )
195                    } else {
196                        return Err(ProtocolError::ContinuationNotStarted);
197                    }
198                }
199                Item::Last(data) => {
200                    if self.flags.contains(Flags::W_CONTINUATION) {
201                        self.flags.remove(Flags::W_CONTINUATION);
202                        Parser::write_message(
203                            dst,
204                            &data[..],
205                            OpCode::Continue,
206                            true,
207                            !self.flags.contains(Flags::SERVER),
208                        )
209                    } else {
210                        return Err(ProtocolError::ContinuationNotStarted);
211                    }
212                }
213            },
214            Message::Nop => {}
215        }
216        Ok(())
217    }
218}
219
220impl Decoder for Codec {
221    type Item = Frame;
222    type Error = ProtocolError;
223
224    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
225        let Some((finished, opcode, payload)) =
226            Parser::parse(src, self.flags.contains(Flags::SERVER), self.max_size)?
227        else {
228            return Ok(None);
229        };
230
231        if !finished {
232            return match opcode {
233                // Continuation frames are exposed individually; aggregation is handled downstream.
234                OpCode::Continue => {
235                    if self.flags.contains(Flags::CONTINUATION) {
236                        Ok(Some(Frame::Continuation(Item::Continue(
237                            payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
238                        ))))
239                    } else {
240                        Err(ProtocolError::ContinuationNotStarted)
241                    }
242                }
243                OpCode::Binary => {
244                    if !self.flags.contains(Flags::CONTINUATION) {
245                        self.flags.insert(Flags::CONTINUATION);
246                        Ok(Some(Frame::Continuation(Item::FirstBinary(
247                            payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
248                        ))))
249                    } else {
250                        Err(ProtocolError::ContinuationStarted)
251                    }
252                }
253                OpCode::Text => {
254                    if !self.flags.contains(Flags::CONTINUATION) {
255                        self.flags.insert(Flags::CONTINUATION);
256                        Ok(Some(Frame::Continuation(Item::FirstText(
257                            payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
258                        ))))
259                    } else {
260                        Err(ProtocolError::ContinuationStarted)
261                    }
262                }
263                _ => {
264                    tracing::error!("Unfinished fragment {opcode:?}");
265                    Err(ProtocolError::ContinuationFragment(opcode))
266                }
267            };
268        }
269
270        match opcode {
271            OpCode::Continue => {
272                if self.flags.contains(Flags::CONTINUATION) {
273                    self.flags.remove(Flags::CONTINUATION);
274                    Ok(Some(Frame::Continuation(Item::Last(
275                        payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
276                    ))))
277                } else {
278                    Err(ProtocolError::ContinuationNotStarted)
279                }
280            }
281            OpCode::Bad => Err(ProtocolError::BadOpCode),
282            OpCode::Close => {
283                if let Some(ref pl) = payload {
284                    let close_reason = Parser::try_parse_close_payload(pl)?;
285                    Ok(Some(Frame::Close(close_reason)))
286                } else {
287                    Ok(Some(Frame::Close(None)))
288                }
289            }
290            OpCode::Ping => Ok(Some(Frame::Ping(
291                payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
292            ))),
293            OpCode::Pong => Ok(Some(Frame::Pong(
294                payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
295            ))),
296            OpCode::Binary => {
297                if self.flags.contains(Flags::CONTINUATION) {
298                    Err(ProtocolError::ContinuationStarted)
299                } else {
300                    Ok(Some(Frame::Binary(
301                        payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
302                    )))
303                }
304            }
305            OpCode::Text => {
306                if self.flags.contains(Flags::CONTINUATION) {
307                    Err(ProtocolError::ContinuationStarted)
308                } else {
309                    Ok(Some(Frame::Text(
310                        payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
311                    )))
312                }
313            }
314        }
315    }
316}
actix_http/ws/codec.rs

actix_http/ws/
codec.rs
