Split functions across modules.

Improve performance.

Signed-off-by: Edu4rdSHL <edu4rdshl@protonmail.com>
This commit is contained in:
Edu4rdSHL 2022-02-13 05:32:25 -05:00
parent 7f974c8414
commit 17776e7b88
No known key found for this signature in database
GPG key ID: 3A574A4009F553E5
7 changed files with 174 additions and 111 deletions

44
Cargo.lock generated
View file

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

View file

@ -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"

110
src/dnslib.rs Normal file
View file

@ -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<String>,
options: ResolverOpts,
) -> AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>> {
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<String>,
resolvers: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
trustable_resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
wildcard_ips: HashSet<String>,
disable_double_check: bool,
mut threads: usize,
show_ip_addresses: bool,
quiet_flag: bool,
) -> HashMap<String, DomainData> {
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::<HashSet<String>>();
} else if let Ok(ip) = trustable_resolver_fut.await {
domain_data.ipv4_addresses = ip
.into_iter()
.map(|x| x.to_string())
.collect::<HashSet<String>>();
}
}
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::<HashMap<String, DomainData>>()
.await
}

View file

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

View file

@ -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<dyn std::error::Error>> {
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<dyn std::error::Error>> {
// 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<dyn std::error::Error>> {
} 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<dyn std::error::Error>> {
let mut stdin = io::stdin();
stdin.read_to_string(&mut buffer).await?;
let hosts: Vec<String> = if matches.is_present("domain") {
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;
buffer
@ -152,38 +151,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
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::<HashSet<String>>();
} else if let Ok(ip) = trustable_resolver_fut.await {
ips = ip
.into_iter()
.map(|x| x.to_string())
.collect::<HashSet<String>>();
}
}
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::<Vec<()>>()
.await;
dnslib::return_hosts_data(
hosts,
resolvers,
trustable_resolver,
wildcard_ips,
disable_double_check,
threads,
show_ip_adress,
false,
)
.await;
Ok(())
}

View file

@ -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<String>,
options: ResolverOpts,
) -> AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>> {
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()
}

View file

@ -0,0 +1,20 @@
use std::collections::HashSet;
#[derive(Clone, Debug)]
pub struct DomainData {
pub ipv4_addresses: HashSet<String>,
pub ipv6_addresses: HashSet<String>,
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,
}
}
}