From 8e7220f704c93e15593998f72394ee37feb6bbdb Mon Sep 17 00:00:00 2001
From: Xnoe <xnoe@xnoe.moe>
Date: Sun, 12 Apr 2026 03:26:25 +0100
Subject: [PATCH] Wildcard support

---
 Cargo.lock     |  2 +-
 Cargo.toml     |  2 +-
 src/cursor.rs  |  8 +++---
 src/ip_pool.rs |  2 +-
 src/main.rs    | 72 +++++++++++++++++++++++++++++++++++++++++++++++---
 5 files changed, 76 insertions(+), 10 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 8d1a9ba..667a87f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -254,7 +254,7 @@ dependencies = [
 ]
 
 [[package]]
-name = "rust-dnsproxy-thing"
+name = "rust-dns-selective-routing"
 version = "0.1.0"
 dependencies = [
  "anyhow",
diff --git a/Cargo.toml b/Cargo.toml
index e19d946..b56555a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,5 +1,5 @@
 [package]
-name = "rust-dnsproxy-thing"
+name = "rust-dns-selective-routing"
 version = "0.1.0"
 edition = "2024"
 
diff --git a/src/cursor.rs b/src/cursor.rs
index 1174e64..52cc2cf 100644
--- a/src/cursor.rs
+++ b/src/cursor.rs
@@ -15,7 +15,7 @@ impl<'a, T> Cursor<'a, T> {
 
     pub fn next(&mut self) -> Option<&T> {
         let next_index = self.index + 1;
-        if next_index >= self.buf.len() {
+        if next_index > self.buf.len() {
             None
         } else {
             let v = &self.buf[self.index];
@@ -25,7 +25,7 @@ impl<'a, T> Cursor<'a, T> {
     }
 
     pub fn seek(&mut self, location: usize) -> Result<(), ()> {
-        if location >= self.buf.len() {
+        if location > self.buf.len() {
             Err(())
         } else {
             self.index = location;
@@ -35,7 +35,7 @@ impl<'a, T> Cursor<'a, T> {
 
     pub fn next_slice(&mut self, amount: usize) -> Option<&'a [T]> {
         let next_index = self.index + amount;
-        if next_index >= self.buf.len() {
+        if next_index > self.buf.len() {
             None
         } else {
             let slice = &self.buf[self.index..next_index];
@@ -50,7 +50,7 @@ impl<'a, T> Cursor<'a, T> {
 
     pub fn forward(&mut self, amount: usize) -> Result<(), ()> {
         let next_index = self.index + amount;
-        if next_index >= self.buf.len() {
+        if next_index > self.buf.len() {
             Err(())
         } else {
             self.index = next_index;
diff --git a/src/ip_pool.rs b/src/ip_pool.rs
index 8a0a572..f544320 100644
--- a/src/ip_pool.rs
+++ b/src/ip_pool.rs
@@ -17,7 +17,7 @@ impl IpPool {
             return None;
         }
         let last_address_number = u32::MAX & !subnet_prefix_mask;
-        let mut pool = VecDeque::with_capacity(last_address_number as usize);
+        let mut pool = VecDeque::with_capacity((last_address_number + 1) as usize);
         for number in 0..=(last_address_number as usize) {
             pool.push_back(Ipv4Addr::from(base_addr_int + (number as u32)));
         }
diff --git a/src/main.rs b/src/main.rs
index a38f77c..49ec772 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -278,7 +278,8 @@ async fn handle_dns_response(buf: &[u8], reply_buf: &mut Vec<u8>, tcp: bool) ->
         if !FwmarkConfigMap
             .lock()
             .await
-            .contains_key(&parts_to_dns_name(&qname_parts))
+            .get(&parts_to_dns_name(&qname_parts))
+            .is_some()
         {
             eprintln!("Not forging due to non-matching qname.");
             forge = false;
@@ -401,8 +402,66 @@ struct Forwarding {
     ttl: u32,
 }
 
+enum FwmarkDomainMatchType {
+    Exact,
+    Suffix
+}
+
+struct FwmarkConfig {
+    entries: Vec<(FwmarkDomainMatchType, Vec<Vec<u8>>, u32)>
+}
+
+impl FwmarkConfig {
+    fn new() -> Self {
+        Self {
+            entries: Vec::new()
+        }
+    }
+
+    fn insert(&mut self, dns_name: Vec<u8>, fwmark: u32) {
+        self.entries.push(
+            (
+                FwmarkDomainMatchType::Exact,
+                dns_name_to_parts(&mut Cursor::from(dns_name.as_slice())).unwrap(),
+                fwmark
+            )
+        );
+    }
+
+    fn insert_wildcard(&mut self, suffix_parts: Vec<Vec<u8>>, fwmark: u32) {
+        self.entries.push(
+            (
+                FwmarkDomainMatchType::Suffix,
+                suffix_parts,
+                fwmark
+            )
+        );
+    }
+
+    fn get(&self, dns_name: &Vec<u8>) -> Option<u32> {
+        let parts = dns_name_to_parts(&mut Cursor::from(dns_name.as_slice())).unwrap();
+
+        for (match_type, _parts, fwmark) in &self.entries {
+            match match_type {
+                FwmarkDomainMatchType::Exact => {
+                    if parts == *_parts {
+                        return Some(*fwmark)
+                    }
+                },
+                FwmarkDomainMatchType::Suffix => {
+                    if parts.ends_with(_parts.as_slice()) {
+                        return Some(*fwmark)
+                    }
+                }
+            }
+        }
+
+        None
+    }
+}
+
 lazy_static! {
-    static ref FwmarkConfigMap: Mutex<HashMap<Vec<u8>, u32>> = Mutex::new(HashMap::new());
+    static ref FwmarkConfigMap: Mutex<FwmarkConfig> = Mutex::new(FwmarkConfig::new());
     static ref ForwardingMap: Mutex<HashMap<Vec<u8>, Vec<Forwarding>>> = Mutex::new(HashMap::new());
     static ref IpAllocator: Mutex<IpPool> =
         Mutex::new(IpPool::new(Ipv4Addr::new(100, 64, 0, 0), 24).unwrap());
@@ -534,8 +593,15 @@ async fn main() -> anyhow::Result<()> {
     if let Some(v) = map.get("domains") {
         let mut fwmark_config = FwmarkConfigMap.lock().await;
         for (domain, fwmark) in v.iter() {
+            let dns_name = string_to_dns_name(domain.to_string());
+            let parts = dns_name_to_parts(&mut Cursor::from(&dns_name)).unwrap();
+
             if let Ok(fwmark) = fwmark.parse::<u32>() {
-                let _ = fwmark_config.insert(string_to_dns_name(domain.to_string()), fwmark);
+                if *parts.get(0).unwrap() == vec![b'*'] {
+                    let _ = fwmark_config.insert_wildcard(parts[1..].to_vec(), fwmark);
+                } else {
+                    let _ = fwmark_config.insert(dns_name, fwmark);
+                }
             }
         }
     }