wowlab_sentinel/utils/
bloom.rs

1#![cfg_attr(
2    test,
3    expect(
4        clippy::cast_possible_truncation,
5        clippy::cast_precision_loss,
6        clippy::cast_sign_loss,
7        reason = "Bloom filter sizing follows the standard floating-point formulas and clamps allocation bounds"
8    )
9)]
10
11use sha2::{Digest, Sha256};
12
13const BITS_PER_BYTE: u64 = 8;
14const HASH_HALF_BYTES: usize = 8;
15const HASH_SECOND_OFFSET: usize = 16;
16
17#[cfg(test)]
18const LN_2: f64 = 2.0_f64;
19
20/// Kirsch-Mitzenmacher SHA-256 bloom filter, interoperable with the TypeScript implementation.
21#[derive(Debug)]
22pub(crate) struct BloomFilter {
23    bits: Vec<u8>,
24    num_bits: u64,
25    num_hashes: u32,
26}
27
28impl BloomFilter {
29    #[cfg(test)]
30    pub(crate) fn new(n: usize, fp_rate: f64) -> Self {
31        assert!(n > 0, "n must be > 0");
32        assert!(fp_rate > 0.0 && fp_rate < 1.0, "fp_rate must be in (0, 1)");
33
34        let raw_bits = optimal_num_bits(n, fp_rate);
35        let num_bytes = (raw_bits as usize).div_ceil(BITS_PER_BYTE as usize);
36        // Byte alignment lets from_bytes reconstruct the filter identically across the TS interop boundary.
37        let num_bits = num_bytes as u64 * BITS_PER_BYTE;
38        let num_hashes = optimal_num_hashes(num_bits, n);
39
40        Self {
41            bits: vec![0u8; num_bytes],
42            num_bits,
43            num_hashes,
44        }
45    }
46
47    #[cfg(test)]
48    pub(crate) fn from_bytes(bytes: Vec<u8>, member_count: usize) -> Self {
49        let num_bits = bytes.len() as u64 * BITS_PER_BYTE;
50        let num_hashes = optimal_num_hashes(num_bits, member_count);
51
52        Self {
53            bits: bytes,
54            num_bits,
55            num_hashes,
56        }
57    }
58
59    #[cfg(test)]
60    pub(crate) fn insert(&mut self, item: &str) {
61        let (h1, h2) = hash_item(item);
62
63        for i in 0..self.num_hashes {
64            let pos = self.position(h1, h2, i);
65
66            self.set_bit(pos);
67        }
68    }
69
70    pub(crate) fn might_contain(&self, item: &str) -> bool {
71        let (h1, h2) = hash_item(item);
72
73        for i in 0..self.num_hashes {
74            let pos = self.position(h1, h2, i);
75
76            if !self.get_bit(pos) {
77                return false;
78            }
79        }
80
81        true
82    }
83
84    #[cfg(test)]
85    pub(crate) fn as_bytes(&self) -> &[u8] {
86        &self.bits
87    }
88
89    #[cfg(test)]
90    pub(crate) fn into_bytes(self) -> Vec<u8> {
91        self.bits
92    }
93
94    fn position(&self, h1: u64, h2: u64, i: u32) -> u64 {
95        h1.wrapping_add(u64::from(i).wrapping_mul(h2)) % self.num_bits
96    }
97
98    // #t(fn: rust_lossy_cast, rust_unchecked_indexing) pos % 8 fits in u8; byte_idx bounded by position modulo num_bits
99    #[cfg(test)]
100    fn set_bit(&mut self, pos: u64) {
101        let byte_idx = (pos / BITS_PER_BYTE) as usize;
102        let bit_idx = (pos % BITS_PER_BYTE) as u8;
103
104        self.bits[byte_idx] |= 1 << bit_idx;
105    }
106
107    // #t(fn: rust_lossy_cast, rust_unchecked_indexing) pos % 8 fits in u8; byte_idx bounded by position modulo num_bits
108    fn get_bit(&self, pos: u64) -> bool {
109        let byte_idx = (pos / BITS_PER_BYTE) as usize;
110        let bit_idx = (pos % BITS_PER_BYTE) as u8;
111
112        (self.bits[byte_idx] >> bit_idx) & 1 == 1
113    }
114}
115
116// #t(fn: rust_unchecked_indexing) SHA-256 digest is always 32 bytes, slicing 0..8 and 8..16 is safe
117fn hash_item(item: &str) -> (u64, u64) {
118    let hash = Sha256::digest(item.as_bytes());
119    let h1 = u64::from_le_bytes(
120        hash[0..HASH_HALF_BYTES]
121            .try_into()
122            .expect("8-byte slice from SHA-256"),
123    );
124    let h2 = u64::from_le_bytes(
125        hash[HASH_HALF_BYTES..HASH_SECOND_OFFSET]
126            .try_into()
127            .expect("8-byte slice from SHA-256"),
128    );
129
130    (h1, h2)
131}
132
133#[cfg(test)]
134fn optimal_num_bits(n: usize, fp_rate: f64) -> u64 {
135    // #t(rust_magic_numbers) ln(2)^2 bloom filter formula constant
136    let m = -(n as f64) * fp_rate.ln() / (LN_2.ln().powi(2));
137
138    m.ceil() as u64
139}
140
141#[cfg(test)]
142fn optimal_num_hashes(num_bits: u64, n: usize) -> u32 {
143    let k = (num_bits as f64 / n as f64) * LN_2.ln();
144    let k = k.round() as u32;
145
146    k.max(1)
147}
148
149#[cfg(test)]
150const FP_RATE: f64 = 0.001;
151
152#[cfg(test)]
153fn create_server_filter(discord_ids: &[String]) -> BloomFilter {
154    let mut filter = BloomFilter::new(discord_ids.len().max(1), FP_RATE);
155
156    for id in discord_ids {
157        filter.insert(id);
158    }
159
160    filter
161}
162
163#[cfg(test)]
164pub(crate) fn filter_hash(bytes: &[u8]) -> String {
165    let hash = Sha256::digest(bytes);
166    // #t(rust_unchecked_indexing) SHA-256 digest is always 32 bytes, 0..8 is safe
167
168    hex::encode(&hash[..HASH_HALF_BYTES])
169}
170
171#[cfg(test)]
172mod tests {
173    use base64::Engine;
174    use googletest::prelude::*;
175
176    use super::*;
177
178    #[gtest]
179    fn test_insert_and_contains() -> Result<()> {
180        let ids: Vec<String> = (0..100)
181            .map(|i| format!("{}", 100_000_000_000_000_000_u64 + i))
182            .collect();
183        let filter = create_server_filter(&ids);
184
185        for id in &ids {
186            verify_true!(filter.might_contain(id))?;
187        }
188
189        Ok(())
190    }
191
192    #[gtest]
193    fn test_false_positive_rate() -> Result<()> {
194        let n = 1000;
195        let ids: Vec<String> = (0..n)
196            .map(|i| format!("{}", 100_000_000_000_000_000_u64 + i))
197            .collect();
198        let filter = create_server_filter(&ids);
199
200        let mut false_positives = 0;
201
202        for i in n..(n + 10000) {
203            let id = format!("{}", 100_000_000_000_000_000_u64 + i);
204
205            if filter.might_contain(&id) {
206                false_positives += 1;
207            }
208        }
209
210        let fp_rate = f64::from(false_positives) / 10000.0;
211
212        verify_true!(fp_rate < 0.005)?;
213
214        Ok(())
215    }
216
217    #[gtest]
218    fn test_roundtrip_from_bytes() -> Result<()> {
219        let ids: Vec<String> = (0..500)
220            .map(|i| format!("{}", 200_000_000_000_000_000_u64 + i))
221            .collect();
222        let filter = create_server_filter(&ids);
223        let bytes = filter.into_bytes();
224
225        let restored = BloomFilter::from_bytes(bytes, 500);
226
227        for id in &ids {
228            verify_true!(restored.might_contain(id))?;
229        }
230
231        Ok(())
232    }
233
234    #[gtest]
235    fn test_interop_vectors() -> Result<()> {
236        let ids = vec![
237            "123456789012345678".to_string(),
238            "987654321098765432".to_string(),
239            "111222333444555666".to_string(),
240        ];
241        let filter = create_server_filter(&ids);
242        let bytes = filter.as_bytes();
243        let b64 = base64::engine::general_purpose::STANDARD.encode(bytes);
244        let hash = filter_hash(bytes);
245
246        eprintln!("interop vector:");
247        eprintln!("  base64: {b64}");
248        eprintln!("  hash: {hash}");
249        eprintln!("  byte_len: {}", bytes.len());
250
251        verify_true!(filter.might_contain("123456789012345678"))?;
252        verify_true!(filter.might_contain("987654321098765432"))?;
253        verify_true!(filter.might_contain("111222333444555666"))?;
254
255        verify_true!(!filter.might_contain("000000000000000000"))?;
256        verify_true!(!filter.might_contain("999999999999999999"))?;
257
258        verify_eq!(bytes.len(), 6)?;
259        verify_eq!(hash, "554721ba80ba8f66")?;
260        verify_eq!(b64, "E5Nn44kd")?;
261
262        Ok(())
263    }
264
265    #[gtest]
266    fn test_filter_sizing() -> Result<()> {
267        let ids: Vec<String> = (0..1000).map(|i| format!("{i}")).collect();
268        let filter = create_server_filter(&ids);
269        let size = filter.as_bytes().len();
270
271        verify_true!(size > 1500 && size < 2500)?;
272
273        Ok(())
274    }
275}
wowlab_sentinel/utils/bloom.rs

wowlab_sentinel/utils/
bloom.rs
