Skip to content

Commit

Permalink
vsock_proxy: Refactor DNS-related functionality
Browse files Browse the repository at this point in the history
- Create a new module 'dns.rs' to encapsulate DNS-related operations.
- Move existing DNS-related functionality into the new 'dns.rs' module.
- Rename 'parse_addr' function to 'resolve' for better semantic clarity.

Signed-off-by: Erdem Meydanli <[email protected]>
  • Loading branch information
meerd committed Mar 13, 2024
1 parent c4e06c2 commit 1dc6107
Show file tree
Hide file tree
Showing 4 changed files with 81 additions and 48 deletions.
62 changes: 62 additions & 0 deletions vsock_proxy/src/dns.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
// Copyright 2019 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 idna::domain_to_ascii;
use std::net::IpAddr;

use crate::{IpAddrType, VsockProxyResult, DnsResolveResult};

/// 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>> {
// 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
}
};

// If there is no restriction, choose randomly
if IpAddrType::IPAddrMixed == ip_addr_type {
return Ok(ips.into_iter().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())
} else if IpAddrType::IPAddrV6Only == ip_addr_type && !ips_v6.is_empty() {
Ok(ips_v6.into_iter().collect())
} else {
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
})
}
15 changes: 14 additions & 1 deletion vsock_proxy/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@
// SPDX-License-Identifier: Apache-2.0

pub mod starter;
pub mod dns;

use std::net::IpAddr;

#[derive(Copy, Clone, PartialEq)]
pub enum IpAddrType {
Expand All @@ -11,4 +14,14 @@ pub enum IpAddrType {
IPAddrV6Only,
/// Allows both IP4 and IP6 addresses
IPAddrMixed
}
}

pub struct DnsResolveResult {
///Resolved address
ip: IpAddr,
///DNS TTL value
ttl: u32
}

/// The most common result type provided by VsockProxy operations.
pub type VsockProxyResult<T> = Result<T, String>;
2 changes: 1 addition & 1 deletion vsock_proxy/src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ use clap::{App, AppSettings, Arg};
use env_logger::init;
use log::info;

use vsock_proxy::{starter::{Proxy, VsockProxyResult}, IpAddrType};
use vsock_proxy::{starter::Proxy, IpAddrType, VsockProxyResult};

fn main() -> VsockProxyResult<()> {
init();
Expand Down
50 changes: 4 additions & 46 deletions vsock_proxy/src/starter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@
/// Contains code for Proxy, a library used for translating vsock traffic to
/// TCP traffic
///
use dns_lookup::lookup_host;
use idna::domain_to_ascii;
use log::info;
use nix::sys::select::{select, FdSet};
use nix::sys::socket::SockType;
Expand All @@ -18,15 +16,12 @@ use threadpool::ThreadPool;
use vsock::{VsockAddr, VsockListener};
use yaml_rust::YamlLoader;

use crate::IpAddrType;
use crate::{dns, IpAddrType, VsockProxyResult};

const BUFF_SIZE: usize = 8192;
pub const VSOCK_PROXY_CID: u32 = 3;
pub const VSOCK_PROXY_PORT: u32 = 8000;

/// The most common result type provided by VsockProxy operations.
pub type VsockProxyResult<T> = Result<T, String>;

/// Checks if the forwarded server is allowed, providing its IP on success.
pub fn check_allowlist(
remote_host: &str,
Expand All @@ -47,9 +42,8 @@ pub fn check_allowlist(
.ok_or("No allowlist field")?;

// Obtain the remote server's IP address.
let mut addrs = Proxy::parse_addr(remote_host, ip_addr_type)
.map_err(|err| format!("Could not parse remote address: {}", err))?;
let remote_addr = *addrs.first().ok_or("No IP address found")?;
let dns_result = dns::resolve_single(remote_host, ip_addr_type)?;
let remote_addr = dns_result.ip;

for raw_service in services {
let addr = raw_service["address"].as_str().ok_or("No address field")?;
Expand All @@ -70,7 +64,7 @@ pub fn check_allowlist(
}

// If hostname matching failed, attempt to match against IPs.
addrs = Proxy::parse_addr(addr, ip_addr_type)?;
let addrs = dns::resolve(addr, ip_addr_type)?;
for addr in addrs.into_iter() {
if addr == remote_addr {
info!("Matched with host IP \"{}\" and port \"{}\"", addr, port);
Expand Down Expand Up @@ -122,42 +116,6 @@ impl Proxy {
})
}

/// Resolve a DNS name (IDNA format) into an IP address (v4 or v6)
pub fn parse_addr(addr: &str, ip_addr_type: IpAddrType) -> VsockProxyResult<Vec<IpAddr>> {
// 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
}
};

// If there is no restriction, choose randomly
if IpAddrType::IPAddrMixed == ip_addr_type {
return Ok(ips.into_iter().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())
} else if IpAddrType::IPAddrV6Only == ip_addr_type && !ips_v6.is_empty() {
Ok(ips_v6.into_iter().collect())
} else {
Err("No accepted IP was found".to_string())
}
}

/// Creates a listening socket
/// Returns the file descriptor for it or the appropriate error
pub fn sock_listen(&self) -> VsockProxyResult<VsockListener> {
Expand Down

0 comments on commit 1dc6107

Please sign in to comment.