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..59d2952a0 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, remote_port: u16, + dns_resolve_date: Option>, + dns_refresh_interval: Option, 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 { 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 = Some(dns_result.ip); - - info!( - "Using IP \"{:?}\" for the given server \"{}\"", - dns_result.ip, remote_host - ); + let remote_addr: Option = None; + let dns_resolve_date: Option> = None; + let dns_refresh_interval: Option = 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, 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 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,