Refactoring proxy threads

Author: raoulstrackx

fortanix-vme/fortanix-vme-runner/src/lib.rs

@@ -8,7 +8,7 @@ use std::thread::{self, JoinHandle};
 use std::io::{self, Error as IoError, ErrorKind as IoErrorKind, Read, Write};
 use std::net::{Shutdown, TcpListener, TcpStream};
 use std::os::unix::io::AsRawFd;
-use std::sync::{Arc, Mutex};
+use std::sync::{Arc, Mutex, RwLock};
 use fortanix_vme_abi::{self, Addr, Response, Request};
 use vsock::{self, SockAddr as VsockAddr, Std, Vsock, VsockListener, VsockStream};
 
@@ -91,23 +91,56 @@ impl Listener {
 
 #[derive(Debug)]
 struct Connection {
-    // Preliminary work for PLAT-367
-    #[allow(dead_code)]
-    remote: Addr,
-    // Preliminary work for PLAT-367
-    #[allow(dead_code)]
-    runner: Addr,
+    tcp_stream: TcpStream,
+    vsock_stream: VsockStream<Std>,
+    remote_name: String,
 }
 
 impl Connection {
-    fn from_tcp_stream(stream: &TcpStream) -> Self {
-        let tcp_remote = stream.peer_addr().unwrap().into();
-        let tcp_runner = stream.local_addr().unwrap().into();
+    fn new(vsock_stream: VsockStream<Std>, tcp_stream: TcpStream, remote_name: String) -> Self {
         Connection {
-            remote: tcp_remote,
-            runner: tcp_runner,
+            tcp_stream,
+            vsock_stream,
+            remote_name,
         }
     }
+
+    fn close(&self) {
+        let _ = self.tcp_stream.shutdown(Shutdown::Both);
+        let _ = self.vsock_stream.shutdown(Shutdown::Both);
+    }
+
+    /// Exchanges messages between the remote server and enclave. Returns `true` when the
+    /// connection should remain active, false otherwise
+    fn proxy(&mut self) -> bool {
+        fn exchange<S: StreamConnection, D: StreamConnection>(src: &mut S, src_name: &str, dst: &mut D, dst_name: &str) -> bool {
+            // According to the `Read` threat documentation, reading 0 bytes
+            // indicates that the connection has been shutdown correctly. So we
+            // close the proxy service
+            // https://doc.rust-lang.org/std/io/trait.Read.html#tymethod.read
+            match Server::transfer_data(src, src_name, dst, dst_name) {
+                Ok(n) if n == 0 => false,
+                Ok(_)           => true,
+                Err(_)          => false,
+            }
+        }
+        let remote = self.tcp_stream.as_raw_fd();
+        let proxy = self.vsock_stream.as_raw_fd();
+
+        let mut read_set = FdSet::new();
+        read_set.insert(remote);
+        read_set.insert(proxy);
+
+        if let Ok(_num) = select(None, Some(&mut read_set), None, None, None) {
+            if read_set.contains(remote) {
+                return exchange(&mut self.tcp_stream, &self.remote_name, &mut self.vsock_stream, "proxy");
+            }
+            if read_set.contains(proxy) {
+                return exchange(&mut self.vsock_stream, "proxy", &mut self.tcp_stream, &self.remote_name);
+            }
+        }
+        true
+    }
 }
 
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
@@ -146,8 +179,8 @@ pub struct Server {
     /// When the enclave instructs to accept a new connection, the runner accepts a new TCP
     /// connection. It then locates the ListenerInfo and finds the information it needs to set up a
     /// new vsock connection to the enclave
-    listeners: Mutex<FnvHashMap<VsockAddr, Arc<Mutex<Listener>>>>,
-    connections: Mutex<FnvHashMap<ConnectionKey, Arc<Mutex<Connection>>>>,
+    listeners: RwLock<FnvHashMap<VsockAddr, Arc<Mutex<Listener>>>>,
+    connections: RwLock<FnvHashMap<ConnectionKey, Arc<Mutex<Connection>>>>,
 }
 
 impl Server {
@@ -230,9 +263,9 @@ impl Server {
      *  [2] remote
      *  [3] proxy
      */
-    fn handle_request_connect(&self, remote_addr: &String, enclave: &mut VsockStream) -> Result<(), IoError> {
+    fn handle_request_connect(server: Arc<Self>, remote_addr: &String, enclave: &mut VsockStream) -> Result<(), IoError> {
         // Connect to remote server
-        let mut remote_socket = TcpStream::connect(remote_addr)?;
+        let remote_socket = TcpStream::connect(remote_addr)?;
         let remote_name = remote_addr.split_terminator(":").next().unwrap_or(remote_addr);
 
         // Create listening socket that the enclave can connect to
@@ -257,50 +290,37 @@ impl Server {
         Self::send(enclave, &response)?;
 
         // Wait for incoming connection from enclave
-        let (mut proxy, _proxy_addr) = proxy_server.accept()?;
+        let (proxy, _proxy_addr) = proxy_server.accept()?;
 
         // Store connection info
-        let k = self.add_connection(&proxy, &remote_socket);
+        server.add_connection(proxy, remote_socket, remote_name.to_string());
 
-        // Pass messages between remote server <-> enclave
-        Self::proxy_connection((&mut remote_socket, remote_name), (&mut proxy, "proxy"));
-
-        // Remove connection info
-        self.remove_connection(&k);
         Ok(())
     }
 
     fn add_listener(&self, addr: VsockAddr, info: Listener) {
-        self.listeners.lock().unwrap().insert(addr, Arc::new(Mutex::new(info)));
+        self.listeners.write().unwrap().insert(addr, Arc::new(Mutex::new(info)));
     }
 
     fn listener(&self, addr: &VsockAddr) -> Option<Arc<Mutex<Listener>>> {
-        self.listeners.lock().unwrap().get(&addr).cloned()
+        self.listeners.read().unwrap().get(&addr).cloned()
     }
 
     // Preliminary work for PLAT-367
     #[allow(dead_code)]
     fn connection(&self, enclave: VsockAddr, runner: VsockAddr) -> Option<Arc<Mutex<Connection>>> {
         let k = ConnectionKey::from_addresses(enclave, runner);
         self.connections
-            .lock()
+            .read()
             .unwrap()
             .get(&k)
             .cloned()
     }
 
-    fn add_connection(&self, runner_enclave: &VsockStream<Std>, runner_remote: &TcpStream) -> ConnectionKey {
-        let k = ConnectionKey::from_vsock_stream(runner_enclave);
-        let info = Connection::from_tcp_stream(runner_remote);
-        self.connections.lock().unwrap().insert(k.clone(), Arc::new(Mutex::new(info)));
-        k
-    }
-
-    fn remove_connection(&self, k: &ConnectionKey) {
-        self.connections
-            .lock()
-            .unwrap()
-            .remove(&k);
+    fn add_connection(&self, runner_enclave: VsockStream<Std>, runner_remote: TcpStream, remote_name: String) {
+        let k = ConnectionKey::from_vsock_stream(&runner_enclave);
+        let info = Connection::new(runner_enclave, runner_remote, remote_name);
+        self.connections.write().unwrap().insert(k.clone(), Arc::new(Mutex::new(info)));
     }
 
     /*
@@ -326,11 +346,11 @@ impl Server {
      *  runner
      *  `enclave`: The runner-enclave vsock connection
      */
-    fn handle_request_bind(&self, addr: &String, enclave_port: u32, enclave: &mut VsockStream) -> Result<(), IoError> {
+    fn handle_request_bind(server: Arc<Self>, addr: &String, enclave_port: u32, enclave: &mut VsockStream) -> Result<(), IoError> {
         let cid: u32 = enclave.peer().unwrap().parse().unwrap_or(vsock::VMADDR_CID_HYPERVISOR);
         let listener = TcpListener::bind(addr)?;
         let local: Addr = listener.local_addr()?.into();
-        self.add_listener(VsockAddr::new(cid, enclave_port), Listener::new(listener));
+        server.add_listener(VsockAddr::new(cid, enclave_port), Listener::new(listener));
         let response = Response::Bound{ local };
         Self::log_communication(
             "runner",
@@ -344,15 +364,15 @@ impl Server {
         Ok(())
     }
 
-    fn handle_request_accept(&self, vsock_listener_port: u32, enclave: &mut VsockStream) -> Result<(), IoError> {
+    fn handle_request_accept(server: Arc<Self>, vsock_listener_port: u32, enclave: &mut VsockStream) -> Result<(), IoError> {
         let enclave_cid: u32 = enclave.peer().unwrap().parse().unwrap_or(vsock::VMADDR_CID_HYPERVISOR);
         let enclave_addr = VsockAddr::new(enclave_cid, vsock_listener_port);
-        let listener = self.listener(&enclave_addr)
+        let listener = server.listener(&enclave_addr)
             .ok_or(IoError::new(IoErrorKind::InvalidInput, "Information about provided file descriptor was not found"))?;
         let listener = listener.lock().unwrap();
 
         match listener.listener.accept() {
-            Ok((mut conn, peer)) => {
+            Ok((conn, peer)) => {
                 let vsock = Vsock::new::<Std>()?;
                 let runner_addr = vsock.addr::<Std>()?;
                 let response = Response::IncomingConnection{
@@ -369,62 +389,55 @@ impl Server {
                     Direction::Right,
                     "vsock");
                 enclave.write(&serde_cbor::ser::to_vec(&response).unwrap())?;
-                let _ = thread::Builder::new().spawn(move || {
-                    let mut proxy = vsock.connect_with_cid_port(enclave_addr.cid(), enclave_addr.port()).unwrap();
-                    //let k = self.add_connection(&proxy, &conn);
-                    Self::proxy_connection((&mut conn, "remote"), (&mut proxy, "proxy"));
-                    //self.remove_connection(&k);
-                });
+
+                let proxy = vsock.connect_with_cid_port(enclave_addr.cid(), enclave_addr.port()).unwrap();
+                server.add_connection(proxy, conn, "remote".to_string());
+
                 Ok(())
             },
             Err(e) => Err(e),
         }
     }
 
-    fn proxy_connection(remote: (&mut TcpStream, &str), proxy: (&mut VsockStream, &str)) {
+    fn proxy_connections(server: Arc<Server>) {
+        let mut closed_connections = Vec::new();
+
         loop {
-            let mut read_set = FdSet::new();
-            read_set.insert(remote.0.as_raw_fd());
-            read_set.insert(proxy.0.as_raw_fd());
-
-            if let Ok(_num) = select(None, Some(&mut read_set), None, None, None) {
-                if read_set.contains(remote.0.as_raw_fd()) {
-                    match Self::transfer_data(remote.0, remote.1, proxy.0, proxy.1) {
-                        Ok(0)  => {
-                            // According to the `Read` threat documentation, reading 0 bytes
-                            // indicates that the connection has been shutdown correctly. So we
-                            // close the proxy service
-                            // https://doc.rust-lang.org/std/io/trait.Read.html#tymethod.read
-                            break
-                        },
-                        Ok(_)  => (),
-                        Err(e) => {
-                            eprintln!("transfer from remote failed: {:?}", e);
-                            break;
-                        }
+            // Exchange messages on every proxy connection
+            // TODO: Store connections as a linked hash map so we don't need to keep a read lock
+            // over the HashMap while every connection is serviced
+            if let Ok(connections) = server.connections.read() {
+                for (key, connection) in connections.iter() {
+                    match connection.try_lock() {
+                        Ok(mut connection) => if !connection.proxy() {
+                                            connection.close();
+                                            closed_connections.push(key.clone());
+                                          }
+                        Err(_) => (),
                     }
                 }
-                if read_set.contains(proxy.0.as_raw_fd()) {
-                    match Self::transfer_data(proxy.0, proxy.1, remote.0, remote.1) {
-                        Ok(0)  => break,
-                        Ok(_)  => (),
-                        Err(e) => {
-                            eprintln!("transfer from proxy failed: {:?}", e);
-                            break;
-                        }
-                    }
+            }
+
+            // Remove closed connections
+            let mut num_connections = None;
+            if let Ok(mut connections) = server.connections.try_write() {
+                while let Some(k) = closed_connections.pop() {
+                    connections.remove(&k);
                 }
+                num_connections = Some(connections.len());
+            }
+
+            if num_connections == Some(0) {
+                thread::yield_now();
             }
         }
-        let _ = proxy.0.shutdown(Shutdown::Both);
-        let _ = remote.0.shutdown(Shutdown::Both);
     }
 
-    fn handle_client(&self, stream: &mut VsockStream) -> Result<(), IoError> {
+    fn handle_client(server: Arc<Self>, stream: &mut VsockStream) -> Result<(), IoError> {
         match Self::read_request(stream) {
-            Ok(Request::Connect{ addr })             => self.handle_request_connect(&addr, stream)?,
-            Ok(Request::Bind{ addr, enclave_port })  => self.handle_request_bind(&addr, enclave_port, stream)?,
-            Ok(Request::Accept{ enclave_port })      => self.handle_request_accept(enclave_port, stream)?,
+            Ok(Request::Connect{ addr })             => Self::handle_request_connect(server, &addr, stream)?,
+            Ok(Request::Bind{ addr, enclave_port })  => Self::handle_request_bind(server, &addr, enclave_port, stream)?,
+            Ok(Request::Accept{ enclave_port })      => Self::handle_request_accept(server, enclave_port, stream)?,
             Err(_e)                                  => return Err(IoError::new(IoErrorKind::InvalidData, "Failed to read request")),
         };
         Ok(())
@@ -434,34 +447,45 @@ impl Server {
         let command_listener = VsockListener::<Std>::bind_with_cid_port(vsock::VMADDR_CID_ANY, port)?;
         Ok(Server {
             command_listener: Mutex::new(command_listener),
-            listeners: Mutex::new(FnvHashMap::default()),
-            connections: Mutex::new(FnvHashMap::default()),
+            listeners: RwLock::new(FnvHashMap::default()),
+            connections: RwLock::new(FnvHashMap::default()),
+        })
+    }
+
+    fn start_proxy_server(server: Arc<Server>) -> Result<JoinHandle<()>, IoError> {
+        thread::Builder::new()
+            .spawn(move || {
+                Self::proxy_connections(server);
+            })
+    }
+
+    fn start_command_server(server: Arc<Server>) -> Result<JoinHandle<()>, IoError> {
+        thread::Builder::new().spawn(move || {
+            let command_listener = server.command_listener.lock().unwrap();
+            for stream in command_listener.incoming() {
+            let server = server.clone();
+            let _ = thread::Builder::new()
+                .spawn(move || {
+                    let mut stream = stream.unwrap();
+                    if let Err(e) = Self::handle_client(server, &mut stream) {
+                        eprintln!("Error handling connection: {}, shutting connection down", e);
+                        let _ = stream.shutdown(Shutdown::Both);
+                    }
+                });
+            }
         })
     }
 
-    pub fn run(port: u32) -> std::io::Result<(JoinHandle<()>, u32)> {
+    pub fn run(port: u32) -> std::io::Result<JoinHandle<()>> {
         println!("Starting enclave runner.");
         let server = Arc::new(Self::bind(port)?);
         let port = server.command_listener.lock().unwrap().local_addr()?.port();
         println!("Listening on vsock port {}...", port);
 
-        let handle = thread::Builder::new().spawn(move || {
-            let server = server;
-            let server = server.clone();
-            let command_listener = server.command_listener.lock().unwrap();
-            for stream in command_listener.incoming() {
-                let server = server.clone();
-                let _ = thread::Builder::new()
-                    .spawn(move || {
-                        let mut stream = stream.unwrap();
-                        if let Err(e) = server.handle_client(&mut stream) {
-                            eprintln!("Error handling connection: {}, shutting connection down", e);
-                            let _ = stream.shutdown(Shutdown::Both);
-                        }
-                    });
-            }
-        })?;
-        Ok((handle, port))
+        Server::start_proxy_server(server.clone())?;
+        let handle = Server::start_command_server(server.clone())?;
+
+        Ok(handle)
     }
 }
 

fortanix-vme/fortanix-vme-runner/src/main.rs

@@ -4,7 +4,7 @@ use std::io::ErrorKind;
 
 fn main() {
     match Server::run(SERVER_PORT) {
-        Ok((server_thread, _port))                   => server_thread.join().expect("Server panicked"),
+        Ok(handle)                                   => { handle.join().unwrap(); },
         Err(e) if e.kind() == ErrorKind::AddrInUse   => println!("Server failed. Do you already have a runner running on vsock port {}? (Error: {:?})", SERVER_PORT, e),
         Err(e)                                       => println!("Server failed. Error: {:?}", e),
     }