Skip to main content

deaddrop_core/protocol/
inventory.rs

1use crate::crypto::{CryptoProvider, DefaultProvider};
2use crate::{HashAlgorithm, ObjectId, Result};
3
4pub struct BloomFilter {
5    pub n: u32,
6    pub k: u8,
7    pub bits: Vec<u8>,
8}
9
10impl BloomFilter {
11    pub fn from_ids(ids: &[ObjectId], bits_per: usize) -> Self {
12        let m_bits = (ids.len().max(1) * bits_per).max(64);
13        let bytes = m_bits.div_ceil(8);
14        let k = 4u8;
15        let mut bits = vec![0u8; bytes];
16        for id in ids {
17            for i in 0..k {
18                let idx = bit_index(id, i, m_bits);
19                bits[idx / 8] |= 1 << (idx % 8);
20            }
21        }
22        Self {
23            n: ids.len() as u32,
24            k,
25            bits,
26        }
27    }
28
29    pub fn may_contain(&self, id: &ObjectId) -> bool {
30        let m_bits = self.bits.len() * 8;
31        if m_bits == 0 {
32            return false;
33        }
34        (0..self.k).all(|i| {
35            let idx = bit_index(id, i, m_bits);
36            self.bits[idx / 8] & (1 << (idx % 8)) != 0
37        })
38    }
39}
40
41fn bit_index(id: &ObjectId, i: u8, m_bits: usize) -> usize {
42    let p = DefaultProvider;
43    let mut buf = Vec::from(id.as_bytes().as_slice());
44    buf.push(i);
45    let h = p.hash(HashAlgorithm::Blake3, &buf);
46    let v = u64::from_le_bytes(h.0[..8].try_into().unwrap());
47    (v as usize) % m_bits
48}
49
50pub enum InventoryAlgo {
51    Sorted,
52    Bloom,
53}
54
55pub fn choose_inventory(local_count: usize, peer_supports_bloom: bool) -> InventoryAlgo {
56    if peer_supports_bloom && local_count > 64 {
57        InventoryAlgo::Bloom
58    } else {
59        InventoryAlgo::Sorted
60    }
61}
62
63pub fn want_from_sorted(peer_ids: &[ObjectId], local: &[ObjectId]) -> Vec<ObjectId> {
64    let set: std::collections::HashSet<_> = local.iter().copied().collect();
65    peer_ids
66        .iter()
67        .copied()
68        .filter(|id| !set.contains(id))
69        .collect()
70}
71
72pub fn want_from_bloom(local: &[ObjectId], bloom: &BloomFilter) -> Vec<ObjectId> {
73    // Receiver of a bloom cannot list peer IDs; this returns local IDs the peer may lack.
74    local
75        .iter()
76        .copied()
77        .filter(|id| !bloom.may_contain(id))
78        .collect()
79}
80
81pub fn _ok() -> Result<()> {
82    Ok(())
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::ObjectId;
89
90    #[test]
91    fn bloom_no_false_negative() {
92        let ids: Vec<_> = (0u8..40)
93            .map(|i| {
94                let mut d = [0u8; 32];
95                d[0] = i;
96                ObjectId::blake3(d)
97            })
98            .collect();
99        let b = BloomFilter::from_ids(&ids, 10);
100        for id in &ids {
101            assert!(b.may_contain(id));
102        }
103    }
104}