feat(dns): allow custom dns server as cli flag (#193)

* feat(dns): allow custom dns server as cli flag

* fix(style): remove commented debugging line (oops)

6 files changed, 48 insertions, 10 deletions

diff --git a/src/main.rs b/src/main.rs
index 93f4c76..39319f7 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -24,6 +24,7 @@ use ::tui::backend::Backend;
 
 use std::process;
 
+use ::std::net::Ipv4Addr;
 use ::std::time::{Duration, Instant};
 use ::tui::backend::CrosstermBackend;
 use std::sync::RwLock;
@@ -48,6 +49,9 @@ pub struct Opt {
     #[structopt(short, long)]
     /// Show DNS queries
     show_dns: bool,
+    #[structopt(short, long)]
+    /// A dns server ip to use instead of the system default
+    dns_server: Option<Ipv4Addr>,
 }
 
 #[derive(StructOpt, Debug, Copy, Clone)]
@@ -76,7 +80,7 @@ fn main() {
 fn try_main() -> Result<(), failure::Error> {
     use os::get_input;
     let opts = Opt::from_args();
-    let os_input = get_input(&opts.interface, !opts.no_resolve)?;
+    let os_input = get_input(&opts.interface, !opts.no_resolve, &opts.dns_server)?;
     let raw_mode = opts.raw;
     if raw_mode {
         let terminal_backend = RawTerminalBackend {};
diff --git a/src/network/dns/resolver.rs b/src/network/dns/resolver.rs
index 074f46e..1bdaf78 100644
--- a/src/network/dns/resolver.rs
+++ b/src/network/dns/resolver.rs
@@ -1,6 +1,9 @@
 use async_trait::async_trait;
 use std::net::IpAddr;
+use std::net::SocketAddr;
+use std::net::{Ipv4Addr, SocketAddrV4};
 use tokio::runtime::Handle;
+use trust_dns_resolver::config::{NameServerConfig, Protocol, ResolverConfig, ResolverOpts};
 use trust_dns_resolver::{error::ResolveErrorKind, TokioAsyncResolver};
 
 #[async_trait]
@@ -11,8 +14,25 @@ pub trait Lookup {
 pub struct Resolver(TokioAsyncResolver);
 
 impl Resolver {
-    pub async fn new(runtime: Handle) -> Result<Self, failure::Error> {
-        let resolver = TokioAsyncResolver::from_system_conf(runtime).await?;
+    pub async fn new(
+        runtime: Handle,
+        dns_server: &Option<Ipv4Addr>,
+    ) -> Result<Self, failure::Error> {
+        let resolver = match dns_server {
+            Some(dns_server_address) => {
+                let mut config = ResolverConfig::new();
+                let options = ResolverOpts::default();
+                let socket = SocketAddr::V4(SocketAddrV4::new(*dns_server_address, 53));
+                let nameserver_config = NameServerConfig {
+                    socket_addr: socket,
+                    protocol: Protocol::Udp,
+                    tls_dns_name: None,
+                };
+                config.add_name_server(nameserver_config);
+                TokioAsyncResolver::new(config, options, runtime).await?
+            }
+            None => TokioAsyncResolver::from_system_conf(runtime).await?,
+        };
         Ok(Self(resolver))
     }
 }
diff --git a/src/os/shared.rs b/src/os/shared.rs
index 3c0fe74..4ef25a9 100644
--- a/src/os/shared.rs
+++ b/src/os/shared.rs
@@ -6,6 +6,7 @@ use ::pnet::datalink::{self, Config, NetworkInterface};
 use ::std::io::{self, ErrorKind, Write};
 use ::tokio::runtime::Runtime;
 
+use ::std::net::Ipv4Addr;
 use ::std::time;
 
 use crate::os::errors::GetInterfaceErrorKind;
@@ -147,6 +148,7 @@ where
 pub fn get_input(
     interface_name: &Option<String>,
     resolve: bool,
+    dns_server: &Option<Ipv4Addr>,
 ) -> Result<OsInputOutput, failure::Error> {
     let network_interfaces = if let Some(name) = interface_name {
         match get_interface(&name) {
@@ -201,13 +203,14 @@ pub fn get_input(
     let write_to_stdout = create_write_to_stdout();
     let dns_client = if resolve {
         let mut runtime = Runtime::new()?;
-        let resolver = match runtime.block_on(dns::Resolver::new(runtime.handle().clone())) {
-            Ok(resolver) => resolver,
-            Err(err) => failure::bail!(
-                "Could not initialize the DNS resolver. Are you offline?\n\nReason: {:?}",
-                err
-            ),
-        };
+        let resolver =
+            match runtime.block_on(dns::Resolver::new(runtime.handle().clone(), dns_server)) {
+                Ok(resolver) => resolver,
+                Err(err) => failure::bail!(
+                    "Could not initialize the DNS resolver. Are you offline?\n\nReason: {:?}",
+                    err
+                ),
+            };
         let dns_client = dns::Client::new(resolver, runtime)?;
         Some(dns_client)
     } else {
diff --git a/src/tests/cases/raw_mode.rs b/src/tests/cases/raw_mode.rs
index 3b4453f..bee763d 100644
--- a/src/tests/cases/raw_mode.rs
+++ b/src/tests/cases/raw_mode.rs
@@ -577,6 +577,7 @@ fn no_resolve_mode() {
         raw: true,
         no_resolve: true,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: false,
             connections: false,
diff --git a/src/tests/cases/test_utils.rs b/src/tests/cases/test_utils.rs
index bd90399..ab8e502 100644
--- a/src/tests/cases/test_utils.rs
+++ b/src/tests/cases/test_utils.rs
@@ -158,6 +158,7 @@ fn opts_factory(raw: bool) -> Opt {
         raw,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: false,
             connections: false,
diff --git a/src/tests/cases/ui.rs b/src/tests/cases/ui.rs
index d526021..14bf861 100644
--- a/src/tests/cases/ui.rs
+++ b/src/tests/cases/ui.rs
@@ -174,6 +174,7 @@ fn basic_only_processes() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: false,
             connections: false,
@@ -200,6 +201,7 @@ fn basic_processes_with_dns_queries() {
         raw: false,
         no_resolve: false,
         show_dns: true,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: false,
             connections: false,
@@ -226,6 +228,7 @@ fn basic_only_connections() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: false,
             connections: true,
@@ -252,6 +255,7 @@ fn basic_only_addresses() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: true,
             connections: false,
@@ -276,6 +280,7 @@ fn two_packets_only_processes() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: false,
             connections: false,
@@ -301,6 +306,7 @@ fn two_packets_only_connections() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: false,
             connections: true,
@@ -326,6 +332,7 @@ fn two_packets_only_addresses() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: true,
             connections: false,
@@ -353,6 +360,7 @@ fn two_windows_split_horizontally() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: true,
             connections: true,
@@ -379,6 +387,7 @@ fn two_windows_split_vertically() {
         raw: false,
         no_resolve: false,
         show_dns: false,
+        dns_server: None,
         render_opts: RenderOpts {
             addresses: true,
             connections: true,