(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.
This commit is contained in:
Eduard Tolosa 2025-07-27 16:19:33 -05:00
parent e06eeffb48
commit 37020360c2
8 changed files with 1130 additions and 610 deletions

1285
Cargo.lock generated

File diff suppressed because it is too large Load diff

View file

@ -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"
# trust-dns-resolver = { git = "https://github.com/Edu4rdSHL/trust-dns", branch = "main" }
clap = "2.34.0"
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" }
hickory-resolver = "0.25.2"
clap = { version = "4.5.41", features = ["derive"] }
rand = "0.9.2"
[profile.release]
lto = true
codegen-units = 1
panic = 'abort'
incremental = false
opt-level = "s"

View file

@ -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,
}

View file

@ -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(
nameserver_ips: HashSet<String>,
pub fn return_tokio_asyncresolver<S: ::std::hash::BuildHasher>(
nameserver_ips: &HashSet<String, S>,
options: ResolverOpts,
) -> AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>> {
let mut name_servers = NameServerConfigGroup::with_capacity(nameserver_ips.len() * 2);
) -> TokioResolver {
let mut name_servers = NameServerConfigGroup::with_capacity(nameserver_ips.len());
name_servers.extend(nameserver_ips.into_iter().flat_map(|server| {
let socket_addr = SocketAddr::V4(match server.parse() {
Ok(a) => a,
Err(e) => unreachable!(
name_servers.extend(nameserver_ips.iter().map(|server| {
let socket_addr = SocketAddr::V4(server.parse().unwrap_or_else(|e| {
panic!(
"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,
})
.chain(std::iter::once(NameServerConfig {
socket_addr,
protocol: Protocol::Tcp,
tls_dns_name: None,
trust_nx_responses: false,
}))
http_endpoint: None,
trust_negative_responses: false,
bind_addr: None,
}
}));
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 resolver_fut = options.resolvers.ipv4_lookup(lookup_host.clone());
let trustable_resolver_fut = options.trustable_resolver.ipv4_lookup(lookup_host);
let lookup_host = format!("{}.", host.trim_end_matches('.'));
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 {
if options.disable_double_check {
ip_lookup = Some(ip);
} else if let Ok(ip) = trustable_resolver_fut.await {
ip_lookup = Some(ip);
let ip_lookup = if options.enable_double_check {
match (
options.resolvers.ipv4_lookup(lookup_host.clone()).await,
options.trustable_resolvers.ipv4_lookup(lookup_host).await,
) {
(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 {
if options.show_ip_address && !domain_data.is_wildcard {
println!("{};{:?}", host, domain_data.ipv4_addresses);
} else if !domain_data.is_wildcard {
println!("{host}");
}
}
print_domain_data(&host, &domain_data, &options);
(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(
hosts: HashSet<String>,
resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
trustable_resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
pub async fn return_cname_data<S: ::std::hash::BuildHasher>(
hosts: HashSet<String, S>,
resolver: TokioResolver,
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 resolver_fut = resolver.lookup(lookup_host.clone(), record_type, request_options);
let trustable_resolver_fut =
trustable_resolver.lookup(lookup_host, record_type, request_options);
let host = host.trim_end_matches('.').to_owned();
let fqdn = format!("{host}.");
let mut domain_data = DomainData::default();
let mut cname_lookup = Option::<Lookup>::None;
let resolver = resolver.clone();
let trustable_resolver = trustable_resolver.clone();
async move {
if let Ok(lookup) = resolver_fut.await {
if disable_double_check {
cname_lookup = Some(lookup);
} else if let Ok(lookup) = trustable_resolver_fut.await {
cname_lookup = Some(lookup);
let cname_lookup = if disable_double_check {
resolver.lookup(fqdn.clone(), RecordType::CNAME).await.ok()
} else {
match (
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;
}
}
}

View file

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

View file

@ -1,12 +1,11 @@
use {
clap::{value_t, App, Arg},
rusolver::{dnslib, structs, utils},
clap::Parser,
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")
.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;
let args = Args::parse();
// Resolver opts
let options = ResolverOpts {
timeout: std::time::Duration::from_secs(timeout),
attempts: retries,
ip_strategy: LookupIpStrategy::Ipv4Only,
num_concurrent_reqs: 1,
shuffle_dns_servers: true,
..Default::default()
};
let options = utils::return_resolver_opts(args.timeout, args.retries);
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 {
nameserver_ips =
utils::return_file_lines(value_t!(matches.value_of("resolvers"), String).unwrap())
.await;
if args.resolvers.is_some() {
nameserver_ips = utils::return_file_lines(&args.resolvers.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 domain = value_t!(matches, "domain", String).unwrap();
wildcard_ips = utils::detect_wildcards(&domain, &trustable_resolver, quiet_flag).await;
let hosts: HashSet<String> = if args.domain.is_some() {
let domain = args.domain.unwrap();
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,
threads,
show_ip_address,
quiet_flag,
enable_double_check: args.enable_double_check,
threads: args.threads,
show_ip_address: args.ip,
quiet_flag: args.quiet_flag,
};
dnslib::return_hosts_data(&options).await;

View file

@ -1,10 +1,4 @@
use {
std::collections::HashSet,
trust_dns_resolver::{
name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime},
AsyncResolver,
},
};
use {hickory_resolver::TokioResolver, std::collections::HashSet};
#[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 trustable_resolver:
AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
pub resolvers: TokioResolver,
pub trustable_resolvers: TokioResolver,
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,

View file

@ -1,18 +1,17 @@
use {
crate::structs::{DomainData, LibOptions},
futures::stream::{self, StreamExt},
rand::{distributions::Alphanumeric, thread_rng as rng, Rng},
std::{collections::HashSet, net::Ipv4Addr},
tokio::{self, fs::File, io::AsyncReadExt},
trust_dns_resolver::{
name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime},
AsyncResolver,
hickory_resolver::{
config::{LookupIpStrategy, ResolverOpts, ServerOrderingStrategy},
TokioResolver,
},
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
// to deal with matches
pub async fn return_file_lines(file: String) -> HashSet<String> {
let mut f = match File::open(&file).await {
pub async fn return_file_lines(file: &str) -> 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 {
Ok(a) => a,
_ => unreachable!("Error reading to string."),
};
buffer.lines().map(|f| format!("{f}:53")).collect()
(f.read_to_string(&mut buffer).await)
.unwrap_or_else(|_| unreachable!("Error reading to string."));
let estimated_lines = buffer.matches('\n').count() + 1;
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!(
"{}.{}.",
rng()
.sample_iter(Alphanumeric)
.take(15)
.map(char::from)
.collect::<String>(),
target
));
let random_subdomain: String = rng()
.sample_iter(Alphanumeric)
.take(15)
.map(char::from)
.collect();
generated_wildcards.insert(format!("{random_subdomain}.{target}."));
}
generated_wilcards = stream::iter(generated_wilcards.clone().into_iter().map(
|host| async move {
if let Ok(ips) = resolvers.ipv4_lookup(host.clone()).await {
ips.into_iter()
.map(|x| x.to_string())
.collect::<Vec<String>>()
} else {
Vec::new()
}
},
))
.buffer_unordered(10)
.map(stream::iter)
.flatten()
.collect()
.await;
let wildcard_ips: HashSet<String> = stream::iter(generated_wildcards.into_iter())
.map(|host| async move {
resolvers.ipv4_lookup(host).await.map_or_else(
|_| Vec::new(),
|ips| {
ips.into_iter()
.filter_map(|ip| {
let ip_str = ip.to_string();
if ip_str.parse::<Ipv4Addr>().is_ok() {
Some(ip_str)
} else {
None
}
})
.collect::<Vec<String>>()
},
)
})
.buffer_unordered(10)
.map(stream::iter)
.flatten()
.collect()
.await;
generated_wilcards.retain(|ip| ip.parse::<Ipv4Addr>().is_ok());
if !generated_wilcards.is_empty() && !quiet_flag {
println!("Wilcards detected for {target} and wildcard's IP saved for furter work.");
println!("Wilcard IPs: {generated_wilcards:?}\n");
if !wildcard_ips.is_empty() && !quiet_flag {
println!("Wildcards detected for {target} and wildcard's IP saved for further work.");
println!("Wildcard IPs: {wildcard_ips:?}\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
}