web_socket/
frame.rs

1#![doc(hidden)]
2
3pub struct Frame<'a> {
4    pub fin: bool,
5    pub opcode: u8,
6    pub data: &'a [u8],
7}
8
9impl<'a> Frame<'a> {
10    #[inline]
11    pub fn encode_without_mask(self) -> Vec<u8> {
12        let mut buf = Vec::<u8>::with_capacity(10 + self.data.len());
13        unsafe {
14            let dist = buf.as_mut_ptr();
15            let head_len = self.encode_header_unchecked(dist, 0);
16            std::ptr::copy_nonoverlapping(self.data.as_ptr(), dist.add(head_len), self.data.len());
17            buf.set_len(head_len + self.data.len());
18        }
19        buf
20    }
21
22    #[inline]
23    pub fn encode_with_mask(self) -> Vec<u8> {
24        let mut buf = Vec::<u8>::with_capacity(14 + self.data.len());
25        let mask = rand::random::<u32>().to_ne_bytes();
26        unsafe {
27            let dist = buf.as_mut_ptr();
28            let head_len = self.encode_header_unchecked(dist, 0x80);
29
30            let [a, b, c, d] = mask;
31            dist.add(head_len).write(a);
32            dist.add(head_len + 1).write(b);
33            dist.add(head_len + 2).write(c);
34            dist.add(head_len + 3).write(d);
35
36            let dist = dist.add(head_len + 4);
37            // TODO: Use SIMD wherever possible for best performance
38            for i in 0..self.data.len() {
39                dist.add(i)
40                    .write(self.data.get_unchecked(i) ^ mask.get_unchecked(i & 3));
41            }
42            buf.set_len(head_len + 4 + self.data.len());
43        }
44        buf
45    }
46
47    /// # SEAFTY
48    ///
49    /// - `dist` must be valid for writes of 10 bytes.
50    pub(crate) unsafe fn encode_header_unchecked(&self, dist: *mut u8, mask_bit: u8) -> usize {
51        dist.write(((self.fin as u8) << 7) | self.opcode);
52        if self.data.len() < 126 {
53            dist.add(1).write(mask_bit | self.data.len() as u8);
54            2
55        } else if self.data.len() < 65536 {
56            let [b2, b3] = (self.data.len() as u16).to_be_bytes();
57            dist.add(1).write(mask_bit | 126);
58            dist.add(2).write(b2);
59            dist.add(3).write(b3);
60            4
61        } else {
62            let [b2, b3, b4, b5, b6, b7, b8, b9] = (self.data.len() as u64).to_be_bytes();
63            dist.add(1).write(mask_bit | 127);
64            dist.add(2).write(b2);
65            dist.add(3).write(b3);
66            dist.add(4).write(b4);
67            dist.add(5).write(b5);
68            dist.add(6).write(b6);
69            dist.add(7).write(b7);
70            dist.add(8).write(b8);
71            dist.add(9).write(b9);
72            10
73        }
74    }
75}
76
77impl<'a> From<&'a str> for Frame<'a> {
78    #[inline]
79    fn from(string: &'a str) -> Self {
80        Self {
81            fin: true,
82            opcode: 1,
83            data: string.as_bytes(),
84        }
85    }
86}
87
88impl<'a> From<&'a [u8]> for Frame<'a> {
89    #[inline]
90    fn from(data: &'a [u8]) -> Self {
91        Self {
92            fin: true,
93            opcode: 2,
94            data,
95        }
96    }
97}
web_socket/frame.rs

web_socket/
frame.rs
