actix_http/ws/
mask.rs

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