From af31c530726928a56611f05d6a79213ba3a27602 Mon Sep 17 00:00:00 2001
From: Peter Cai <peter@typeblog.net>
Date: Mon, 5 Apr 2021 16:21:52 +0800
Subject: [PATCH] override: implement suffix matching

---
 src/lib.rs      |  1 +
 src/override.rs | 30 +++++++++++++---
 src/trie_map.rs | 91 +++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 117 insertions(+), 5 deletions(-)
 create mode 100644 src/trie_map.rs

diff --git a/src/lib.rs b/src/lib.rs
index 7603693..ea062bc 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -3,6 +3,7 @@ mod client;
 mod kv;
 mod r#override;
 mod server;
+mod trie_map;
 mod util;
 
 use cfg_if::cfg_if;
diff --git a/src/override.rs b/src/override.rs
index c9ecfde..eda3be9 100644
--- a/src/override.rs
+++ b/src/override.rs
@@ -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
         }
diff --git a/src/trie_map.rs b/src/trie_map.rs
new file mode 100644
index 0000000..4b375ba
--- /dev/null
+++ b/src/trie_map.rs
@@ -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
+    }
+}