From 1a79e0aa41787169b0e2ef51c25cd10f6f66ced2 Mon Sep 17 00:00:00 2001
From: Erdem Meydanli <meydanli@amazon.com>
Date: Wed, 13 Mar 2024 14:13:41 +0000
Subject: [PATCH] vsock_proxy: Perform DNS resolution after the expiration of
 the TTL

The proxy used to resolve the server hostname only once during its
startup. This behavior was creating issues where the DNS addresses changed
frequently. This commit changes the behavior to resolve DNS each time after
the TTL expires.

Signed-off-by: Erdem Meydanli <meydanli@amazon.com>
---
 Cargo.lock                           |  1 +
 vsock_proxy/Cargo.toml               |  1 +
 vsock_proxy/src/main.rs              |  2 +-
 vsock_proxy/src/proxy.rs             | 38 ++++++++++++++++++++--------
 vsock_proxy/tests/connection_test.rs |  2 +-
 5 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index e36991745..6e1b40153 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1974,6 +1974,7 @@ dependencies = [
 name = "vsock-proxy"
 version = "0.1.0"
 dependencies = [
+ "chrono",
  "clap",
  "dns-lookup",
  "env_logger",
diff --git a/vsock_proxy/Cargo.toml b/vsock_proxy/Cargo.toml
index 85ca11e6b..e47a4a7a7 100644
--- a/vsock_proxy/Cargo.toml
+++ b/vsock_proxy/Cargo.toml
@@ -8,6 +8,7 @@ rust-version = "1.68"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
+chrono = "0.4"
 clap = "3.2"
 dns-lookup = "2.0.3"
 env_logger = "0.10"
diff --git a/vsock_proxy/src/main.rs b/vsock_proxy/src/main.rs
index 588aa47cb..9fd78cf2b 100644
--- a/vsock_proxy/src/main.rs
+++ b/vsock_proxy/src/main.rs
@@ -114,7 +114,7 @@ fn main() -> VsockProxyResult<()> {
     let _ = check_allowlist(&remote_host, remote_port, config_file, ip_addr_type)
         .map_err(|err| format!("Error at checking the allowlist: {}", err))?;
 
-    let proxy = Proxy::new(
+    let mut proxy = Proxy::new(
         local_port,
         remote_host,
         remote_port,
diff --git a/vsock_proxy/src/proxy.rs b/vsock_proxy/src/proxy.rs
index 5e1194485..11cb6eb31 100644
--- a/vsock_proxy/src/proxy.rs
+++ b/vsock_proxy/src/proxy.rs
@@ -4,7 +4,7 @@
 
 /// Contains code for Proxy, a library used for translating vsock traffic to
 /// TCP traffic
-///
+use chrono::{DateTime, Duration, Utc};
 use log::info;
 use nix::sys::select::{select, FdSet};
 use nix::sys::socket::SockType;
@@ -82,8 +82,11 @@ pub struct Proxy {
     remote_host: String,
     remote_addr: Option<IpAddr>,
     remote_port: u16,
+    dns_resolve_date: Option<DateTime<Utc>>,
+    dns_refresh_interval: Option<Duration>,
     pool: ThreadPool,
     sock_type: SockType,
+    ip_addr_type: IpAddrType
 }
 
 impl Proxy {
@@ -96,22 +99,20 @@ impl Proxy {
     ) -> VsockProxyResult<Self> {
         let pool = ThreadPool::new(num_workers);
         let sock_type = SockType::Stream;
-
-        let dns_result = dns::resolve_single(&remote_host, ip_addr_type)?;
-        let remote_addr: Option<IpAddr> = Some(dns_result.ip);
-
-        info!(
-            "Using IP \"{:?}\" for the given server \"{}\"",
-            dns_result.ip, remote_host
-        );
+        let remote_addr: Option<IpAddr> = None;
+        let dns_resolve_date: Option<DateTime<Utc>> = None;
+        let dns_refresh_interval: Option<Duration> = None;
 
         Ok(Proxy {
             local_port,
             remote_host,
             remote_addr,
             remote_port,
+            dns_resolve_date,
+            dns_refresh_interval,
             pool,
             sock_type,
+            ip_addr_type
         })
     }
 
@@ -129,12 +130,29 @@ impl Proxy {
     /// Accepts an incoming connection coming on listener and handles it on a
     /// different thread
     /// Returns the handle for the new thread or the appropriate error
-    pub fn sock_accept(&self, listener: &VsockListener) -> VsockProxyResult<()> {
+    pub fn sock_accept(&mut self, listener: &VsockListener) -> VsockProxyResult<()> {
         let (mut client, client_addr) = listener
             .accept()
             .map_err(|_| "Could not accept connection")?;
         info!("Accepted connection on {:?}", client_addr);
 
+        let needs_resolve = |d: DateTime<Utc>, i: Duration| {
+            (Utc::now() - d + Duration::seconds(2)) > i
+        };
+
+        if self.dns_resolve_date == None || needs_resolve(self.dns_resolve_date.unwrap(), self.dns_refresh_interval.unwrap()) {
+            info!("Resolving hostname: {}.", self.remote_host);
+            let result = dns::resolve_single(&self.remote_host, self.ip_addr_type)?;
+            self.dns_resolve_date = Some(Utc::now());
+            self.dns_refresh_interval = Some(Duration::seconds(result.ttl as i64));
+            self.remote_addr = Some(result.ip);
+
+            info!(
+                "Using IP \"{:?}\" for the given server \"{}\". (TTL: {} secs)",
+                result.ip, self.remote_host, result.ttl
+            );
+        }
+
         let sockaddr = SocketAddr::new(self.remote_addr.unwrap(), self.remote_port);
         let sock_type = self.sock_type;
         self.pool.execute(move || {
diff --git a/vsock_proxy/tests/connection_test.rs b/vsock_proxy/tests/connection_test.rs
index 303750b9a..8614c878f 100644
--- a/vsock_proxy/tests/connection_test.rs
+++ b/vsock_proxy/tests/connection_test.rs
@@ -29,7 +29,7 @@ fn test_tcp_connection() {
             - {address: 127.0.0.1, port: 9000}",
     )
     .unwrap();
-    let proxy = Proxy::new(
+    let mut proxy = Proxy::new(
         vsock_proxy::proxy::VSOCK_PROXY_PORT,
         addr,
         9000,