From 75669f6d1eee2112de3f33c246a59a97d36c9881 Mon Sep 17 00:00:00 2001
From: Erdem Meydanli <meydanli@amazon.com>
Date: Wed, 13 Mar 2024 13:17:37 +0000
Subject: [PATCH] vsock_proxy: Handle allowlisting out of Proxy

Perform allowlisted hosts' check before creating a Proxy instance and
terminate the application if necessary.

Signed-off-by: Erdem Meydanli <meydanli@amazon.com>
---
 vsock_proxy/src/lib.rs               |  4 ++--
 vsock_proxy/src/main.rs              | 13 ++++++++++---
 vsock_proxy/src/proxy.rs             | 23 +++++++++++------------
 vsock_proxy/tests/connection_test.rs |  3 +--
 4 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/vsock_proxy/src/lib.rs b/vsock_proxy/src/lib.rs
index 88d18a9c4..f9c8c0a8d 100644
--- a/vsock_proxy/src/lib.rs
+++ b/vsock_proxy/src/lib.rs
@@ -18,9 +18,9 @@ pub enum IpAddrType {
 
 pub struct DnsResolveResult {
     ///Resolved address
-    ip: IpAddr,
+    pub ip: IpAddr,
     ///DNS TTL value
-    ttl: u32
+    pub ttl: u32
 }
 
 /// The most common result type provided by VsockProxy operations.
diff --git a/vsock_proxy/src/main.rs b/vsock_proxy/src/main.rs
index 44f524ac0..588aa47cb 100644
--- a/vsock_proxy/src/main.rs
+++ b/vsock_proxy/src/main.rs
@@ -10,7 +10,7 @@ use clap::{App, AppSettings, Arg};
 use env_logger::init;
 use log::info;
 
-use vsock_proxy::{proxy::Proxy, IpAddrType, VsockProxyResult};
+use vsock_proxy::{proxy::{check_allowlist, Proxy}, IpAddrType, VsockProxyResult};
 
 fn main() -> VsockProxyResult<()> {
     init();
@@ -104,14 +104,21 @@ fn main() -> VsockProxyResult<()> {
         .parse::<usize>()
         .map_err(|_| "Number of workers is not valid")?;
 
+    if num_workers == 0 {
+        return Err("Number of workers must not be 0".to_string());
+    }
+
+    info!("Checking allowlist configuration");
     let config_file = matches.value_of("config_file");
+    let remote_host = String::from(remote_addr);
+    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(
         local_port,
-        remote_addr,
+        remote_host,
         remote_port,
         num_workers,
-        config_file,
         ip_addr_type
     )
     .map_err(|err| format!("Could not create proxy: {}", err))?;
diff --git a/vsock_proxy/src/proxy.rs b/vsock_proxy/src/proxy.rs
index ed9ea8ae9..5e1194485 100644
--- a/vsock_proxy/src/proxy.rs
+++ b/vsock_proxy/src/proxy.rs
@@ -1,6 +1,6 @@
 // Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
-#![deny(warnings)]
+//#![deny(warnings)]
 
 /// Contains code for Proxy, a library used for translating vsock traffic to
 /// TCP traffic
@@ -79,7 +79,8 @@ pub fn check_allowlist(
 /// Configuration parameters for port listening and remote destination
 pub struct Proxy {
     local_port: u32,
-    remote_addr: IpAddr,
+    remote_host: String,
+    remote_addr: Option<IpAddr>,
     remote_port: u16,
     pool: ThreadPool,
     sock_type: SockType,
@@ -88,27 +89,25 @@ pub struct Proxy {
 impl Proxy {
     pub fn new(
         local_port: u32,
-        remote_host: &str,
+        remote_host: String,
         remote_port: u16,
         num_workers: usize,
-        config_file: Option<&str>,
         ip_addr_type: IpAddrType
     ) -> VsockProxyResult<Self> {
-        if num_workers == 0 {
-            return Err("Number of workers must not be 0".to_string());
-        }
-        info!("Checking allowlist configuration");
-        let remote_addr = check_allowlist(remote_host, remote_port, config_file, ip_addr_type)
-            .map_err(|err| format!("Error at checking the allowlist: {}", err))?;
         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 \"{}\"",
-            remote_addr, remote_host
+            dns_result.ip, remote_host
         );
+
         Ok(Proxy {
             local_port,
+            remote_host,
             remote_addr,
             remote_port,
             pool,
@@ -136,7 +135,7 @@ impl Proxy {
             .map_err(|_| "Could not accept connection")?;
         info!("Accepted connection on {:?}", client_addr);
 
-        let sockaddr = SocketAddr::new(self.remote_addr, self.remote_port);
+        let sockaddr = SocketAddr::new(self.remote_addr.unwrap(), self.remote_port);
         let sock_type = self.sock_type;
         self.pool.execute(move || {
             let mut server = match sock_type {
diff --git a/vsock_proxy/tests/connection_test.rs b/vsock_proxy/tests/connection_test.rs
index f333cc35d..303750b9a 100644
--- a/vsock_proxy/tests/connection_test.rs
+++ b/vsock_proxy/tests/connection_test.rs
@@ -31,10 +31,9 @@ fn test_tcp_connection() {
     .unwrap();
     let proxy = Proxy::new(
         vsock_proxy::proxy::VSOCK_PROXY_PORT,
-        &addr,
+        addr,
         9000,
         2,
-        file.path().to_str(),
         IpAddrType::IPAddrMixed,
     )
     .unwrap();