websocket_base/ws/util/
mask.rs

1//! Utility functions for masking data frame payload data
2use rand;
3use std::io::Result as IoResult;
4use std::io::Write;
5
6/// Struct to pipe data into another writer,
7/// while masking the data being written
8pub struct Masker<'w> {
9	key: [u8; 4],
10	pos: usize,
11	end: &'w mut dyn Write,
12}
13
14impl<'w> Masker<'w> {
15	/// Create a new Masker with the key and the endpoint
16	/// to be writer to.
17	pub fn new(key: [u8; 4], endpoint: &'w mut dyn Write) -> Self {
18		Masker {
19			key,
20			pos: 0,
21			end: endpoint,
22		}
23	}
24}
25
26impl<'w> Write for Masker<'w> {
27	fn write(&mut self, data: &[u8]) -> IoResult<usize> {
28		let mut buf = Vec::with_capacity(data.len());
29		for &byte in data.iter() {
30			buf.push(byte ^ self.key[self.pos]);
31			self.pos = (self.pos + 1) % self.key.len();
32		}
33		self.end.write(&buf)
34	}
35
36	fn flush(&mut self) -> IoResult<()> {
37		self.end.flush()
38	}
39}
40
41/// Generates a random masking key
42pub fn gen_mask() -> [u8; 4] {
43	rand::random()
44}
45
46/// Masks data to send to a server and writes
47pub fn mask_data(mask: [u8; 4], data: &[u8]) -> Vec<u8> {
48	let mut out = Vec::with_capacity(data.len());
49	let zip_iter = data.iter().zip(mask.iter().cycle());
50	for (&buf_item, &key_item) in zip_iter {
51		out.push(buf_item ^ key_item);
52	}
53	out
54}
55
56#[cfg(all(feature = "nightly", test))]
57mod tests {
58	use super::*;
59	use test;
60
61	#[test]
62	fn test_mask_data() {
63		let key = [1u8, 2u8, 3u8, 4u8];
64		let original = vec![10u8, 11u8, 12u8, 13u8, 14u8, 15u8, 16u8, 17u8];
65		let expected = vec![11u8, 9u8, 15u8, 9u8, 15u8, 13u8, 19u8, 21u8];
66		let obtained = mask_data(key, &original[..]);
67		let reversed = mask_data(key, &obtained[..]);
68
69		assert_eq!(original, reversed);
70		assert_eq!(obtained, expected);
71	}
72
73	#[bench]
74	fn bench_mask_data(b: &mut test::Bencher) {
75		let buffer = b"The quick brown fox jumps over the lazy dog";
76		let key = gen_mask();
77		b.iter(|| {
78			let mut output = mask_data(key, buffer);
79			test::black_box(&mut output);
80		});
81	}
82
83	#[bench]
84	fn bench_gen_mask(b: &mut test::Bencher) {
85		b.iter(|| {
86			let mut key = gen_mask();
87			test::black_box(&mut key);
88		});
89	}
90}
websocket_base/ws/util/mask.rs

websocket_base/ws/util/
mask.rs
