From 4dea54773d407105178a482bea45f4e4d566c8d5 Mon Sep 17 00:00:00 2001 From: Edu4rdSHL Date: Wed, 12 May 2021 20:41:18 -0500 Subject: [PATCH] Big performance improvement. Add --retries option. Signed-off-by: Edu4rdSHL --- src/main.rs | 142 +++++++++++++++++++++++++++------------------------- 1 file changed, 75 insertions(+), 67 deletions(-) diff --git a/src/main.rs b/src/main.rs index 0616ad3..3f52120 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,10 +1,10 @@ use { clap::{value_t, App, Arg}, futures::stream::{self, StreamExt}, - rand::{distributions::Alphanumeric, prelude::SliceRandom, thread_rng as rng, Rng}, + rand::{distributions::Alphanumeric, thread_rng as rng, Rng}, std::{ collections::HashSet, - net::{IpAddr, Ipv4Addr}, + net::{Ipv4Addr, SocketAddr}, }, tokio::{ self, @@ -12,7 +12,10 @@ use { io::{self, AsyncReadExt}, }, trust_dns_resolver::{ - config::{NameServerConfigGroup, ResolverConfig, ResolverOpts}, + config::{ + LookupIpStrategy, NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig, + ResolverOpts, + }, name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime}, AsyncResolver, TokioAsyncResolver, }, @@ -31,7 +34,14 @@ async fn main() -> Result<(), Box> { .long("threads") .takes_value(true) .help("Number of threads. Default: 100"), - ).arg( + ) + .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") @@ -71,48 +81,52 @@ async fn main() -> Result<(), Box> { let show_ip_adress = 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(|_| 1); + let retries = value_t!(matches.value_of("retries"), usize).unwrap_or_else(|_| 0); let quiet_flag = matches.is_present("quiet"); // Resolver opts let options = ResolverOpts { timeout: std::time::Duration::from_secs(timeout), + attempts: retries, + ip_strategy: LookupIpStrategy::Ipv4Only, ..Default::default() }; // Create resolvers - let mut dns_ips = HashSet::new(); + let mut nameserver_ips = HashSet::new(); if matches.is_present("resolvers") { - dns_ips = return_file_lines(value_t!(matches.value_of("resolvers"), String).unwrap()).await; + nameserver_ips = + return_file_lines(value_t!(matches.value_of("resolvers"), String).unwrap()).await; } else { - let built_in_dns = vec![ + let built_in_nameservers = vec![ // Cloudflare - "1.1.1.1", - "1.0.0.1", + "1.1.1.1:53", + "1.0.0.1:53", // Google - "8.8.8.8", - "8.8.4.4", + "8.8.8.8:53", + "8.8.4.4:53", // Quad9 - "9.9.9.9", - "149.112.112.112", + "9.9.9.9:53", + "149.112.112.112:53", // OpenDNS - "208.67.222.222", - "208.67.220.220", + "208.67.222.222:53", + "208.67.220.220:53", // Verisign - "64.6.64.6", - "64.6.65.6", + "64.6.64.6:53", + "64.6.65.6:53", // UncensoredDNS - "91.239.100.100", - "89.233.43.71", + "91.239.100.100:53", + "89.233.43.71:53", // dns.watch - "84.200.69.80", - "84.200.70.40", + "84.200.69.80:53", + "84.200.70.40:53", ]; - for ip in built_in_dns { - dns_ips.insert(ip.to_string()); + for ip in built_in_nameservers { + nameserver_ips.insert(ip.to_string()); } } - let resolvers = return_tokio_dns(dns_ips, options).await; + let resolvers = return_tokio_asyncresolver(nameserver_ips, options); let mut wildcard_ips = HashSet::new(); // Read stdin @@ -132,10 +146,7 @@ async fn main() -> Result<(), Box> { }; futures::stream::iter(hosts.into_iter().map(|host| { - let resolver_fut = resolvers - .choose(&mut rng()) - .expect("failed to retrieve DNS resolver") - .ipv4_lookup(host.clone()); + let resolver_fut = resolvers.ipv4_lookup(host.clone()); let wildcard_ips = wildcard_ips.clone(); async move { if let Ok(ip) = resolver_fut.await { @@ -173,46 +184,48 @@ async fn return_file_lines(file: String) -> HashSet { Ok(a) => a, _ => unreachable!("Error reading to string."), }; - buffer.lines().map(str::to_owned).collect() + buffer.lines().map(|f| format!("{}:53", f)).collect() } -async fn return_tokio_dns( - dns_ips: HashSet, +fn return_tokio_asyncresolver( + nameserver_ips: HashSet, options: ResolverOpts, -) -> Vec>> { - let mut resolvers = Vec::new(); - for ip in dns_ips { - resolvers.push( - TokioAsyncResolver::tokio( - ResolverConfig::from_parts( - None, - vec![], - NameServerConfigGroup::from_ips_clear( - &[IpAddr::V4(match ip.parse() { - Ok(a) => a, - Err(e) => { - eprintln!( - "Error adding {} to the list of resolvers, only IPv4 addresses are allowed. Please fix the problem and try again. Error: {}", - ip, e - ); - std::process::exit(1) - } - })], - 53, - false, - ), - ), - options, - ) - .unwrap(), - ) - } - resolvers +) -> AsyncResolver> { + let mut name_servers = NameServerConfigGroup::with_capacity(nameserver_ips.len() * 2); + + name_servers.extend(nameserver_ips.into_iter().flat_map(|server| { + let socket_addr = SocketAddr::V4(match server.parse() { + Ok(a) => a, + Err(e) => unreachable!( + "Error parsing the server {}, only IPv4 are allowed. Error: {}", + server, e + ), + }); + + std::iter::once(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, + })) + })); + + TokioAsyncResolver::tokio( + ResolverConfig::from_parts(None, vec![], name_servers), + options, + ) + .unwrap() } async fn detect_wildcards( target: &str, - resolvers: &[AsyncResolver>], + resolvers: &AsyncResolver>, quiet_flag: bool, ) -> HashSet { if !quiet_flag { @@ -233,12 +246,7 @@ async fn detect_wildcards( generated_wilcards = stream::iter(generated_wilcards.clone().into_iter().map( |host| async move { - if let Ok(ips) = resolvers - .choose(&mut rng()) - .expect("failed to retrieve DNS resolver") - .ipv4_lookup(host.clone()) - .await - { + if let Ok(ips) = resolvers.ipv4_lookup(host.clone()).await { ips.into_iter() .map(|x| x.to_string()) .collect::>()