Use a Resolver with multiple nameserver ips instead of creating multiple Resolver instances.

Signed-off-by: Edu4rdSHL <edu4rdshl@protonmail.com>
This commit is contained in:
Edu4rdSHL 2021-05-16 07:40:29 -05:00
parent fb6dbe11d0
commit ab00d977bd
No known key found for this signature in database
GPG key ID: 3A574A4009F553E5
2 changed files with 67 additions and 36 deletions

View file

@ -1,9 +1,8 @@
use { use {
log::error, crate::{files, structs::Args},
rand::{seq::SliceRandom, thread_rng as rng}, std::{collections::HashSet, net::SocketAddr},
std::net::{IpAddr, Ipv4Addr},
trust_dns_resolver::{ trust_dns_resolver::{
config::{NameServerConfigGroup, ResolverConfig, ResolverOpts}, config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig, ResolverOpts},
Resolver, Resolver,
}, },
}; };
@ -19,29 +18,51 @@ pub fn get_records(resolver: &Resolver, domain: &str) -> String {
} }
} }
pub fn get_resolver(resolvers_ips: &[Ipv4Addr], opts: &ResolverOpts) -> Resolver { pub fn get_resolver(nameserver_ips: HashSet<SocketAddr>, opts: ResolverOpts) -> Resolver {
match Resolver::new( let mut name_servers = NameServerConfigGroup::with_capacity(nameserver_ips.len() * 2);
ResolverConfig::from_parts( name_servers.extend(nameserver_ips.into_iter().flat_map(|socket_addr| {
None, std::iter::once(NameServerConfig {
vec![], socket_addr,
NameServerConfigGroup::from_ips_clear( protocol: Protocol::Udp,
&[IpAddr::V4( tls_dns_name: None,
resolvers_ips trust_nx_responses: false,
.choose(&mut rng()) })
.expect("failed to read ipv4 string") .chain(std::iter::once(NameServerConfig {
.to_owned(), socket_addr,
)], protocol: Protocol::Tcp,
53, tls_dns_name: None,
false, trust_nx_responses: false,
), }))
), }));
*opts, Resolver::new(ResolverConfig::from_parts(None, vec![], name_servers), opts).unwrap()
) { }
Ok(resolver) => resolver,
Err(e) => { pub fn return_socket_address(args: &Args) -> HashSet<SocketAddr> {
error!("Failed to create the resolver. Error: {}\n", e); let mut resolver_ips = HashSet::new();
std::process::exit(1) if args.custom_resolvers {
for r in &files::return_file_targets(&args, args.resolvers.clone()) {
let server = r.to_owned() + ":53";
let socket_addr = SocketAddr::V4(match server.parse() {
Ok(a) => a,
Err(e) => unreachable!(
"Error parsing the server {}, only IPv4 are allowed. Error: {}",
r, e
),
});
resolver_ips.insert(socket_addr);
}
} else {
for r in &args.resolvers {
let server = r.to_owned() + ":53";
let socket_addr = SocketAddr::V4(match server.parse() {
Ok(a) => a,
Err(e) => unreachable!(
"Error parsing the server {}, only IPv4 are allowed. Error: {}",
r, e
),
});
resolver_ips.insert(socket_addr);
} }
} }
resolver_ips
} }

View file

@ -15,7 +15,10 @@ use {
net::Ipv4Addr, net::Ipv4Addr,
time::Duration, time::Duration,
}, },
trust_dns_resolver::config::ResolverOpts, trust_dns_resolver::{
config::{LookupIpStrategy, ResolverOpts},
Resolver,
},
}; };
lazy_static! { lazy_static! {
@ -57,7 +60,16 @@ pub fn parallel_resolver_all(args: &mut Args) -> Result<()> {
) )
} }
let data = parallel_resolver_engine(&args, args.targets.clone()); let opts = ResolverOpts {
timeout: Duration::from_secs(1),
ip_strategy: LookupIpStrategy::Ipv4Only,
num_concurrent_reqs: 1,
..Default::default()
};
let resolver = networking::get_resolver(networking::return_socket_address(args), opts);
let data = parallel_resolver_engine(&args, args.targets.clone(), resolver);
let mut table = Table::new(); let mut table = Table::new();
table.set_titles(row![ table.set_titles(row![
@ -186,19 +198,17 @@ pub fn parallel_resolver_all(args: &mut Args) -> Result<()> {
Ok(()) Ok(())
} }
fn parallel_resolver_engine(args: &Args, targets: HashSet<String>) -> HashMap<String, ResolvData> { fn parallel_resolver_engine(
let opts = ResolverOpts { args: &Args,
timeout: Duration::from_secs(2), targets: HashSet<String>,
..Default::default() resolver: Resolver,
}; ) -> HashMap<String, ResolvData> {
let resolv_data: HashMap<String, ResolvData> = targets let resolv_data: HashMap<String, ResolvData> = targets
.par_iter() .par_iter()
.map(|target| { .map(|target| {
let fqdn_target = format!("{}.", target); let fqdn_target = format!("{}.", target);
let mut resolv_data = ResolvData::default(); let mut resolv_data = ResolvData::default();
resolv_data.ip = resolv_data.ip = networking::get_records(&resolver, &fqdn_target);
networking::get_records(&networking::get_resolver(&RESOLVERS, &opts), &fqdn_target);
(target.to_owned(), resolv_data) (target.to_owned(), resolv_data)
}) })
.collect(); .collect();