tungstenite/protocol/frame/
mask.rs

Help
1/// Generate a random frame mask.
2#[inline]
3pub fn generate_mask() -> [u8; 4] {
4    rand::random()
5}
6
7/// Mask/unmask a frame.
8#[inline]
9pub 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]
15fn 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]
23pub 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)]
45mod tests {
46    use super::*;
47
48    #[test]
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}
tungstenite/protocol/frame/mask.rs

tungstenite/protocol/frame/
mask.rs
