zerodds_websocket_bridge/
masking.rs1use core::sync::atomic::{AtomicU64, Ordering};
7
8pub fn apply_mask(data: &mut [u8], key: [u8; 4]) {
16 for (i, b) in data.iter_mut().enumerate() {
17 *b ^= key[i & 0x3];
18 }
19}
20
21#[must_use]
34pub fn generate_masking_key() -> [u8; 4] {
35 static COUNTER: AtomicU64 = AtomicU64::new(0xCAFE_F00D_DEAD_BEEF);
36 let n = COUNTER.fetch_add(0x9E37_79B9_7F4A_7C15, Ordering::SeqCst);
37 let mut x = n;
38 x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
39 x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
40 x ^= x >> 31;
41 let bytes = x.to_le_bytes();
42 [bytes[0], bytes[1], bytes[2], bytes[3]]
43}
44
45pub trait MaskingKeyProvider {
53 fn next_key(&mut self) -> [u8; 4];
55}
56
57#[derive(Debug, Default)]
61pub struct InsecureSplitmixProvider;
62
63impl MaskingKeyProvider for InsecureSplitmixProvider {
64 fn next_key(&mut self) -> [u8; 4] {
65 generate_masking_key()
66 }
67}
68
69pub struct ClosureMaskingKeyProvider<F: FnMut() -> [u8; 4]>(pub F);
73
74impl<F: FnMut() -> [u8; 4]> MaskingKeyProvider for ClosureMaskingKeyProvider<F> {
75 fn next_key(&mut self) -> [u8; 4] {
76 (self.0)()
77 }
78}
79
80#[cfg(test)]
81mod tests {
82 use super::*;
83
84 #[test]
85 fn apply_mask_is_symmetric() {
86 let key = [0x01, 0x02, 0x03, 0x04];
88 let original: alloc::vec::Vec<u8> = (0..16).collect();
89 let mut masked = original.clone();
90 apply_mask(&mut masked, key);
91 assert_ne!(masked, original);
92 apply_mask(&mut masked, key);
93 assert_eq!(masked, original);
94 }
95
96 #[test]
97 fn apply_mask_xors_with_key_modulo_4() {
98 let key = [0xFF, 0x00, 0xFF, 0x00];
100 let mut data = [0xAA; 8];
101 apply_mask(&mut data, key);
102 assert_eq!(data, [0x55, 0xAA, 0x55, 0xAA, 0x55, 0xAA, 0x55, 0xAA]);
103 }
104
105 #[test]
106 fn apply_mask_handles_partial_key_alignment() {
107 let key = [0x01, 0x02, 0x03, 0x04];
110 let mut data = [0x10, 0x20, 0x30, 0x40, 0x50];
111 apply_mask(&mut data, key);
112 assert_eq!(data[0], 0x10 ^ 0x01);
114 assert_eq!(data[1], 0x20 ^ 0x02);
115 assert_eq!(data[2], 0x30 ^ 0x03);
116 assert_eq!(data[3], 0x40 ^ 0x04);
117 assert_eq!(data[4], 0x50 ^ 0x01);
118 }
119
120 #[test]
121 fn empty_payload_is_unchanged() {
122 let mut data: alloc::vec::Vec<u8> = alloc::vec::Vec::new();
123 apply_mask(&mut data, [1, 2, 3, 4]);
124 assert!(data.is_empty());
125 }
126
127 #[test]
128 fn generate_masking_key_returns_4_bytes() {
129 let k = generate_masking_key();
130 assert_eq!(k.len(), 4);
131 }
132
133 #[test]
134 fn generate_masking_key_returns_distinct_values_across_calls() {
135 let k1 = generate_masking_key();
138 let k2 = generate_masking_key();
139 assert_ne!(k1, k2);
140 }
141
142 #[test]
143 fn insecure_splitmix_provider_implements_trait() {
144 let mut p = InsecureSplitmixProvider;
145 let k = p.next_key();
146 assert_eq!(k.len(), 4);
147 }
148
149 #[test]
150 fn closure_provider_calls_user_supplied_fn() {
151 let mut p = ClosureMaskingKeyProvider(|| [9, 9, 9, 9]);
152 assert_eq!(p.next_key(), [9, 9, 9, 9]);
153 }
154}