✨ Add support for wildcard domain match

justfile

@@ -25,5 +25,6 @@ apply-patch: init
 # Initialize all tools needed
 init:
   @cargo patch-crate --version || cargo install patch-crate
-  # @cargo patch-crate
+  @cargo set-version --version || cargo install cargo-edit
+  @#@cargo patch-crate
 

src/collections.rs

@@ -1,38 +1,71 @@
 use std::{collections::HashMap, fmt::Debug};
 
-use crate::libdns::proto::rr::Name;
+use crate::{config::WildcardName, libdns::proto::rr::Name};
 
 #[derive(Debug)]
-pub struct DomainMap<T: Debug>(HashMap<Name, T>);
+pub struct DomainMap<T: Debug>(HashMap<WildcardName, T>);
 
 impl<T: Debug> DomainMap<T> {
     #[inline]
     pub fn new() -> Self {
         Self::default()
     }
 
-    pub fn find(&self, domain: &Name) -> Option<&T> {
-        let mut domain = domain.to_owned();
+    pub fn find(&self, name: &Name) -> Option<&T> {
+        if self.0.is_empty() {
+            return None;
+        }
+        let mut name = if name.is_wildcard() {
+            name.base_name()
+        } else {
+            name.to_owned()
+        };
 
+        let mut lvl = 0;
         loop {
-            if let Some(v) = self.0.get(&domain) {
-                return Some(v);
+            {
+                if lvl == 0 {
+                    let wildcard_name = WildcardName::Full(name);
+                    if let Some(v) = self.0.get(&wildcard_name) {
+                        return Some(v);
+                    }
+                    name = wildcard_name.into();
+                }
+
+                if lvl == 1 {
+                    let wildcard_name = WildcardName::Sub(name);
+                    if let Some(v) = self.0.get(&wildcard_name) {
+                        return Some(v);
+                    }
+                    name = wildcard_name.into();
+                }
+
+                if lvl >= 1 {
+                    let wildcard_name = WildcardName::Suffix(name);
+                    if let Some(v) = self.0.get(&wildcard_name) {
+                        return Some(v);
+                    }
+                    name = wildcard_name.into();
+                }
+
+                let wildcard_name = WildcardName::Default(name);
+                if let Some(v) = self.0.get(&wildcard_name) {
+                    return Some(v);
+                }
+                name = wildcard_name.into();
             }
 
-            if !domain.is_fqdn() {
-                domain.set_fqdn(true);
+            if !name.is_fqdn() {
+                name.set_fqdn(true);
                 continue;
             }
 
-            if domain.is_root() {
+            if name.is_root() {
                 break;
             }
 
-            if domain.is_wildcard() {
-                domain = domain.base_name();
-            } else {
-                domain = domain.into_wildcard();
-            }
+            name = name.base_name();
+            lvl += 1;
         }
 
         None
@@ -54,14 +87,15 @@ impl<T: Debug> DomainMap<T> {
     }
 
     #[inline]
-    pub fn insert(&mut self, mut domain: Name, value: T) -> Option<T> {
-        domain.set_fqdn(true);
-        self.0.insert(domain, value)
+    pub fn insert(&mut self, name: impl Into<WildcardName>, value: T) -> Option<T> {
+        let mut name = name.into();
+        name.set_fqdn(true);
+        self.0.insert(name, value)
     }
 
     #[inline]
-    pub fn remove(&mut self, domain: &Name) -> Option<T> {
-        self.0.remove(domain)
+    pub fn remove(&mut self, name: &WildcardName) -> Option<T> {
+        self.0.remove(name)
     }
 }
 
@@ -72,9 +106,9 @@ impl<T: Debug> Default for DomainMap<T> {
     }
 }
 
-impl<T: Debug> From<HashMap<Name, T>> for DomainMap<T> {
+impl<T: Debug> From<HashMap<WildcardName, T>> for DomainMap<T> {
     #[inline]
-    fn from(value: HashMap<Name, T>) -> Self {
+    fn from(value: HashMap<WildcardName, T>) -> Self {
         Self(value)
     }
 }
@@ -103,12 +137,12 @@ impl DomainSet {
         self.0.is_empty()
     }
 
-    pub fn insert(&mut self, domain: Name) -> bool {
+    pub fn insert(&mut self, domain: impl Into<WildcardName>) -> bool {
         self.0.insert(domain, ()).is_none()
     }
 
     #[inline]
-    pub fn remove(&mut self, domain: &Name) -> bool {
+    pub fn remove(&mut self, domain: &WildcardName) -> bool {
         self.0.remove(domain).is_some()
     }
 }
@@ -365,15 +399,14 @@ mod phf {
 mod tests {
 
     use super::*;
-    use std::str::FromStr;
 
     #[test]
     fn test_set_contains() {
         let mut set = DomainSet::new();
 
-        set.insert(Name::from_str("example.com").unwrap());
-        assert!(set.contains(&Name::from_str("example.com").unwrap()));
-        assert!(set.contains(&Name::from_str("www.example.com").unwrap()));
+        set.insert(WildcardName::Default("example.com".parse().unwrap()));
+        assert!(set.contains(&"example.com".parse().unwrap()));
+        assert!(set.contains(&"www.example.com".parse().unwrap()));
     }
 
     #[cfg(feature = "experimental-trie")]

src/config/domain.rs

@@ -0,0 +1,188 @@
+use std::ops::Deref;
+
+use crate::libdns::proto::rr::Name;
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum WildcardName {
+    /// Match domain suffix, include itself.
+    ///
+    /// examples:
+    ///   /example.com/      => match: a.example.com, b.a.example.com;
+    Default(Name),
+
+    /// The + wildcard can match domain suffix.
+    ///
+    /// examples:
+    ///   /+.example.com/    => match: a.example.com, b.a.example.com; but not match: example.com
+    ///   /.example.com/
+    Suffix(Name),
+
+    /// The * wildcard can only match one-level domains.
+    /// examples:
+    ///   /*.example.com/    => match: a.example.com, b.example.com; but not match: example.com, b.a.example.com
+    Sub(Name),
+
+    /// Match full domain.
+    ///
+    /// examples:
+    /// /-.example.com/     => match: example.com
+    Full(Name),
+}
+
+impl WildcardName {
+    pub fn is_match(&self, name: &Name) -> bool {
+        match self {
+            WildcardName::Default(n) => n.zone_of(name),
+            WildcardName::Suffix(n) => n != name && n.zone_of(name),
+            WildcardName::Sub(n) => n == &name.base_name(),
+            WildcardName::Full(n) => n == name,
+        }
+    }
+}
+
+impl std::cmp::PartialOrd for WildcardName {
+    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl std::cmp::Ord for WildcardName {
+    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
+        use std::cmp::Ordering::*;
+        match self.deref().cmp(other.deref()) {
+            Less => Less,
+            Greater => Greater,
+            Equal => {
+                let ord = |w: &WildcardName| match w {
+                    WildcardName::Default(_) => 3u8,
+                    WildcardName::Suffix(_) => 2,
+                    WildcardName::Sub(_) => 1,
+                    WildcardName::Full(_) => 0,
+                };
+                let a = ord(self);
+                let b = ord(other);
+                a.cmp(&b)
+            }
+        }
+    }
+}
+
+impl std::ops::Deref for WildcardName {
+    type Target = Name;
+
+    fn deref(&self) -> &Self::Target {
+        match self {
+            WildcardName::Default(n) => n,
+            WildcardName::Suffix(n) => n,
+            WildcardName::Sub(n) => n,
+            WildcardName::Full(n) => n,
+        }
+    }
+}
+
+impl std::ops::DerefMut for WildcardName {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        match self {
+            WildcardName::Default(n) => n,
+            WildcardName::Suffix(n) => n,
+            WildcardName::Sub(n) => n,
+            WildcardName::Full(n) => n,
+        }
+    }
+}
+
+impl std::fmt::Display for WildcardName {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            WildcardName::Default(n) => write!(f, "{n}"),
+            WildcardName::Suffix(n) => write!(f, "+.{n}"),
+            WildcardName::Sub(n) => write!(f, "*.{n}"),
+            WildcardName::Full(n) => write!(f, "-.{n}"),
+        }
+    }
+}
+
+impl std::convert::From<Name> for WildcardName {
+    fn from(value: Name) -> Self {
+        if value.is_wildcard() {
+            Self::Sub(value.base_name())
+        } else {
+            Self::Default(value)
+        }
+    }
+}
+
+impl From<WildcardName> for Name {
+    fn from(val: WildcardName) -> Self {
+        match val {
+            WildcardName::Default(n) => n,
+            WildcardName::Suffix(n) => n,
+            WildcardName::Sub(n) => n,
+            WildcardName::Full(n) => n,
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub enum Domain {
+    Name(WildcardName),
+    Set(String),
+}
+
+impl std::fmt::Display for Domain {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Domain::Name(n) => write!(f, "{n}"),
+            Domain::Set(n) => write!(f, "domain-set:{n}"),
+        }
+    }
+}
+
+impl From<Domain> for String {
+    fn from(value: Domain) -> Self {
+        value.to_string()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::str::FromStr;
+
+    use super::WildcardName;
+
+    #[test]
+    fn test_wildcard_name_default() {
+        let wildcard_name = WildcardName::from_str("example.com").unwrap();
+
+        assert!(wildcard_name.is_match(&"example.com".parse().unwrap()));
+        assert!(wildcard_name.is_match(&"a.example.com".parse().unwrap()));
+        assert!(wildcard_name.is_match(&"b.a.example.com".parse().unwrap()));
+    }
+
+    #[test]
+    fn test_wildcard_name_suffix() {
+        let wildcard_name = WildcardName::from_str("+.example.com").unwrap();
+
+        assert!(!wildcard_name.is_match(&"example.com".parse().unwrap()));
+        assert!(wildcard_name.is_match(&"a.example.com".parse().unwrap()));
+        assert!(wildcard_name.is_match(&"b.a.example.com".parse().unwrap()));
+    }
+
+    #[test]
+    fn test_wildcard_name_sub() {
+        let wildcard_name = WildcardName::from_str("*.example.com").unwrap();
+
+        assert!(!wildcard_name.is_match(&"example.com".parse().unwrap()));
+        assert!(wildcard_name.is_match(&"a.example.com".parse().unwrap()));
+        assert!(!wildcard_name.is_match(&"b.a.example.com".parse().unwrap()));
+    }
+
+    #[test]
+    fn test_wildcard_name_full() {
+        let wildcard_name = WildcardName::from_str("-.example.com").unwrap();
+
+        assert!(wildcard_name.is_match(&"example.com".parse().unwrap()));
+        assert!(!wildcard_name.is_match(&"a.example.com".parse().unwrap()));
+        assert!(!wildcard_name.is_match(&"b.a.example.com".parse().unwrap()));
+    }
+}