nitro: Refactor, separate different functions

Signed-off-by: Tyler Fanelli <tfanelli@redhat.com>

Author: tylerfanelli

src/nitro/src/enclaves.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::error::NitroError;
+use libc::c_int;
 use nitro_enclaves::{
     launch::{ImageType, Launcher, MemoryInfo, PollTimeout, StartFlags},
     Device,
@@ -12,7 +13,7 @@ use std::{
     os::fd::AsFd,
     path::PathBuf,
 };
-use vsock::{VsockAddr, VsockListener};
+use vsock::{VsockAddr, VsockListener, VsockStream};
 
 type Result<T> = std::result::Result<T, NitroError>;
 
@@ -41,6 +42,22 @@ impl NitroEnclave {
     pub fn run(&mut self) -> Result<u32> {
         let _rootfs_archive = self.rootfs_archive()?;
 
+        let (cid, _stream) = self.start()?;
+
+        Ok(cid)
+    }
+
+    fn rootfs_archive(&self) -> Result<Vec<u8>> {
+        let mut builder = tar::Builder::new(Vec::new());
+
+        builder
+            .append_dir_all("rootfs", self.rootfs.clone())
+            .unwrap();
+
+        Ok(builder.into_inner().unwrap())
+    }
+
+    fn launch(&mut self) -> Result<u32> {
         let device = Device::open().map_err(NitroError::DeviceOpen)?;
 
         let mut launcher = Launcher::new(&device).map_err(NitroError::VmCreate)?;
@@ -52,61 +69,59 @@ impl NitroEnclave {
             launcher.add_vcpu(None).map_err(NitroError::VcpuAdd)?;
         }
 
-        let sockaddr = VsockAddr::new(VMADDR_CID_PARENT, ENCLAVE_READY_VSOCK_PORT);
-        let listener = VsockListener::bind(&sockaddr).map_err(NitroError::HeartbeatBind)?;
-
         let cid = launcher
             .start(self.start_flags, None)
             .map_err(NitroError::VmStart)?;
 
-        // Safe to unwrap.
-        let cid: u32 = cid.try_into().unwrap();
+        Ok(cid.try_into().unwrap()) // Safe to unwrap.
+    }
 
+    fn poll(&self, listener: &VsockListener) -> Result<()> {
         let poll_timeout = PollTimeout::try_from((&self.image, self.mem_size_mib << 20))
             .map_err(NitroError::PollTimeoutCalculate)?;
 
-        enclave_check(listener, poll_timeout.into(), cid)?;
+        let mut poll_fds = [PollFd::new(listener.as_fd(), PollFlags::POLLIN)];
+        let result = poll(
+            &mut poll_fds,
+            NixPollTimeout::from(c_int::from(poll_timeout) as u16),
+        );
 
-        Ok(cid)
-    }
-
-    fn rootfs_archive(&self) -> Result<Vec<u8>> {
-        let mut builder = tar::Builder::new(Vec::new());
+        if result == Ok(0) {
+            return Err(NitroError::PollNoSelectedEvents);
+        }
 
-        builder
-            .append_dir_all("rootfs", self.rootfs.clone())
-            .unwrap();
+        if result != Ok(1) {
+            return Err(NitroError::PollMoreThanOneSelectedEvent);
+        }
 
-        Ok(builder.into_inner().unwrap())
+        Ok(())
     }
-}
 
-fn enclave_check(listener: VsockListener, poll_timeout_ms: libc::c_int, cid: u32) -> Result<()> {
-    let mut poll_fds = [PollFd::new(listener.as_fd(), PollFlags::POLLIN)];
-    let result = poll(&mut poll_fds, NixPollTimeout::from(poll_timeout_ms as u16));
-    if result == Ok(0) {
-        return Err(NitroError::PollNoSelectedEvents);
-    } else if result != Ok(1) {
-        return Err(NitroError::PollMoreThanOneSelectedEvent);
-    }
+    fn start(&mut self) -> Result<(u32, VsockStream)> {
+        let sockaddr = VsockAddr::new(VMADDR_CID_PARENT, ENCLAVE_READY_VSOCK_PORT);
+        let listener = VsockListener::bind(&sockaddr).map_err(NitroError::HeartbeatBind)?;
 
-    let mut stream = listener.accept().map_err(NitroError::HeartbeatAccept)?;
+        let cid = self.launch()?;
+        self.poll(&listener)?;
 
-    let mut buf = [0u8];
-    let bytes = stream.0.read(&mut buf).map_err(NitroError::HeartbeatRead)?;
+        let mut stream = listener.accept().map_err(NitroError::HeartbeatAccept)?;
 
-    if bytes != 1 || buf[0] != HEART_BEAT {
-        return Err(NitroError::EnclaveHeartbeatNotDetected);
-    }
+        if stream.1.cid() != cid {
+            return Err(NitroError::HeartbeatCidMismatch);
+        }
 
-    stream
-        .0
-        .write_all(&buf)
-        .map_err(NitroError::HeartbeatWrite)?;
+        let mut buf = [0u8];
+        let bytes = stream.0.read(&mut buf).map_err(NitroError::HeartbeatRead)?;
 
-    if stream.1.cid() != cid {
-        return Err(NitroError::HeartbeatCidMismatch);
-    }
+        if bytes != 1 || buf[0] != HEART_BEAT {
+            return Err(NitroError::EnclaveHeartbeatNotDetected);
+        }
 
-    Ok(())
+        stream
+            .0
+            .write_all(&buf)
+            .map_err(NitroError::HeartbeatWrite)?;
+
+        Ok((cid, stream.0))
+    }
 }