126 lines
3.3 KiB
Rust
126 lines
3.3 KiB
Rust
use std::{sync::Arc, collections::HashMap};
|
|
|
|
use tokio::{
|
|
io::{BufReader, AsyncBufReadExt},
|
|
sync::RwLock,
|
|
};
|
|
|
|
#[derive(Debug)]
|
|
enum Leaf {
|
|
Blocked,
|
|
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<Leaf>>,
|
|
}
|
|
|
|
impl BlockList {
|
|
pub async fn new(path: &str) -> BlockList {
|
|
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();
|
|
while let Ok(_) = file.read_line(&mut line).await {
|
|
if line == "" {
|
|
break
|
|
}
|
|
|
|
root.insert(line.trim_end());
|
|
|
|
line = String::new();
|
|
}
|
|
|
|
match &root {
|
|
Leaf::Blocked =>
|
|
panic!("the whole internet has been blocked"),
|
|
Leaf::Tree(_) => {}
|
|
}
|
|
root
|
|
}).await.unwrap();
|
|
|
|
|
|
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);
|
|
}
|