vsock_proxy: change function's signature

Modify the signature of `dns::resolve` function to return a `DnsResolveResult`
struct instead of a plain vector of IP addresses.

Signed-off-by: Erdem Meydanli <meydanli@amazon.com>

vsock_proxy/src/dns.rs

@@ -12,48 +12,50 @@ use std::net::IpAddr;
 use crate::{DnsResolveResult, IpAddrType, VsockProxyResult};
 
 /// Resolve a DNS name (IDNA format) into multiple IP addresses (v4 or v6)
-pub fn resolve(addr: &str, ip_addr_type: IpAddrType) -> VsockProxyResult<Vec<IpAddr>> {
+pub fn resolve(addr: &str, ip_addr_type: IpAddrType) -> VsockProxyResult<Vec<DnsResolveResult>> {
     // IDNA parsing
     let addr = domain_to_ascii(addr).map_err(|_| "Could not parse domain name")?;
 
     // DNS lookup
     // It results in a vector of IPs (V4 and V6)
-    let ips = match lookup_host(&addr) {
-        Err(_) => {
-            return Ok(vec![]);
-        }
-        Ok(v) => {
-            if v.is_empty() {
-                return Ok(v);
-            }
-            v
-        }
-    };
+    let ips = lookup_host(&addr).map_err(|_| "DNS lookup failed!")?;
+
+    if ips.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().collect());
+        return Ok(ips
+            .into_iter()
+            .map(|ip| DnsResolveResult { ip, ttl })
+            .collect());
     }
 
     // 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().collect())
+        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().collect())
+        Ok(ips_v6
+            .into_iter()
+            .map(|ip| DnsResolveResult { ip, ttl })
+            .collect())
     } else {
-        Err("No accepted IP was found".to_string())
+        Err("No accepted IP was found.".to_string())
     }
 }
 
 /// Resolve a DNS name (IDNA format) into a single address with a TTL value
 pub fn resolve_single(addr: &str, ip_addr_type: IpAddrType) -> VsockProxyResult<DnsResolveResult> {
-    let addrs = resolve(addr, ip_addr_type)
-        .map_err(|err| format!("Could not parse remote address: {}", err))?;
-
-    let ip = *addrs.first().ok_or("No IP address found")?;
-    let ttl = 60; // IMPORTANT TODO: Obtain this value dynamically
-
-    Ok(DnsResolveResult { ip, ttl })
+    let rresults = resolve(addr, ip_addr_type)?;
+    // Return the first resolved IP address and its TTL value.
+    // We already know that rresults containt at least one item.
+    Ok(rresults.first().unwrap().clone())
 }

vsock_proxy/src/lib.rs

@@ -16,6 +16,7 @@ pub enum IpAddrType {
     IPAddrMixed,
 }
 
+#[derive(Copy, Clone)]
 pub struct DnsResolveResult {
     ///Resolved address
     pub ip: IpAddr,

vsock_proxy/src/proxy.rs

@@ -64,9 +64,9 @@ pub fn check_allowlist(
             }
 
             // If hostname matching failed, attempt to match against IPs.
-            let addrs = dns::resolve(addr, ip_addr_type)?;
-            for addr in addrs.into_iter() {
-                if addr == remote_addr {
+            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);
                 }