From 8c28f2e0004f66db5a172fb786d3bd82ef602d68 Mon Sep 17 00:00:00 2001
From: "B. Blechschmidt" <git@blechschmidt.io>
Date: Wed, 28 Aug 2024 23:06:37 +0200
Subject: [PATCH] Implement --virtual-dns-pool

---
 src/args.rs        |  6 ++++++
 src/lib.rs         |  2 +-
 src/virtual_dns.rs | 51 ++++++----------------------------------------
 3 files changed, 13 insertions(+), 46 deletions(-)

diff --git a/src/args.rs b/src/args.rs
index 7ecf8ed..ef6e636 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -6,6 +6,7 @@ use tproxy_config::IpCidr;
 use std::ffi::OsString;
 
 use std::net::{IpAddr, SocketAddr, ToSocketAddrs};
+use std::str::FromStr;
 
 #[derive(Debug, Clone, clap::Parser)]
 #[command(author, version, about = "Tunnel interface to proxy.", long_about = None)]
@@ -76,6 +77,10 @@ pub struct Args {
     #[arg(long, value_name = "IP", default_value = "8.8.8.8")]
     pub dns_addr: IpAddr,
 
+    /// IP address pool to be used by virtual DNS in CIDR notation.
+    #[arg(long, value_name = "CIDR", default_value = "198.18.0.0/15")]
+    pub virtual_dns_pool: IpCidr,
+
     /// IPs used in routing setup which should bypass the tunnel,
     /// in the form of IP or IP/CIDR. Multiple IPs can be specified,
     /// e.g. --bypass 3.4.5.0/24 --bypass 5.6.7.8
@@ -132,6 +137,7 @@ impl Default for Args {
             tcp_timeout: 600,
             udp_timeout: 10,
             verbosity: ArgVerbosity::Info,
+            virtual_dns_pool: IpCidr::from_str("198.18.0.0/15").unwrap(),
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index 3df7040..033c542 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -161,7 +161,7 @@ where
     let dns_addr = args.dns_addr;
     let ipv6_enabled = args.ipv6_enabled;
     let virtual_dns = if args.dns == ArgDns::Virtual {
-        Some(Arc::new(Mutex::new(VirtualDns::new())))
+        Some(Arc::new(Mutex::new(VirtualDns::new(args.virtual_dns_pool))))
     } else {
         None
     };
diff --git a/src/virtual_dns.rs b/src/virtual_dns.rs
index 966ae5a..8dae2a7 100644
--- a/src/virtual_dns.rs
+++ b/src/virtual_dns.rs
@@ -4,9 +4,9 @@ use std::{
     collections::HashMap,
     convert::TryInto,
     net::{IpAddr, Ipv4Addr, Ipv6Addr},
-    str::FromStr,
     time::{Duration, Instant},
 };
+use tproxy_config::IpCidr;
 
 const MAPPING_TIMEOUT: u64 = 60; // Mapping timeout in seconds
 
@@ -27,29 +27,17 @@ pub struct VirtualDns {
     next_addr: IpAddr,
 }
 
-impl Default for VirtualDns {
-    fn default() -> Self {
-        let start_addr = Ipv4Addr::from_str("198.18.0.0").unwrap();
-        let prefix_len = 15;
-
-        let network_addr = calculate_network_addr(start_addr, prefix_len);
-        let broadcast_addr = calculate_broadcast_addr(start_addr, prefix_len);
-
+impl VirtualDns {
+    pub fn new(ip_pool: IpCidr) -> Self {
         Self {
             trailing_dot: false,
-            next_addr: start_addr.into(),
+            next_addr: ip_pool.first_address(),
             name_to_ip: HashMap::default(),
-            network_addr: IpAddr::from(network_addr),
-            broadcast_addr: IpAddr::from(broadcast_addr),
+            network_addr: ip_pool.first_address(),
+            broadcast_addr: ip_pool.last_address(),
             lru_cache: LruCache::new_unbounded(),
         }
     }
-}
-
-impl VirtualDns {
-    pub fn new() -> Self {
-        VirtualDns::default()
-    }
 
     /// Returns the DNS response to send back to the client.
     pub fn generate_query(&mut self, data: &[u8]) -> Result<(Vec<u8>, String, IpAddr)> {
@@ -160,30 +148,3 @@ impl VirtualDns {
         }
     }
 }
-
-fn calculate_network_addr(ip: std::net::Ipv4Addr, prefix_len: u8) -> std::net::Ipv4Addr {
-    let mask = (!0u32) << (32 - prefix_len);
-    let ip_u32 = u32::from_be_bytes(ip.octets());
-    std::net::Ipv4Addr::from((ip_u32 & mask).to_be_bytes())
-}
-
-fn calculate_broadcast_addr(ip: std::net::Ipv4Addr, prefix_len: u8) -> std::net::Ipv4Addr {
-    let mask = (!0u32) >> prefix_len;
-    let ip_u32 = u32::from_be_bytes(ip.octets());
-    std::net::Ipv4Addr::from((ip_u32 | mask).to_be_bytes())
-}
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn test_cidr_addr() {
-        let start_addr = Ipv4Addr::from_str("198.18.0.0").unwrap();
-        let prefix_len = 15;
-        let network_addr = calculate_network_addr(start_addr, prefix_len);
-        let broadcast_addr = calculate_broadcast_addr(start_addr, prefix_len);
-        assert_eq!(network_addr, Ipv4Addr::from_str("198.18.0.0").unwrap());
-        assert_eq!(broadcast_addr, Ipv4Addr::from_str("198.19.255.255").unwrap());
-    }
-}