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