hunter/block_list: integrate

cave/src/live_file.rs

@@ -1,7 +1,7 @@
 use std::ops::Deref;
-use std::{sync::Arc, path::Path};
+use std::sync::Arc;
 use futures::{Future, StreamExt};
-use tokio::{fs::File, io::AsyncReadExt};
+use tokio::fs::File;
 use tokio::sync::RwLock;
 
 use inotify::{Inotify, WatchMask};
@@ -34,7 +34,6 @@ where
         inotify.event_stream([0; 1024])
             .unwrap()
             .for_each(|result| {
-                dbg!(&result);
                 let path = path.clone();
                 let dir = dir.clone();
                 let f = f.clone();

hunter/src/block_list.rs

@@ -11,13 +11,53 @@ enum Leaf {
     Tree(HashMap<String, Leaf>),
 }
 
+impl Leaf {
+    pub fn insert(&mut self, host: &str) {
+        match self {
+            Leaf::Blocked => {
+                return;
+            }
+            &mut Leaf::Tree(ref mut tree) => {
+                if let Some(pos) = host.rfind('.') {
+                    tree.entry(host[pos + 1..].to_string())
+                        .or_insert(Leaf::Tree(HashMap::new()))
+                        .insert(&host[..pos]);
+                } else {
+                    tree.insert(host.to_string(), Leaf::Blocked);
+                }
+            }
+        }
+    }
+
+    pub fn is_blocked(&self, host: &str) -> bool {
+        match self {
+            Leaf::Blocked =>
+                true,
+            Leaf::Tree(ref tree) if host.len() > 0 => {
+                let (label, new_host) = if let Some(pos) = host.rfind('.') {
+                    (&host[pos + 1..], &host[..pos])
+                } else {
+                    (host, "")
+                };
+                if let Some(leaf) = tree.get(label) {
+                    leaf.is_blocked(new_host)
+                } else {
+                    false
+                }
+            }
+            _ =>
+                false,
+        }
+    }
+}
+
 pub struct BlockList {
-    tree: Arc<RwLock<HashMap<String, Leaf>>>,
+    tree: Arc<RwLock<Leaf>>,
 }
 
 impl BlockList {
     pub async fn new(path: &str) -> BlockList {
-        let tree = cave::live_file::load(path, |file| async move {
+        let root = cave::live_file::load(path, |file| async move {
             let mut root = Leaf::Tree(HashMap::new());
             let mut file = BufReader::new(file);
             let mut line = String::new();
@@ -26,43 +66,60 @@ impl BlockList {
                     break
                 }
 
-                fn insert(leaf: &mut Leaf, mut iter: impl Iterator<Item = (usize, String)>) {
-                    match leaf {
-                        Leaf::Blocked => {
-                            return;
-                        }
-                        &mut Leaf::Tree(ref mut tree) => {
-                            if let Some((i, label)) = iter.next() {
-                                let leaf = tree.entry(label)
-                                    .or_insert(if i > 0 {
-                                        Leaf::Tree(HashMap::new())
-                                    } else {
-                                        Leaf::Blocked
-                                    });
-                                insert(leaf, iter);
-                            }
-                        }
-                    }
-                }
-                insert(&mut root, line.trim_end().split('.')
-                       .collect::<Vec<_>>()
-                       .into_iter()
-                       .map(|s| s.to_string())
-                       .enumerate()
-                       .rev());
+                root.insert(line.trim_end());
 
                 line = String::new();
             }
 
-            match root {
+            match &root {
                 Leaf::Blocked =>
                     panic!("the whole internet has been blocked"),
-                Leaf::Tree(tree) =>
-                    tree,
+                Leaf::Tree(_) => {}
             }
+            root
         }).await.unwrap();
 
 
-        BlockList { tree }
+        BlockList { tree: root }
+    }
+
+    pub async fn is_blocked(&self, host: &str) -> bool {
+        if self.tree.read().await
+            .is_blocked(host)
+        {
+            tracing::warn!("host {} is blocked", host);
+            true
+        } else {
+            false
+        }
     }
 }
+
+#[cfg(test)]
+#[tokio::test]
+async fn test_blocklist() {
+    let bl = BlockList {
+        tree: Arc::new(RwLock::new(Leaf::Tree(HashMap::new()))),
+    };
+    let mut root = bl.tree.write().await;
+    root.insert("bad.actor");
+    root.insert("evil.actor");
+    drop(root);
+    assert_eq!(true, bl.is_blocked("bad.actor").await);
+    assert_eq!(true, bl.is_blocked("evil.actor").await);
+    assert_eq!(false, bl.is_blocked("good.actor").await);
+    assert_eq!(false, bl.is_blocked("not-bad.actor").await);
+    assert_eq!(false, bl.is_blocked("actor").await);
+}
+
+#[cfg(test)]
+#[tokio::test]
+async fn test_blocklist_subdomain() {
+    let bl = BlockList {
+        tree: Arc::new(RwLock::new(Leaf::Tree(HashMap::new()))),
+    };
+    bl.tree.write().await
+        .insert("bad.actor");
+    assert_eq!(true, bl.is_blocked("bad.actor").await);
+    assert_eq!(true, bl.is_blocked("very.bad.actor").await);
+}

hunter/src/main.rs

@@ -37,10 +37,9 @@ async fn main() {
     let posts_cache = posts_cache::PostsCache::new(65536);
 
     let block_list = block_list::BlockList::new(&config.blocklist).await;
-    panic!("block_list");
 
     cave::systemd::status("Starting scheduler");
-    let mut scheduler = scheduler::Scheduler::new();
+    let mut scheduler = scheduler::Scheduler::new(block_list);
     cave::systemd::status("Loading known hosts from config");
     for host in config.hosts.into_iter() {
         scheduler.introduce(host).await;

hunter/src/scheduler.rs

@@ -4,6 +4,8 @@ use std::time::Duration;
 use rand::{thread_rng, Rng};
 use tokio::time::Instant;
 
+use crate::block_list::BlockList;
+
 const MIN_INTERVAL: Duration = Duration::from_secs(30);
 const MAX_INTERVAL: Duration = Duration::from_secs(7200);
 const DEFAULT_INTERVAL: Duration = Duration::from_secs(120);
@@ -19,13 +21,15 @@ pub struct Instance {
 pub struct Scheduler {
     instances: HashMap<Host, Instance>,
     queue: BTreeMap<Instant, Host>,
+    block_list: BlockList,
 }
 
 impl Scheduler {
-    pub fn new() -> Self {
+    pub fn new(block_list: BlockList) -> Self {
         Scheduler {
             instances: HashMap::new(),
             queue: BTreeMap::new(),
+            block_list,
         }
     }
 
@@ -38,12 +42,9 @@ impl Scheduler {
     }
 
     pub async fn introduce(&mut self, host: String) -> bool {
-        // if EVIL_DOMAINS.iter().any(|domain| {
-        //     let o = host.len().saturating_sub(domain.len());
-        //     &host[o..] == *domain
-        // } ) {
-        //     return false;
-        // }
+        if self.block_list.is_blocked(&host).await {
+            return false;
+        }
 
         let now = Instant::now();
         let host = Arc::new(host);