reifydb_network/websocket/
frame.rs1#[derive(Debug, Clone)]
5pub struct WebSocketFrame {
6 pub fin: bool,
7 pub opcode: WebSocketOpcode,
8 pub payload: Vec<u8>,
9}
10
11#[derive(Debug, Clone, Copy, PartialEq)]
12pub enum WebSocketOpcode {
13 Continuation,
14 Text,
15 Binary,
16 Close,
17 Ping,
18 Pong,
19 Unknown(u8),
20}
21
22impl From<u8> for WebSocketOpcode {
23 fn from(value: u8) -> Self {
24 match value {
25 0x0 => Self::Continuation,
26 0x1 => Self::Text,
27 0x2 => Self::Binary,
28 0x8 => Self::Close,
29 0x9 => Self::Ping,
30 0xA => Self::Pong,
31 other => Self::Unknown(other),
32 }
33 }
34}
35
36pub fn parse_ws_frame(data: &[u8]) -> Result<Option<(u8, Vec<u8>)>, Box<dyn std::error::Error>> {
37 if data.len() < 2 {
38 return Ok(None);
39 }
40
41 let first_byte = data[0];
42 let second_byte = data[1];
43
44 let fin = (first_byte & 0x80) != 0;
45 let opcode = first_byte & 0x0F;
46 let masked = (second_byte & 0x80) != 0;
47 let mut payload_len = (second_byte & 0x7F) as usize;
48
49 let mut pos = 2;
50
51 if payload_len == 126 {
53 if data.len() < pos + 2 {
54 return Ok(None);
55 }
56 payload_len = u16::from_be_bytes([data[pos], data[pos + 1]]) as usize;
57 pos += 2;
58 } else if payload_len == 127 {
59 if data.len() < pos + 8 {
60 return Ok(None);
61 }
62 payload_len = u64::from_be_bytes([
63 data[pos],
64 data[pos + 1],
65 data[pos + 2],
66 data[pos + 3],
67 data[pos + 4],
68 data[pos + 5],
69 data[pos + 6],
70 data[pos + 7],
71 ]) as usize;
72 pos += 8;
73 }
74
75 let mask_key = if masked {
77 if data.len() < pos + 4 {
78 return Ok(None);
79 }
80 let key = [data[pos], data[pos + 1], data[pos + 2], data[pos + 3]];
81 pos += 4;
82 Some(key)
83 } else {
84 None
85 };
86
87 if data.len() < pos + payload_len {
89 return Ok(None);
90 }
91
92 let mut payload = data[pos..pos + payload_len].to_vec();
93
94 if let Some(mask) = mask_key {
96 for (i, byte) in payload.iter_mut().enumerate() {
97 *byte ^= mask[i % 4];
98 }
99 }
100
101 if !fin {
102 return Ok(Some((opcode, payload)));
104 }
105
106 Ok(Some((opcode, payload)))
107}
108
109pub fn build_ws_frame(opcode: u8, payload: &[u8]) -> Vec<u8> {
110 let mut frame = Vec::new();
111
112 frame.push(0x80 | opcode);
114
115 let payload_len = payload.len();
117 if payload_len < 126 {
118 frame.push(payload_len as u8);
119 } else if payload_len < 65536 {
120 frame.push(126);
121 frame.extend_from_slice(&(payload_len as u16).to_be_bytes());
122 } else {
123 frame.push(127);
124 frame.extend_from_slice(&(payload_len as u64).to_be_bytes());
125 }
126
127 frame.extend_from_slice(payload);
129
130 frame
131}