actix_http/ws/
frame.rs

1use std::cmp::min;
2
3use bytes::{Buf, BufMut, BytesMut};
4use tracing::debug;
5
6use super::{
7    mask::apply_mask,
8    proto::{CloseCode, CloseReason, OpCode},
9    ProtocolError,
10};
11
12/// A struct representing a WebSocket frame.
13#[derive(Debug)]
14pub struct Parser;
15
16impl Parser {
17    fn parse_metadata(
18        src: &[u8],
19        server: bool,
20    ) -> Result<Option<(usize, bool, OpCode, usize, Option<[u8; 4]>)>, ProtocolError> {
21        let chunk_len = src.len();
22
23        let mut idx = 2;
24        if chunk_len < 2 {
25            return Ok(None);
26        }
27
28        let first = src[0];
29        let second = src[1];
30        let finished = first & 0x80 != 0;
31
32        // check masking
33        let masked = second & 0x80 != 0;
34        if !masked && server {
35            return Err(ProtocolError::UnmaskedFrame);
36        } else if masked && !server {
37            return Err(ProtocolError::MaskedFrame);
38        }
39
40        // Op code
41        let opcode = OpCode::from(first & 0x0F);
42
43        if let OpCode::Bad = opcode {
44            return Err(ProtocolError::InvalidOpcode(first & 0x0F));
45        }
46
47        let len = second & 0x7F;
48        let length = if len == 126 {
49            if chunk_len < 4 {
50                return Ok(None);
51            }
52            let len = usize::from(u16::from_be_bytes(
53                TryFrom::try_from(&src[idx..idx + 2]).unwrap(),
54            ));
55            idx += 2;
56            len
57        } else if len == 127 {
58            if chunk_len < 10 {
59                return Ok(None);
60            }
61            let len = u64::from_be_bytes(TryFrom::try_from(&src[idx..idx + 8]).unwrap());
62            idx += 8;
63            len as usize
64        } else {
65            len as usize
66        };
67
68        let mask = if server {
69            if chunk_len < idx + 4 {
70                return Ok(None);
71            }
72
73            let mask = TryFrom::try_from(&src[idx..idx + 4]).unwrap();
74
75            idx += 4;
76
77            Some(mask)
78        } else {
79            None
80        };
81
82        Ok(Some((idx, finished, opcode, length, mask)))
83    }
84
85    /// Parse the input stream into a frame.
86    pub fn parse(
87        src: &mut BytesMut,
88        server: bool,
89        max_size: usize,
90    ) -> Result<Option<(bool, OpCode, Option<BytesMut>)>, ProtocolError> {
91        // try to parse ws frame metadata
92        let (idx, finished, opcode, length, mask) = match Parser::parse_metadata(src, server)? {
93            None => return Ok(None),
94            Some(res) => res,
95        };
96
97        // not enough data
98        if src.len() < idx + length {
99            let min_length = min(length, max_size);
100            if src.capacity() < idx + min_length {
101                src.reserve(idx + min_length - src.capacity());
102            }
103            return Ok(None);
104        }
105
106        // remove prefix
107        src.advance(idx);
108
109        // check for max allowed size
110        if length > max_size {
111            // drop the payload
112            src.advance(length);
113            return Err(ProtocolError::Overflow);
114        }
115
116        // no need for body
117        if length == 0 {
118            return Ok(Some((finished, opcode, None)));
119        }
120
121        let mut data = src.split_to(length);
122
123        // control frames must have length <= 125
124        match opcode {
125            OpCode::Ping | OpCode::Pong if length > 125 => {
126                return Err(ProtocolError::InvalidLength(length));
127            }
128            OpCode::Close if length > 125 => {
129                debug!("Received close frame with payload length exceeding 125. Morphing to protocol close frame.");
130                return Ok(Some((true, OpCode::Close, None)));
131            }
132            _ => {}
133        }
134
135        // unmask
136        if let Some(mask) = mask {
137            apply_mask(&mut data, mask);
138        }
139
140        Ok(Some((finished, opcode, Some(data))))
141    }
142
143    /// Parse the payload of a close frame.
144    pub fn parse_close_payload(payload: &[u8]) -> Option<CloseReason> {
145        if payload.len() >= 2 {
146            let raw_code = u16::from_be_bytes(TryFrom::try_from(&payload[..2]).unwrap());
147            let code = CloseCode::from(raw_code);
148            let description = if payload.len() > 2 {
149                Some(String::from_utf8_lossy(&payload[2..]).into())
150            } else {
151                None
152            };
153            Some(CloseReason { code, description })
154        } else {
155            None
156        }
157    }
158
159    /// Generate binary representation
160    pub fn write_message<B: AsRef<[u8]>>(
161        dst: &mut BytesMut,
162        pl: B,
163        op: OpCode,
164        fin: bool,
165        mask: bool,
166    ) {
167        let payload = pl.as_ref();
168        let one: u8 = if fin {
169            0x80 | Into::<u8>::into(op)
170        } else {
171            op.into()
172        };
173        let payload_len = payload.len();
174        let (two, p_len) = if mask {
175            (0x80, payload_len + 4)
176        } else {
177            (0, payload_len)
178        };
179
180        if payload_len < 126 {
181            dst.reserve(p_len + 2);
182            dst.put_slice(&[one, two | payload_len as u8]);
183        } else if payload_len <= 65_535 {
184            dst.reserve(p_len + 4);
185            dst.put_slice(&[one, two | 126]);
186            dst.put_u16(payload_len as u16);
187        } else {
188            dst.reserve(p_len + 10);
189            dst.put_slice(&[one, two | 127]);
190            dst.put_u64(payload_len as u64);
191        };
192
193        if mask {
194            let mask = rand::random::<[u8; 4]>();
195            dst.put_slice(mask.as_ref());
196            dst.put_slice(payload.as_ref());
197            let pos = dst.len() - payload_len;
198            apply_mask(&mut dst[pos..], mask);
199        } else {
200            dst.put_slice(payload.as_ref());
201        }
202    }
203
204    /// Create a new Close control frame.
205    #[inline]
206    pub fn write_close(dst: &mut BytesMut, reason: Option<CloseReason>, mask: bool) {
207        let payload = match reason {
208            None => Vec::new(),
209            Some(reason) => {
210                let mut payload = Into::<u16>::into(reason.code).to_be_bytes().to_vec();
211                if let Some(description) = reason.description {
212                    payload.extend(description.as_bytes());
213                }
214                payload
215            }
216        };
217
218        Parser::write_message(dst, payload, OpCode::Close, true, mask)
219    }
220}
221
222#[cfg(test)]
223mod tests {
224    use bytes::Bytes;
225
226    use super::*;
227
228    struct F {
229        finished: bool,
230        opcode: OpCode,
231        payload: Bytes,
232    }
233
234    fn is_none(frm: &Result<Option<(bool, OpCode, Option<BytesMut>)>, ProtocolError>) -> bool {
235        matches!(*frm, Ok(None))
236    }
237
238    fn extract(frm: Result<Option<(bool, OpCode, Option<BytesMut>)>, ProtocolError>) -> F {
239        match frm {
240            Ok(Some((finished, opcode, payload))) => F {
241                finished,
242                opcode,
243                payload: payload
244                    .map(|b| b.freeze())
245                    .unwrap_or_else(|| Bytes::from("")),
246            },
247            _ => unreachable!("error"),
248        }
249    }
250
251    #[test]
252    fn test_parse() {
253        let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0001u8][..]);
254        assert!(is_none(&Parser::parse(&mut buf, false, 1024)));
255
256        let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0001u8][..]);
257        buf.extend(b"1");
258
259        let frame = extract(Parser::parse(&mut buf, false, 1024));
260        assert!(!frame.finished);
261        assert_eq!(frame.opcode, OpCode::Text);
262        assert_eq!(frame.payload.as_ref(), &b"1"[..]);
263    }
264
265    #[test]
266    fn test_parse_length0() {
267        let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0000u8][..]);
268        let frame = extract(Parser::parse(&mut buf, false, 1024));
269        assert!(!frame.finished);
270        assert_eq!(frame.opcode, OpCode::Text);
271        assert!(frame.payload.is_empty());
272    }
273
274    #[test]
275    fn test_parse_length2() {
276        let mut buf = BytesMut::from(&[0b0000_0001u8, 126u8][..]);
277        assert!(is_none(&Parser::parse(&mut buf, false, 1024)));
278
279        let mut buf = BytesMut::from(&[0b0000_0001u8, 126u8][..]);
280        buf.extend(&[0u8, 4u8][..]);
281        buf.extend(b"1234");
282
283        let frame = extract(Parser::parse(&mut buf, false, 1024));
284        assert!(!frame.finished);
285        assert_eq!(frame.opcode, OpCode::Text);
286        assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
287    }
288
289    #[test]
290    fn test_parse_length4() {
291        let mut buf = BytesMut::from(&[0b0000_0001u8, 127u8][..]);
292        assert!(is_none(&Parser::parse(&mut buf, false, 1024)));
293
294        let mut buf = BytesMut::from(&[0b0000_0001u8, 127u8][..]);
295        buf.extend(&[0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 0u8, 4u8][..]);
296        buf.extend(b"1234");
297
298        let frame = extract(Parser::parse(&mut buf, false, 1024));
299        assert!(!frame.finished);
300        assert_eq!(frame.opcode, OpCode::Text);
301        assert_eq!(frame.payload.as_ref(), &b"1234"[..]);
302    }
303
304    #[test]
305    fn test_parse_frame_mask() {
306        let mut buf = BytesMut::from(&[0b0000_0001u8, 0b1000_0001u8][..]);
307        buf.extend(b"0001");
308        buf.extend(b"1");
309
310        assert!(Parser::parse(&mut buf, false, 1024).is_err());
311
312        let frame = extract(Parser::parse(&mut buf, true, 1024));
313        assert!(!frame.finished);
314        assert_eq!(frame.opcode, OpCode::Text);
315        assert_eq!(frame.payload, Bytes::from(vec![1u8]));
316    }
317
318    #[test]
319    fn test_parse_frame_no_mask() {
320        let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0001u8][..]);
321        buf.extend([1u8]);
322
323        assert!(Parser::parse(&mut buf, true, 1024).is_err());
324
325        let frame = extract(Parser::parse(&mut buf, false, 1024));
326        assert!(!frame.finished);
327        assert_eq!(frame.opcode, OpCode::Text);
328        assert_eq!(frame.payload, Bytes::from(vec![1u8]));
329    }
330
331    #[test]
332    fn test_parse_frame_max_size() {
333        let mut buf = BytesMut::from(&[0b0000_0001u8, 0b0000_0010u8][..]);
334        buf.extend([1u8, 1u8]);
335
336        assert!(Parser::parse(&mut buf, true, 1).is_err());
337
338        if let Err(ProtocolError::Overflow) = Parser::parse(&mut buf, false, 0) {
339        } else {
340            unreachable!("error");
341        }
342    }
343
344    #[test]
345    fn test_parse_frame_max_size_recoverability() {
346        let mut buf = BytesMut::new();
347        // The first text frame with length == 2, payload doesn't matter.
348        buf.extend([0b0000_0001u8, 0b0000_0010u8, 0b0000_0000u8, 0b0000_0000u8]);
349        // Next binary frame with length == 2 and payload == `[0x1111_1111u8, 0x1111_1111u8]`.
350        buf.extend([0b0000_0010u8, 0b0000_0010u8, 0b1111_1111u8, 0b1111_1111u8]);
351
352        assert_eq!(buf.len(), 8);
353        assert!(matches!(
354            Parser::parse(&mut buf, false, 1),
355            Err(ProtocolError::Overflow)
356        ));
357        assert_eq!(buf.len(), 4);
358        let frame = extract(Parser::parse(&mut buf, false, 2));
359        assert!(!frame.finished);
360        assert_eq!(frame.opcode, OpCode::Binary);
361        assert_eq!(
362            frame.payload,
363            Bytes::from(vec![0b1111_1111u8, 0b1111_1111u8])
364        );
365        assert_eq!(buf.len(), 0);
366    }
367
368    #[test]
369    fn test_ping_frame() {
370        let mut buf = BytesMut::new();
371        Parser::write_message(&mut buf, Vec::from("data"), OpCode::Ping, true, false);
372
373        let mut v = vec![137u8, 4u8];
374        v.extend(b"data");
375        assert_eq!(&buf[..], &v[..]);
376    }
377
378    #[test]
379    fn test_pong_frame() {
380        let mut buf = BytesMut::new();
381        Parser::write_message(&mut buf, Vec::from("data"), OpCode::Pong, true, false);
382
383        let mut v = vec![138u8, 4u8];
384        v.extend(b"data");
385        assert_eq!(&buf[..], &v[..]);
386    }
387
388    #[test]
389    fn test_close_frame() {
390        let mut buf = BytesMut::new();
391        let reason = (CloseCode::Normal, "data");
392        Parser::write_close(&mut buf, Some(reason.into()), false);
393
394        let mut v = vec![136u8, 6u8, 3u8, 232u8];
395        v.extend(b"data");
396        assert_eq!(&buf[..], &v[..]);
397    }
398
399    #[test]
400    fn test_empty_close_frame() {
401        let mut buf = BytesMut::new();
402        Parser::write_close(&mut buf, None, false);
403        assert_eq!(&buf[..], &vec![0x88, 0x00][..]);
404    }
405}
actix_http/ws/frame.rs

actix_http/ws/
frame.rs
