From d4d796d3b19511ac2d177b8cf6e380aea8ccdd07 Mon Sep 17 00:00:00 2001 From: wyhaya Date: Tue, 31 Dec 2019 20:07:51 +0800 Subject: [PATCH] optimize domain name matching --- Cargo.lock | 9 +-- Cargo.toml | 4 +- README.md | 32 ++++----- src/config.rs | 106 +++------------------------- src/main.rs | 7 +- src/matcher.rs | 184 +++++++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 217 insertions(+), 125 deletions(-) create mode 100644 src/matcher.rs diff --git a/Cargo.lock b/Cargo.lock index d212ccd..eb9e04c 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 141ecc9..07d1216 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "updns" -version = "0.1.1" +version = "0.1.2" edition = "2018" authors = ["wyhaya "] @@ -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"]} \ No newline at end of file +tokio = {version = "0.2.6", features = ["fs", "io-util", "macros", "net", "stream", "time"]} \ No newline at end of file diff --git a/README.md b/README.md index df51b91..2983135 100644 --- a/README.md +++ b/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. - -Now change your local DNS server to `127.0.0.1` 🚀 +You may use `sudo` to run this command because you will use the `53` port ## Running in docker @@ -65,31 +63,29 @@ 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 update the config file at any time, updns will listen for file changes +You can specify standard domains, or utilize [regular expressions](https://rustexp.lpil.uk "rustexp") for dynamic matching + +> Regular expression starts with `~` ```ini -bind 0.0.0.0:53 # Binding address -proxy 8.8.8.8:53 # Proxy address -timeout 2000 # Proxy timeout (ms) +bind 0.0.0.0:53 # Binding address +proxy 8.8.8.8:53 # Proxy address +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 +example.com 1.1.1.1 +*.example.com 2.2.2.2 +~^\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) diff --git a/src/config.rs b/src/config.rs index 2ad7827..5d35105 100644 --- a/src/config.rs +++ b/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 { - // 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, @@ -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")); - } -} diff --git a/src/main.rs b/src/main.rs index e106c34..86512cd 100644 --- a/src/main.rs +++ b/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 ); } diff --git a/src/matcher.rs b/src/matcher.rs new file mode 100644 index 0000000..68b86b5 --- /dev/null +++ b/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 { + // Use regex: ~^example\.com$ + if raw.starts_with(REGEX_WORD) { + let reg = raw.replacen(REGEX_WORD, "", 1); + let mode = MatchMode::Regex(Regex::new(®)?); + 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, +} + +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() {} +}