Skip to main content

oxihuman_core/
websocket_frame.rs

1// Copyright (C) 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3#![allow(dead_code)]
4
5//! WebSocket frame encode/decode stub (RFC 6455).
6
7/// WebSocket opcode.
8#[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    /// Parse a raw opcode byte.
21    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/// A decoded WebSocket frame.
35#[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/// WebSocket frame error.
45#[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
66/// Encode a WebSocket frame (server-side: unmasked).
67pub 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
83/// Decode a WebSocket frame from a byte slice (unmasked only).
84pub 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
129/// Apply (or remove) a WebSocket masking key to a payload in place.
130pub 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
136/// Return `true` if the frame is a control frame.
137pub fn is_control_frame(frame: &WsFrame) -> bool {
138    matches!(
139        frame.opcode,
140        WsOpcode::Close | WsOpcode::Ping | WsOpcode::Pong
141    )
142}
143
144/// Build a simple text frame.
145pub 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        /* small payload encodes to 2 + payload bytes */
162        let f = text_frame("hi");
163        let mut buf = vec![];
164        encode_frame(&f, &mut buf);
165        assert_eq!(buf.len(), 4); /* 2 header + 2 payload */
166    }
167
168    #[test]
169    fn test_decode_text_frame() {
170        /* round-trip encode/decode */
171        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        /* Ping is a control frame */
182        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        /* Text is not a control frame */
195        let f = text_frame("x");
196        assert!(!is_control_frame(&f));
197    }
198
199    #[test]
200    fn test_apply_mask_roundtrip() {
201        /* applying mask twice restores original */
202        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        /* opcode 1 is Text */
213        assert_eq!(WsOpcode::from_u8(1), Some(WsOpcode::Text));
214    }
215
216    #[test]
217    fn test_opcode_unknown() {
218        /* opcode 3 is unknown */
219        assert!(WsOpcode::from_u8(3).is_none());
220    }
221
222    #[test]
223    fn test_insufficient_data() {
224        /* single byte returns error */
225        assert!(decode_frame(&[0x81]).is_err());
226    }
227
228    #[test]
229    fn test_fin_bit() {
230        /* FIN bit preserved */
231        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        /* binary opcode round-trip */
241        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}