http_type/websocket_frame/
impl.rs

1use crate::*;
2
3/// Implements the `Default` trait for `WebSocketFrame`.
4///
5/// Provides a default `WebSocketFrame` with `fin: false`, `opcode: WebSocketOpcode::Text`,
6/// `mask: false`, and an empty `payload_data`.
7impl Default for WebSocketFrame {
8    /// Returns the default `WebSocketFrame`.
9    ///
10    /// # Returns
11    ///
12    /// A default `WebSocketFrame` instance.
13    fn default() -> Self {
14        Self {
15            fin: false,
16            opcode: WebSocketOpcode::Text,
17            mask: false,
18            payload_data: Vec::new(),
19        }
20    }
21}
22
23impl WebSocketOpcode {
24    /// Creates a `WebSocketOpcode` from a raw u8 value.
25    ///
26    /// # Arguments
27    ///
28    /// - `opcode`: The raw opcode value.
29    ///
30    /// # Returns
31    ///
32    /// A `WebSocketOpcode` enum variant corresponding to the raw value.
33    #[inline]
34    pub fn from_u8(opcode: u8) -> Self {
35        match opcode {
36            0x0 => Self::Continuation,
37            0x1 => Self::Text,
38            0x2 => Self::Binary,
39            0x8 => Self::Close,
40            0x9 => Self::Ping,
41            0xA => Self::Pong,
42            _ => Self::Reserved(opcode),
43        }
44    }
45
46    /// Converts the `WebSocketOpcode` to its raw u8 value.
47    ///
48    /// # Returns
49    ///
50    /// The raw u8 value of the opcode.
51    #[inline]
52    pub fn to_u8(&self) -> u8 {
53        match self {
54            Self::Continuation => 0x0,
55            Self::Text => 0x1,
56            Self::Binary => 0x2,
57            Self::Close => 0x8,
58            Self::Ping => 0x9,
59            Self::Pong => 0xA,
60            Self::Reserved(code) => *code,
61        }
62    }
63
64    /// Checks if the opcode is a control frame.
65    ///
66    /// # Returns
67    ///
68    /// `true` if the opcode represents a control frame (Close, Ping, Pong), otherwise `false`.
69    #[inline]
70    pub fn is_control(&self) -> bool {
71        matches!(self, Self::Close | Self::Ping | Self::Pong)
72    }
73
74    /// Checks if the opcode is a data frame.
75    ///
76    /// # Returns
77    ///
78    /// `true` if the opcode represents a data frame (Text, Binary, Continuation), otherwise `false`.
79    #[inline]
80    pub fn is_data(&self) -> bool {
81        matches!(self, Self::Text | Self::Binary | Self::Continuation)
82    }
83
84    /// Checks if the opcode is a continuation frame.
85    ///
86    /// # Returns
87    ///
88    /// `true` if the opcode is `Continuation`, otherwise `false`.
89    #[inline]
90    pub fn is_continuation(&self) -> bool {
91        matches!(self, Self::Continuation)
92    }
93
94    /// Checks if the opcode is a text frame.
95    ///
96    /// # Returns
97    ///
98    /// `true` if the opcode is `Text`, otherwise `false`.
99    #[inline]
100    pub fn is_text(&self) -> bool {
101        matches!(self, Self::Text)
102    }
103
104    /// Checks if the opcode is a binary frame.
105    ///
106    /// # Returns
107    ///
108    /// `true` if the opcode is `Binary`, otherwise `false`.
109    #[inline]
110    pub fn is_binary(&self) -> bool {
111        matches!(self, Self::Binary)
112    }
113
114    /// Checks if the opcode is a close frame.
115    ///
116    /// # Returns
117    ///
118    /// `true` if the opcode is `Close`, otherwise `false`.
119    #[inline]
120    pub fn is_close(&self) -> bool {
121        matches!(self, Self::Close)
122    }
123
124    /// Checks if the opcode is a ping frame.
125    ///
126    /// # Returns
127    ///
128    /// `true` if the opcode is `Ping`, otherwise `false`.
129    #[inline]
130    pub fn is_ping(&self) -> bool {
131        matches!(self, Self::Ping)
132    }
133
134    /// Checks if the opcode is a pong frame.
135    ///
136    /// # Returns
137    ///
138    /// `true` if the opcode is `Pong`, otherwise `false`.
139    #[inline]
140    pub fn is_pong(&self) -> bool {
141        matches!(self, Self::Pong)
142    }
143
144    /// Checks if the opcode is a reserved frame.
145    ///
146    /// # Returns
147    ///
148    /// `true` if the opcode is `Reserved(_)`, otherwise `false`.
149    #[inline]
150    pub fn is_reserved(&self) -> bool {
151        matches!(self, Self::Reserved(_))
152    }
153}
154
155impl WebSocketFrame {
156    /// Decodes a WebSocket frame from the provided data slice.
157    ///
158    /// This function parses the raw bytes from a WebSocket stream according to the WebSocket protocol
159    /// specification to reconstruct a `WebSocketFrame`. It handles FIN bit, opcode, mask bit,
160    /// payload length (including extended lengths), mask key, and the payload data itself.
161    ///
162    /// # Arguments
163    ///
164    /// - `AsRef<[u8]>` - The raw data to decode into a WebSocket frame.
165    ///
166    /// # Returns
167    ///
168    /// - `Some((WebSocketFrame, usize))`: If the frame is successfully decoded, returns the decoded frame
169    ///   and the number of bytes consumed from the input slice.
170    /// - `None`: If the frame is incomplete or malformed.
171    pub fn decode_ws_frame<D>(data: D) -> WebsocketFrameWithLengthOption
172    where
173        D: AsRef<[u8]>,
174    {
175        let data_ref: &[u8] = data.as_ref();
176        if data_ref.len() < 2 {
177            return None;
178        }
179        let mut index: usize = 0;
180        let fin: bool = (data_ref[index] & 0b1000_0000) != 0;
181        let opcode: WebSocketOpcode = WebSocketOpcode::from_u8(data_ref[index] & 0b0000_1111);
182        index += 1;
183        let mask: bool = (data_ref[index] & 0b1000_0000) != 0;
184        let mut payload_len: usize = (data_ref[index] & 0b0111_1111) as usize;
185        index += 1;
186        if payload_len == 126 {
187            if data_ref.len() < index + 2 {
188                return None;
189            }
190            payload_len = u16::from_be_bytes(data_ref[index..index + 2].try_into().ok()?) as usize;
191            index += 2;
192        } else if payload_len == 127 {
193            if data_ref.len() < index + 8 {
194                return None;
195            }
196            payload_len = u64::from_be_bytes(data_ref[index..index + 8].try_into().ok()?) as usize;
197            index += 8;
198        }
199        let mask_key: Option<[u8; 4]> = if mask {
200            if data_ref.len() < index + 4 {
201                return None;
202            }
203            let key: [u8; 4] = data_ref[index..index + 4].try_into().ok()?;
204            index += 4;
205            Some(key)
206        } else {
207            None
208        };
209        if data_ref.len() < index + payload_len {
210            return None;
211        }
212        let mut payload: Vec<u8> = data_ref[index..index + payload_len].to_vec();
213        if let Some(mask_key) = mask_key {
214            for (i, byte) in payload.iter_mut().enumerate() {
215                *byte ^= mask_key[i % 4];
216            }
217        }
218        index += payload_len;
219        let frame: WebSocketFrame = WebSocketFrame {
220            fin,
221            opcode,
222            mask,
223            payload_data: payload,
224        };
225        Some((frame, index))
226    }
227
228    /// Creates a list of response frames from the provided body.
229    ///
230    /// This method segments the response body into WebSocket frames, respecting the maximum frame size
231    /// and handling UTF-8 character boundaries for text frames. It determines the appropriate opcode
232    /// (Text or Binary) based on the body's content.
233    ///
234    /// # Arguments
235    ///
236    /// - `AsRef<[u8]>` - A reference to a response body (payload) as a byte slice.
237    ///
238    /// # Returns
239    ///
240    /// - A vector of `ResponseBody` (byte vectors), where each element represents a framed WebSocket message.
241    pub fn create_frame_list<D>(data: D) -> Vec<ResponseBody>
242    where
243        D: AsRef<[u8]>,
244    {
245        let data_ref: &[u8] = data.as_ref();
246        let total_len: usize = data_ref.len();
247        let mut offset: usize = 0;
248        let mut frames_list: Vec<ResponseBody> =
249            Vec::with_capacity((total_len / MAX_FRAME_SIZE) + 1);
250        let mut is_first_frame: bool = true;
251        let is_valid_utf8: bool = std::str::from_utf8(data_ref).is_ok();
252        let base_opcode: WebSocketOpcode = if is_valid_utf8 {
253            WebSocketOpcode::Text
254        } else {
255            WebSocketOpcode::Binary
256        };
257        while offset < total_len {
258            let remaining: usize = total_len - offset;
259            let mut frame_size: usize = remaining.min(MAX_FRAME_SIZE);
260            if is_valid_utf8 && frame_size < remaining {
261                while frame_size > 0 && (data_ref[offset + frame_size] & 0xC0) == 0x80 {
262                    frame_size -= 1;
263                }
264                if frame_size == 0 {
265                    frame_size = remaining.min(MAX_FRAME_SIZE);
266                }
267            }
268            let mut frame: ResponseBody = Vec::with_capacity(frame_size + 10);
269            let opcode: WebSocketOpcode = if is_first_frame {
270                base_opcode
271            } else {
272                WebSocketOpcode::Continuation
273            };
274            let fin: u8 = if remaining > frame_size { 0x00 } else { 0x80 };
275            let opcode_byte: u8 = opcode.to_u8() & 0x0F;
276            frame.push(fin | opcode_byte);
277            if frame_size < 126 {
278                frame.push(frame_size as u8);
279            } else if frame_size <= MAX_FRAME_SIZE {
280                frame.push(126);
281                frame.extend_from_slice(&(frame_size as u16).to_be_bytes());
282            } else {
283                frame.push(127);
284                frame.extend_from_slice(&(frame_size as u16).to_be_bytes());
285            }
286            let end: usize = offset + frame_size;
287            frame.extend_from_slice(&data_ref[offset..end]);
288            frames_list.push(frame);
289            offset = end;
290            is_first_frame = false;
291        }
292        frames_list
293    }
294
295    /// Calculates the SHA-1 hash of the input data.
296    ///
297    /// This function implements the SHA-1 cryptographic hash algorithm according to RFC 3174.
298    /// It processes the input data in 512-bit (64-byte) blocks and produces a 160-bit (20-byte) hash.
299    ///
300    /// # Arguments
301    ///
302    /// - `AsRef<[u8]>` - The input data to be hashed.
303    ///
304    /// # Returns
305    ///
306    /// - A 20-byte array representing the SHA-1 hash of the input data.
307    pub fn sha1<D>(data: D) -> [u8; 20]
308    where
309        D: AsRef<[u8]>,
310    {
311        let data_ref: &[u8] = data.as_ref();
312        let mut hash_state: [u32; 5] = HASH_STATE;
313        let mut padded_data: Vec<u8> = Vec::from(data_ref);
314        let original_length_bits: u64 = (padded_data.len() * 8) as u64;
315        padded_data.push(0x80);
316        while !(padded_data.len() + 8).is_multiple_of(64) {
317            padded_data.push(0);
318        }
319        padded_data.extend_from_slice(&original_length_bits.to_be_bytes());
320        for block in padded_data.chunks_exact(64) {
321            let mut message_schedule: [u32; 80] = [0u32; 80];
322            for (i, block_chunk) in block.chunks_exact(4).enumerate().take(16) {
323                message_schedule[i] = u32::from_be_bytes([
324                    block_chunk[0],
325                    block_chunk[1],
326                    block_chunk[2],
327                    block_chunk[3],
328                ]);
329            }
330            for i in 16..80 {
331                message_schedule[i] = (message_schedule[i - 3]
332                    ^ message_schedule[i - 8]
333                    ^ message_schedule[i - 14]
334                    ^ message_schedule[i - 16])
335                    .rotate_left(1);
336            }
337            let [mut a, mut b, mut c, mut d, mut e] = hash_state;
338            for (i, &word) in message_schedule.iter().enumerate() {
339                let (f, k) = match i {
340                    0..=19 => ((b & c) | (!b & d), 0x5A827999),
341                    20..=39 => (b ^ c ^ d, 0x6ED9EBA1),
342                    40..=59 => ((b & c) | (b & d) | (c & d), 0x8F1BBCDC),
343                    _ => (b ^ c ^ d, 0xCA62C1D6),
344                };
345                let temp: u32 = a
346                    .rotate_left(5)
347                    .wrapping_add(f)
348                    .wrapping_add(e)
349                    .wrapping_add(k)
350                    .wrapping_add(word);
351                e = d;
352                d = c;
353                c = b.rotate_left(30);
354                b = a;
355                a = temp;
356            }
357            hash_state[0] = hash_state[0].wrapping_add(a);
358            hash_state[1] = hash_state[1].wrapping_add(b);
359            hash_state[2] = hash_state[2].wrapping_add(c);
360            hash_state[3] = hash_state[3].wrapping_add(d);
361            hash_state[4] = hash_state[4].wrapping_add(e);
362        }
363        let mut result: [u8; 20] = [0u8; 20];
364        for (i, &val) in hash_state.iter().enumerate() {
365            result[i * 4..(i + 1) * 4].copy_from_slice(&val.to_be_bytes());
366        }
367        result
368    }
369
370    /// Generates a WebSocket accept key from the client-provided key.
371    ///
372    /// This function is used during the WebSocket handshake to validate the client's request.
373    /// It concatenates the client's key with a specific GUID, calculates the SHA-1 hash of the result,
374    /// and then encodes the hash in base64.
375    ///
376    /// # Arguments
377    ///
378    /// - `AsRef<str>` - The client-provided key (typically from the `Sec-WebSocket-Key` header).
379    ///
380    /// # Returns
381    ///
382    /// - A string representing the generated WebSocket accept key (typically for the `Sec-WebSocket-Accept` header).
383    pub fn generate_accept_key<K>(key: K) -> String
384    where
385        K: AsRef<str>,
386    {
387        let key_ref: &str = key.as_ref();
388        let mut data: [u8; 60] = [0u8; 60];
389        data[..24].copy_from_slice(&key_ref.as_bytes()[..24.min(key_ref.len())]);
390        data[24..].copy_from_slice(GUID);
391        let hash: [u8; 20] = Self::sha1(data);
392        Self::base64_encode(hash)
393    }
394
395    /// Encodes the input data as a base64 string.
396    ///
397    /// This function implements the Base64 encoding scheme, converting binary data into an ASCII string format.
398    /// It processes the input data in chunks of 3 bytes and encodes them into 4 base64 characters.
399    /// Padding with '=' characters is applied if necessary.
400    ///
401    /// # Arguments
402    ///
403    /// - `AsRef<[u8]>` - The data to encode in base64.
404    ///
405    /// # Returns
406    ///
407    /// - A string with the base64 encoded representation of the input data.
408    pub fn base64_encode<D>(data: D) -> String
409    where
410        D: AsRef<[u8]>,
411    {
412        let data_ref: &[u8] = data.as_ref();
413        let mut encoded_data: Vec<u8> = Vec::with_capacity(data_ref.len().div_ceil(3) * 4);
414        for chunk in data_ref.chunks(3) {
415            let mut buffer: [u8; 3] = [0u8; 3];
416            buffer[..chunk.len()].copy_from_slice(chunk);
417            let indices: [u8; 4] = [
418                buffer[0] >> 2,
419                ((buffer[0] & 0b11) << 4) | (buffer[1] >> 4),
420                ((buffer[1] & 0b1111) << 2) | (buffer[2] >> 6),
421                buffer[2] & 0b111111,
422            ];
423            for &idx in &indices[..chunk.len() + 1] {
424                encoded_data.push(BASE64_CHARSET_TABLE[idx as usize]);
425            }
426            while !encoded_data.len().is_multiple_of(4) {
427                encoded_data.push(EQUAL_BYTES[0]);
428            }
429        }
430        String::from_utf8(encoded_data).unwrap()
431    }
432
433    /// Checks if the opcode is a continuation frame.
434    ///
435    /// # Returns
436    ///
437    /// `true` if the opcode is `Continuation`, otherwise `false`.
438    #[inline]
439    pub fn is_continuation_opcode(&self) -> bool {
440        self.opcode.is_continuation()
441    }
442
443    /// Checks if the opcode is a text frame.
444    ///
445    /// # Returns
446    ///
447    /// `true` if the opcode is `Text`, otherwise `false`.
448    #[inline]
449    pub fn is_text_opcode(&self) -> bool {
450        self.opcode.is_text()
451    }
452
453    /// Checks if the opcode is a binary frame.
454    ///
455    /// # Returns
456    ///
457    /// `true` if the opcode is `Binary`, otherwise `false`.
458    #[inline]
459    pub fn is_binary_opcode(&self) -> bool {
460        self.opcode.is_binary()
461    }
462
463    /// Checks if the opcode is a close frame.
464    ///
465    /// # Returns
466    ///
467    /// `true` if the opcode is `Close`, otherwise `false`.
468    #[inline]
469    pub fn is_close_opcode(&self) -> bool {
470        self.opcode.is_close()
471    }
472
473    /// Checks if the opcode is a ping frame.
474    ///
475    /// # Returns
476    ///
477    /// `true` if the opcode is `Ping`, otherwise `false`.
478    #[inline]
479    pub fn is_ping_opcode(&self) -> bool {
480        self.opcode.is_ping()
481    }
482
483    /// Checks if the opcode is a pong frame.
484    ///
485    /// # Returns
486    ///
487    /// `true` if the opcode is `Pong`, otherwise `false`.
488    #[inline]
489    pub fn is_pong_opcode(&self) -> bool {
490        self.opcode.is_pong()
491    }
492
493    /// Checks if the opcode is a reserved frame.
494    ///
495    /// # Returns
496    ///
497    /// `true` if the opcode is `Reserved(_)`, otherwise `false`.
498    #[inline]
499    pub fn is_reserved_opcode(&self) -> bool {
500        self.opcode.is_reserved()
501    }
502}