Use LibOptions for library options.

Signed-off-by: Eduard Tolosa <edu4rdshl@protonmail.com>

src/dnslib.rs

@@ -1,17 +1,15 @@
-use trust_dns_resolver::proto::{rr::RecordType, xfer::DnsRequestOptions};
-
 use {
-    crate::structs::DomainData,
+    crate::structs::{DomainData, LibOptions},
     futures::stream::{self, StreamExt},
-    std::collections::{HashMap, HashSet},
+    std::{
-    trust_dns_resolver::config::ResolverOpts,
+        collections::{HashMap, HashSet},
-};
+        net::SocketAddr,
-
+    },
-use {
-    std::net::SocketAddr,
     trust_dns_resolver::{
-        config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig},
+        config::{NameServerConfig, NameServerConfigGroup, Protocol, ResolverConfig, ResolverOpts},
+        lookup::{Ipv4Lookup, Lookup},
         name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime},
+        proto::{rr::RecordType, xfer::DnsRequestOptions},
         AsyncResolver, TokioAsyncResolver,
     },
 };
@@ -52,51 +50,46 @@ pub fn return_tokio_asyncresolver(
     .unwrap()
 }
 
-#[allow(clippy::too_many_arguments)]
+pub async fn return_hosts_data(options: &LibOptions) -> HashMap<String, DomainData> {
-pub async fn return_hosts_data(
+    let threads = if options.hosts.len() < options.threads {
-    hosts: HashSet<String>,
+        options.hosts.len()
-    resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
+    } else {
-    trustable_resolver: AsyncResolver<GenericConnection, GenericConnectionProvider<TokioRuntime>>,
+        options.threads
-    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)
+    stream::iter(options.hosts.clone().into_iter().map(|host| {
-        .map(|host| {
+        let lookup_host = host.trim_end_matches('.').to_owned() + ".";
-            let resolver_fut = resolver.ipv4_lookup(host.trim_end_matches('.').to_owned() + ".");
+
-            let trustable_resolver_fut =
+        let resolver_fut = options.resolvers.ipv4_lookup(lookup_host.clone());
-                trustable_resolver.ipv4_lookup(host.trim_end_matches('.').to_owned() + ".");
+        let trustable_resolver_fut = options.trustable_resolver.ipv4_lookup(lookup_host);
-            let wildcard_ips = wildcard_ips.clone();
+        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 disable_double_check {
+                if options.disable_double_check {
-                        domain_data.ipv4_addresses = ip
+                    ip_lookup = Some(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
+                    ip_lookup = Some(ip);
-                            .into_iter()
-                            .map(|x| x.to_string())
-                            .collect::<HashSet<String>>();
                 }
             }
+
+            if let Some(ip_lookup) = ip_lookup {
+                for ip in ip_lookup.iter() {
+                    domain_data.ipv4_addresses.insert(ip.to_string());
+                }
+            }
+
             domain_data.is_wildcard = domain_data
                 .ipv4_addresses
                 .iter()
                 .all(|ip| wildcard_ips.contains(ip));
 
-                if !quiet_flag {
+            if !options.quiet_flag {
-                    if show_ip_addresses && !domain_data.is_wildcard {
+                if options.show_ip_address && !domain_data.is_wildcard {
                     println!("{};{:?}", host, domain_data.ipv4_addresses);
                 } else if !domain_data.is_wildcard {
                     println!("{}", host)
@@ -105,7 +98,7 @@ pub async fn return_hosts_data(
 
             (host, domain_data)
         }
-        })
+    }))
     .buffer_unordered(threads)
     .collect::<HashMap<String, DomainData>>()
     .await
@@ -128,41 +121,32 @@ pub async fn return_cname_data(
 
     stream::iter(hosts)
         .map(|host| {
-            let resolver_fut = resolver.lookup(
+            let lookup_host = host.trim_end_matches('.').to_owned() + ".";
-                host.trim_end_matches('.').to_owned() + ".",
+            let resolver_fut = resolver.lookup(lookup_host.clone(), record_type, request_options);
-                record_type,
+            let trustable_resolver_fut =
-                request_options,
+                trustable_resolver.lookup(lookup_host, record_type, request_options);
-            );
-            let trustable_resolver_fut = trustable_resolver.lookup(
-                host.trim_end_matches('.').to_owned() + ".",
-                record_type,
-                request_options,
-            );
 
             let mut domain_data = DomainData::default();
 
+            let mut cname_lookup = Option::<Lookup>::None;
+
             async move {
                 if let Ok(lookup) = resolver_fut.await {
                     if disable_double_check {
-                        domain_data.cname = lookup
+                        cname_lookup = Some(lookup);
-                            .iter()
-                            .filter_map(|rdata| rdata.as_cname())
-                            .map(|name| {
-                                let name = name.to_string();
-                                name[..name.len() - 1].to_owned()
-                            })
-                            .collect();
                     } else if let Ok(lookup) = trustable_resolver_fut.await {
-                        domain_data.cname = lookup
+                        cname_lookup = Some(lookup);
-                            .iter()
-                            .filter_map(|rdata| rdata.as_cname())
-                            .map(|name| {
-                                let name = name.to_string();
-                                name[..name.len() - 1].to_owned()
-                            })
-                            .collect();
                     }
                 }
+
+                if let Some(lookup) = cname_lookup {
+                    for record in lookup.iter() {
+                        if let Some(cname) = record.as_cname() {
+                            domain_data.cname = cname.to_string();
+                        }
+                    }
+                }
+
                 (host, domain_data)
             }
         })

src/main.rs

@@ -1,6 +1,6 @@
 use {
     clap::{value_t, App, Arg},
-    rusolver::{dnslib, utils},
+    rusolver::{dnslib, structs, utils},
     std::collections::HashSet,
     tokio::{
         self,
@@ -9,7 +9,7 @@ use {
     trust_dns_resolver::config::{LookupIpStrategy, ResolverOpts},
 };
 
-// Please add support for AAAA, TXT, SRV, NAPTR, PTR, CNAME, DNAME, MX, NS, SOA, LOC, SVCB, HTTPS, SPF, CAA and AVC resource records.
+// WIP: add support for AAAA, TXT, SRV, NAPTR, PTR, CNAME, DNAME, MX, NS, SOA, LOC, SVCB, HTTPS, SPF, CAA and AVC resource records.
 // This could use a new command line option such as -t, e.g. echo www.example.com | rusolver -i -t AAAA. It might also make sense
 // to change -i/--ip to -d/--data with the text Display the record data.
 
@@ -76,7 +76,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .get_matches();
 
     // Assign values or use defaults
-    let show_ip_adress = matches.is_present("ip");
+    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);
@@ -118,7 +118,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         "84.200.70.40:53",
     ]
     .iter()
-    .map(|x| x.to_string())
+    .map(ToString::to_string)
     .collect();
 
     // Create resolvers
@@ -152,17 +152,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         buffer.lines().map(str::to_owned).collect()
     };
 
-    dnslib::return_hosts_data(
+    let options = structs::LibOptions {
         hosts,
         resolvers,
         trustable_resolver,
         wildcard_ips,
         disable_double_check,
         threads,
-        show_ip_adress,
+        show_ip_address,
-        false,
+        quiet_flag,
-    )
+    };
-    .await;
+
+    dnslib::return_hosts_data(&options).await;
 
     Ok(())
 }

src/structs.rs

@@ -1,4 +1,10 @@
-use std::collections::HashSet;
+use {
+    std::collections::HashSet,
+    trust_dns_resolver::{
+        name_server::{GenericConnection, GenericConnectionProvider, TokioRuntime},
+        AsyncResolver,
+    },
+};
 
 #[derive(Clone, Debug)]
 pub struct DomainData {
@@ -18,3 +24,16 @@ impl Default for 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 wildcard_ips: HashSet<String>,
+    pub disable_double_check: bool,
+    pub threads: usize,
+    pub show_ip_address: bool,
+    pub quiet_flag: bool,
+}