override: implement suffix matching
This commit is contained in:
parent
4adeae71cc
commit
af31c53072
|
@ -3,6 +3,7 @@ mod client;
|
|||
mod kv;
|
||||
mod r#override;
|
||||
mod server;
|
||||
mod trie_map;
|
||||
mod util;
|
||||
|
||||
use cfg_if::cfg_if;
|
||||
|
|
|
@ -1,3 +1,4 @@
|
|||
use crate::trie_map::TrieMap;
|
||||
use domain::base::{rdata::UnknownRecordData, Compose, Dname, Question, Record, Rtype};
|
||||
use domain::rdata::{Aaaa, AllRecordData, A};
|
||||
use std::collections::HashMap;
|
||||
|
@ -5,29 +6,43 @@ use std::net::IpAddr;
|
|||
|
||||
pub struct OverrideResolver {
|
||||
simple_matches: HashMap<String, IpAddr>,
|
||||
suffix_matches: TrieMap<IpAddr>,
|
||||
override_ttl: u32,
|
||||
}
|
||||
|
||||
impl OverrideResolver {
|
||||
pub fn new(overrides: HashMap<String, String>, override_ttl: u32) -> OverrideResolver {
|
||||
let (simple_matches, suffix_matches) = Self::build_match_tables(overrides);
|
||||
OverrideResolver {
|
||||
simple_matches: Self::build_simple_match_table(overrides),
|
||||
suffix_matches,
|
||||
simple_matches,
|
||||
override_ttl,
|
||||
}
|
||||
}
|
||||
|
||||
fn build_simple_match_table(overrides: HashMap<String, String>) -> HashMap<String, IpAddr> {
|
||||
let mut ret = HashMap::new();
|
||||
fn build_match_tables(
|
||||
overrides: HashMap<String, String>,
|
||||
) -> (HashMap<String, IpAddr>, TrieMap<IpAddr>) {
|
||||
let mut simple = HashMap::new();
|
||||
let mut suffix = TrieMap::new();
|
||||
for (k, v) in overrides.into_iter() {
|
||||
match v.parse::<IpAddr>() {
|
||||
Ok(addr) => {
|
||||
ret.insert(k, addr);
|
||||
if k.starts_with("*.") {
|
||||
// Anything starting with a wildcard character is a suffix match
|
||||
// we convert it to a prefix match by reversing the domain
|
||||
// Note that we get rid of the wildcard but keep the dot, i.e.
|
||||
// we don't allow suffix match in the middle of a part of a domain
|
||||
suffix.put_prefix(k[1..].chars().rev().collect::<String>(), addr);
|
||||
} else {
|
||||
simple.insert(k, addr);
|
||||
}
|
||||
}
|
||||
// Ignore malformed IP addresses
|
||||
Err(_) => continue,
|
||||
}
|
||||
}
|
||||
return ret;
|
||||
(simple, suffix)
|
||||
}
|
||||
|
||||
pub fn try_resolve(
|
||||
|
@ -44,6 +59,11 @@ impl OverrideResolver {
|
|||
let name = question.qname().to_string();
|
||||
if let Some(addr) = self.simple_matches.get(&name) {
|
||||
self.respond_with_addr(question, addr)
|
||||
} else if let Some(addr) = self
|
||||
.suffix_matches
|
||||
.get_prefix(name.chars().rev().collect::<String>())
|
||||
{
|
||||
self.respond_with_addr(question, addr)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
|
|
91
src/trie_map.rs
Normal file
91
src/trie_map.rs
Normal file
|
@ -0,0 +1,91 @@
|
|||
struct TrieMapNode<T> {
|
||||
label: u8,
|
||||
value: Option<T>,
|
||||
children: Vec<TrieMapNode<T>>,
|
||||
}
|
||||
|
||||
impl<T> TrieMapNode<T> {
|
||||
fn find_child(&self, label: u8) -> Option<usize> {
|
||||
for (idx, child) in self.children.iter().enumerate() {
|
||||
if child.label == label {
|
||||
return Some(idx);
|
||||
}
|
||||
}
|
||||
return None;
|
||||
}
|
||||
|
||||
fn traverse_trie_mut<'a, 'b>(
|
||||
&'a mut self,
|
||||
prefix: &'b [u8],
|
||||
) -> (&'a mut TrieMapNode<T>, &'b [u8]) {
|
||||
if prefix.len() == 0 {
|
||||
return (self, prefix);
|
||||
}
|
||||
|
||||
if let Some(idx) = self.find_child(prefix[0]) {
|
||||
self.children[idx].traverse_trie_mut(&prefix[1..])
|
||||
} else {
|
||||
(self, prefix)
|
||||
}
|
||||
}
|
||||
|
||||
fn traverse_trie_for_value<'a, 'b>(
|
||||
&'a self,
|
||||
prefix: &'b [u8],
|
||||
mut last_value: Option<&'a T>,
|
||||
) -> (&'a TrieMapNode<T>, Option<&'a T>, &'b [u8]) {
|
||||
if self.value.is_some() {
|
||||
last_value = self.value.as_ref();
|
||||
}
|
||||
|
||||
if prefix.len() == 0 {
|
||||
return (self, last_value, prefix);
|
||||
}
|
||||
|
||||
if let Some(idx) = self.find_child(prefix[0]) {
|
||||
self.children[idx].traverse_trie_for_value(&prefix[1..], last_value)
|
||||
} else {
|
||||
(self, last_value, prefix)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// A Map implemented with a trie, so that when a (K, V) pair is
|
||||
// inserted into the map, any key whose prefix matches K will also
|
||||
// be mapped to V.
|
||||
// The prefix match is greedy, i.e. if multiple key prefixes match
|
||||
// one key, then the mapped value is the value of the longest prefix
|
||||
pub struct TrieMap<T> {
|
||||
root: TrieMapNode<T>,
|
||||
}
|
||||
|
||||
impl<T> TrieMap<T> {
|
||||
pub fn new() -> TrieMap<T> {
|
||||
TrieMap {
|
||||
root: TrieMapNode {
|
||||
label: 0,
|
||||
value: None,
|
||||
children: Vec::new(),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub fn put_prefix(&mut self, prefix: impl AsRef<[u8]>, value: impl Into<T>) {
|
||||
let (mut node, remaining_prefix) = self.root.traverse_trie_mut(prefix.as_ref());
|
||||
for b in remaining_prefix {
|
||||
let new_node = TrieMapNode {
|
||||
label: *b,
|
||||
value: None,
|
||||
children: Vec::new(),
|
||||
};
|
||||
node.children.push(new_node);
|
||||
node = node.children.last_mut().unwrap();
|
||||
}
|
||||
node.value = Some(value.into());
|
||||
}
|
||||
|
||||
pub fn get_prefix(&self, prefix: impl AsRef<[u8]>) -> Option<&T> {
|
||||
let (_, value, _) = self.root.traverse_trie_for_value(prefix.as_ref(), None);
|
||||
value
|
||||
}
|
||||
}
|
Loading…
Reference in a new issue