From 17776e7b88ca9c8a31c5f3b7e50a962b8e5b6f70 Mon Sep 17 00:00:00 2001 From: Edu4rdSHL Date: Sun, 13 Feb 2022 05:32:25 -0500 Subject: [PATCH] Split functions across modules. Improve performance. Signed-off-by: Edu4rdSHL --- Cargo.lock | 44 +++++++++---------- Cargo.toml | 6 +-- src/dnslib.rs | 110 ++++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 5 +-- src/main.rs | 56 ++++++++---------------- src/resolver.rs | 44 ------------------- src/structs.rs | 20 +++++++++ 7 files changed, 174 insertions(+), 111 deletions(-) create mode 100644 src/dnslib.rs delete mode 100644 src/resolver.rs diff --git a/Cargo.lock b/Cargo.lock index a216954..201f80e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -96,9 +96,9 @@ dependencies = [ [[package]] name = "futures" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28560757fe2bb34e79f907794bb6b22ae8b0e5c669b638a1132f2592b19035b4" +checksum = "f73fe65f54d1e12b726f517d3e2135ca3125a437b6d998caf1962961f7172d9e" dependencies = [ "futures-channel", "futures-core", @@ -111,9 +111,9 @@ dependencies = [ [[package]] name = "futures-channel" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba3dda0b6588335f360afc675d0564c17a77a2bda81ca178a4b6081bd86c7f0b" +checksum = "c3083ce4b914124575708913bca19bfe887522d6e2e6d0952943f5eac4a74010" dependencies = [ "futures-core", "futures-sink", @@ -121,15 +121,15 @@ dependencies = [ [[package]] name = "futures-core" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d0c8ff0461b82559810cdccfde3215c3f373807f5e5232b71479bff7bb2583d7" +checksum = "0c09fd04b7e4073ac7156a9539b57a484a8ea920f79c7c675d05d289ab6110d3" [[package]] name = "futures-executor" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29d6d2ff5bb10fb95c85b8ce46538a2e5f5e7fdc755623a7d4529ab8a4ed9d2a" +checksum = "9420b90cfa29e327d0429f19be13e7ddb68fa1cccb09d65e5706b8c7a749b8a6" dependencies = [ "futures-core", "futures-task", @@ -138,15 +138,15 @@ dependencies = [ [[package]] name = "futures-io" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b1f9d34af5a1aac6fb380f735fe510746c38067c5bf16c7fd250280503c971b2" +checksum = "fc4045962a5a5e935ee2fdedaa4e08284547402885ab326734432bed5d12966b" [[package]] name = "futures-macro" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6dbd947adfffb0efc70599b3ddcf7b5597bb5fa9e245eb99f62b3a5f7bb8bd3c" +checksum = "33c1e13800337f4d4d7a316bf45a567dbcb6ffe087f16424852d97e97a91f512" dependencies = [ "proc-macro2", "quote", @@ -155,21 +155,21 @@ dependencies = [ [[package]] name = "futures-sink" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e3055baccb68d74ff6480350f8d6eb8fcfa3aa11bdc1a1ae3afdd0514617d508" +checksum = "21163e139fa306126e6eedaf49ecdb4588f939600f0b1e770f4205ee4b7fa868" [[package]] name = "futures-task" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ee7c6485c30167ce4dfb83ac568a849fe53274c831081476ee13e0dce1aad72" +checksum = "57c66a976bf5909d801bbef33416c41372779507e7a6b3a5e25e4749c58f776a" [[package]] name = "futures-util" -version = "0.3.19" +version = "0.3.21" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9b5cf40b47a271f77a8b1bec03ca09044d99d2372c0de244e66430761127164" +checksum = "d8b7abd5d659d9b90c8cba917f6ec750a74e2dc23902ef9cd4cc8c8b22e6036a" dependencies = [ "futures-channel", "futures-core", @@ -645,9 +645,9 @@ dependencies = [ [[package]] name = "trust-dns-proto" -version = "0.20.3" +version = "0.20.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ad0d7f5db438199a6e2609debe3f69f808d074e0a2888ee0bccb45fe234d03f4" +checksum = "ca94d4e9feb6a181c690c4040d7a24ef34018d8313ac5044a61d21222ae24e31" dependencies = [ "async-trait", "cfg-if", @@ -670,9 +670,9 @@ dependencies = [ [[package]] name = "trust-dns-resolver" -version = "0.20.3" +version = "0.20.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6ad17b608a64bd0735e67bde16b0636f8aa8591f831a25d18443ed00a699770" +checksum = "ecae383baad9995efaa34ce8e57d12c3f305e545887472a492b838f4b5cfb77a" dependencies = [ "cfg-if", "futures-util", diff --git a/Cargo.toml b/Cargo.toml index 0b47865..e714d97 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,8 +13,8 @@ readme = "README.md" [dependencies] tokio = { version = "1.16.1", features = ["full"] } -futures = "0.3.19" -trust-dns-resolver = "0.20.3" +futures = "0.3.21" +trust-dns-resolver = "0.20.4" # trust-dns-resolver = { git = "https://github.com/Edu4rdSHL/trust-dns", branch = "main" } -clap = "2.33.3" +clap = "2.33.4" rand = "0.8.4" diff --git a/src/dnslib.rs b/src/dnslib.rs new file mode 100644 index 0000000..899c883 --- /dev/null +++ b/src/dnslib.rs @@ -0,0 +1,110 @@ +use { + crate::structs::DomainData, + futures::stream::{self, StreamExt}, + std::collections::{HashMap, HashSet}, + trust_dns_resolver::config::ResolverOpts, +}; + +use { + std::net::SocketAddr, + trust_dns_resolver::{ + config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig}, + name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime}, + AsyncResolver, TokioAsyncResolver, + }, +}; + +pub fn return_tokio_asyncresolver( + nameserver_ips: HashSet, + options: ResolverOpts, +) -> 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() +} + +#[allow(clippy::too_many_arguments)] +pub async fn return_hosts_data( + hosts: HashSet, + resolvers: AsyncResolver>, + trustable_resolver: AsyncResolver>, + wildcard_ips: HashSet, + disable_double_check: bool, + mut threads: usize, + show_ip_addresses: bool, + quiet_flag: bool, +) -> HashMap { + if hosts.len() < threads { + threads = hosts.len(); + } + + stream::iter(hosts) + .map(|host| { + let resolver_fut = resolvers.ipv4_lookup(host.trim_end_matches('.').to_owned() + "."); + let trustable_resolver_fut = + trustable_resolver.ipv4_lookup(host.trim_end_matches('.').to_owned() + "."); + let wildcard_ips = wildcard_ips.clone(); + + async move { + let mut domain_data = DomainData::default(); + + if let Ok(ip) = resolver_fut.await { + if disable_double_check { + domain_data.ipv4_addresses = ip + .into_iter() + .map(|x| x.to_string()) + .collect::>(); + } else if let Ok(ip) = trustable_resolver_fut.await { + domain_data.ipv4_addresses = ip + .into_iter() + .map(|x| x.to_string()) + .collect::>(); + } + } + domain_data.is_wildcard = domain_data + .ipv4_addresses + .iter() + .all(|ip| wildcard_ips.contains(ip)); + + if !quiet_flag { + if show_ip_addresses && !domain_data.is_wildcard { + println!("{};{:?}", host, domain_data.ipv4_addresses); + } else if !domain_data.is_wildcard { + println!("{}", host) + } + } + + (host, domain_data) + } + }) + .buffer_unordered(threads) + .collect::>() + .await +} diff --git a/src/lib.rs b/src/lib.rs index cdacf92..de23ad6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,3 @@ -pub mod resolver; +pub mod dnslib; +pub mod structs; pub mod utils; - -mod structs; diff --git a/src/main.rs b/src/main.rs index fb83357..53ca1ec 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,6 @@ use { clap::{value_t, App, Arg}, - futures::stream::{self, StreamExt}, - rusolver::{resolver, utils}, + rusolver::{dnslib, utils}, std::collections::HashSet, tokio::{ self, @@ -52,7 +51,7 @@ async fn main() -> Result<(), Box> { Arg::with_name("timeout") .long("timeout") .takes_value(true) - .help("Timeout in seconds. Default: 1"), + .help("Timeout in seconds. Default: 3"), ) .arg( Arg::with_name("ip") @@ -79,7 +78,7 @@ async fn main() -> Result<(), Box> { // Assign values or use defaults 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 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"); @@ -132,8 +131,8 @@ async fn main() -> Result<(), Box> { } else { nameserver_ips = built_in_nameservers.clone(); } - let resolvers = resolver::return_tokio_asyncresolver(nameserver_ips, options); - let trustable_resolver = resolver::return_tokio_asyncresolver(built_in_nameservers, options); + let resolvers = dnslib::return_tokio_asyncresolver(nameserver_ips, options); + let trustable_resolver = dnslib::return_tokio_asyncresolver(built_in_nameservers, options); let mut wildcard_ips = HashSet::new(); // Read stdin @@ -141,7 +140,7 @@ async fn main() -> Result<(), Box> { let mut stdin = io::stdin(); stdin.read_to_string(&mut buffer).await?; - let hosts: Vec = if matches.is_present("domain") { + let hosts: HashSet = if matches.is_present("domain") { let domain = value_t!(matches, "domain", String).unwrap(); wildcard_ips = utils::detect_wildcards(&domain, &trustable_resolver, quiet_flag).await; buffer @@ -152,38 +151,17 @@ async fn main() -> Result<(), Box> { buffer.lines().map(str::to_owned).collect() }; - stream::iter(hosts) - .map(|host| { - let resolver_fut = resolvers.ipv4_lookup(host.trim_end_matches('.').to_owned() + "."); - let trustable_resolver_fut = - trustable_resolver.ipv4_lookup(host.trim_end_matches('.').to_owned() + "."); - let wildcard_ips = wildcard_ips.clone(); - - async move { - let mut ips = HashSet::new(); - if let Ok(ip) = resolver_fut.await { - if disable_double_check { - ips = ip - .into_iter() - .map(|x| x.to_string()) - .collect::>(); - } else if let Ok(ip) = trustable_resolver_fut.await { - ips = ip - .into_iter() - .map(|x| x.to_string()) - .collect::>(); - } - } - if show_ip_adress && !ips.iter().all(|ip| wildcard_ips.contains(ip)) { - println!("{};{:?}", host, ips) - } else if !ips.iter().all(|ip| wildcard_ips.contains(ip)) { - println!("{}", host) - } - } - }) - .buffer_unordered(threads) - .collect::>() - .await; + dnslib::return_hosts_data( + hosts, + resolvers, + trustable_resolver, + wildcard_ips, + disable_double_check, + threads, + show_ip_adress, + false, + ) + .await; Ok(()) } diff --git a/src/resolver.rs b/src/resolver.rs deleted file mode 100644 index 582ed07..0000000 --- a/src/resolver.rs +++ /dev/null @@ -1,44 +0,0 @@ -use { - std::{collections::HashSet, net::SocketAddr}, - trust_dns_resolver::{ - config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig, ResolverOpts}, - name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime}, - AsyncResolver, TokioAsyncResolver, - }, -}; - -pub fn return_tokio_asyncresolver( - nameserver_ips: HashSet, - options: ResolverOpts, -) -> 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() -} diff --git a/src/structs.rs b/src/structs.rs index e69de29..f281670 100644 --- a/src/structs.rs +++ b/src/structs.rs @@ -0,0 +1,20 @@ +use std::collections::HashSet; + +#[derive(Clone, Debug)] +pub struct DomainData { + pub ipv4_addresses: HashSet, + pub ipv6_addresses: HashSet, + pub cname: String, + pub is_wildcard: bool, +} + +impl Default for DomainData { + fn default() -> Self { + DomainData { + ipv4_addresses: HashSet::new(), + ipv6_addresses: HashSet::new(), + cname: String::from(""), + is_wildcard: false, + } + } +}