1#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum OpCode {
8 Continuation,
9 Text,
10 Binary,
11 Close,
12 Ping,
13 Pong,
14}
15
16impl OpCode {
17 pub fn from_u8(b: u8) -> Option<OpCode> {
18 Some(match b {
19 0x0 => OpCode::Continuation,
20 0x1 => OpCode::Text,
21 0x2 => OpCode::Binary,
22 0x8 => OpCode::Close,
23 0x9 => OpCode::Ping,
24 0xA => OpCode::Pong,
25 _ => return None,
26 })
27 }
28
29 pub fn as_u8(self) -> u8 {
30 match self {
31 OpCode::Continuation => 0x0,
32 OpCode::Text => 0x1,
33 OpCode::Binary => 0x2,
34 OpCode::Close => 0x8,
35 OpCode::Ping => 0x9,
36 OpCode::Pong => 0xA,
37 }
38 }
39
40 pub fn is_control(self) -> bool {
41 matches!(self, OpCode::Close | OpCode::Ping | OpCode::Pong)
42 }
43}
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq)]
46pub struct Header {
47 pub fin: bool,
48 pub rsv1: bool,
49 pub rsv2: bool,
50 pub rsv3: bool,
51 pub opcode: u8,
52 pub mask: Option<[u8; 4]>,
53 pub payload_len: u64,
54 pub header_len: usize,
55}
56
57pub fn parse_header(buf: &[u8]) -> Result<Option<Header>, &'static str> {
59 if buf.len() < 2 {
60 return Ok(None);
61 }
62 let b0 = buf[0];
63 let b1 = buf[1];
64 let masked = b1 & 0x80 != 0;
65 let len7 = (b1 & 0x7F) as u64;
66 let mut pos = 2;
67 let payload_len = match len7 {
68 126 => {
69 if buf.len() < pos + 2 {
70 return Ok(None);
71 }
72 let l = u16::from_be_bytes([buf[2], buf[3]]) as u64;
73 pos += 2;
74 if l < 126 {
75 return Err("non-minimal payload length encoding");
76 }
77 l
78 }
79 127 => {
80 if buf.len() < pos + 8 {
81 return Ok(None);
82 }
83 let l = u64::from_be_bytes(buf[2..10].try_into().expect("8 bytes"));
84 pos += 8;
85 if l >> 63 != 0 {
86 return Err("payload length has most significant bit set");
87 }
88 if l <= 0xFFFF {
89 return Err("non-minimal payload length encoding");
90 }
91 l
92 }
93 l => l,
94 };
95 let mask = if masked {
96 if buf.len() < pos + 4 {
97 return Ok(None);
98 }
99 let m = [buf[pos], buf[pos + 1], buf[pos + 2], buf[pos + 3]];
100 pos += 4;
101 Some(m)
102 } else {
103 None
104 };
105 Ok(Some(Header {
106 fin: b0 & 0x80 != 0,
107 rsv1: b0 & 0x40 != 0,
108 rsv2: b0 & 0x20 != 0,
109 rsv3: b0 & 0x10 != 0,
110 opcode: b0 & 0x0F,
111 mask,
112 payload_len,
113 header_len: pos,
114 }))
115}
116
117pub fn apply_mask(data: &mut [u8], mask: [u8; 4]) {
119 for (i, b) in data.iter_mut().enumerate() {
120 *b ^= mask[i & 3];
121 }
122}
123
124pub fn encode(
126 out: &mut Vec<u8>,
127 opcode: OpCode,
128 rsv1: bool,
129 payload: &[u8],
130 mask: Option<[u8; 4]>,
131) {
132 out.push(0x80 | if rsv1 { 0x40 } else { 0 } | opcode.as_u8());
133 let mask_bit = if mask.is_some() { 0x80 } else { 0 };
134 match payload.len() {
135 n if n < 126 => out.push(mask_bit | n as u8),
136 n if n <= 0xFFFF => {
137 out.push(mask_bit | 126);
138 out.extend_from_slice(&(n as u16).to_be_bytes());
139 }
140 n => {
141 out.push(mask_bit | 127);
142 out.extend_from_slice(&(n as u64).to_be_bytes());
143 }
144 }
145 match mask {
146 Some(m) => {
147 out.extend_from_slice(&m);
148 let start = out.len();
149 out.extend_from_slice(payload);
150 apply_mask(&mut out[start..], m);
151 }
152 None => out.extend_from_slice(payload),
153 }
154}
155
156#[cfg(test)]
157mod tests {
158 use super::*;
159
160 #[test]
161 fn rfc6455_examples() {
162 let h = parse_header(&[0x81, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f])
164 .unwrap()
165 .unwrap();
166 assert!(
167 h.fin && h.mask.is_none() && h.opcode == 1 && h.payload_len == 5 && h.header_len == 2
168 );
169 let mut f = [
171 0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58,
172 ];
173 let h = parse_header(&f).unwrap().unwrap();
174 let m = h.mask.unwrap();
175 apply_mask(&mut f[6..], m);
176 assert_eq!(&f[6..], b"Hello");
177 let mut out = Vec::new();
179 encode(&mut out, OpCode::Binary, false, &[0u8; 256], None);
180 assert_eq!(&out[..4], &[0x82, 0x7E, 0x01, 0x00]);
181 out.clear();
183 encode(&mut out, OpCode::Binary, false, &vec![0u8; 65536], None);
184 assert_eq!(&out[..10], &[0x82, 0x7F, 0, 0, 0, 0, 0, 1, 0, 0]);
185 assert_eq!(parse_header(&out).unwrap().unwrap().payload_len, 65536);
186 assert_eq!(parse_header(&[0x81]).unwrap(), None);
188 assert_eq!(parse_header(&[0x82, 0x7F, 0]).unwrap(), None);
189 }
190}