Refactoring keeping track of TcpListeners

Author: raoulstrackx

Cargo.lock

@@ -19,11 +19,13 @@ pub enum Request {
     Bind {
         /// The address the listen to in the parent VM
         addr: String,
-        /// The port the enclave is listening on to receive connections from the parent VM
+        /// The port the enclave is listening on to receive connections from the parent VM. This
+        /// port will also be used to reference the connection
         enclave_port: u32,
     },
     Accept {
-        fd: i32,
+        /// The Vsock port the enclave is listening on
+        enclave_port: u32,
     }
 }
 
@@ -76,9 +78,6 @@ pub enum Response {
     Bound {
         /// The local TCP address the parent VM is listening on
         local: Addr,
-        /// The id used to identify the listener. It can be used for subsequent calls (e.g., to
-        /// accept new incoming connections)
-        fd: i32,
     },
     IncomingConnection {
         /// The local address (as used by the runner)

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

@@ -8,10 +8,9 @@ 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::os::unix::prelude::RawFd;
 use std::sync::{Arc, Mutex};
-use fortanix_vme_abi::{self, Response, Request};
-use vsock::{self, Std, Vsock, VsockListener, VsockStream};
+use fortanix_vme_abi::{self, Addr, Response, Request};
+use vsock::{self, SockAddr as VsockAddr, Std, Vsock, VsockListener, VsockStream};
 
 const PROXY_BUFF_SIZE: usize = 4192;
 
@@ -80,10 +79,14 @@ impl StreamConnection for VsockStream {
     }
 }
 
-struct ListenerInfo {
+struct Listener {
     listener: TcpListener,
-    enclave_cid: u32,
-    enclave_port: u32,
+}
+
+impl Listener {
+    fn new(listener: TcpListener) -> Self {
+        Listener{ listener }
+    }
 }
 
 pub struct Server {
@@ -93,7 +96,7 @@ 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<RawFd, Arc<Mutex<ListenerInfo>>>>,
+    listeners: Mutex<FnvHashMap<VsockAddr, Arc<Mutex<Listener>>>>,
 }
 
 impl Server {
@@ -210,14 +213,12 @@ impl Server {
         Ok(())
     }
 
-    fn add_listener_info(&self, info: ListenerInfo) -> RawFd {
-        let fd = info.listener.as_raw_fd();
-        self.listeners.lock().unwrap().insert(fd, Arc::new(Mutex::new(info)));
-        fd
+    fn add_listener(&self, addr: VsockAddr, info: Listener) {
+        self.listeners.lock().unwrap().insert(addr, Arc::new(Mutex::new(info)));
     }
 
-    fn listener_info(&self, fd: &RawFd) -> Option<Arc<Mutex<ListenerInfo>>> {
-        self.listeners.lock().unwrap().get(&fd).cloned()
+    fn listener(&self, addr: &VsockAddr) -> Option<Arc<Mutex<Listener>>> {
+        self.listeners.lock().unwrap().get(&addr).cloned()
     }
 
     /*
@@ -246,9 +247,9 @@ impl Server {
     fn handle_request_bind(&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 = listener.local_addr()?.into();
-        let fd = self.add_listener_info(ListenerInfo{ listener, enclave_cid: cid, enclave_port });
-        let response = Response::Bound{ local, fd };
+        let local: Addr = listener.local_addr()?.into();
+        self.add_listener(VsockAddr::new(cid, enclave_port), Listener::new(listener));
+        let response = Response::Bound{ local };
         Self::log_communication(
             "runner",
             enclave.local_port().unwrap_or_default(),
@@ -261,12 +262,14 @@ impl Server {
         Ok(())
     }
 
-    fn handle_request_accept(&self, fd: RawFd, enclave: &mut VsockStream) -> Result<(), IoError> {
-        let listener_info = self.listener_info(&fd)
+    fn handle_request_accept(&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)
             .ok_or(IoError::new(IoErrorKind::InvalidInput, "Information about provided file descriptor was not found"))?;
-        let listener_info = listener_info.lock().unwrap();
-        let (cid, port) = (listener_info.enclave_cid, listener_info.enclave_port);
-        match listener_info.listener.accept() {
+        let listener = listener.lock().unwrap();
+
+        match listener.listener.accept() {
             Ok((mut conn, peer)) => {
                 let vsock = Vsock::new::<Std>()?;
                 let response = Response::IncomingConnection{
@@ -284,7 +287,7 @@ impl Server {
                     "vsock");
                 enclave.write(&serde_cbor::ser::to_vec(&response).unwrap())?;
                 let _ = thread::Builder::new().spawn(move || {
-                    let mut proxy = vsock.connect_with_cid_port(cid, port).unwrap();
+                    let mut proxy = vsock.connect_with_cid_port(enclave_addr.cid(), enclave_addr.port()).unwrap();
                     Self::proxy_connection((&mut conn, "remote"), (&mut proxy, "proxy"));
                 });
                 Ok(())
@@ -336,7 +339,7 @@ impl Server {
         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{ fd })                => self.handle_request_accept(fd, stream)?,
+            Ok(Request::Accept{ enclave_port })      => self.handle_request_accept(enclave_port, stream)?,
             Err(_e)                                  => return Err(IoError::new(IoErrorKind::InvalidData, "Failed to read request")),
         };
         Ok(())