Merge pull request #32 from bwks/issue-31-dns

Issue 31 dns

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "netkraken"
-version = "0.1.4"
+version = "0.1.5"
 edition = "2021"
 
 [[bin]]

src/tcp/client.rs

@@ -1,9 +1,8 @@
-use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
-use anyhow::Result;
+use anyhow::{bail, Result};
 use futures::StreamExt;
 use tokio::net::TcpSocket;
 use tokio::signal;
@@ -14,9 +13,13 @@ use crate::core::common::{
     OutputOptions, PingOptions,
 };
 use crate::core::konst::{BIND_ADDR, BIND_PORT, BUFFER_SIZE};
+use crate::util::dns::resolve_host;
 use crate::util::handler::{io_error_switch_handler, loop_handler, output_handler2};
-use crate::util::message::{client_result_msg, client_summary_msg, ping_header_msg};
+use crate::util::message::{
+    client_result_msg, client_summary_msg, ping_header_msg, resolved_ips_msg,
+};
 use crate::util::parser::parse_ipaddr;
+use crate::util::result::get_results_map;
 use crate::util::time::{calc_connect_ms, time_now_us};
 
 #[derive(Debug)]
@@ -57,8 +60,6 @@ impl TcpClient {
     }
 
     pub async fn connect(&self) -> Result<()> {
-        let mut results_map: HashMap<String, HashMap<String, Vec<f64>>> = HashMap::new();
-
         let src_ip_port = IpPort {
             ip: parse_ipaddr(&self.src_ip)?,
             port: self.src_port,
@@ -68,49 +69,20 @@ impl TcpClient {
 
         let hosts = vec![host_records.clone()];
 
-        let lookup_data: Vec<HostRecord> = futures::stream::iter(hosts)
-            .map(|host| {
-                async move {
-                    //
-                    HostRecord::new(&host.host, host.port).await
-                }
-            })
-            .buffer_unordered(BUFFER_SIZE)
-            .collect()
-            .await;
+        let resolved_hosts = resolve_host(hosts).await;
 
-        let mut resolved_hosts: Vec<HostRecord> = vec![];
-        for lookup in lookup_data.clone() {
-            if lookup.ipv4_sockets.is_empty() && lookup.ipv6_sockets.is_empty() {
-                println!("{} returned no IPs", lookup.host);
-                continue;
-            }
-            resolved_hosts.push(lookup.clone());
-            println!(
-                "{} resolves to {} IPs",
-                lookup.host,
-                lookup.ipv4_sockets.len() + lookup.ipv6_sockets.len()
-            );
-            results_map.insert(lookup.host.to_owned(), HashMap::new());
-            for addr in lookup.ipv4_sockets {
-                println!(" - {}", addr.ip());
-                results_map
-                    .get_mut(&lookup.host)
-                    // this should never fail because we just inserted lookup.host
-                    .unwrap()
-                    .insert(addr.to_string(), vec![]);
-            }
-            for addr in lookup.ipv6_sockets {
-                println!(" - {}", addr.ip());
-                results_map
-                    .get_mut(&lookup.host)
-                    // this should never fail because we just inserted lookup.host
-                    .unwrap()
-                    .insert(addr.to_string(), vec![]);
+        for record in &resolved_hosts {
+            match record.ipv4_sockets.is_empty() && record.ipv6_sockets.is_empty() {
+                true => bail!("{} did not resolve to an IP address", record.host),
+                false => {
+                    let resolved_host_msg = resolved_ips_msg(record);
+                    println!("{resolved_host_msg}");
+                }
             }
-            println!();
         }
 
+        let mut results_map = get_results_map(&resolved_hosts);
+
         let mut count: u16 = 0;
         let mut send_count: u16 = 0;
 

src/udp/client.rs

@@ -1,9 +1,8 @@
-use std::collections::HashMap;
 use std::net::SocketAddr;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::Arc;
 
-use anyhow::Result;
+use anyhow::{bail, Result};
 use futures::StreamExt;
 use tokio::net::UdpSocket;
 use tokio::signal;
@@ -14,9 +13,13 @@ use crate::core::common::{
 };
 use crate::core::common::{ConnectMethod, ConnectRecord, ConnectResult};
 use crate::core::konst::{BIND_ADDR, BIND_PORT, BUFFER_SIZE, MAX_PACKET_SIZE, PING_MSG};
+use crate::util::dns::resolve_host;
 use crate::util::handler::{io_error_switch_handler, loop_handler, output_handler2};
-use crate::util::message::{client_result_msg, client_summary_msg, ping_header_msg};
+use crate::util::message::{
+    client_result_msg, client_summary_msg, ping_header_msg, resolved_ips_msg,
+};
 use crate::util::parser::parse_ipaddr;
+use crate::util::result::get_results_map;
 use crate::util::time::{calc_connect_ms, time_now_us};
 
 pub struct UdpClient {
@@ -48,8 +51,6 @@ impl UdpClient {
     }
 
     pub async fn connect(&self) -> Result<()> {
-        let mut results_map: HashMap<String, HashMap<String, Vec<f64>>> = HashMap::new();
-
         let src_ip_port = IpPort {
             ip: parse_ipaddr(&self.src_ip)?,
             port: self.src_port,
@@ -59,49 +60,20 @@ impl UdpClient {
 
         let hosts = vec![host_records.clone()];
 
-        let lookup_data: Vec<HostRecord> = futures::stream::iter(hosts)
-            .map(|host| {
-                async move {
-                    //
-                    HostRecord::new(&host.host, host.port).await
+        let resolved_hosts = resolve_host(hosts).await;
+
+        for record in &resolved_hosts {
+            match record.ipv4_sockets.is_empty() && record.ipv6_sockets.is_empty() {
+                true => bail!("{} did not resolve to an IP address", record.host),
+                false => {
+                    let resolved_host_msg = resolved_ips_msg(record);
+                    println!("{resolved_host_msg}");
                 }
-            })
-            .buffer_unordered(BUFFER_SIZE)
-            .collect()
-            .await;
-
-        let mut resolved_hosts: Vec<HostRecord> = vec![];
-        for lookup in lookup_data.clone() {
-            if lookup.ipv4_sockets.is_empty() && lookup.ipv6_sockets.is_empty() {
-                println!("{} returned no IPs", lookup.host);
-                continue;
             }
-            resolved_hosts.push(lookup.clone());
-            println!(
-                "{} resolves to {} IPs",
-                lookup.host,
-                lookup.ipv4_sockets.len() + lookup.ipv6_sockets.len()
-            );
-            results_map.insert(lookup.host.to_owned(), HashMap::new());
-            for addr in lookup.ipv4_sockets {
-                println!(" - {}", addr.ip());
-                results_map
-                    .get_mut(&lookup.host)
-                    // this should never fail because we just inserted lookup.host
-                    .unwrap()
-                    .insert(addr.to_string(), vec![]);
-            }
-            for addr in lookup.ipv6_sockets {
-                println!(" - {}", addr.ip());
-                results_map
-                    .get_mut(&lookup.host)
-                    // this should never fail because we just inserted lookup.host
-                    .unwrap()
-                    .insert(addr.to_string(), vec![]);
-            }
-            println!();
         }
 
+        let mut results_map = get_results_map(&resolved_hosts);
+
         let mut count: u16 = 0;
         let mut send_count: u16 = 0;
 

src/util/dns.rs

@@ -0,0 +1,19 @@
+use futures::StreamExt;
+
+use crate::core::common::HostRecord;
+use crate::core::konst::BUFFER_SIZE;
+
+pub async fn resolve_host(hosts: Vec<HostRecord>) -> Vec<HostRecord> {
+    let lookup_data: Vec<HostRecord> = futures::stream::iter(hosts)
+        .map(|host| {
+            async move {
+                //
+                HostRecord::new(&host.host, host.port).await
+            }
+        })
+        .buffer_unordered(BUFFER_SIZE)
+        .collect()
+        .await;
+
+    lookup_data
+}

src/util/message.rs

@@ -1,4 +1,6 @@
-use crate::core::common::{ClientSummary, ConnectMethod, ConnectRecord, ConnectResult};
+use std::net::SocketAddr;
+
+use crate::core::common::{ClientSummary, ConnectMethod, ConnectRecord, ConnectResult, HostRecord};
 
 /// Return the CLI header message
 pub fn cli_header_msg() -> String {
@@ -16,6 +18,30 @@ Press CRTL+C to exit
     )
 }
 
+/// Return a list of resolved IPs from a hostname
+pub fn resolved_ips_msg(host_record: &HostRecord) -> String {
+    let ip_records: Vec<&SocketAddr> = host_record
+        .ipv4_sockets
+        .iter()
+        .chain(host_record.ipv6_sockets.iter())
+        .collect();
+
+    let ip_str = ip_records
+        .iter()
+        .map(|x| format!(" {}", x.ip()))
+        .collect::<Vec<String>>()
+        .join("\n");
+
+    format!(
+        "{} resolves to {} IPs
+{}
+",
+        host_record.host,
+        host_record.ipv4_sockets.len() + host_record.ipv6_sockets.len(),
+        ip_str,
+    )
+}
+
 /// Return a ping header message
 pub fn ping_header_msg(destination: &String, port: u16, protocol: ConnectMethod) -> String {
     format!(
@@ -139,8 +165,75 @@ pub fn calc_loss_percent(sent: u16, received: u16) -> f64 {
 
 #[cfg(test)]
 mod tests {
+    use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
+
+    use crate::core::common::HostRecord;
     use crate::util::message::*;
 
+    #[test]
+    fn resolved_ips_msg_with_no_ips_is_expected() {
+        let host_record = HostRecord {
+            host: "blah.bleh".to_owned(),
+            port: 443,
+            ipv4_sockets: vec![],
+            ipv6_sockets: vec![],
+        };
+        let msg = resolved_ips_msg(&host_record);
+
+        assert_eq!(msg, "blah.bleh resolves to 0 IPs\n\n");
+    }
+
+    #[test]
+    fn resolved_ips_msg_with_only_ip4s_is_expected() {
+        let host_record = HostRecord {
+            host: "blah.bleh".to_owned(),
+            port: 443,
+            ipv4_sockets: vec![SocketAddr::new(
+                IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
+                443,
+            )],
+            ipv6_sockets: vec![],
+        };
+        let msg = resolved_ips_msg(&host_record);
+
+        assert_eq!(msg, "blah.bleh resolves to 1 IPs\n 127.0.0.1\n");
+    }
+
+    #[test]
+    fn resolved_ips_msg_with_only_ip6s_is_expected() {
+        let host_record = HostRecord {
+            host: "blah.bleh".to_owned(),
+            port: 443,
+            ipv4_sockets: vec![],
+            ipv6_sockets: vec![SocketAddr::new(
+                IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)),
+                443,
+            )],
+        };
+        let msg = resolved_ips_msg(&host_record);
+
+        assert_eq!(msg, "blah.bleh resolves to 1 IPs\n ::1\n");
+    }
+
+    #[test]
+    fn resolved_ips_msg_with_ipv4_and_ip6s_is_expected() {
+        let host_record = HostRecord {
+            host: "blah.bleh".to_owned(),
+            port: 443,
+            ipv4_sockets: vec![SocketAddr::new(
+                IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)),
+                443,
+            )],
+            ipv6_sockets: vec![SocketAddr::new(
+                IpAddr::V6(Ipv6Addr::new(0, 0, 0, 0, 0, 0, 0, 1)),
+                443,
+            )],
+        };
+        let msg = resolved_ips_msg(&host_record);
+
+        assert_eq!(msg, "blah.bleh resolves to 2 IPs\n 127.0.0.1\n ::1\n");
+    }
+
     #[test]
     fn ping_header_msg_is_expected() {
         let msg = ping_header_msg(&"198.51.100.1".to_owned(), 443, ConnectMethod::TCP);

src/util/mod.rs

@@ -1,4 +1,6 @@
+pub mod dns;
 pub mod handler;
 pub mod message;
 pub mod parser;
+pub mod result;
 pub mod time;