vsock_proxy: Use hickory_resolver for DNS resolution

Replace the existing DNS lookup implementation with hickory_resolver.
This allows accessing additional information from DNS records, such as
TTL values, and enhances the functionality of the dns module.

Signed-off-by: Erdem Meydanli <[email protected]>

vsock_proxy/Cargo.toml

@@ -12,6 +12,7 @@ chrono = "0.4"
 clap = "3.2"
 dns-lookup = "2.0.3"
 env_logger = "0.10"
+hickory-resolver = "0.24"
 idna = "0.3.0"
 log = "0.4"
 nix = "0.26"

vsock_proxy/src/dns.rs

@@ -1,13 +1,13 @@
-// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2019-2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 #![deny(warnings)]
 
 /// Contains code for Proxy, a library used for translating vsock traffic to
 /// TCP traffic
 ///
-use dns_lookup::lookup_host;
+use hickory_resolver::config::*;
+use hickory_resolver::Resolver;
 use idna::domain_to_ascii;
-use std::net::IpAddr;
 
 use crate::{DnsResolveResult, IpAddrType, VsockProxyResult};
 
@@ -18,35 +18,42 @@ pub fn resolve(addr: &str, ip_addr_type: IpAddrType) -> VsockProxyResult<Vec<Dns
 
     // DNS lookup
     // It results in a vector of IPs (V4 and V6)
-    let ips = lookup_host(&addr).map_err(|_| "DNS lookup failed!")?;
+    let resolver = Resolver::new(ResolverConfig::default(), ResolverOpts::default()).unwrap();
+
+    let rresults: Vec<DnsResolveResult> = resolver
+        .lookup_ip(&addr)
+        .map_err(|error| format!("DNS lookup failed! {:?}", error))?
+        .as_lookup()
+        .records()
+        .iter()
+        .filter_map(|record| {
+            if let Some(rdata) = record.data() {
+                if let Some(ip) = rdata.ip_addr() {
+                    let ttl = record.ttl();
+                    return Some(DnsResolveResult { ip, ttl });
+                }
+            }
+            None
+        })
+        .collect();
 
-    if ips.is_empty() {
+    if rresults.is_empty() {
         return Err("DNS lookup returned no IP addresses!".into());
     }
 
-    let ttl = 60; //TODO: update hardcoded value
-
     // If there is no restriction, choose randomly
     if IpAddrType::IPAddrMixed == ip_addr_type {
-        return Ok(ips
-            .into_iter()
-            .map(|ip| DnsResolveResult { ip, ttl })
-            .collect());
+        return Ok(rresults);
     }
 
-    // Split the IPs in v4 and v6
-    let (ips_v4, ips_v6): (Vec<_>, Vec<_>) = ips.into_iter().partition(IpAddr::is_ipv4);
-
-    if IpAddrType::IPAddrV4Only == ip_addr_type && !ips_v4.is_empty() {
-        Ok(ips_v4
-            .into_iter()
-            .map(|ip| DnsResolveResult { ip, ttl })
-            .collect())
-    } else if IpAddrType::IPAddrV6Only == ip_addr_type && !ips_v6.is_empty() {
-        Ok(ips_v6
-            .into_iter()
-            .map(|ip| DnsResolveResult { ip, ttl })
-            .collect())
+    //Partition the resolution results into groups that use IPv4 or IPv6 addresses.
+    let (rresults_with_ipv4, rresults_with_ipv6): (Vec<_>, Vec<_>) =
+        rresults.into_iter().partition(|result| result.ip.is_ipv4());
+
+    if IpAddrType::IPAddrV4Only == ip_addr_type && !rresults_with_ipv4.is_empty() {
+        Ok(rresults_with_ipv4)
+    } else if IpAddrType::IPAddrV6Only == ip_addr_type && !rresults_with_ipv6.is_empty() {
+        Ok(rresults_with_ipv6)
     } else {
         Err("No accepted IP was found.".to_string())
     }
@@ -65,8 +72,8 @@ mod tests {
     use super::*;
 
     use ctor::ctor;
-    use std::sync::Once;
     use std::env;
+    use std::sync::Once;
 
     static TEST_INIT: Once = Once::new();
 

vsock_proxy/src/lib.rs

@@ -1,4 +1,4 @@
-// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2019-2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
 pub mod dns;
@@ -16,7 +16,7 @@ pub enum IpAddrType {
     IPAddrMixed,
 }
 
-#[derive(Copy, Clone)]
+#[derive(Copy, Clone, Debug)]
 pub struct DnsResolveResult {
     ///Resolved address
     pub ip: IpAddr,

vsock_proxy/src/proxy.rs

@@ -1,11 +1,11 @@
-// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// Copyright 2019-2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 //#![deny(warnings)]
 
 /// Contains code for Proxy, a library used for translating vsock traffic to
 /// TCP traffic
 use chrono::{DateTime, Duration, Utc};
-use log::info;
+use log::{info, warn};
 use nix::sys::select::{select, FdSet};
 use nix::sys::socket::SockType;
 use std::fs::File;
@@ -64,11 +64,23 @@ pub fn check_allowlist(
             }
 
             // If hostname matching failed, attempt to match against IPs.
-            let rresults = dns::resolve(addr, ip_addr_type)?;
-            for rresult in rresults.into_iter() {
-                if rresult.ip == remote_addr {
-                    info!("Matched with host IP \"{}\" and port \"{}\"", addr, port);
-                    return Ok(remote_addr);
+            let rresults = dns::resolve(addr, ip_addr_type);
+            let remote_addr_matched = rresults.clone().ok().and_then(|rresults| {
+                rresults
+                    .into_iter()
+                    .find(|rresult| rresult.ip == remote_addr)
+                    .map(|_| {
+                        info!("Matched with host IP \"{}\" and port \"{}\"", addr, port);
+                        remote_addr
+                    })
+            });
+
+            match remote_addr_matched {
+                Some(matched_addr) => return Ok(matched_addr),
+                None => {
+                    if rresults.is_err() {
+                        warn!("Unable to resolve allow listed host: {:?}.", addr);
+                    }
                 }
             }
         }