Refactor for better error handling and entropy seeding prep

Makefile

@@ -255,7 +255,7 @@ $(OUT_DIR)/init:
 		RUSTFLAGS='-C target-feature=+crt-static' cargo build \
 			--target x86_64-unknown-linux-gnu \
 			--release && \
-		cp /src/init/target/release/init /out/init \
+		cp /src/init/target/x86_64-unknown-linux-gnu/release/init /out/init \
 	")
 
 

src/init/init.rs

@@ -1,49 +1,57 @@
-extern crate libc;
+use libc::{
-use libc::c_ulong;
+    c_ulong,
-use libc::c_int;
+    c_int,
-use libc::c_void;
+    c_void,
-use libc::MS_NOSUID;
+    MS_NOSUID,
-use libc::MS_NOEXEC;
+    MS_NOEXEC,
-use libc::MS_NODEV;
+    MS_NODEV,
-use std::mem::zeroed;
+};
-use std::mem::size_of;
+use std::{
-use std::ffi::CString;
+    mem::zeroed,
-use std::fs::File;
+    mem::size_of,
-use std::os::unix::io::AsRawFd;
+    ffi::CString,
+    fs::{File, read_to_string},
+    os::unix::io::AsRawFd,
+    fmt,
+};
 
-// Log errors to console
+struct SystemError {
-pub fn error(message: String){
+    message: String,
-    eprintln!("{} {}", boot_time(), message);
 }
-
+impl fmt::Display for SystemError {
-// Log info to console
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-pub fn info(message: String){
+        write!(f, "{} {}", boot_time(), self.message)
-    println!("{} {}", boot_time(), message);
-}
-
-pub fn reboot(){
-    use libc::{reboot, RB_AUTOBOOT};
-    unsafe {
-        reboot(RB_AUTOBOOT);
     }
 }
 
+// Log dmesg formatted log to console
+fn dmesg(message: String){
+    println!("{} {}", boot_time(), message);
+}
+
 // Dmesg formatted seconds since boot
-pub fn boot_time() -> String {
+fn boot_time() -> String {
     use libc::{clock_gettime, timespec, CLOCK_BOOTTIME};
     let mut t = timespec { tv_sec: 0, tv_nsec: 0 };
     unsafe { clock_gettime(CLOCK_BOOTTIME, &mut t as *mut timespec); }
     format!("[ {: >4}.{}]", t.tv_sec, t.tv_nsec / 1000).to_string()
 }
 
+fn reboot(){
+    use libc::{reboot, RB_AUTOBOOT};
+    unsafe {
+        reboot(RB_AUTOBOOT);
+    }
+}
+
 // libc::mount casting/error wrapper
-pub fn mount(
+fn mount(
     src: &str,
     target: &str,
     fstype: &str,
     flags: c_ulong,
     data: &str,
-) {
+) -> Result<(), SystemError> {
     use libc::mount;
     let src_cs = CString::new(src).unwrap();
     let fstype_cs = CString::new(fstype).unwrap();
@@ -58,18 +66,18 @@ pub fn mount(
             data_cs.as_ptr() as *const c_void
         )
     } != 0 {
-        error(format!("Failed to mount: {}", target));
+        Err(SystemError { message: format!("Failed to mount: {}", target) })
     } else {
-        info(format!("Mounted: {}", target));
+        Ok(())
     }
 }
 
 // libc::freopen casting/error wrapper
-pub fn freopen(
+fn freopen(
     filename: &str,
     mode: &str,
     file: c_int,
-) {
+) -> Result<(), SystemError> {
     use libc::{freopen, fdopen};
     let filename_cs = CString::new(filename).unwrap();
     let mode_cs = CString::new(mode).unwrap();
@@ -80,91 +88,173 @@ pub fn freopen(
             fdopen(file, mode_cs.as_ptr() as *const i8)
         )
     }.is_null() {
-        error(format!("Failed to freopen: {}", filename));
+        Err(SystemError { message: format!("Failed to freopen: {}", filename) })
+    } else {
+        Ok(())
     }
 }
 
 // Insert kernel module into memory
-pub fn insmod(path: &str){
+fn insmod(path: &str) -> Result<(), SystemError> {
     use libc::{syscall, SYS_finit_module};
     let file = File::open(path).unwrap();
     let fd = file.as_raw_fd();
     if unsafe { syscall(SYS_finit_module, fd, &[0u8; 1], 0) } < 0 {
-        error(format!("Failed to insert kernel module: {}", path));
+        Err(SystemError {
+            message: format!("Failed to insert kernel module: {}", path)
+        })
     } else {
-        info(format!("Loaded kernel module: {}", path));
+        Ok(())
     }
 }
 
-// Signal to Nitro hypervisor that booting was successful
+// Instantiate a socket
-pub fn nitro_heartbeat(){
+fn socket_connect(
-    use libc::{connect, socket, write, read, close, sockaddr, sockaddr_vm, SOCK_STREAM, AF_VSOCK};
+    family: c_int,
-    let mut buf: [u8; 1] = [0; 1];
+    port: u32,
-    buf[0] = 0xB7; // AWS Nitro heartbeat value
+    cid: u32,
-    unsafe {
+) -> Result<c_int, SystemError> {
+    use libc::{connect, socket, sockaddr, sockaddr_vm, SOCK_STREAM};
+    let fd = unsafe { socket(family, SOCK_STREAM, 0) };
+    if unsafe {
         let mut sa: sockaddr_vm = zeroed();
-        sa.svm_family = AF_VSOCK as _;
+        sa.svm_family = family as _;
-        sa.svm_port = 9000;
+        sa.svm_port = port;
-        sa.svm_cid = 3;
+        sa.svm_cid = cid;
-        let fd = socket(AF_VSOCK, SOCK_STREAM, 0);
         connect(
             fd,
             &sa as *const _ as *mut sockaddr,
             size_of::<sockaddr_vm>() as _,
-        );
+        )
+    } < 0 {
+        Err(SystemError {
+            message: format!("Failed to connect to socket: {}", family)
+        })
+    } else {
+        Ok(fd)
+    }
+}
+
+// Signal to Nitro hypervisor that booting was successful
+fn nitro_heartbeat() {
+    use libc::{write, read, close, AF_VSOCK};
+    let mut buf: [u8; 1] = [0; 1];
+    buf[0] = 0xB7; // AWS Nitro heartbeat value
+    let fd = match socket_connect(AF_VSOCK, 9000, 3) {
+        Ok(f)=> f,
+        Err(e)=> {
+            eprintln!("{}", e);
+            return
+        },
+    };
+    unsafe {
         write(fd, buf.as_ptr() as _, 1);
         read(fd, buf.as_ptr() as _, 1);
         close(fd);
     }
-    info(format!("Sent NSM heartbeat"));
+    dmesg(format!("Sent NSM heartbeat"));
 }
 
 // Get entropy sample from Nitro device
-pub fn nitro_get_entropy() -> u8 {
+fn nitro_get_entropy() -> Result<[u8; 256], SystemError> {
-	use nsm_lib::{nsm_lib_init, nsm_get_random};
+    use nsm_api::api::ErrorCode;
-	use nsm_api::api::ErrorCode;
+    use nsm_lib::{nsm_get_random, nsm_lib_init};
+
     let nsm_fd = nsm_lib_init();
     if nsm_fd < 0 {
-    	error(format!("Failed to connect to NSM device"));
+        return Err(SystemError {
+            message: String::from("Failed to connect to NSM device")
+        });
     };
-    let mut dest: [u8; 256] = [0; 256];
+
+    let mut dest = [0u8; 256];
     let mut dest_len = dest.len();
+
     let status = unsafe {
-		nsm_get_random(nsm_fd, dest.as_mut_ptr(), &mut dest_len)
+        nsm_get_random(nsm_fd, dest.as_mut_ptr(), &mut dest_len)
     };
     match status {
-        ErrorCode::Success => info(format!("Entropy seeding success")),
+        ErrorCode::Success => {
-    	_ => error(format!("Failed to get entropy from NSM device")),
+            Ok(dest)
+        },
+        _ => Err(SystemError {
+            message: String::from("Failed to get entropy from NSM device")
+        })
     }
 }
 
-pub fn init_nitro(){
+fn get_random_poolsize() -> Result<usize, SystemError> {
+    let ps_path = "/proc/sys/kernel/random/poolsize";
+    let size_s = read_to_string(ps_path).unwrap_or_else(|_| String::new());
+    if size_s.is_empty(){
+		return Err(SystemError {
+            message: String::from("Failed to read kernel random poolsize"),
+        })
+    };
+	match size_s.parse::<usize>() {
+		Ok(size) => Ok(size),
+		Err(_) => Err(SystemError {
+            message: String::from("Failed to parse kernel random poolsize"),
+        }),
+    }
+}
+
+// Initialize nitro device
+fn init_nitro(){
+    // TODO: error handling
     nitro_heartbeat();
-    insmod("/nsm.ko");
+
-    nitro_seed_entropy();
+	match insmod("/nsm.ko") {
+        Ok(())=> dmesg(format!("Loaded nsm.ko")),
+        Err(e)=> eprintln!("{}", e)
+    };
+	match get_random_poolsize() {
+        Ok(size)=> dmesg(format!("Kernel entropy pool size: {}", size)),
+        Err(e)=> eprintln!("{}", e)
+    };
+    match nitro_get_entropy() {
+        Ok(_)=> dmesg(format!("Got NSM Entropy sample")),
+        Err(e)=> eprintln!("{}", e)
+    };
 }
 
 // Initialize console with stdin/stdout/stderr
-pub fn init_console() {
+fn init_console() {
-    freopen("/dev/console", "r", 0);
+    let args = [
-    freopen("/dev/console", "w", 1);
+        ("/dev/console", "r", 0),
-    freopen("/dev/console", "w", 2);
+        ("/dev/console", "w", 1),
-    info(format!("Initialized console"));
+        ("/dev/console", "w", 2),
+    ];
+    for (filename, mode, file) in args {
+        match freopen(filename, mode, file) {
+            Ok(())=> {},
+            Err(e)=> eprintln!("{}", e),
+        }
+    }
 }
 
 // Mount common filesystems with conservative permissions
-pub fn init_rootfs() {
+fn init_rootfs() {
-    mount("devtmpfs", "/dev", "devtmpfs", MS_NOSUID | MS_NOEXEC, "mode=0755");
+    let args = [
-    mount("proc", "/proc", "proc", MS_NODEV | MS_NOSUID | MS_NOEXEC, "hidepid=2");
+        ("devtmpfs", "/dev", "devtmpfs", MS_NOSUID | MS_NOEXEC, "mode=0755"),
-    mount("tmpfs", "/run", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "mode=0755");
+        ("devtmpfs", "/dev", "devtmpfs", MS_NOSUID | MS_NOEXEC, "mode=0755"),
-    mount("tmpfs", "/tmp", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "");
+        ("proc", "/proc", "proc", MS_NODEV | MS_NOSUID | MS_NOEXEC, "hidepid=2"),
-    mount("shm", "/dev/shm", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "mode=0755");
+        ("tmpfs", "/run", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "mode=0755"),
-    mount("devpts", "/dev/pts", "devpts", MS_NOSUID | MS_NOEXEC, "");
+        ("tmpfs", "/tmp", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, ""),
-    mount("sysfs", "/sys", "sysfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "");
+        ("shm", "/dev/shm", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "mode=0755"),
-    mount("cgroup_root", "/sys/fs/cgroup", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "mode=0755");
+        ("devpts", "/dev/pts", "devpts", MS_NOSUID | MS_NOEXEC, ""),
+        ("sysfs", "/sys", "sysfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, ""),
+        ("cgroup_root", "/sys/fs/cgroup", "tmpfs", MS_NODEV | MS_NOSUID | MS_NOEXEC, "mode=0755"),
+    ];
+    for (src, target, fstype, flags, data) in args {
+        match mount(src, target, fstype, flags, data) {
+            Ok(())=> dmesg(format!("Mounted {}", target)),
+            Err(e)=> eprintln!("{}", e),
+        }
+    }
 }
 
-pub fn boot(){
+fn boot(){
     init_rootfs();
     init_console();
     init_nitro();
@@ -172,6 +262,6 @@ pub fn boot(){
 
 fn main() {
     boot();
-    info("EnclaveOS Booted".to_string());
+    dmesg("EnclaveOS Booted".to_string());
     reboot();
 }