optimize domain name matching

Cargo.lock

@@ -540,11 +540,12 @@ dependencies = [
 
 [[package]]
 name = "tokio"
-version = "0.2.2"
+version = "0.2.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 dependencies = [
  "bytes 0.5.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "futures-core 0.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
+ "iovec 0.1.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "libc 0.2.62 (registry+https://github.com/rust-lang/crates.io-index)",
  "memchr 2.2.1 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -576,7 +577,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 
 [[package]]
 name = "updns"
-version = "0.1.1"
+version = "0.1.2"
 dependencies = [
  "ace 0.0.3 (registry+https://github.com/rust-lang/crates.io-index)",
  "dirs 2.0.2 (registry+https://github.com/rust-lang/crates.io-index)",
@@ -584,7 +585,7 @@ dependencies = [
  "lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "regex 1.3.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "time 0.1.42 (registry+https://github.com/rust-lang/crates.io-index)",
- "tokio 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "tokio 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)",
 ]
 
 [[package]]
@@ -693,7 +694,7 @@ dependencies = [
 "checksum synstructure 0.10.2 (registry+https://github.com/rust-lang/crates.io-index)" = "02353edf96d6e4dc81aea2d8490a7e9db177bf8acb0e951c24940bf866cb313f"
 "checksum thread_local 0.3.6 (registry+https://github.com/rust-lang/crates.io-index)" = "c6b53e329000edc2b34dbe8545fd20e55a333362d0a321909685a19bd28c3f1b"
 "checksum time 0.1.42 (registry+https://github.com/rust-lang/crates.io-index)" = "db8dcfca086c1143c9270ac42a2bbd8a7ee477b78ac8e45b19abfb0cbede4b6f"
-"checksum tokio 0.2.2 (registry+https://github.com/rust-lang/crates.io-index)" = "2e765bf9f550bd9b8a970633ca3b56b8120c4b6c5dcbe26a93744cb02fee4b17"
+"checksum tokio 0.2.6 (registry+https://github.com/rust-lang/crates.io-index)" = "0e1bef565a52394086ecac0a6fa3b8ace4cb3a138ee1d96bd2b93283b56824e3"
 "checksum tokio-macros 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "d5795a71419535c6dcecc9b6ca95bdd3c2d6142f7e8343d7beb9923f129aa87e"
 "checksum unicode-xid 0.1.0 (registry+https://github.com/rust-lang/crates.io-index)" = "fc72304796d0818e357ead4e000d19c9c174ab23dc11093ac919054d20a6a7fc"
 "checksum unicode-xid 0.2.0 (registry+https://github.com/rust-lang/crates.io-index)" = "826e7639553986605ec5979c7dd957c7895e93eabed50ab2ffa7f6128a75097c"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "updns"
-version = "0.1.1"
+version = "0.1.2"
 edition = "2018"
 
 authors = ["wyhaya <wyhaya@gmail.com>"]
@@ -24,4 +24,4 @@ futures = "0.3.1"
 lazy_static = "1.4.0"
 regex = "1.3.1"
 time = "0.1.42"
-tokio = {version = "0.2.2", features = ["fs", "io-util",  "macros", "net", "stream", "time"]}
+tokio = {version = "0.2.6", features = ["fs", "io-util",  "macros", "net", "stream", "time"]}

README.md

@@ -8,7 +8,7 @@
 
 ---
 
-updns is a simple DNS proxy server developed using `Rust`. You can intercept any domain name and return the ip you need.
+updns is a simple DNS proxy server developed using `Rust`. You can intercept any domain name and return the ip you need
 
 ## Install
 
@@ -20,7 +20,7 @@ Or use `cargo` to install
 cargo install updns
 ```
 
-## Start to use
+## Start to use 🚀
 
 ```bash
 updns
@@ -28,9 +28,7 @@ updns
 updns -c /your/hosts
 ```
 
-You may use `sudo` to run this command because you will use the `53` port, make sure you have sufficient permissions.
+You may use `sudo` to run this command because you will use the `53` port
-
-Now change your local DNS server to `127.0.0.1` 🚀
 
 ## Running in docker
 
@@ -65,10 +63,11 @@ Option:
 
 ## Config
 
-You can use `updns config` command and then call `vim` quick edit, or use `updns path` find the updns's installation directory and edit the `config` file
+You can use `updns config` command and then call `vim` edit, or find `~/.updns/config`  edit
 
-You can specify standard domains, or utilize [regular expressions](https://rustexp.lpil.uk "rustexp") for dynamic matching, 
+You can specify standard domains, or utilize [regular expressions](https://rustexp.lpil.uk "rustexp") for dynamic matching
-You can update the config file at any time, updns will listen for file changes
+
+> Regular expression starts with `~`
 
 ```ini
 bind     0.0.0.0:53      # Binding address
@@ -78,18 +77,15 @@ timeout  2000           # Proxy timeout (ms)
 # Domain matching
 example.com              1.1.1.1
 *.example.com            2.2.2.2
-^\w+\.example\.[a-z]+$  3.3.3.3
+~^\w+\.example\.[a-z]+$  3.3.3.3
 
+# IPv6
 test.com                ::
 
 # Import from other file
 import /other/hosts
 ```
 
-## Todo
-
-* Dynamically update port bindings
-
 ## Reference
 
 [Building a DNS server in Rust](https://github.com/EmilHernvall/dnsguide)

src/config.rs

@@ -1,3 +1,4 @@
+use crate::matcher::Matcher;
 use futures::future::{BoxFuture, FutureExt};
 use regex::Regex;
 use std::{
@@ -34,7 +35,7 @@ pub enum InvalidType {
 }
 
 impl InvalidType {
-    pub fn as_str(&self) -> &str {
+    pub fn description(&self) -> &str {
         match self {
             InvalidType::SocketAddr => "Cannot parse socket address",
             InvalidType::IpAddr => "Cannot parse ip address",
@@ -47,7 +48,7 @@ impl InvalidType {
 
 #[derive(Debug)]
 pub struct Hosts {
-    record: Vec<(Host, IpAddr)>,
+    record: Vec<(Matcher, IpAddr)>,
 }
 
 impl Hosts {
@@ -55,7 +56,7 @@ impl Hosts {
         Hosts { record: Vec::new() }
     }
 
-    fn push(&mut self, record: (Host, IpAddr)) {
+    fn push(&mut self, record: (Matcher, IpAddr)) {
         self.record.push(record);
     }
 
@@ -65,7 +66,7 @@ impl Hosts {
         }
     }
 
-    pub fn iter(&mut self) -> Iter<(Host, IpAddr)> {
+    pub fn iter(&mut self) -> Iter<(Matcher, IpAddr)> {
         self.record.iter()
     }
 
@@ -79,61 +80,6 @@ impl Hosts {
     }
 }
 
-// domain match
-const TEXT: &str = "abcdefghijklmnopqrstuvwxyz0123456789-.";
-const WILDCARD: &str = "abcdefghijklmnopqrstuvwxyz0123456789-.*";
-
-#[derive(Debug)]
-pub struct Host(MatchMode);
-
-#[derive(Debug)]
-enum MatchMode {
-    Text(String),
-    Regex(Regex),
-}
-
-impl Host {
-    fn new(domain: &str) -> result::Result<Host, regex::Error> {
-        // example.com
-        if Self::is_text(domain) {
-            return Ok(Host(MatchMode::Text(domain.to_string())));
-        }
-
-        // *.example.com
-        if Self::is_wildcard(domain) {
-            let s = format!("^{}$", domain.replace(".", r"\.").replace("*", r"[^.]+"));
-            return Ok(Host(MatchMode::Regex(Regex::new(&s)?)));
-        }
-
-        // use regex
-        Ok(Host(MatchMode::Regex(Regex::new(domain)?)))
-    }
-
-    fn is_text(domain: &str) -> bool {
-        domain.chars().all(|item| TEXT.chars().any(|c| item == c))
-    }
-
-    fn is_wildcard(domain: &str) -> bool {
-        domain
-            .chars()
-            .all(|item| WILDCARD.chars().any(|c| item == c))
-    }
-
-    pub fn is_match(&self, domain: &str) -> bool {
-        match &self.0 {
-            MatchMode::Text(text) => text == domain,
-            MatchMode::Regex(reg) => reg.is_match(domain),
-        }
-    }
-
-    pub fn as_str(&self) -> &str {
-        match &self.0 {
-            MatchMode::Text(text) => text,
-            MatchMode::Regex(reg) => reg.as_str(),
-        }
-    }
-}
-
 #[derive(Debug)]
 pub struct Config {
     pub bind: Vec<SocketAddr>,
@@ -223,17 +169,17 @@ impl Parser {
     // match host
     // example.com 0.0.0.0
     // 0.0.0.0 example.com
-    fn record(left: &str, right: &str) -> result::Result<(Host, IpAddr), InvalidType> {
+    fn record(left: &str, right: &str) -> result::Result<(Matcher, IpAddr), InvalidType> {
         // ip domain
         if let Ok(ip) = right.parse() {
-            return Host::new(left)
+            return Matcher::new(left)
                 .map(|host| (host, ip))
                 .map_err(|_| InvalidType::Regex);
         }
 
         // domain ip
         if let Ok(ip) = left.parse() {
-            return Host::new(right)
+            return Matcher::new(right)
                 .map(|host| (host, ip))
                 .map_err(|_| InvalidType::Regex);
         }
@@ -307,39 +253,3 @@ impl Parser {
             .boxed()
     }
 }
-
-#[cfg(test)]
-mod test_host {
-    use super::*;
-
-    #[test]
-    fn test_create() {}
-
-    #[test]
-    fn test_text() {
-        let host = Host::new("example.com").unwrap();
-        assert!(host.is_match("example.com"));
-        assert!(!host.is_match("-example.com"));
-        assert!(!host.is_match("example.com.cn"));
-    }
-
-    #[test]
-    fn test_wildcard() {
-        let host = Host::new("*.example.com").unwrap();
-        assert!(host.is_match("test.example.com"));
-        assert!(!host.is_match("test.example.test"));
-        assert!(!host.is_match("test.test.com"));
-
-        let host = Host::new("*.example.*").unwrap();
-        assert!(host.is_match("test.example.test"));
-        assert!(!host.is_match("example.com"));
-        assert!(!host.is_match("test.test.test"));
-    }
-
-    #[test]
-    fn test_regex() {
-        let host = Host::new("^example.com$").unwrap();
-        assert!(host.is_match("example.com"));
-        assert!(!host.is_match("test.example.com"));
-    }
-}

src/main.rs

@@ -3,6 +3,7 @@ extern crate lazy_static;
 
 mod config;
 mod lib;
+mod matcher;
 mod watch;
 
 use ace::App;
@@ -136,11 +137,11 @@ async fn main() {
                 let n = config
                     .hosts
                     .iter()
-                    .map(|(r, _)| r.as_str().len())
+                    .map(|(m, _)| m.to_string().len())
                     .fold(0, |a, b| a.max(b));
 
                 for (host, ip) in config.hosts.iter() {
-                    println!("{:domain$}    {}", host.as_str(), ip, domain = n);
+                    println!("{:domain$}    {}", host.to_string(), ip, domain = n);
                 }
             }
             "config" => {
@@ -227,7 +228,7 @@ fn output_invalid(errors: &[Invalid]) {
         error!(
             "[line:{}] {} `{}`",
             invalid.line,
-            invalid.kind.as_str(),
+            invalid.kind.description(),
             invalid.source
         );
     }

src/matcher.rs

@@ -0,0 +1,184 @@
+use regex::{Error, Regex};
+use std::fmt;
+
+#[derive(Debug)]
+pub struct Matcher(MatchMode);
+
+#[derive(Debug)]
+enum MatchMode {
+    Static(String),
+    Wildcard(WildcardMatch),
+    Regex(Regex),
+}
+
+const REGEX_WORD: char = '~';
+const WILDCARD: char = '*';
+
+impl Matcher {
+    pub fn new(raw: &str) -> Result<Self, Error> {
+        // Use regex: ~^example\.com$
+        if raw.starts_with(REGEX_WORD) {
+            let reg = raw.replacen(REGEX_WORD, "", 1);
+            let mode = MatchMode::Regex(Regex::new(&reg)?);
+            return Ok(Matcher(mode));
+        }
+
+        // Use wildcard match: *.example.com
+        let find = raw.chars().any(|c| c == WILDCARD);
+        if find {
+            let mode = MatchMode::Wildcard(WildcardMatch::new(raw));
+            return Ok(Matcher(mode));
+        }
+
+        // Plain Text: example.com
+        Ok(Matcher(MatchMode::Static(raw.to_string())))
+    }
+
+    pub fn is_match(&self, domain: &str) -> bool {
+        match &self.0 {
+            MatchMode::Static(raw) => raw == domain,
+            MatchMode::Wildcard(raw) => raw.is_match(domain),
+            MatchMode::Regex(raw) => raw.is_match(domain),
+        }
+    }
+}
+
+#[derive(Debug)]
+struct WildcardMatch {
+    chars: Vec<char>,
+}
+
+impl WildcardMatch {
+    fn new(raw: &str) -> Self {
+        let mut chars = Vec::with_capacity(raw.len());
+        for c in raw.chars() {
+            chars.push(c);
+        }
+        Self { chars }
+    }
+
+    fn is_match(&self, text: &str) -> bool {
+        let mut chars = text.chars();
+        let mut dot = false;
+
+        for cur in &self.chars {
+            match cur {
+                '*' => {
+                    match chars.next() {
+                        Some(c) => {
+                            if c == '.' {
+                                return false;
+                            }
+                        }
+                        None => return false,
+                    }
+                    while let Some(n) = chars.next() {
+                        if n == '.' {
+                            dot = true;
+                            break;
+                        }
+                    }
+                }
+                word => {
+                    if dot {
+                        if word == &'.' {
+                            dot = false;
+                            continue;
+                        } else {
+                            return false;
+                        }
+                    }
+                    match chars.next() {
+                        Some(c) => {
+                            if word != &c {
+                                return false;
+                            }
+                        }
+                        None => return false,
+                    }
+                }
+            }
+        }
+        if dot {
+            return false;
+        }
+        chars.next().is_none()
+    }
+}
+
+impl fmt::Display for Matcher {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match &self.0 {
+            MatchMode::Static(raw) => write!(f, "{}", raw),
+            MatchMode::Wildcard(raw) => {
+                let mut s = String::new();
+                for ch in raw.chars.clone() {
+                    s.push(ch);
+                }
+                write!(f, "{}", s)
+            }
+            MatchMode::Regex(raw) => write!(f, "~{}", raw.as_str()),
+        }
+    }
+}
+
+#[cfg(test)]
+mod test_matcher {
+    use super::*;
+
+    #[test]
+    fn test_create() {}
+
+    #[test]
+    fn test_text() {
+        let matcher = Matcher::new("example.com").unwrap();
+        assert!(matcher.is_match("example.com"));
+        assert!(!matcher.is_match("-example.com"));
+        assert!(!matcher.is_match("example.com.cn"));
+    }
+
+    #[test]
+    fn test_wildcard() {
+        let matcher = Matcher::new("*").unwrap();
+        assert!(matcher.is_match("localhost"));
+        assert!(!matcher.is_match(".localhost"));
+        assert!(!matcher.is_match("localhost."));
+        assert!(!matcher.is_match("local.host"));
+
+        let matcher = Matcher::new("*.com").unwrap();
+        assert!(matcher.is_match("test.com"));
+        assert!(matcher.is_match("example.com"));
+        assert!(!matcher.is_match("test.test"));
+        assert!(!matcher.is_match(".test.com"));
+        assert!(!matcher.is_match("test.com."));
+        assert!(!matcher.is_match("test.test.com"));
+
+        let matcher = Matcher::new("*.*").unwrap();
+        assert!(matcher.is_match("test.test"));
+        assert!(!matcher.is_match(".test.test"));
+        assert!(!matcher.is_match("test.test."));
+        assert!(!matcher.is_match("test.test.test"));
+
+        let matcher = Matcher::new("*.example.com").unwrap();
+        assert!(matcher.is_match("test.example.com"));
+        assert!(matcher.is_match("example.example.com"));
+        assert!(!matcher.is_match("test.example.com.com"));
+        assert!(!matcher.is_match("test.test.example.com"));
+
+        let matcher = Matcher::new("*.example.*").unwrap();
+        assert!(matcher.is_match("test.example.com"));
+        assert!(matcher.is_match("example.example.com"));
+        assert!(!matcher.is_match("test.test.example.test"));
+        assert!(!matcher.is_match("test.example.test.test"));
+    }
+
+    #[test]
+    fn test_regex() {
+        let matcher = Matcher::new("~^example.com$").unwrap();
+        assert!(matcher.is_match("example.com"));
+        assert!(!matcher.is_match("test.example.com"));
+    }
+
+    #[test]
+    fn test_to_string() {}
+}