1extern crate rand;
2extern crate sha1;
3
4pub mod base64;
5
6pub use base64::B64Encode;
7pub use base64::B64;
8use rand::Rng;
9
10macro_rules! get_fin {
11 ($a: expr) => {
12 ($a >> 7) & 0b1
13 };
14}
15
16macro_rules! get_fin_enum {
17 ($a: expr) => {
18 match $a {
19 0 => FIN::NotFinal,
20 1 => FIN::Final,
21 _ => panic!(),
22 }
23 };
24}
25
26macro_rules! get_fin_enum_value {
27 ($a: expr) => {
28 match $a {
29 FIN::NotFinal => 0,
30 FIN::Final => 1,
31 _ => panic!(),
32 }
33 };
34}
35
36macro_rules! get_opcode {
37 ($a: expr) => {
38 $a & 0b1111
39 };
40}
41
42macro_rules! get_opcode_enum {
43 ($a: expr) => {
44 match $a {
45 0x0 => OPCODE::Fragment,
46 0x1 => OPCODE::Txt,
47 0x2 => OPCODE::Bin,
48 0x3 => OPCODE::Rsv10,
49 0x4 => OPCODE::Rsv11,
50 0x5 => OPCODE::Rsv12,
51 0x6 => OPCODE::Rsv13,
52 0x7 => OPCODE::Rsv14,
53 0x8 => OPCODE::Close,
54 0x9 => OPCODE::Ping,
55 0xA => OPCODE::Pong,
56 0xB => OPCODE::Rsv20,
57 0xC => OPCODE::Rsv21,
58 0xD => OPCODE::Rsv22,
59 0xE => OPCODE::Rsv23,
60 0xF => OPCODE::Rsv24,
61 _ => panic!(),
62 }
63 };
64}
65
66macro_rules! get_opcode_enum_value {
67 ($a: expr) => {
68 match $a {
69 OPCODE::Fragment => 0x0,
70 OPCODE::Txt => 0x1,
71 OPCODE::Bin => 0x2,
72 OPCODE::Rsv10 => 0x3,
73 OPCODE::Rsv11 => 0x4,
74 OPCODE::Rsv12 => 0x5,
75 OPCODE::Rsv13 => 0x6,
76 OPCODE::Rsv14 => 0x7,
77 OPCODE::Close => 0x8,
78 OPCODE::Ping => 0x9,
79 OPCODE::Pong => 0xA,
80 OPCODE::Rsv20 => 0xB,
81 OPCODE::Rsv21 => 0xC,
82 OPCODE::Rsv22 => 0xD,
83 OPCODE::Rsv23 => 0xE,
84 OPCODE::Rsv24 => 0xF,
85 _ => panic!(),
86 }
87 };
88}
89
90macro_rules! get_payload_length {
91 ($a: expr) => {
92 ($a & 0b_0111_1111) as usize
93 };
94}
95
96macro_rules! decode_payload_length_16 {
97 ($a: expr, $b: expr) => {
98 (($a as usize) << 8 & ($b as usize))
99 };
100}
101
102macro_rules! decode_payload_length_64 {
103 ($a: expr, $b: expr, $c: expr, $d: expr, $e: expr, $f: expr, $g: expr, $h: expr) => { (
104 ($a as usize) << 56 &
105 ($b as usize) << 48 &
106 ($c as usize) << 40 &
107 ($d as usize) << 32 &
108 ($e as usize) << 24 &
109 ($f as usize) << 16 &
110 ($g as usize) << 8 &
111 ($h as usize)
112 ) }
113}
114
115macro_rules! is_masked {
116 ($a: expr) => {
117 $a >> 7 == 1
118 };
119}
120
121#[derive(Debug, PartialEq)]
122pub enum OPCODE {
123 Fragment,
124 Txt,
125 Bin,
126 Rsv10,
127 Rsv11,
128 Rsv12,
129 Rsv13,
130 Rsv14,
131 Close,
132 Ping,
133 Pong,
134 Rsv20,
135 Rsv21,
136 Rsv22,
137 Rsv23,
138 Rsv24,
139}
140
141#[derive(Debug, PartialEq)]
142pub enum FIN {
143 NotFinal,
144 Final,
145}
146
147pub fn gen_mask() -> [u8; 4] {
148 let mut rng = rand::thread_rng();
149 let mut vec = [0u8; 4];
150 for x in vec.iter_mut() {
151 *x = rng.gen();
152 }
153 vec
154}
155
156pub fn gen_sec_websocket_key() -> String {
157 let mut rng = rand::thread_rng();
158 let mut vec = [0u8; 16];
159 for x in vec.iter_mut() {
160 *x = rng.gen();
161 }
162 B64::<String>::encode(&vec)
163}
164
165pub fn get_sec_websocket_accept(sec_web_socket_key: &str) -> String {
166 let mut hash = sha1::Sha1::new();
167 hash.update(
168 (sec_web_socket_key.to_owned() + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").as_bytes(),
169 );
170 B64::<String>::encode(&hash.digest().bytes())
171}
172
173#[derive(Debug)]
174pub struct Frame {
175 fin_rsv_opcode: u8,
176 mask_payload_length: u8,
177 payload_length: Vec<u8>,
178 mask: Option<[u8; 4]>,
179 payload: Vec<u8>,
180}
181
182impl Frame {
183 pub fn new(fin: FIN, opcode: OPCODE, data: &[u8]) -> Frame {
184 let fin_rsv_opcode: u8 = (get_fin_enum_value!(fin) << 7) | (get_opcode_enum_value!(opcode));
185 let payload_length: Vec<u8>;
186
187 let mask_payload_length: u8 = match data.len() {
188 len if len > 125 && len <= (u16::MAX as usize) => {
189 payload_length = vec![(len >> 8 & 0b_1111_1111) as u8, (len & 0b_1111_1111) as u8];
190 126u8
191 }
192 len if len > (u16::MAX as usize) => {
193 payload_length = vec![
194 (len >> 56 & 0b_1111_1111) as u8,
195 (len >> 48 & 0b_1111_1111) as u8,
196 (len >> 40 & 0b_1111_1111) as u8,
197 (len >> 32 & 0b_1111_1111) as u8,
198 (len >> 24 & 0b_1111_1111) as u8,
199 (len >> 16 & 0b_1111_1111) as u8,
200 (len >> 8 & 0b_1111_1111) as u8,
201 (len & 0b_1111_1111) as u8,
202 ];
203 127u8
204 }
205 len => {
206 payload_length = vec![];
207 len as u8
208 }
209 };
210 let payload = data.to_vec();
211
212 Frame {
213 fin_rsv_opcode,
214 mask_payload_length,
215 payload_length,
216 mask: None,
217 payload,
218 }
219 }
220 pub fn fin(&self) -> FIN {
221 get_fin_enum!(get_fin!(self.fin_rsv_opcode))
222 }
223 pub fn opcode(&self) -> OPCODE {
224 get_opcode_enum!(get_opcode!(self.fin_rsv_opcode))
225 }
226 pub fn masking(&mut self, mask: &[u8; 4]) {
227 if let Some(mask_) = self.mask {
228 if mask_ != *mask {
229 self.unmasking();
230 self.masking(mask);
231 }
232 } else {
233 for i in 0..self.payload.len() {
234 self.payload[i] ^= mask[i % 4];
235 }
236 self.mask = Some(*mask);
237 self.mask_payload_length ^= 0b_1000_0000;
238 }
239 }
240 pub fn unmasking(&mut self) {
241 if let Some(mask) = self.mask {
242 for i in 0..self.payload.len() {
243 self.payload[i] ^= mask[i % 4];
244 }
245 self.mask = None;
246 self.mask_payload_length &= 0b_0111_1111;
247 }
248 }
249 pub fn to_vec(&self) -> Vec<u8> {
250 let mut vec = Vec::<u8>::with_capacity(self.payload.len() + 14);
251
252 vec.push(self.fin_rsv_opcode);
253 vec.push(self.mask_payload_length);
254
255 vec.extend_from_slice(self.payload_length.as_slice());
256
257 if let Some(mask) = self.mask {
258 vec.extend_from_slice(&mask);
259 }
260
261 vec.extend_from_slice(self.payload.as_slice());
262
263 vec
264 }
265}
266
267pub fn decode_frame(v: &[u8], mut index: usize) -> Frame {
268 let fin_rsv_opcode = v[index];
269 let mask_payload_length = v[index + 1];
270
271 index += 2;
272
273 let mut temp_payload_length = get_payload_length!(mask_payload_length);
274 let payload_length = match temp_payload_length {
275 126 => {
276 index += 2;
277 temp_payload_length = decode_payload_length_16!(v[2], v[3]);
278 vec![v[2], v[3]]
279 }
280 127 => {
281 index += 8;
282 temp_payload_length =
283 decode_payload_length_64!(v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9]);
284 vec![v[2], v[3], v[4], v[5], v[6], v[7], v[8], v[9]]
285 }
286 _ => {
287 vec![]
288 }
289 };
290
291 let mask: Option<[u8; 4]> = if is_masked!(mask_payload_length) {
292 index += 4;
293 Some([v[index - 4], v[index - 3], v[index - 2], v[index - 1]])
294 } else {
295 None
296 };
297
298 let payload = v[index..(index + temp_payload_length)].to_vec();
299
300 index += temp_payload_length;
301
302 Frame {
303 fin_rsv_opcode,
304 mask_payload_length,
305 payload_length,
306 mask,
307 payload,
308 }
309}