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

http_type/websocket_frame/
impl.rs
