1 use crate::errors::*; 2 3 use rustc_hash::FxHashMap; 4 use std::fs::File; 5 use std::io::{BufRead, BufReader}; 6 use std::path::Path; 7 use std::sync::Arc; 8 9 const MAX_ITERATIONS: usize = 5; 10 11 #[derive(Debug)] 12 struct BlackListInner { 13 map: FxHashMap<Vec<u8>, ()>, 14 } 15 16 #[derive(Clone, Debug)] 17 pub struct BlackList { 18 inner: Arc<BlackListInner>, 19 max_iterations: usize, 20 } 21 22 impl BlackList { new(map: FxHashMap<Vec<u8>, ()>, max_iterations: usize) -> Self23 pub fn new(map: FxHashMap<Vec<u8>, ()>, max_iterations: usize) -> Self { 24 let inner = Arc::new(BlackListInner { map }); 25 BlackList { 26 inner, 27 max_iterations, 28 } 29 } 30 load(path: impl AsRef<Path>) -> Result<Self, Error>31 pub fn load(path: impl AsRef<Path>) -> Result<Self, Error> { 32 let mut map = FxHashMap::default(); 33 let fp = BufReader::new(File::open(path)?); 34 for (line_nb, line) in fp.lines().enumerate() { 35 let line = line?; 36 let mut line = line.trim(); 37 if line.is_empty() || line.starts_with('#') { 38 continue; 39 } 40 while line.starts_with("*.") { 41 line = &line[2..]; 42 } 43 while line.ends_with('.') { 44 line = &line[..line.len() - 1]; 45 } 46 let qname = line.as_bytes().to_ascii_lowercase(); 47 if qname.is_empty() { 48 bail!("Unexpected blacklist rule at line {}", line_nb) 49 } 50 map.insert(qname, ()); 51 } 52 Ok(BlackList::new(map, MAX_ITERATIONS)) 53 } 54 find(&self, qname: &[u8]) -> bool55 pub fn find(&self, qname: &[u8]) -> bool { 56 let qname = qname.to_ascii_lowercase(); 57 let mut qname = qname.as_slice(); 58 let map = &self.inner.map; 59 let mut iterations = self.max_iterations; 60 while qname.len() >= 4 && iterations > 0 { 61 if map.contains_key(qname) { 62 return true; 63 } 64 let mut it = qname.splitn(2, |x| *x == b'.'); 65 if it.next().is_none() { 66 break; 67 } 68 qname = match it.next() { 69 None => break, 70 Some(qname) => qname, 71 }; 72 iterations -= 1; 73 } 74 false 75 } 76 } 77