1 /// Generate a random frame mask.
2 #[inline]
generate_mask() -> [u8; 4]3 pub fn generate_mask() -> [u8; 4] {
4     rand::random()
5 }
6 
7 /// Mask/unmask a frame.
8 #[inline]
apply_mask(buf: &mut [u8], mask: [u8; 4])9 pub fn apply_mask(buf: &mut [u8], mask: [u8; 4]) {
10     apply_mask_fast32(buf, mask)
11 }
12 
13 /// A safe unoptimized mask application.
14 #[inline]
apply_mask_fallback(buf: &mut [u8], mask: [u8; 4])15 fn apply_mask_fallback(buf: &mut [u8], mask: [u8; 4]) {
16     for (i, byte) in buf.iter_mut().enumerate() {
17         *byte ^= mask[i & 3];
18     }
19 }
20 
21 /// Faster version of `apply_mask()` which operates on 4-byte blocks.
22 #[inline]
apply_mask_fast32(buf: &mut [u8], mask: [u8; 4])23 pub fn apply_mask_fast32(buf: &mut [u8], mask: [u8; 4]) {
24     let mask_u32 = u32::from_ne_bytes(mask);
25 
26     let (prefix, words, suffix) = unsafe { buf.align_to_mut::<u32>() };
27     apply_mask_fallback(prefix, mask);
28     let head = prefix.len() & 3;
29     let mask_u32 = if head > 0 {
30         if cfg!(target_endian = "big") {
31             mask_u32.rotate_left(8 * head as u32)
32         } else {
33             mask_u32.rotate_right(8 * head as u32)
34         }
35     } else {
36         mask_u32
37     };
38     for word in words.iter_mut() {
39         *word ^= mask_u32;
40     }
41     apply_mask_fallback(suffix, mask_u32.to_ne_bytes());
42 }
43 
44 #[cfg(test)]
45 mod tests {
46     use super::*;
47 
48     #[test]
test_apply_mask()49     fn test_apply_mask() {
50         let mask = [0x6d, 0xb6, 0xb2, 0x80];
51         let unmasked = [
52             0xf3, 0x00, 0x01, 0x02, 0x03, 0x80, 0x81, 0x82, 0xff, 0xfe, 0x00, 0x17, 0x74, 0xf9,
53             0x12, 0x03,
54         ];
55 
56         for data_len in 0..=unmasked.len() {
57             let unmasked = &unmasked[0..data_len];
58             // Check masking with different alignment.
59             for off in 0..=3 {
60                 if unmasked.len() < off {
61                     continue;
62                 }
63                 let mut masked = unmasked.to_vec();
64                 apply_mask_fallback(&mut masked[off..], mask);
65 
66                 let mut masked_fast = unmasked.to_vec();
67                 apply_mask_fast32(&mut masked_fast[off..], mask);
68 
69                 assert_eq!(masked, masked_fast);
70             }
71         }
72     }
73 }
74