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