reifydb_network/websocket/
frame.rs

1// Copyright (c) reifydb.com 2025
2// This file is licensed under the AGPL-3.0-or-later, see license.md file
3
4#[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	// Extended payload length
52	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	// Masking key
76	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	// Payload
88	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	// Unmask payload if necessary
95	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		// Handle fragmented frames - for now, just return what we have
103		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	// First byte: FIN = 1, RSV = 0, OPCODE = opcode
113	frame.push(0x80 | opcode);
114
115	// Payload length
116	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	// Payload (no masking for server->client frames)
128	frame.extend_from_slice(payload);
129
130	frame
131}