From e813362e5cfaa8452bd3df3adb65acfd9d76a3c3 Mon Sep 17 00:00:00 2001 From: Aram Drevekenin Date: Thu, 1 Oct 2020 15:02:55 +0200 Subject: 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) --- src/main.rs | 6 +++++- src/network/dns/resolver.rs | 24 ++++++++++++++++++++++-- src/os/shared.rs | 17 ++++++++++------- src/tests/cases/raw_mode.rs | 1 + src/tests/cases/test_utils.rs | 1 + src/tests/cases/ui.rs | 9 +++++++++ 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, } #[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 { - let resolver = TokioAsyncResolver::from_system_conf(runtime).await?; + pub async fn new( + runtime: Handle, + dns_server: &Option, + ) -> Result { + 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, resolve: bool, + dns_server: &Option, ) -> Result { 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, -- cgit v1.2.3