vinteumorg · Davidson-Souza · May 22, 2024 · Apr 30, 2024
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/crates/floresta-chain/src/pruned_utreexo/chainstore.rs b/crates/floresta-chain/src/pruned_utreexo/chainstore.rs
@@ -83,7 +83,7 @@ impl Encodable for DiskBlockHeader {
         &self,
         writer: &mut W,
     ) -> core::result::Result<usize, ioError> {
-        let mut len = 80 + 1; // Header + tag + height
+        let mut len = 80 + 1; // Header + tag
         match self {
             DiskBlockHeader::FullyValid(header, height) => {
                 0x00_u8.consensus_encode(writer)?;

diff --git a/crates/floresta-compact-filters/src/kv_filter_database.rs b/crates/floresta-compact-filters/src/kv_filter_database.rs
@@ -1,3 +1,6 @@
+use core::fmt;
+use std::fmt::Debug;
+use std::fmt::Formatter;
 use std::path::PathBuf;
 
 use kv::Bucket;
@@ -13,6 +16,12 @@ pub struct KvFilterStore {
     bucket: Bucket<'static, Integer, Vec<u8>>,
 }
 
+impl Debug for KvFilterStore {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("KvFilterStore").finish()
+    }
+}
+
 impl KvFilterStore {
     /// Creates a new [KvFilterStore] that stores it's content in `datadir`.
     ///
@@ -33,6 +42,7 @@ impl KvFilterStore {
         let bucket = store.bucket(Some("cfilters")).unwrap();
         KvFilterStore { bucket }
     }
+
     /// Creates a new [KvFilterStore] that stores it's content with a given config
     pub fn with_config(config: Config) -> Self {
         let store = kv::Store::new(config).expect("Could not open database");

diff --git a/crates/floresta-compact-filters/src/lib.rs b/crates/floresta-compact-filters/src/lib.rs
@@ -24,7 +24,9 @@ use log::error;
 
 mod bip158;
 pub mod kv_filter_database;
+pub mod network_filters;
 
+pub use bip158::BlockFilter;
 use bip158::*;
 
 /// A database that stores our compact filters
@@ -91,9 +93,6 @@ pub struct BlockFilterBackend {
     key: [u8; 32],
 }
 
-unsafe impl Sync for BlockFilterBackend {}
-unsafe impl Send for BlockFilterBackend {}
-
 /// Keeps track of a unfinnished BIP-158 block filter. We use this to add new elements
 /// to the filter, until there's nothing more to add
 struct FilterBuilder<'a> {
@@ -255,7 +254,7 @@ impl BlockConsumer for BlockFilterBackend {
 /// Fields have the same meaning as in the backend itself.
 #[derive(Default)]
 pub struct FilterBackendBuilder {
-    /// Were we should store our filter.
+    /// Where we should store our filter.
     storage: Option<Box<dyn BlockFilterStore>>,
     /// What types of outputs should we store.
     whitelisted_outputs: u8,

diff --git a/crates/floresta-compact-filters/src/network_filters.rs b/crates/floresta-compact-filters/src/network_filters.rs
@@ -0,0 +1,67 @@
+use std::sync::PoisonError;
+use std::sync::RwLock;
+use std::sync::RwLockWriteGuard;
+
+use bitcoin::BlockHash;
+use floresta_chain::pruned_utreexo::BlockchainInterface;
+
+use crate::BlockFilterStore;
+
+#[derive(Debug)]
+pub struct NetworkFilters<Storage: BlockFilterStore + Send + Sync> {
+    filters: Storage,
+    height: RwLock<u32>,
+}
+
+impl<Storage: BlockFilterStore + Send + Sync> NetworkFilters<Storage> {
+    pub fn new(filters: Storage, height: u32) -> Self {
+        Self {
+            filters,
+            height: RwLock::new(height),
+        }
+    }
+
+    pub fn get_filter(&self, height: u32) -> Option<crate::BlockFilter> {
+        self.filters.get_filter(height as u64)
+    }
+
+    pub fn match_any(
+        &self,
+        query: Vec<&[u8]>,
+        start_height: u32,
+        end_height: u32,
+        chain: impl BlockchainInterface,
+    ) -> Vec<BlockHash> {
+        let mut blocks = Vec::new();
+        for height in start_height..end_height {
+            let Some(filter) = self.filters.get_filter(height as u64) else {
+                continue;
+            };
+
+            let mut query = query.clone().into_iter();
+            let hash = chain.get_block_hash(height).unwrap();
+
+            if filter.match_any(&hash, &mut query).unwrap() {
+                let block_hash = chain.get_block_hash(height).unwrap();
+                blocks.push(block_hash);
+            }
+        }
+
+        blocks
+    }
+
+    pub fn push_filter(
+        &self,
+        height: u32,
+        filter: crate::BlockFilter,
+    ) -> Result<(), PoisonError<RwLockWriteGuard<u32>>> {
+        self.filters.put_filter(height as u64, filter);
+        self.height.write().map(|mut self_height| {
+            *self_height = height;
+        })
+    }
+
+    pub fn get_height(&self) -> u32 {
+        self.height.read().map(|height| *height).unwrap_or(0)
+    }
+}
diff --git a/crates/floresta-electrum/src/electrum_protocol.rs b/crates/floresta-electrum/src/electrum_protocol.rs
@@ -1,3 +1,4 @@
+use std::borrow::Borrow;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::sync::Arc;
@@ -14,7 +15,6 @@ use bitcoin::consensus::deserialize;
 use bitcoin::consensus::encode::serialize_hex;
 use bitcoin::hashes::hex::FromHex;
 use bitcoin::hashes::sha256;
-use bitcoin::hashes::Hash;
 use bitcoin::ScriptBuf;
 use bitcoin::Transaction;
 use bitcoin::TxOut;
@@ -23,8 +23,8 @@ use floresta_chain::pruned_utreexo::BlockchainInterface;
 use floresta_common::get_hash_from_u8;
 use floresta_common::get_spk_hash;
 use floresta_common::spsc::Channel;
-use floresta_compact_filters::BlockFilterBackend;
-use floresta_compact_filters::QueryType;
+use floresta_compact_filters::kv_filter_database::KvFilterStore;
+use floresta_compact_filters::network_filters::NetworkFilters;
 use floresta_watch_only::kv_database::KvDatabase;
 use floresta_watch_only::AddressCache;
 use floresta_watch_only::CachedTransaction;
@@ -102,7 +102,7 @@ pub struct ElectrumServer<Blockchain: BlockchainInterface> {
     pub client_addresses: HashMap<sha256::Hash, Arc<Client>>,
     /// A Arc-ed copy of the block filters backend that we can use to check if a
     /// block contains a transaction that we are interested in.
-    pub block_filters: Option<Arc<BlockFilterBackend>>,
+    pub block_filters: Option<Arc<NetworkFilters<KvFilterStore>>>,
     /// An interface to a running node, used to broadcast transactions and request
     /// blocks.
     pub node_interface: Arc<NodeInterface>,
@@ -120,7 +120,7 @@ impl<Blockchain: BlockchainInterface> ElectrumServer<Blockchain> {
         address: String,
         address_cache: Arc<RwLock<AddressCache<KvDatabase>>>,
         chain: Arc<Blockchain>,
-        block_filters: Option<Arc<BlockFilterBackend>>,
+        block_filters: Option<Arc<NetworkFilters<KvFilterStore>>>,
         node_interface: Arc<NodeInterface>,
     ) -> Result<ElectrumServer<Blockchain>, Box<dyn std::error::Error>> {
         let listener = Arc::new(TcpListener::bind(address).await?);
@@ -442,40 +442,32 @@ impl<Blockchain: BlockchainInterface> ElectrumServer<Blockchain> {
     /// transactions, once a new address is added by subscription.
     async fn rescan_with_block_filters(
         &self,
-        cfilters: &BlockFilterBackend,
+        cfilters: &Arc<NetworkFilters<KvFilterStore>>,
         addresses: Vec<sha256::Hash>,
     ) -> Result<(), super::error::Error> {
         // By default, we look from 1..tip
         let height = self.chain.get_height().unwrap_or(0) as u64;
-
-        let addresses = addresses
-            .into_iter()
-            .map(|a| QueryType::ScriptHash(a.to_byte_array()))
+        let mut _addresses = addresses
+            .iter()
+            .map(|hash| hash.borrow())
             .collect::<Vec<_>>();
 
         // TODO (Davidson): Let users select what the starting and end height is
         let blocks: Vec<_> = cfilters
-            .match_any(1, height, &addresses)
-            .unwrap_or_default()
+            .match_any(_addresses, 1, height as u32, self.chain.clone())
             .into_iter()
-            .flat_map(|height| {
-                self.chain
-                    .get_block_hash(height as u32)
-                    .into_iter()
-                    .zip(Some(height))
-            })
-            .flat_map(|(hash, height)| {
-                self.node_interface
-                    .get_block(hash)
-                    .ok()
-                    .flatten()
-                    .map(|block| (block, height))
-            })
+            .flat_map(|hash| self.node_interface.get_block(hash).ok().flatten())
             .collect();
 
         // Tells users about the transactions we found
-        for (block, height) in blocks {
-            self.handle_block(block, height as u32).await;
+        for block in blocks {
+            let height = self
+                .chain
+                .get_block_height(&block.block_hash())
+                .ok()
+                .flatten()
+                .unwrap();
+            self.handle_block(block, height).await;
         }
 
         Ok(())

diff --git a/crates/floresta-wire/Cargo.toml b/crates/floresta-wire/Cargo.toml
@@ -29,6 +29,7 @@ rand = "0.8.5"
 bitcoin = { version = "0.31", features = ["serde", "std", "bitcoinconsensus"] }
 dns-lookup = "1.0.8"
 floresta-chain = { path = "../floresta-chain" }
+floresta-compact-filters = { path = "../floresta-compact-filters" }
 thiserror = "1.0"
 floresta-common = { path = "../floresta-common" }
 oneshot = "0.1.5"

diff --git a/crates/floresta-wire/src/p2p_wire/address_man.rs b/crates/floresta-wire/src/p2p_wire/address_man.rs
@@ -374,6 +374,10 @@ impl AddressMan {
     }
     /// Updates the state of an address
     pub fn update_set_state(&mut self, idx: usize, state: AddressState) -> &mut Self {
+        if matches!(state, AddressState::Banned(_)) {
+            return self;
+        }
+
         match state {
             AddressState::Banned(_) => {
                 self.good_addresses.retain(|&x| x != idx);

diff --git a/crates/floresta-wire/src/p2p_wire/chain_selector.rs b/crates/floresta-wire/src/p2p_wire/chain_selector.rs
@@ -113,7 +113,7 @@ impl NodeContext for ChainSelector {
     const REQUEST_TIMEOUT: u64 = 10; // Ban peers stalling our IBD
     const TRY_NEW_CONNECTION: u64 = 10; // Try creating connections more aggressively
 
-    fn get_required_services(&self, _utreexo_peers: usize) -> ServiceFlags {
+    fn get_required_services(&self) -> ServiceFlags {
         ServiceFlags::NETWORK | ServiceFlags::UTREEXO
     }
 }