Skip to content

Commit

Permalink
vsock_proxy: Use hickory_resolver for DNS resolution
Browse files Browse the repository at this point in the history
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]>
  • Loading branch information
meerd committed Apr 10, 2024
1 parent 4740118 commit 6161520
Show file tree
Hide file tree
Showing 5 changed files with 226 additions and 36 deletions.
174 changes: 172 additions & 2 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions vsock_proxy/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
57 changes: 32 additions & 25 deletions vsock_proxy/src/dns.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand All @@ -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())
}
Expand All @@ -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();

Expand Down
4 changes: 2 additions & 2 deletions vsock_proxy/src/lib.rs
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -16,7 +16,7 @@ pub enum IpAddrType {
IPAddrMixed,
}

#[derive(Copy, Clone)]
#[derive(Copy, Clone, Debug)]
pub struct DnsResolveResult {
///Resolved address
pub ip: IpAddr,
Expand Down
Loading

0 comments on commit 6161520

Please sign in to comment.