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 <[email protected]>
meerd · Apr 9, 2024 · dfaac56 · dfaac56
1 parent 75669f6
commit dfaac56
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 13 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git 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
@@ -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
@@ -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 {
@@ -92,26 +95,24 @@ impl Proxy {
         remote_host: String,
         remote_port: u16,
         num_workers: usize,
-        ip_addr_type: IpAddrType
+        ip_addr_type: IpAddrType,
     ) -> 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,33 @@ 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.is_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
@@ -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,