aws-nitro: Separate guest kernel module loading, updates

Makefile

@@ -23,7 +23,8 @@ AWS_NITRO_INIT_SRC = \
         init/aws-nitro/archive.c				\
         init/aws-nitro/args_reader.c			\
         init/aws-nitro/fs.c				\
-        init/aws-nitro/device/include/*			\
+        init/aws-nitro/mod.c					\
+		init/aws-nitro/device/include/*			\
 		init/aws-nitro/device/app_stdio_output.c	\
 		init/aws-nitro/device/device.c              \
 		init/aws-nitro/device/net_tap_afvsock.c	\
@@ -140,7 +141,7 @@ endif
 
 AWS_NITRO_INIT_BINARY= init/aws-nitro/init
 $(AWS_NITRO_INIT_BINARY): $(AWS_NITRO_INIT_SRC)
-	$(CC) -O2 -static -Wall $(AWS_NITRO_INIT_LD_FLAGS) -o $@ $(AWS_NITRO_INIT_SRC) $(AWS_NITRO_INIT_LD_FLAGS)
+	$(CC) -O2 -static -s -Wall $(AWS_NITRO_INIT_LD_FLAGS) -o $@ $(AWS_NITRO_INIT_SRC) $(AWS_NITRO_INIT_LD_FLAGS)
 
 # Sysroot preparation rules for cross-compilation on macOS
 DEBIAN_PACKAGES = libc6 libc6-dev libgcc-12-dev linux-libc-dev

init/aws-nitro/args_reader.c

@@ -25,7 +25,7 @@ enum {
     ENCLAVE_ARG_ID_EXEC_ARGV,
     ENCLAVE_ARG_ID_EXEC_ENVP,
     ENCLAVE_ARG_ID_NETWORK_PROXY,
-    ENCLAVE_ARG_ID_DEBUG,
+    ENCLAVE_ARG_ID_APP_OUTPUT,
 
     ENCLAVE_ARGS_FINISHED = 255,
 };
@@ -255,8 +255,8 @@ static int __args_reader_read(int sock_fd, struct enclave_args *args)
         case ENCLAVE_ARG_ID_NETWORK_PROXY:
             args->network_proxy = true;
             break;
-        case ENCLAVE_ARG_ID_DEBUG:
-            args->debug = true;
+        case ENCLAVE_ARG_ID_APP_OUTPUT:
+            args->app_output = true;
             break;
 
         /*

init/aws-nitro/device/app_stdio_output.c

@@ -40,7 +40,7 @@ int app_stdio_output(unsigned int vsock_port)
     ret = setsockopt(sock_fd, AF_VSOCK, SO_VM_SOCKETS_CONNECT_TIMEOUT,
                      (void *)&timeval, sizeof(struct timeval));
     if (ret < 0) {
-        perror("unable to connect to host socket");
+        perror("unable to set application output vsock timeout");
         close(sock_fd);
         return -errno;
     }
@@ -81,4 +81,4 @@ void app_stdio_close(void)
         close(APP_STDIO_OUTPUT_VSOCK_FD);
         APP_STDIO_OUTPUT_VSOCK_FD = -1;
     }
-}
+}

init/aws-nitro/include/args_reader.h

@@ -16,7 +16,7 @@ struct enclave_args {
     char **exec_argv;             // Execution argument vector.
     char **exec_envp;             // Execution environment pointer.
     bool network_proxy;           // Indicate if networking is configured.
-    bool debug;                   // Indicate if running in debug mode.
+    bool app_output;              // Indicate if running in non-debug mode.
 };
 
 int args_reader_read(struct enclave_args *, unsigned int);

init/aws-nitro/include/mod.h

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#ifndef _KRUN_MOD_LOADING_H
+#define _KRUN_MOD_LOADING_H
+
+int mods_load(void);
+
+#endif // _KRUN_MOD_LOADING_H

init/aws-nitro/main.c

@@ -16,7 +16,6 @@
 #include <sys/reboot.h>
 #include <sys/socket.h>
 #include <sys/stat.h>
-#include <sys/syscall.h>
 #include <sys/types.h>
 #include <sys/wait.h>
 #include <unistd.h>
@@ -25,9 +24,7 @@
 #include "include/archive.h"
 #include "include/args_reader.h"
 #include "include/fs.h"
-
-#define finit_module(fd, param_values, flags)                                  \
-    (int)syscall(__NR_finit_module, fd, param_values, flags)
+#include "include/mod.h"
 
 #define NSM_PCR_EXEC_DATA 17
 
@@ -41,47 +38,6 @@ enum {
     VSOCK_PORT_OFFSET_SIGNAL_HANDLER = 5,
 };
 
-/*
- * Load the NSM kernel module.
- */
-static int nsm_load(void)
-{
-    const char *file_name = "nsm.ko";
-    int fd, ret;
-
-    // Open and load the kernel module.
-    fd = open(file_name, O_RDONLY | O_CLOEXEC);
-    if (fd < 0 && errno == ENOENT)
-        return 0;
-    else if (fd < 0) {
-        perror("nsm.ko open");
-        return -errno;
-    }
-
-    ret = finit_module(fd, "", 0);
-    if (ret < 0) {
-        close(fd);
-        perror("nsm.ko finit_module");
-        return -errno;
-    }
-
-    // Close the file descriptor.
-    ret = close(fd);
-    if (ret < 0) {
-        perror("nsm.ko close");
-        return -errno;
-    }
-
-    // The NSM module file is no longer needed, remove it.
-    ret = unlink(file_name);
-    if (ret < 0) {
-        perror("nsm.ko unlink");
-        return -errno;
-    }
-
-    return 0;
-}
-
 /*
  * Mount the extracted rootfs and switch the root directory to it.
  */
@@ -351,10 +307,10 @@ static int proxies_init(int cid, struct enclave_args *args, int shutdown_fd)
 
     /*
      * If not running in debug mode, initialize the application output proxy.
-     * In debug mode, the enclave uses the console (which is already connected)
-     * for output.
+     * Otherwise, the enclave uses the console (which is already connected) for
+     * output.
      */
-    if (!args->debug) {
+    if (args->app_output) {
         ret = device_init(KRUN_NE_DEV_APP_OUTPUT_STDIO,
                           cid + VSOCK_PORT_OFFSET_OUTPUT, shutdown_fd);
     }
@@ -398,7 +354,7 @@ static int proxies_exit(struct enclave_args *args, int shutdown_fd)
     }
 
     // If not in debug mode, close the application output vsock.
-    if (!args->debug)
+    if (args->app_output)
         app_stdio_close();
 
     return ret;
@@ -447,6 +403,15 @@ int main(int argc, char *argv[])
         return -errno;
     }
 
+    /*
+     * Some linux modules (virtio-mmio, for example) may be required for console
+     * output. Load these modules immediately to ensure they are available to
+     * the initrd.
+     */
+    ret = mods_load();
+    if (ret < 0)
+        goto out;
+
     // Initialize early debug output with /dev/console.
     ret = console_init();
     if (ret < 0)
@@ -457,11 +422,6 @@ int main(int argc, char *argv[])
     if (cid == 0)
         goto out;
 
-    // Initialize the NSM kernel module.
-    ret = nsm_load();
-    if (ret < 0)
-        goto out;
-
     // Read the enclave arguments from the host.
     ret = args_reader_read(&args, cid + VSOCK_PORT_OFFSET_ARGS_READER);
     if (ret < 0)

init/aws-nitro/mod.c

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: Apache-2.0
+
+#include <dirent.h>
+#include <errno.h>
+#include <fcntl.h>
+#include <stddef.h>
+#include <stdio.h>
+#include <string.h>
+#include <sys/syscall.h>
+#include <unistd.h>
+
+#include "include/mod.h"
+
+#define KRUN_LINUX_MODS_DIR_NAME "/krun_linux_mods"
+#define MOD_FILE_NAME_BUF_SIZE 256
+
+#define finit_module(fd, param_values, flags)                                  \
+    (int)syscall(__NR_finit_module, fd, param_values, flags)
+
+/*
+ * Load a kernel module.
+ */
+static int mod_load(const char *path)
+{
+    int fd, ret;
+
+    // Open and load the kernel module.
+    fd = open(path, O_RDONLY | O_CLOEXEC);
+    if (fd < 0) {
+        if (errno == ENOENT)
+            return 0;
+
+        fprintf(stderr, "open module %s (errno %d)\n", path, errno);
+        return -errno;
+    }
+
+    ret = finit_module(fd, "", 0);
+    if (ret < 0) {
+        close(fd);
+        fprintf(stderr, "init module %s (errno %d)\n", path, errno);
+        return -errno;
+    }
+
+    // Close the file descriptor and remove the module file.
+    ret = close(fd);
+    if (ret < 0) {
+        fprintf(stderr, "close module %s (errno %d)\n", path, errno);
+        return -errno;
+    }
+
+    ret = unlink(path);
+    if (ret < 0) {
+        fprintf(stderr, "unlink module %s (errno %d)\n", path, errno);
+        return -errno;
+    }
+
+    return 0;
+}
+
+/*
+ * Load the configured kernel modules.
+ */
+int mods_load(void)
+{
+    char path[MOD_FILE_NAME_BUF_SIZE + sizeof(KRUN_LINUX_MODS_DIR_NAME) + 1];
+    struct dirent *entry;
+    int ret;
+    DIR *dir;
+
+    ret = 0;
+
+    dir = opendir(KRUN_LINUX_MODS_DIR_NAME);
+    if (dir != NULL) {
+        while ((entry = readdir(dir)) != NULL) {
+            /*
+             * Ignore the "." and ".." directory entries, as they are not kernel
+             * modules.
+             */
+            if (strcmp(entry->d_name, ".") == 0 ||
+                strcmp(entry->d_name, "..") == 0)
+                continue;
+
+            // Copy the full path of the module file.
+            snprintf(path, sizeof(path), "%s/%s", KRUN_LINUX_MODS_DIR_NAME,
+                     entry->d_name);
+            ret = mod_load(path);
+            if (ret < 0)
+                break;
+        }
+        closedir(dir);
+    } else if (errno != ENOENT) {
+        ret = -errno;
+        perror("unable to open kernel module configuration directory");
+    }
+
+    return ret;
+}

src/aws_nitro/src/enclave/args_writer.rs

@@ -139,8 +139,8 @@ pub enum EnclaveArg<'a> {
     ExecEnvp(Vec<String>),
     // Network proxy.
     NetworkProxy,
-    // Debug logs.
-    Debug,
+    // Application output.
+    AppOutput,
 
     // Placeholder argument where libkrun notifies the initramfs that all arguments have been
     // written and it can now close the vsock connection.
@@ -157,7 +157,7 @@ impl From<&EnclaveArg<'_>> for u8 {
             EnclaveArg::ExecArgv(_) => 2,
             EnclaveArg::ExecEnvp(_) => 3,
             EnclaveArg::NetworkProxy => 4,
-            EnclaveArg::Debug => 5,
+            EnclaveArg::AppOutput => 5,
 
             EnclaveArg::Finished => 255,
         }

src/aws_nitro/src/enclave/mod.rs

@@ -19,6 +19,7 @@ use std::{
     io::{self, Read, Write},
     os::fd::RawFd,
     path::{Path, PathBuf},
+    thread::{self, JoinHandle},
 };
 use tar::HeaderMode;
 use vsock::{VsockAddr, VsockListener, VMADDR_CID_ANY};
@@ -81,6 +82,12 @@ impl NitroEnclave {
         // Launch the enclave and write the configured launch parameters to the initramfs.
         let (cid, timeout) = self.start().map_err(Error::Start)?;
 
+        // If debug mode is enabled, attach to the serial console immediately after the enclave VM
+        // is started to get logs.
+        if self.debug {
+            self.start_console_debug(cid).map_err(Error::DeviceProxy)?;
+        }
+
         writer.write_args(cid, timeout).map_err(Error::ArgsWrite)?;
 
         // Establish the vsock listener for the application's return code upon termination.
@@ -154,9 +161,11 @@ impl NitroEnclave {
     fn proxies(&self) -> Result<DeviceProxyList, proxy::Error> {
         let mut proxies: Vec<Box<dyn Send + DeviceProxy>> = vec![];
 
-        // All enclaves will include a proxy for debug/application output.
-        let output = OutputProxy::new(&self.output_path, self.debug)?;
-        proxies.push(Box::new(output));
+        // Only specify application output proxy if not running in debug mode.
+        if !self.debug {
+            let output = OutputProxy::new(&self.output_path, false)?;
+            proxies.push(Box::new(output));
+        }
 
         if let Some(fd) = self.net_unixfd {
             let net = NetProxy::try_from(fd)?;
@@ -239,6 +248,19 @@ impl NitroEnclave {
             libc::pthread_sigmask(sig, &set, std::ptr::null_mut());
         }
     }
+
+    /// Start the debug output thread, read from the serial console. Run this proxy separate
+    /// from the others, as it is not easily controlled by the user due to the connection to
+    /// the enclave VM's serial console.
+    fn start_console_debug(&self, cid: u32) -> Result<(), proxy::Error> {
+        let mut output = OutputProxy::new(&self.output_path, true)?;
+        let mut vsock_rcv = output.vsock(cid)?;
+        let _: JoinHandle<Result<(), proxy::Error>> = thread::spawn(move || loop {
+            output.rcv(&mut vsock_rcv)?;
+        });
+
+        Ok(())
+    }
 }
 
 /// Each service provided to an enclave is done so via vsock. Each service has a designated port

src/aws_nitro/src/enclave/proxy/proxies/output.rs

@@ -48,10 +48,9 @@ impl OutputProxy {
 impl DeviceProxy for OutputProxy {
     /// Enclave argument of the proxy.
     fn arg(&self) -> Option<EnclaveArg<'_>> {
-        // The enclave only needs to be made aware that it is to be run in debug mode.
         match self.debug {
-            true => Some(EnclaveArg::Debug),
-            false => None,
+            true => None,
+            false => Some(EnclaveArg::AppOutput),
         }
     }