nitro: Write rootfs tar archive to enclave

Author: tylerfanelli

init/nitro/main.c

@@ -7,6 +7,7 @@
 #include <stdint.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 #include <unistd.h>
 #include <sys/mount.h>
 #include <sys/reboot.h>
@@ -177,23 +178,59 @@ krun_signal()
     // communication is established.
     ret = write(sock_fd, buf, 1);
     if (ret != 1) {
+        close(sock_fd);
         perror("vsock write");
         return -errno;
     }
 
     ret = read(sock_fd, buf, 1);
     if (ret != 1) {
+        close(sock_fd);
         perror("vsock read");
         return -errno;
     }
 
-    if (buf[0] != HEART_BEAT)
+    if (buf[0] != HEART_BEAT) {
+        close(sock_fd);
         return -1;
+    }
 
-    ret = close(sock_fd);
-    if (ret < 0) {
-        perror("vsock close");
-        return -errno;
+    return sock_fd;
+}
+
+int
+rootfs_rcv(int sock_fd)
+{
+    uint8_t len_b[8], *rootfs_archive;
+    ssize_t read_len;
+    uint64_t len, idx = 0;
+    int ret;
+
+    read_len = read(sock_fd, len_b, 8);
+    if (read_len < 8) {
+        close(sock_fd);
+        perror("vsock rootfs archive length read");
+        return -1;
+    }
+
+    memcpy(&len, len_b, sizeof(uint64_t));
+
+    rootfs_archive = (uint8_t *) malloc(len);
+    if (rootfs_archive == NULL) {
+        close(sock_fd);
+        perror("rootfs archive malloc");
+        return -1;
+    }
+
+    while (len) {
+        read_len = read(sock_fd, &rootfs_archive[idx], len);
+        if (read_len <= 0) {
+            close(sock_fd);
+            perror("vsock rootfs archive read");
+            return -1;
+        }
+        idx += read_len;
+        len -= read_len;
     }
 
     return 0;
@@ -202,7 +239,7 @@ krun_signal()
 int
 main(int argc, char *argv[])
 {
-    int ret;
+    int ret, sock_fd;
 
     // Block all signals.
     ret = sig_mask(SIG_BLOCK);
@@ -219,7 +256,11 @@ main(int argc, char *argv[])
     if (ret < 0)
         exit(ret);
 
-    ret = krun_signal();
+    sock_fd = krun_signal();
+    if (sock_fd < 0)
+        exit(ret);
+
+    ret = rootfs_rcv(sock_fd);
     if (ret < 0)
         exit(ret);
 

src/nitro/src/enclaves.rs

@@ -40,9 +40,11 @@ pub struct NitroEnclave {
 impl NitroEnclave {
     /// Run the enclave.
     pub fn run(&mut self) -> Result<u32> {
-        let _rootfs_archive = self.rootfs_archive()?;
+        let rootfs_archive = self.rootfs_archive()?;
 
-        let (cid, _stream) = self.start()?;
+        let (cid, stream) = self.start()?;
+
+        self.write_rootfs(rootfs_archive, stream)?;
 
         Ok(cid)
     }
@@ -124,4 +126,21 @@ impl NitroEnclave {
 
         Ok((cid, stream.0))
     }
+
+    fn write_rootfs(&self, archive: Vec<u8>, mut stream: VsockStream) -> Result<()> {
+        let len: u64 = archive
+            .len()
+            .try_into()
+            .or(Err(NitroError::RootFsTooLarge))?;
+
+        stream
+            .write_all(&len.to_ne_bytes())
+            .map_err(NitroError::RootFsLenWrite)?;
+
+        stream
+            .write_all(&archive)
+            .map_err(NitroError::RootFsWrite)?;
+
+        Ok(())
+    }
 }

src/nitro/src/error.rs

@@ -23,6 +23,9 @@ pub enum NitroError {
     VsockSetTimeout,
     VsockConnect,
     IpcWrite(io::Error),
+    RootFsLenWrite(io::Error),
+    RootFsWrite(io::Error),
+    RootFsTooLarge,
 }
 
 impl fmt::Display for NitroError {
@@ -64,6 +67,13 @@ impl fmt::Display for NitroError {
             NitroError::IpcWrite(e) => {
                 format!("unable to write enclave vsock data to UNIX IPC socket: {e}")
             }
+            NitroError::RootFsLenWrite(e) => {
+                format!("unable to write rootfs archive length to enclave: {e}")
+            }
+            NitroError::RootFsWrite(e) => {
+                format!("unable to write rootfs archive to enclave: {e}")
+            }
+            NitroError::RootFsTooLarge => "rootfs size is larger than 64 bytes".to_string(),
         };
 
         write!(f, "{}", msg)