1#![allow(dead_code)]
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9#[repr(u8)]
10pub enum WsOpcode {
11 Continuation = 0x0,
12 Text = 0x1,
13 Binary = 0x2,
14 Close = 0x8,
15 Ping = 0x9,
16 Pong = 0xA,
17}
18
19impl WsOpcode {
20 pub fn from_u8(b: u8) -> Option<Self> {
22 match b & 0x0F {
23 0x0 => Some(Self::Continuation),
24 0x1 => Some(Self::Text),
25 0x2 => Some(Self::Binary),
26 0x8 => Some(Self::Close),
27 0x9 => Some(Self::Ping),
28 0xA => Some(Self::Pong),
29 _ => None,
30 }
31 }
32}
33
34#[derive(Debug, Clone, PartialEq)]
36pub struct WsFrame {
37 pub fin: bool,
38 pub opcode: WsOpcode,
39 pub masked: bool,
40 pub masking_key: Option<[u8; 4]>,
41 pub payload: Vec<u8>,
42}
43
44#[derive(Debug, Clone, PartialEq)]
46pub enum WsError {
47 InsufficientData,
48 UnknownOpcode(u8),
49 Rsv1Set,
50 Rsv2Set,
51 Rsv3Set,
52}
53
54impl std::fmt::Display for WsError {
55 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56 match self {
57 Self::InsufficientData => write!(f, "insufficient data for WebSocket frame"),
58 Self::UnknownOpcode(b) => write!(f, "unknown WebSocket opcode: 0x{b:x}"),
59 Self::Rsv1Set => write!(f, "RSV1 bit set without negotiated extension"),
60 Self::Rsv2Set => write!(f, "RSV2 bit set without negotiated extension"),
61 Self::Rsv3Set => write!(f, "RSV3 bit set without negotiated extension"),
62 }
63 }
64}
65
66pub fn encode_frame(frame: &WsFrame, buf: &mut Vec<u8>) {
68 let b0 = (if frame.fin { 0x80 } else { 0 }) | (frame.opcode as u8);
69 buf.push(b0);
70 let payload_len = frame.payload.len();
71 if payload_len <= 125 {
72 buf.push(payload_len as u8);
73 } else if payload_len <= 65535 {
74 buf.push(126);
75 buf.extend_from_slice(&(payload_len as u16).to_be_bytes());
76 } else {
77 buf.push(127);
78 buf.extend_from_slice(&(payload_len as u64).to_be_bytes());
79 }
80 buf.extend_from_slice(&frame.payload);
81}
82
83pub fn decode_frame(buf: &[u8]) -> Result<WsFrame, WsError> {
85 if buf.len() < 2 {
86 return Err(WsError::InsufficientData);
87 }
88 let b0 = buf[0];
89 let b1 = buf[1];
90 let fin = (b0 & 0x80) != 0;
91 let opcode_raw = b0 & 0x0F;
92 let opcode = WsOpcode::from_u8(opcode_raw).ok_or(WsError::UnknownOpcode(opcode_raw))?;
93 let masked = (b1 & 0x80) != 0;
94 let len_field = (b1 & 0x7F) as usize;
95 let payload_start = 2;
96 let payload_len = if len_field <= 125 {
97 len_field
98 } else if len_field == 126 {
99 if buf.len() < 4 {
100 return Err(WsError::InsufficientData);
101 }
102 u16::from_be_bytes([buf[2], buf[3]]) as usize
103 } else {
104 if buf.len() < 10 {
105 return Err(WsError::InsufficientData);
106 }
107 u64::from_be_bytes(buf[2..10].try_into().unwrap_or_default()) as usize
108 };
109 let offset = if len_field <= 125 {
110 payload_start
111 } else if len_field == 126 {
112 4
113 } else {
114 10
115 };
116 if buf.len() < offset + payload_len {
117 return Err(WsError::InsufficientData);
118 }
119 let payload = buf[offset..offset + payload_len].to_vec();
120 Ok(WsFrame {
121 fin,
122 opcode,
123 masked,
124 masking_key: None,
125 payload,
126 })
127}
128
129pub fn apply_mask(payload: &mut [u8], key: [u8; 4]) {
131 for (i, byte) in payload.iter_mut().enumerate() {
132 *byte ^= key[i % 4];
133 }
134}
135
136pub fn is_control_frame(frame: &WsFrame) -> bool {
138 matches!(
139 frame.opcode,
140 WsOpcode::Close | WsOpcode::Ping | WsOpcode::Pong
141 )
142}
143
144pub fn text_frame(payload: &str) -> WsFrame {
146 WsFrame {
147 fin: true,
148 opcode: WsOpcode::Text,
149 masked: false,
150 masking_key: None,
151 payload: payload.as_bytes().to_vec(),
152 }
153}
154
155#[cfg(test)]
156mod tests {
157 use super::*;
158
159 #[test]
160 fn test_encode_small_frame() {
161 let f = text_frame("hi");
163 let mut buf = vec![];
164 encode_frame(&f, &mut buf);
165 assert_eq!(buf.len(), 4); }
167
168 #[test]
169 fn test_decode_text_frame() {
170 let f = text_frame("hello");
172 let mut buf = vec![];
173 encode_frame(&f, &mut buf);
174 let decoded = decode_frame(&buf).expect("should succeed");
175 assert_eq!(decoded.opcode, WsOpcode::Text);
176 assert_eq!(decoded.payload, b"hello");
177 }
178
179 #[test]
180 fn test_is_control_ping() {
181 let f = WsFrame {
183 fin: true,
184 opcode: WsOpcode::Ping,
185 masked: false,
186 masking_key: None,
187 payload: vec![],
188 };
189 assert!(is_control_frame(&f));
190 }
191
192 #[test]
193 fn test_is_control_text_false() {
194 let f = text_frame("x");
196 assert!(!is_control_frame(&f));
197 }
198
199 #[test]
200 fn test_apply_mask_roundtrip() {
201 let key = [0xAB, 0xCD, 0xEF, 0x12];
203 let original = vec![1u8, 2, 3, 4];
204 let mut data = original.clone();
205 apply_mask(&mut data, key);
206 apply_mask(&mut data, key);
207 assert_eq!(data, original);
208 }
209
210 #[test]
211 fn test_opcode_from_u8_text() {
212 assert_eq!(WsOpcode::from_u8(1), Some(WsOpcode::Text));
214 }
215
216 #[test]
217 fn test_opcode_unknown() {
218 assert!(WsOpcode::from_u8(3).is_none());
220 }
221
222 #[test]
223 fn test_insufficient_data() {
224 assert!(decode_frame(&[0x81]).is_err());
226 }
227
228 #[test]
229 fn test_fin_bit() {
230 let f = text_frame("x");
232 let mut buf = vec![];
233 encode_frame(&f, &mut buf);
234 let d = decode_frame(&buf).expect("should succeed");
235 assert!(d.fin);
236 }
237
238 #[test]
239 fn test_binary_opcode() {
240 let f = WsFrame {
242 fin: true,
243 opcode: WsOpcode::Binary,
244 masked: false,
245 masking_key: None,
246 payload: vec![0xFF],
247 };
248 let mut buf = vec![];
249 encode_frame(&f, &mut buf);
250 let d = decode_frame(&buf).expect("should succeed");
251 assert_eq!(d.opcode, WsOpcode::Binary);
252 }
253}