(feat): refactor code

This refactor is aimed to provide better readability for future maintenance, at the same time that provides several performance improvements reducing the number of calls that we do to resolvers.

Cargo.toml

@@ -1,8 +1,8 @@
 [package]
 name = "rusolver"
-version = "0.9.1"
+version = "0.10.0"
 authors = ["Eduard Tolosa <edu4rdshl@protonmail.com>"]
-edition = "2018"
+edition = "2021"
 description = "Fast DNS resolver written in Rust."
 repository = "https://github.com/edu4rdshl/rusolver"
 license = "GPL-3.0-or-later"
@@ -12,20 +12,13 @@ readme = "README.md"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-tokio = { version = "1.44.2", features = ["full"] }
+tokio = { version = "1.47.0", features = ["full"] }
 futures = "0.3.31"
-trust-dns-resolver = "0.20.4"
+hickory-resolver = "0.25.2"
-# trust-dns-resolver = { git = "https://github.com/Edu4rdSHL/trust-dns", branch = "main" }
+clap = { version = "4.5.41", features = ["derive"] }
-clap = "2.34.0"
+rand = "0.9.2"
-rand = "0.8.5"
-
-# https://github.com/bluejekyll/trust-dns/pull/1632
-[patch.crates-io]
-trust-dns-resolver = { git = "https://github.com/Findomain/trust-dns", package = "trust-dns-resolver", branch = "custombranch" }
 
 [profile.release]
 lto = true
 codegen-units = 1
 panic = 'abort'
-incremental = false
-opt-level = "s"

src/args.rs

@@ -0,0 +1,46 @@
+use clap::{arg, Parser};
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+pub struct Args {
+    #[arg(
+        short,
+        long,
+        default_value_t = 100,
+        help = "Number of threads. Default: 100"
+    )]
+    pub threads: usize,
+
+    #[arg(
+        long,
+        default_value_t = 0,
+        help = "Number of retries after lookup failure before giving up. Defaults to 0"
+    )]
+    pub retries: usize,
+
+    #[arg(
+        short,
+        long,
+        help = "Target domain. When it's specified, a wordlist can be used from stdin for bruteforcing."
+    )]
+    pub domain: Option<String>,
+
+    #[arg(short, long, help = "File with DNS ips.")]
+    pub resolvers: Option<String>,
+
+    #[arg(long, default_value_t = 3, help = "Timeout in seconds. Default: 3")]
+    pub timeout: u64,
+
+    #[arg(short, long, help = "Display the record data.")]
+    pub ip: bool,
+
+    #[arg(
+        short,
+        long,
+        help = "Enable the double verification algorithm for subdomains. Default: false"
+    )]
+    pub enable_double_check: bool,
+
+    #[arg(short, long, help = "Quiet mode, no output except errors.")]
+    pub quiet_flag: bool,
+}

src/dnslib.rs

@@ -1,54 +1,52 @@
 use {
-    crate::structs::{DomainData, LibOptions},
+    crate::{
+        structs::{DomainData, LibOptions},
+        utils::print_domain_data,
+    },
     futures::stream::{self, StreamExt},
+    hickory_resolver::{
+        config::{NameServerConfig, NameServerConfigGroup, ResolverConfig, ResolverOpts},
+        name_server::TokioConnectionProvider,
+        proto::{rr::RecordType, xfer::Protocol},
+        TokioResolver,
+    },
     std::{
         collections::{HashMap, HashSet},
         net::SocketAddr,
     },
-    trust_dns_resolver::{
-        config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig, ResolverOpts},
-        lookup::{Ipv4Lookup, Lookup},
-        name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime},
-        proto::{rr::RecordType, xfer::DnsRequestOptions},
-        AsyncResolver, TokioAsyncResolver,
-    },
 };
 
 #[must_use]
-pub fn return_tokio_asyncresolver(
+pub fn return_tokio_asyncresolver<S: ::std::hash::BuildHasher>(
-    nameserver_ips: HashSet<String>,
+    nameserver_ips: &HashSet<String, S>,
     options: ResolverOpts,
-) -> AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>> {
+) -> TokioResolver {
-    let mut name_servers = NameServerConfigGroup::with_capacity(nameserver_ips.len() * 2);
+    let mut name_servers = NameServerConfigGroup::with_capacity(nameserver_ips.len());
 
-    name_servers.extend(nameserver_ips.into_iter().flat_map(|server| {
+    name_servers.extend(nameserver_ips.iter().map(|server| {
-        let socket_addr = SocketAddr::V4(match server.parse() {
+        let socket_addr = SocketAddr::V4(server.parse().unwrap_or_else(|e| {
-            Ok(a) => a,
+            panic!(
-            Err(e) => unreachable!(
                 "Error parsing the server {}, only IPv4 are allowed. Error: {}",
                 server, e
-            ),
+            )
-        });
+        }));
 
-        std::iter::once(NameServerConfig {
+        NameServerConfig {
             socket_addr,
             protocol: Protocol::Udp,
             tls_dns_name: None,
-            trust_nx_responses: false,
+            http_endpoint: None,
-        })
+            trust_negative_responses: false,
-        .chain(std::iter::once(NameServerConfig {
+            bind_addr: None,
-            socket_addr,
+        }
-            protocol: Protocol::Tcp,
-            tls_dns_name: None,
-            trust_nx_responses: false,
-        }))
     }));
 
-    TokioAsyncResolver::tokio(
+    TokioResolver::builder_with_config(
         ResolverConfig::from_parts(None, vec![], name_servers),
-        options,
+        TokioConnectionProvider::default(),
     )
-    .unwrap()
+    .with_options(options)
+    .build()
 }
 
 pub async fn return_hosts_data(options: &LibOptions) -> HashMap<String, DomainData> {
@@ -59,24 +57,22 @@ pub async fn return_hosts_data(options: &LibOptions) -> HashMap<String, DomainDa
     };
 
     stream::iter(options.hosts.clone().into_iter().map(|host| {
-        let lookup_host = host.trim_end_matches('.').to_owned() + ".";
+        let lookup_host = format!("{}.", host.trim_end_matches('.'));
-
-        let resolver_fut = options.resolvers.ipv4_lookup(lookup_host.clone());
-        let trustable_resolver_fut = options.trustable_resolver.ipv4_lookup(lookup_host);
         let wildcard_ips = options.wildcard_ips.clone();
-
         let mut domain_data = DomainData::default();
 
-        let mut ip_lookup = Option::<Ipv4Lookup>::None;
-
         async move {
-            if let Ok(ip) = resolver_fut.await {
+            let ip_lookup = if options.enable_double_check {
-                if options.disable_double_check {
+                match (
-                    ip_lookup = Some(ip);
+                    options.resolvers.ipv4_lookup(lookup_host.clone()).await,
-                } else if let Ok(ip) = trustable_resolver_fut.await {
+                    options.trustable_resolvers.ipv4_lookup(lookup_host).await,
-                    ip_lookup = Some(ip);
+                ) {
+                    (Ok(_), Ok(ip)) => Some(ip),
+                    _ => None,
                 }
-            }
+            } else {
+                options.resolvers.ipv4_lookup(lookup_host).await.ok()
+            };
 
             if let Some(ip_lookup) = ip_lookup {
                 for ip in ip_lookup.iter() {
@@ -89,13 +85,7 @@ pub async fn return_hosts_data(options: &LibOptions) -> HashMap<String, DomainDa
                 .iter()
                 .all(|ip| wildcard_ips.contains(ip));
 
-            if !options.quiet_flag {
+            print_domain_data(&host, &domain_data, &options);
-                if options.show_ip_address && !domain_data.is_wildcard {
-                    println!("{};{:?}", host, domain_data.ipv4_addresses);
-                } else if !domain_data.is_wildcard {
-                    println!("{host}");
-                }
-            }
 
             (host, domain_data)
         }
@@ -106,10 +96,10 @@ pub async fn return_hosts_data(options: &LibOptions) -> HashMap<String, DomainDa
 }
 
 // Used internally for now
-pub async fn return_cname_data(
+pub async fn return_cname_data<S: ::std::hash::BuildHasher>(
-    hosts: HashSet<String>,
+    hosts: HashSet<String, S>,
-    resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
+    resolver: TokioResolver,
-    trustable_resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
+    trustable_resolver: TokioResolver,
     disable_double_check: bool,
     mut threads: usize,
 ) -> HashMap<String, DomainData> {
@@ -117,33 +107,34 @@ pub async fn return_cname_data(
         threads = hosts.len();
     }
 
-    let request_options = DnsRequestOptions::default();
-    let record_type = RecordType::CNAME;
-
     stream::iter(hosts)
         .map(|host| {
-            let lookup_host = host.trim_end_matches('.').to_owned() + ".";
+            let host = host.trim_end_matches('.').to_owned();
-            let resolver_fut = resolver.lookup(lookup_host.clone(), record_type, request_options);
+            let fqdn = format!("{host}.");
-            let trustable_resolver_fut =
-                trustable_resolver.lookup(lookup_host, record_type, request_options);
 
-            let mut domain_data = DomainData::default();
+            let resolver = resolver.clone();
-
+            let trustable_resolver = trustable_resolver.clone();
-            let mut cname_lookup = Option::<Lookup>::None;
 
             async move {
-                if let Ok(lookup) = resolver_fut.await {
+                let cname_lookup = if disable_double_check {
-                    if disable_double_check {
+                    resolver.lookup(fqdn.clone(), RecordType::CNAME).await.ok()
-                        cname_lookup = Some(lookup);
+                } else {
-                    } else if let Ok(lookup) = trustable_resolver_fut.await {
+                    match (
-                        cname_lookup = Some(lookup);
+                        resolver.lookup(fqdn.clone(), RecordType::CNAME).await,
+                        trustable_resolver.lookup(fqdn, RecordType::CNAME).await,
+                    ) {
+                        (Ok(_), Ok(lookup)) => Some(lookup),
+                        _ => None,
                     }
-                }
+                };
+
+                let mut domain_data = DomainData::default();
 
                 if let Some(lookup) = cname_lookup {
                     for record in lookup.iter() {
                         if let Some(cname) = record.as_cname() {
                             domain_data.cname = cname.to_string();
+                            break;
                         }
                     }
                 }

src/lib.rs

@@ -1,3 +1,4 @@
+pub mod args;
 pub mod dnslib;
 pub mod structs;
 pub mod utils;

src/main.rs

@@ -1,12 +1,11 @@
 use {
-    clap::{value_t, App, Arg},
+    clap::Parser,
-    rusolver::{dnslib, structs, utils},
+    rusolver::{args::Args, dnslib, structs, utils},
     std::collections::HashSet,
     tokio::{
         self,
         io::{self, AsyncReadExt},
     },
-    trust_dns_resolver::config::{LookupIpStrategy, ResolverOpts},
 };
 
 // WIP: add support for AAAA, TXT, SRV, NAPTR, PTR, CNAME, DNAME, MX, NS, SOA, LOC, SVCB, HTTPS, SPF, CAA and AVC resource records.
@@ -16,83 +15,10 @@ use {
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Eval args
-    let matches = App::new("Rusolver")
+    let args = Args::parse();
-        .version(clap::crate_version!())
-        .author("Eduard Tolosa <edu4rdshl@protonmail.com>")
-        .about("Fast DNS resolver written in Rust.")
-        .arg(
-            Arg::with_name("threads")
-                .short("t")
-                .long("threads")
-                .takes_value(true)
-                .help("Number of threads. Default: 100"),
-        )
-        .arg(
-            Arg::with_name("retries")
-                .long("retries")
-                .takes_value(true)
-                .help("Number of retries after lookup failure before giving up. Defaults to 0"),
-        )
-        .arg(
-            Arg::with_name("domain")
-                .short("d")
-                .long("domain")
-                .takes_value(true)
-                .help("Target domain. When it's specified, a wordlist can be used from stdin for bruteforcing."),
-        )
-        .arg(
-            Arg::with_name("resolvers")
-                .short("r")
-                .long("resolvers")
-                .takes_value(true)
-                .help("File with DNS ips."),
-        )
-        .arg(
-            Arg::with_name("timeout")
-                .long("timeout")
-                .takes_value(true)
-                .help("Timeout in seconds. Default: 3"),
-        )
-        .arg(
-            Arg::with_name("ip")
-                .short("i")
-                .long("ip")
-                .takes_value(false)
-                .help("Show the discovered IP addresses. Default: false"),
-        )
-        .arg(
-            Arg::with_name("no-verify")
-                .long("no-verify")
-                .takes_value(false)
-                .help("Disables the double verification algorithm for valid subdomains -NOT RECOMMENDED-. Default: false"),
-        )
-        .arg(
-            Arg::with_name("quiet")
-                .short("q")
-                .long("quiet")
-                .takes_value(false)
-                .help("Enables quiet mode."),
-        )
-        .get_matches();
-
-    // Assign values or use defaults
-    let show_ip_address = matches.is_present("ip");
-    let threads = value_t!(matches.value_of("threads"), usize).unwrap_or_else(|_| 100);
-    let timeout = value_t!(matches.value_of("timeout"), u64).unwrap_or_else(|_| 3);
-    let retries = value_t!(matches.value_of("retries"), usize).unwrap_or_else(|_| 0);
-    let quiet_flag = matches.is_present("quiet");
-    let custom_resolvers = matches.is_present("resolvers");
-    let disable_double_check = matches.is_present("no-verify") && custom_resolvers;
 
     // Resolver opts
-    let options = ResolverOpts {
+    let options = utils::return_resolver_opts(args.timeout, args.retries);
-        timeout: std::time::Duration::from_secs(timeout),
-        attempts: retries,
-        ip_strategy: LookupIpStrategy::Ipv4Only,
-        num_concurrent_reqs: 1,
-        shuffle_dns_servers: true,
-        ..Default::default()
-    };
 
     let built_in_nameservers: HashSet<String> = vec![
         // Cloudflare
@@ -110,12 +36,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         // Verisign
         "64.6.64.6:53",
         "64.6.65.6:53",
-        // UncensoredDNS
-        "91.239.100.100:53",
-        "89.233.43.71:53",
-        // dns.watch
-        "84.200.69.80:53",
-        "84.200.70.40:53",
     ]
     .iter()
     .map(ToString::to_string)
@@ -124,16 +44,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // Create resolvers
     let mut nameserver_ips;
 
-    if custom_resolvers {
+    if args.resolvers.is_some() {
-        nameserver_ips =
+        nameserver_ips = utils::return_file_lines(&args.resolvers.unwrap()).await;
-            utils::return_file_lines(value_t!(matches.value_of("resolvers"), String).unwrap())
-                .await;
         nameserver_ips.retain(|ip| !ip.is_empty());
     } else {
         nameserver_ips = built_in_nameservers.clone();
     }
-    let resolvers = dnslib::return_tokio_asyncresolver(nameserver_ips, options);
+
-    let trustable_resolver = dnslib::return_tokio_asyncresolver(built_in_nameservers, options);
+    let resolvers = dnslib::return_tokio_asyncresolver(&nameserver_ips, options.clone());
+    let trustable_resolvers = dnslib::return_tokio_asyncresolver(&built_in_nameservers, options);
     let mut wildcard_ips = HashSet::new();
 
     // Read stdin
@@ -141,9 +60,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut stdin = io::stdin();
     stdin.read_to_string(&mut buffer).await?;
 
-    let hosts: HashSet<String> = if matches.is_present("domain") {
+    let hosts: HashSet<String> = if args.domain.is_some() {
-        let domain = value_t!(matches, "domain", String).unwrap();
+        let domain = args.domain.unwrap();
-        wildcard_ips = utils::detect_wildcards(&domain, &trustable_resolver, quiet_flag).await;
+        wildcard_ips =
+            utils::detect_wildcards(&domain, &trustable_resolvers, args.quiet_flag).await;
         buffer
             .lines()
             .map(|word| format!("{word}.{domain}"))
@@ -155,12 +75,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let options = structs::LibOptions {
         hosts,
         resolvers,
-        trustable_resolver,
+        trustable_resolvers,
         wildcard_ips,
-        disable_double_check,
+        enable_double_check: args.enable_double_check,
-        threads,
+        threads: args.threads,
-        show_ip_address,
+        show_ip_address: args.ip,
-        quiet_flag,
+        quiet_flag: args.quiet_flag,
     };
 
     dnslib::return_hosts_data(&options).await;

src/structs.rs

@@ -1,10 +1,4 @@
-use {
+use {hickory_resolver::TokioResolver, std::collections::HashSet};
-    std::collections::HashSet,
-    trust_dns_resolver::{
-        name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime},
-        AsyncResolver,
-    },
-};
 
 #[derive(Clone, Debug, Default)]
 pub struct DomainData {
@@ -17,11 +11,10 @@ pub struct DomainData {
 #[derive(Clone, Debug)]
 pub struct LibOptions {
     pub hosts: HashSet<String>,
-    pub resolvers: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
+    pub resolvers: TokioResolver,
-    pub trustable_resolver:
+    pub trustable_resolvers: TokioResolver,
-        AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
     pub wildcard_ips: HashSet<String>,
-    pub disable_double_check: bool,
+    pub enable_double_check: bool,
     pub threads: usize,
     pub show_ip_address: bool,
     pub quiet_flag: bool,

src/utils.rs

@@ -1,18 +1,17 @@
 use {
+    crate::structs::{DomainData, LibOptions},
     futures::stream::{self, StreamExt},
-    rand::{distributions::Alphanumeric, thread_rng as rng, Rng},
+    hickory_resolver::{
-    std::{collections::HashSet, net::Ipv4Addr},
+        config::{LookupIpStrategy, ResolverOpts, ServerOrderingStrategy},
-    tokio::{self, fs::File, io::AsyncReadExt},
+        TokioResolver,
-    trust_dns_resolver::{
-        name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime},
-        AsyncResolver,
     },
+    rand::{distr::Alphanumeric, rng, Rng},
+    std::{collections::HashSet, net::Ipv4Addr},
+    tokio::{fs::File, io::AsyncReadExt},
 };
 
-// In the future I may need to implement error propagation, but for now it's fine
+pub async fn return_file_lines(file: &str) -> HashSet<String> {
-// to deal with matches
+    let mut f = match File::open(file).await {
-pub async fn return_file_lines(file: String) -> HashSet<String> {
-    let mut f = match File::open(&file).await {
         Ok(file) => file,
         Err(e) => {
             eprintln!("Error opening resolvers file. Error: {e}");
@@ -21,58 +20,88 @@ pub async fn return_file_lines(file: String) -> HashSet<String> {
     };
     let mut buffer = String::new();
 
-    match f.read_to_string(&mut buffer).await {
+    (f.read_to_string(&mut buffer).await)
-        Ok(a) => a,
+        .unwrap_or_else(|_| unreachable!("Error reading to string."));
-        _ => unreachable!("Error reading to string."),
+
-    };
+    let estimated_lines = buffer.matches('\n').count() + 1;
-    buffer.lines().map(|f| format!("{f}:53")).collect()
+    let mut result = HashSet::with_capacity(estimated_lines);
+
+    for line in buffer.lines() {
+        if !line.is_empty() {
+            result.insert(format!("{line}:53"));
+        }
+    }
+    result
 }
 
 pub async fn detect_wildcards(
     target: &str,
-    resolvers: &AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
+    resolvers: &TokioResolver,
     quiet_flag: bool,
 ) -> HashSet<String> {
     if !quiet_flag {
         println!("Running wildcards detection for {target}...\n");
     }
-    let mut generated_wilcards: HashSet<String> = HashSet::new();
+
+    let mut generated_wildcards = HashSet::with_capacity(19);
+
+    // Generate random subdomains for wildcard detection
     for _ in 1..20 {
-        generated_wilcards.insert(format!(
+        let random_subdomain: String = rng()
-            "{}.{}.",
+            .sample_iter(Alphanumeric)
-            rng()
+            .take(15)
-                .sample_iter(Alphanumeric)
+            .map(char::from)
-                .take(15)
+            .collect();
-                .map(char::from)
+        generated_wildcards.insert(format!("{random_subdomain}.{target}."));
-                .collect::<String>(),
-            target
-        ));
     }
 
-    generated_wilcards = stream::iter(generated_wilcards.clone().into_iter().map(
+    let wildcard_ips: HashSet<String> = stream::iter(generated_wildcards.into_iter())
-        |host| async move {
+        .map(|host| async move {
-            if let Ok(ips) = resolvers.ipv4_lookup(host.clone()).await {
+            resolvers.ipv4_lookup(host).await.map_or_else(
-                ips.into_iter()
+                |_| Vec::new(),
-                    .map(|x| x.to_string())
+                |ips| {
-                    .collect::<Vec<String>>()
+                    ips.into_iter()
-            } else {
+                        .filter_map(|ip| {
-                Vec::new()
+                            let ip_str = ip.to_string();
-            }
+                            if ip_str.parse::<Ipv4Addr>().is_ok() {
-        },
+                                Some(ip_str)
-    ))
+                            } else {
-    .buffer_unordered(10)
+                                None
-    .map(stream::iter)
+                            }
-    .flatten()
+                        })
-    .collect()
+                        .collect::<Vec<String>>()
-    .await;
+                },
+            )
+        })
+        .buffer_unordered(10)
+        .map(stream::iter)
+        .flatten()
+        .collect()
+        .await;
 
-    generated_wilcards.retain(|ip| ip.parse::<Ipv4Addr>().is_ok());
+    if !wildcard_ips.is_empty() && !quiet_flag {
-
+        println!("Wildcards detected for {target} and wildcard's IP saved for further work.");
-    if !generated_wilcards.is_empty() && !quiet_flag {
+        println!("Wildcard IPs: {wildcard_ips:?}\n");
-        println!("Wilcards detected for {target} and wildcard's IP saved for furter work.");
-        println!("Wilcard IPs: {generated_wilcards:?}\n");
     } else if !quiet_flag {
-        println!("No wilcards detected for {target}, nice!\n");
+        println!("No wildcards detected for {target}, nice!\n");
     }
-    generated_wilcards
+    wildcard_ips
+}
+
+pub fn print_domain_data(host: &str, domain_data: &DomainData, options: &LibOptions) {
+    if options.show_ip_address {
+        println!("{}: {:?}", host, domain_data.ipv4_addresses);
+    } else {
+        println!("{}", host);
+    }
+}
+
+pub fn return_resolver_opts(timeout: u64, retries: usize) -> ResolverOpts {
+    let mut options = ResolverOpts::default();
+    options.timeout = std::time::Duration::from_secs(timeout);
+    options.attempts = retries;
+    options.ip_strategy = LookupIpStrategy::Ipv4Only;
+    options.num_concurrent_reqs = 1;
+    options.server_ordering_strategy = ServerOrderingStrategy::RoundRobin;
+    options
 }