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
84impl Codec {
85    /// Create new WebSocket frames decoder.
86    pub const fn new() -> Codec {
87        Codec {
88            max_size: 65_536,
89            flags: Flags::SERVER,
90        }
91    }
92
93    /// Set max frame size.
94    ///
95    /// By default max size is set to 64KiB.
96    #[must_use = "This returns the a new Codec, without modifying the original."]
97    pub fn max_size(mut self, size: usize) -> Self {
98        self.max_size = size;
99        self
100    }
101
102    /// Set decoder to client mode.
103    ///
104    /// By default decoder works in server mode.
105    #[must_use = "This returns the a new Codec, without modifying the original."]
106    pub fn client_mode(mut self) -> Self {
107        self.flags.remove(Flags::SERVER);
108        self
109    }
110}
111
112impl Default for Codec {
113    fn default() -> Self {
114        Self::new()
115    }
116}
117
118impl Encoder<Message> for Codec {
119    type Error = ProtocolError;
120
121    fn encode(&mut self, item: Message, dst: &mut BytesMut) -> Result<(), Self::Error> {
122        match item {
123            Message::Text(txt) => Parser::write_message(
124                dst,
125                txt,
126                OpCode::Text,
127                true,
128                !self.flags.contains(Flags::SERVER),
129            ),
130            Message::Binary(bin) => Parser::write_message(
131                dst,
132                bin,
133                OpCode::Binary,
134                true,
135                !self.flags.contains(Flags::SERVER),
136            ),
137            Message::Ping(txt) => Parser::write_message(
138                dst,
139                txt,
140                OpCode::Ping,
141                true,
142                !self.flags.contains(Flags::SERVER),
143            ),
144            Message::Pong(txt) => Parser::write_message(
145                dst,
146                txt,
147                OpCode::Pong,
148                true,
149                !self.flags.contains(Flags::SERVER),
150            ),
151            Message::Close(reason) => {
152                Parser::write_close(dst, reason, !self.flags.contains(Flags::SERVER))
153            }
154            Message::Continuation(cont) => match cont {
155                Item::FirstText(data) => {
156                    if self.flags.contains(Flags::W_CONTINUATION) {
157                        return Err(ProtocolError::ContinuationStarted);
158                    } else {
159                        self.flags.insert(Flags::W_CONTINUATION);
160                        Parser::write_message(
161                            dst,
162                            &data[..],
163                            OpCode::Text,
164                            false,
165                            !self.flags.contains(Flags::SERVER),
166                        )
167                    }
168                }
169                Item::FirstBinary(data) => {
170                    if self.flags.contains(Flags::W_CONTINUATION) {
171                        return Err(ProtocolError::ContinuationStarted);
172                    } else {
173                        self.flags.insert(Flags::W_CONTINUATION);
174                        Parser::write_message(
175                            dst,
176                            &data[..],
177                            OpCode::Binary,
178                            false,
179                            !self.flags.contains(Flags::SERVER),
180                        )
181                    }
182                }
183                Item::Continue(data) => {
184                    if self.flags.contains(Flags::W_CONTINUATION) {
185                        Parser::write_message(
186                            dst,
187                            &data[..],
188                            OpCode::Continue,
189                            false,
190                            !self.flags.contains(Flags::SERVER),
191                        )
192                    } else {
193                        return Err(ProtocolError::ContinuationNotStarted);
194                    }
195                }
196                Item::Last(data) => {
197                    if self.flags.contains(Flags::W_CONTINUATION) {
198                        self.flags.remove(Flags::W_CONTINUATION);
199                        Parser::write_message(
200                            dst,
201                            &data[..],
202                            OpCode::Continue,
203                            true,
204                            !self.flags.contains(Flags::SERVER),
205                        )
206                    } else {
207                        return Err(ProtocolError::ContinuationNotStarted);
208                    }
209                }
210            },
211            Message::Nop => {}
212        }
213        Ok(())
214    }
215}
216
217impl Decoder for Codec {
218    type Item = Frame;
219    type Error = ProtocolError;
220
221    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
222        let Some((finished, opcode, payload)) =
223            Parser::parse(src, self.flags.contains(Flags::SERVER), self.max_size)?
224        else {
225            return Ok(None);
226        };
227
228        if !finished {
229            return match opcode {
230                // Continuation frames are exposed individually; aggregation is handled downstream.
231                OpCode::Continue => {
232                    if self.flags.contains(Flags::CONTINUATION) {
233                        Ok(Some(Frame::Continuation(Item::Continue(
234                            payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
235                        ))))
236                    } else {
237                        Err(ProtocolError::ContinuationNotStarted)
238                    }
239                }
240                OpCode::Binary => {
241                    if !self.flags.contains(Flags::CONTINUATION) {
242                        self.flags.insert(Flags::CONTINUATION);
243                        Ok(Some(Frame::Continuation(Item::FirstBinary(
244                            payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
245                        ))))
246                    } else {
247                        Err(ProtocolError::ContinuationStarted)
248                    }
249                }
250                OpCode::Text => {
251                    if !self.flags.contains(Flags::CONTINUATION) {
252                        self.flags.insert(Flags::CONTINUATION);
253                        Ok(Some(Frame::Continuation(Item::FirstText(
254                            payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
255                        ))))
256                    } else {
257                        Err(ProtocolError::ContinuationStarted)
258                    }
259                }
260                _ => {
261                    tracing::error!("Unfinished fragment {opcode:?}");
262                    Err(ProtocolError::ContinuationFragment(opcode))
263                }
264            };
265        }
266
267        match opcode {
268            OpCode::Continue => {
269                if self.flags.contains(Flags::CONTINUATION) {
270                    self.flags.remove(Flags::CONTINUATION);
271                    Ok(Some(Frame::Continuation(Item::Last(
272                        payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
273                    ))))
274                } else {
275                    Err(ProtocolError::ContinuationNotStarted)
276                }
277            }
278            OpCode::Bad => Err(ProtocolError::BadOpCode),
279            OpCode::Close => {
280                if let Some(ref pl) = payload {
281                    let close_reason = Parser::try_parse_close_payload(pl)?;
282                    Ok(Some(Frame::Close(close_reason)))
283                } else {
284                    Ok(Some(Frame::Close(None)))
285                }
286            }
287            OpCode::Ping => Ok(Some(Frame::Ping(
288                payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
289            ))),
290            OpCode::Pong => Ok(Some(Frame::Pong(
291                payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
292            ))),
293            OpCode::Binary => {
294                if self.flags.contains(Flags::CONTINUATION) {
295                    Err(ProtocolError::ContinuationStarted)
296                } else {
297                    Ok(Some(Frame::Binary(
298                        payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
299                    )))
300                }
301            }
302            OpCode::Text => {
303                if self.flags.contains(Flags::CONTINUATION) {
304                    Err(ProtocolError::ContinuationStarted)
305                } else {
306                    Ok(Some(Frame::Text(
307                        payload.map(|pl| pl.freeze()).unwrap_or_else(Bytes::new),
308                    )))
309                }
310            }
311        }
312    }
313}
actix_http/ws/codec.rs

actix_http/ws/
codec.rs
