http_type/websocket_frame/
impl.rs

1use crate::*;
2
3impl Default for WebSocketFrame {
4    fn default() -> Self {
5        Self {
6            fin: false,
7            opcode: WebSocketOpcode::Text,
8            mask: false,
9            payload_data: Vec::new(),
10        }
11    }
12}
13
14impl WebSocketOpcode {
15    /// Create a WebSocketOpcode from a raw u8 value
16    ///
17    /// # Parameters
18    /// - `opcode`: The raw opcode value
19    ///
20    /// # Returns
21    /// - A WebSocketOpcode enum variant corresponding to the raw value
22    pub fn from_u8(opcode: u8) -> Self {
23        match opcode {
24            0x0 => Self::Continuation,
25            0x1 => Self::Text,
26            0x2 => Self::Binary,
27            0x8 => Self::Close,
28            0x9 => Self::Ping,
29            0xA => Self::Pong,
30            _ => Self::Reserved(opcode),
31        }
32    }
33
34    /// Convert the WebSocketOpcode to its raw u8 value
35    ///
36    /// # Returns
37    /// - The raw u8 value of the opcode
38    pub fn to_u8(&self) -> u8 {
39        match self {
40            Self::Continuation => 0x0,
41            Self::Text => 0x1,
42            Self::Binary => 0x2,
43            Self::Close => 0x8,
44            Self::Ping => 0x9,
45            Self::Pong => 0xA,
46            Self::Reserved(code) => *code,
47        }
48    }
49
50    /// Check if the opcode is a control frame
51    ///
52    /// # Returns
53    /// - true if the opcode represents a control frame (Close, Ping, Pong)
54    /// - false otherwise
55    pub fn is_control(&self) -> bool {
56        matches!(self, Self::Close | Self::Ping | Self::Pong)
57    }
58
59    /// Check if the opcode is a data frame
60    ///
61    /// # Returns
62    /// - true if the opcode represents a data frame (Text, Binary, Continuation)
63    /// - false otherwise
64    pub fn is_data(&self) -> bool {
65        matches!(self, Self::Text | Self::Binary | Self::Continuation)
66    }
67
68    /// Checks if the frame is a continuation frame.
69    ///
70    /// # Parameters
71    /// - `self`: The current frame.
72    ///
73    /// # Returns
74    /// - `true` if the frame is `Continuation`, otherwise `false`.
75    pub fn is_continuation(&self) -> bool {
76        matches!(self, Self::Continuation)
77    }
78
79    /// Checks if the frame is a text frame.
80    ///
81    /// # Parameters
82    /// - `self`: The current frame.
83    ///
84    /// # Returns
85    /// - `true` if the frame is `Text`, otherwise `false`.
86    pub fn is_text(&self) -> bool {
87        matches!(self, Self::Text)
88    }
89
90    /// Checks if the frame is a binary frame.
91    ///
92    /// # Parameters
93    /// - `self`: The current frame.
94    ///
95    /// # Returns
96    /// - `true` if the frame is `Binary`, otherwise `false`.
97    pub fn is_binary(&self) -> bool {
98        matches!(self, Self::Binary)
99    }
100
101    /// Checks if the frame is a close frame.
102    ///
103    /// # Parameters
104    /// - `self`: The current frame.
105    ///
106    /// # Returns
107    /// - `true` if the frame is `Close`, otherwise `false`.
108    pub fn is_close(&self) -> bool {
109        matches!(self, Self::Close)
110    }
111
112    /// Checks if the frame is a ping frame.
113    ///
114    /// # Parameters
115    /// - `self`: The current frame.
116    ///
117    /// # Returns
118    /// - `true` if the frame is `Ping`, otherwise `false`.
119    pub fn is_ping(&self) -> bool {
120        matches!(self, Self::Ping)
121    }
122
123    /// Checks if the frame is a pong frame.
124    ///
125    /// # Parameters
126    /// - `self`: The current frame.
127    ///
128    /// # Returns
129    /// - `true` if the frame is `Pong`, otherwise `false`.
130    pub fn is_pong(&self) -> bool {
131        matches!(self, Self::Pong)
132    }
133
134    /// Checks if the frame is a reserved frame.
135    ///
136    /// # Parameters
137    /// - `self`: The current frame.
138    ///
139    /// # Returns
140    /// - `true` if the frame is `Reserved(_)`, otherwise `false`.
141    pub fn is_reserved(&self) -> bool {
142        matches!(self, Self::Reserved(_))
143    }
144}
145
146impl WebSocketFrame {
147    /// decode_ws_frame
148    ///
149    /// # Parameters
150    /// - `data`: The raw data slice from the WebSocket stream.
151    ///
152    /// # Returns
153    /// - An Option containing a tuple (WebSocketFrame, usize), where the WebSocketFrame is the decoded frame and usize is the number of bytes consumed.
154    ///   Returns None if the frame is incomplete.
155    pub fn decode_ws_frame(data: &[u8]) -> WebsocketFrameWithLengthOption {
156        if data.len() < 2 {
157            return None;
158        }
159        let mut index: usize = 0;
160        let fin: bool = (data[index] & 0b1000_0000) != 0;
161        let opcode: WebSocketOpcode = WebSocketOpcode::from_u8(data[index] & 0b0000_1111);
162        index += 1;
163        let mask: bool = (data[index] & 0b1000_0000) != 0;
164        let mut payload_len: usize = (data[index] & 0b0111_1111) as usize;
165        index += 1;
166        if payload_len == 126 {
167            if data.len() < index + 2 {
168                return None;
169            }
170            payload_len = u16::from_be_bytes(data[index..index + 2].try_into().ok()?) as usize;
171            index += 2;
172        } else if payload_len == 127 {
173            if data.len() < index + 8 {
174                return None;
175            }
176            payload_len = u64::from_be_bytes(data[index..index + 8].try_into().ok()?) as usize;
177            index += 8;
178        }
179        let mask_key: Option<[u8; 4]> = if mask {
180            if data.len() < index + 4 {
181                return None;
182            }
183            let key: [u8; 4] = data[index..index + 4].try_into().ok()?;
184            index += 4;
185            Some(key)
186        } else {
187            None
188        };
189        if data.len() < index + payload_len {
190            return None;
191        }
192        let mut payload: Vec<u8> = data[index..index + payload_len].to_vec();
193        if let Some(mask_key) = mask_key {
194            for (i, byte) in payload.iter_mut().enumerate() {
195                *byte ^= mask_key[i % 4];
196            }
197        }
198        index += payload_len;
199        let frame: WebSocketFrame = WebSocketFrame {
200            fin,
201            opcode,
202            mask,
203            payload_data: payload,
204        };
205        Some((frame, index))
206    }
207
208    /// create_response_frame_list
209    ///
210    /// - `body`: A reference to a response body (payload) as a byte slice.
211    /// - Returns: A vector of response bodies (frames) representing the framed data.
212    pub fn create_response_frame_list(body: &ResponseBody) -> Vec<ResponseBody> {
213        let total_len: usize = body.len();
214        let mut offset: usize = 0;
215        let mut frames_list: Vec<ResponseBody> = Vec::new();
216        let mut is_first_frame: bool = true;
217        let is_valid_utf8: bool = std::str::from_utf8(body).is_ok();
218        let base_opcode: WebSocketOpcode = if is_valid_utf8 {
219            WebSocketOpcode::Text
220        } else {
221            WebSocketOpcode::Binary
222        };
223        while offset < total_len {
224            let remaining: usize = total_len - offset;
225            let mut frame_size: usize = remaining.min(MAX_FRAME_SIZE);
226            if is_valid_utf8 && frame_size < remaining {
227                while frame_size > 0 && (body[offset + frame_size] & 0xC0) == 0x80 {
228                    frame_size -= 1;
229                }
230                if frame_size == 0 {
231                    frame_size = remaining.min(MAX_FRAME_SIZE);
232                }
233            }
234            let mut frame: ResponseBody = Vec::with_capacity(frame_size + 10);
235            let opcode: WebSocketOpcode = if is_first_frame {
236                base_opcode
237            } else {
238                WebSocketOpcode::Continuation
239            };
240            let fin_bit: u8 = if remaining > frame_size { 0x00 } else { 0x80 };
241            let opcode_byte: u8 = opcode.to_u8() & 0x0F;
242            frame.push(fin_bit | opcode_byte);
243            if frame_size < 126 {
244                frame.push(frame_size as u8);
245            } else if frame_size <= MAX_FRAME_SIZE {
246                frame.push(126);
247                frame.extend_from_slice(&(frame_size as u16).to_be_bytes());
248            } else {
249                frame.push(127);
250                frame.extend_from_slice(&(frame_size as u64).to_be_bytes());
251            }
252            frame.extend_from_slice(&body[offset..offset + frame_size]);
253            frames_list.push(frame);
254            offset += frame_size;
255            is_first_frame = false;
256        }
257        frames_list
258    }
259
260    /// sha1
261    ///
262    /// - `data`: A byte slice containing the input data to be hashed.
263    /// - Returns: A 20-byte array representing the SHA-1 hash of the input data.
264    pub fn sha1(data: &[u8]) -> [u8; 20] {
265        let mut hash_state: [u32; 5] = HASH_STATE;
266        let mut padded_data: Vec<u8> = Vec::from(data);
267        let original_length_bits: u64 = (padded_data.len() * 8) as u64;
268        padded_data.push(0x80);
269        while (padded_data.len() + 8) % 64 != 0 {
270            padded_data.push(0);
271        }
272        padded_data.extend_from_slice(&original_length_bits.to_be_bytes());
273        for block in padded_data.chunks_exact(64) {
274            let mut message_schedule: [u32; 80] = [0u32; 80];
275            for (i, block_chunk) in block.chunks_exact(4).enumerate().take(16) {
276                message_schedule[i] = u32::from_be_bytes([
277                    block_chunk[0],
278                    block_chunk[1],
279                    block_chunk[2],
280                    block_chunk[3],
281                ]);
282            }
283            for i in 16..80 {
284                message_schedule[i] = (message_schedule[i - 3]
285                    ^ message_schedule[i - 8]
286                    ^ message_schedule[i - 14]
287                    ^ message_schedule[i - 16])
288                    .rotate_left(1);
289            }
290            let [mut a, mut b, mut c, mut d, mut e] = hash_state;
291            for (i, &word) in message_schedule.iter().enumerate() {
292                let (f, k) = match i {
293                    0..=19 => ((b & c) | (!b & d), 0x5A827999),
294                    20..=39 => (b ^ c ^ d, 0x6ED9EBA1),
295                    40..=59 => ((b & c) | (b & d) | (c & d), 0x8F1BBCDC),
296                    _ => (b ^ c ^ d, 0xCA62C1D6),
297                };
298                let temp: u32 = a
299                    .rotate_left(5)
300                    .wrapping_add(f)
301                    .wrapping_add(e)
302                    .wrapping_add(k)
303                    .wrapping_add(word);
304                e = d;
305                d = c;
306                c = b.rotate_left(30);
307                b = a;
308                a = temp;
309            }
310            hash_state[0] = hash_state[0].wrapping_add(a);
311            hash_state[1] = hash_state[1].wrapping_add(b);
312            hash_state[2] = hash_state[2].wrapping_add(c);
313            hash_state[3] = hash_state[3].wrapping_add(d);
314            hash_state[4] = hash_state[4].wrapping_add(e);
315        }
316        let mut result: [u8; 20] = [0u8; 20];
317        for (i, &val) in hash_state.iter().enumerate() {
318            result[i * 4..(i + 1) * 4].copy_from_slice(&val.to_be_bytes());
319        }
320        result
321    }
322
323    /// generate_accept_key
324    ///
325    /// - `key`: A string slice containing the client-provided key.
326    /// - Returns: A string representing the generated WebSocket accept key.
327    pub fn generate_accept_key(key: &str) -> String {
328        let mut data: [u8; 60] = [0u8; 60];
329        data[..24].copy_from_slice(&key.as_bytes()[..24.min(key.len())]);
330        data[24..].copy_from_slice(GUID);
331        let hash: [u8; 20] = Self::sha1(&data);
332        Self::base64_encode(&hash)
333    }
334
335    /// base64_encode
336    ///
337    /// - `data`: A byte slice containing the data to encode in base64.
338    /// - Returns: A string with the base64 encoded representation of the input data.
339    pub fn base64_encode(data: &[u8]) -> String {
340        let mut encoded_data: Vec<u8> = Vec::with_capacity((data.len() + 2) / 3 * 4);
341        for chunk in data.chunks(3) {
342            let mut buffer: [u8; 3] = [0u8; 3];
343            buffer[..chunk.len()].copy_from_slice(chunk);
344            let indices: [u8; 4] = [
345                buffer[0] >> 2,
346                ((buffer[0] & 0b11) << 4) | (buffer[1] >> 4),
347                ((buffer[1] & 0b1111) << 2) | (buffer[2] >> 6),
348                buffer[2] & 0b111111,
349            ];
350            for &idx in &indices[..chunk.len() + 1] {
351                encoded_data.push(BASE64_CHARSET_TABLE[idx as usize]);
352            }
353            while encoded_data.len() % 4 != 0 {
354                encoded_data.push(EQUAL_BYTES[0]);
355            }
356        }
357        String::from_utf8(encoded_data).unwrap()
358    }
359
360    /// Checks if the opcode is a continuation frame.
361    ///
362    /// # Parameters
363    /// - `self`: The current frame.
364    ///
365    /// # Returns
366    /// - `true` if the opcode is `Continuation`, otherwise `false`.
367    pub fn is_continuation_opcode(&self) -> bool {
368        self.opcode.is_continuation()
369    }
370
371    /// Checks if the opcode is a text frame.
372    ///
373    /// # Parameters
374    /// - `self`: The current frame.
375    ///
376    /// # Returns
377    /// - `true` if the opcode is `Text`, otherwise `false`.
378    pub fn is_text_opcode(&self) -> bool {
379        self.opcode.is_text()
380    }
381
382    /// Checks if the opcode is a binary frame.
383    ///
384    /// # Parameters
385    /// - `self`: The current frame.
386    ///
387    /// # Returns
388    /// - `true` if the opcode is `Binary`, otherwise `false`.
389    pub fn is_binary_opcode(&self) -> bool {
390        self.opcode.is_binary()
391    }
392
393    /// Checks if the opcode is a close frame.
394    ///
395    /// # Parameters
396    /// - `self`: The current frame.
397    ///
398    /// # Returns
399    /// - `true` if the opcode is `Close`, otherwise `false`.
400    pub fn is_close_opcode(&self) -> bool {
401        self.opcode.is_close()
402    }
403
404    /// Checks if the opcode is a ping frame.
405    ///
406    /// # Parameters
407    /// - `self`: The current frame.
408    ///
409    /// # Returns
410    /// - `true` if the opcode is `Ping`, otherwise `false`.
411    pub fn is_ping_opcode(&self) -> bool {
412        self.opcode.is_ping()
413    }
414
415    /// Checks if the opcode is a pong frame.
416    ///
417    /// # Parameters
418    /// - `self`: The current frame.
419    ///
420    /// # Returns
421    /// - `true` if the opcode is `Pong`, otherwise `false`.
422    pub fn is_pong_opcode(&self) -> bool {
423        self.opcode.is_pong()
424    }
425
426    /// Checks if the opcode is a reserved frame.
427    ///
428    /// # Parameters
429    /// - `self`: The current frame.
430    ///
431    /// # Returns
432    /// - `true` if the opcode is `Reserved(_)`, otherwise `false`.
433    pub fn is_reserved_opcode(&self) -> bool {
434        self.opcode.is_reserved()
435    }
436}
http_type/websocket_frame/impl.rs

http_type/websocket_frame/
impl.rs
