fixup! wasmtime: first custom syscall demo

wasmtime-custom-syscall/src/app.rs

@@ -16,14 +16,40 @@ mod lib;
 
 use lib::{enarx_syscalls, CustomNumbers};
 
-fn main() {
+fn main() -> ! {
+    udp_receive_loop();
+}
+
+fn simple_print() {
     let (a, b) = (5, 7);
     let result = a + b;
 
     enarx_syscalls::println(&format!("{} + {} = {}", a, b, result));
+}
 
+fn mutating_syscalls() {
     enarx_syscalls::println(&format!(
         "{:?}",
         enarx_syscalls::mutate_custom_numbers(&mut CustomNumbers(0, 1))
     ));
 }
+
+fn udp_receive_loop() -> ! {
+    enarx_syscalls::socket_udp_bind("127.0.0.1:34254");
+    loop {
+        let mut buf = [0; 1024];
+        enarx_syscalls::println(&format!(
+            "received {} bytes",
+            enarx_syscalls::socket_udp_receive(&mut buf),
+        ));
+
+        enarx_syscalls::println(&format!(
+            "{}",
+            std::str::from_utf8(&buf)
+                .unwrap()
+                .split("\n")
+                .nth(0)
+                .unwrap()
+        ));
+    }
+}

wasmtime-custom-syscall/src/lib.rs

@@ -1,3 +1,5 @@
+#![allow(dead_code)]
+
 #[derive(Debug)]
 #[repr(C)]
 pub struct CustomNumbers(pub u64, pub u64);
@@ -7,11 +9,14 @@ pub mod enarx_syscalls {
     extern "C" {
         fn __print_str(ptr: *const u8, len: u64);
         fn __mutate_custom_numbers(custom: &mut super::CustomNumbers) -> &mut super::CustomNumbers;
+
+        fn __socket_udp_bind(ptr: *const u8, len: u64);
+        fn __socket_udp_receive(buf_ptr: *mut u8, len: u64) -> u64;
     }
 
-    fn _print_str(s: &str) {
-        use std::convert::TryInto;
+    use std::convert::TryInto;
 
+    fn _print_str(s: &str) {
         unsafe { __print_str(s.as_ptr(), s.len().try_into().unwrap()) }
     }
 
@@ -29,4 +34,11 @@ pub mod enarx_syscalls {
     pub fn mutate_custom_numbers(custom: &mut super::CustomNumbers) -> &mut super::CustomNumbers {
         unsafe { __mutate_custom_numbers(custom) }
     }
+
+    pub fn socket_udp_bind(addr: &str) {
+        unsafe { __socket_udp_bind(addr.as_ptr(), addr.len().try_into().unwrap()) }
+    }
+    pub fn socket_udp_receive(buf: &mut [u8]) -> u64 {
+        unsafe { __socket_udp_receive(buf.as_mut_ptr(), buf.len().try_into().unwrap()) }
+    }
 }

wasmtime-custom-syscall/src/main.rs

@@ -14,6 +14,7 @@
 
 mod lib;
 
+use std::convert::TryInto;
 use std::error::Error;
 
 // The basic WASM demo itself.
@@ -29,7 +30,7 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     // Instantiate the syscalls.
     let enarx_syscalls_module = HostRef::new(wrap_wasmtime_module!(
-        &store, |_imports| EnarxSyscalls; module(enarx_syscalls_mod)
+        &store, |_imports| EnarxSyscalls::default(); module(enarx_syscalls_mod)
     )?);
     let enarx_syscalls_instance = Instance::new(
         &store,
@@ -46,48 +47,59 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     // Create the app instance with the enarx_syscalls imported
     let app_instance = {
-        // Collect all the potential imports
+        // Prepare a lookup-map of module imports and their index
         let app_module_imports = app_module
             .borrow()
             .imports()
             .iter()
-            .map(|import| {
-                format!(
-                    "{}/{}",
-                    import.module().to_string(),
-                    import.name().to_string(),
+            .enumerate()
+            .map(|(i, import)| {
+                (
+                    format!(
+                        "{}/{}",
+                        import.module().to_string(),
+                        import.name().to_string(),
+                    ),
+                    i,
                 )
             })
-            .collect::<std::collections::HashSet<_>>();
+            .collect::<std::collections::HashMap<_, _>>();
+        println!("[RUNTIME] App imports: {:?}", app_module_imports);
 
-        println!("App imports: {:?}", app_module_imports);
+        // Instance imports for the app must only include the ones which the module actually references.
+        // They need to be in the same order they are referenced.
+        let app_instance_imports = enarx_syscalls_instance
+            .exports()
+            .iter()
+            .cloned()
+            // TODO: verify the assumption that instance exports are in the same order as the module exports
+            .zip(enarx_syscalls_module.borrow().exports().iter())
+            .filter_map(|(ext, exp)| match (&ext, exp) {
+                (wasmtime_api::Extern::Func(_), exp) => {
+                    // instance imports may only include references which are imported by the module
+                    let lookup =
+                        format!("{}/{}", SYSCALLS_WASM_MODULE_NAME, exp.name().to_string());
+                    if let Some(index) = app_module_imports.get(&lookup) {
+                        println!("[RUNTIME] Including export '{}' at index {}", lookup, index);
+                        Some((index, ext))
+                    } else {
+                        println!("[RUNTIME] Skipping export '{}'", lookup);
+                        None
+                    }
+                }
+                // TODO: figure out what to do with non-Func externs
+                _ => None,
+            })
+            // ensure the same order as observed in the imports
+            .collect::<std::collections::BTreeMap<_, _>>()
+            .values()
+            .cloned()
+            .collect::<Vec<wasmtime_api::Extern>>();
 
         HostRef::new(Instance::new(
             &store,
             &app_module,
-            enarx_syscalls_instance
-                .exports()
-                .iter()
-                .cloned()
-                // TODO: verifty the assumption that instance exports are in the same order as the module exports
-                .zip(enarx_syscalls_module.borrow().exports().iter())
-                .filter_map(|(ext, exp)| match (&ext, exp) {
-                    // TODO: verify why the following is the case
-                    // imports may only include exports which are actually used and they have to be in the correct order
-                    (wasmtime_api::Extern::Func(_), exp) => {
-                        let lookup = format!("enarx_syscalls/{}", exp.name().to_string());
-                        if app_module_imports.contains(&lookup) {
-                            println!("Including export '{}'", lookup);
-                            Some(ext)
-                        } else {
-                            println!("Skipping export '{}'", lookup);
-                            None
-                        }
-                    }
-                    _ => Some(ext),
-                })
-                .collect::<Vec<_>>()
-                .as_slice(),
+            app_instance_imports.as_slice(),
         )?)
     };
 
@@ -106,7 +118,17 @@ trait App {
     fn main(&self);
 }
 
-pub struct EnarxSyscalls;
+pub struct EnarxSyscalls {
+    pub udp_socket: Option<std::net::UdpSocket>,
+}
+
+impl Default for EnarxSyscalls {
+    fn default() -> Self {
+        EnarxSyscalls { udp_socket: None }
+    }
+}
+
+static SYSCALLS_WASM_MODULE_NAME: &str = "enarx_syscalls";
 
 use wasmtime_bindings_macro::wasmtime_impl;
 #[wasmtime_impl(module(enarx_syscalls_mod),
@@ -115,18 +137,12 @@ use wasmtime_bindings_macro::wasmtime_impl;
 ]
 impl EnarxSyscalls {
     fn __print_char(&self, c: u8) {
-        print!("{}", c as char);
+        print!("[WASM] {}", c as char);
     }
 
     fn __print_str(&self, ptr: *const u8, len: u64) {
-        use std::convert::TryInto;
-
-        let s = std::str::from_utf8(unsafe {
-            std::slice::from_raw_parts(ptr, len.try_into().unwrap())
-        })
-        .unwrap();
-
-        print!("{}", s);
+        let wasm_str = WasmString(ptr, len);
+        print!("[WASM] {}", wasm_str.to_str());
     }
 
     // TODO: investigate why we can't return `&mut` here
@@ -137,4 +153,39 @@ impl EnarxSyscalls {
 
         custom
     }
+
+    fn __socket_udp_bind(&mut self, ptr: *const u8, len: u64) {
+        let addr_str = WasmString(ptr, len);
+        self.udp_socket = Some(std::net::UdpSocket::bind(addr_str.to_str()).unwrap());
+    }
+
+    fn __socket_udp_receive(&mut self, buf_ptr: *mut u8, len: u64) -> u64 {
+        let mut buf = unsafe { std::slice::from_raw_parts_mut(buf_ptr, len.try_into().unwrap()) };
+
+        if let Some(socket) = self.udp_socket.as_ref() {
+            println!("[RUNTIME] receiving from {:?}", socket);
+            let (amt, src) = socket.recv_from(&mut buf).unwrap();
+            println!("[RUNTIME] received {} bytes from {}", amt, src);
+            amt.try_into().unwrap()
+        } else {
+            0
+        }
+    }
+}
+
+pub struct WasmString(*const u8, u64);
+
+impl WasmString {
+    fn to_str(&self) -> &str {
+        self.into()
+    }
+}
+
+impl std::convert::From<&WasmString> for &str {
+    fn from(wasm_string: &WasmString) -> Self {
+        std::str::from_utf8(unsafe {
+            std::slice::from_raw_parts(wasm_string.0, wasm_string.1.try_into().unwrap())
+        })
+        .unwrap()
+    }
 }