Making the library code safer

crates/water/src/runtime/core.rs

@@ -54,6 +54,10 @@ impl H2O<Host> {
         // store.data_mut().preview1_ctx = Some(WasiCtxBuilder::new().inherit_stdio().preopened_dir(path, ".")?.build());
         store.data_mut().preview1_ctx = Some(WasiCtxBuilder::new().inherit_stdio().build());
 
+        if store.data().preview1_ctx.is_none() {
+            return Err(anyhow::anyhow!("Failed to retrieve preview1_ctx from Host"));
+        }
+
         wasmtime_wasi::add_to_linker(&mut linker, |h: &mut Host| h.preview1_ctx.as_mut().unwrap())?;
 
         // initializing stuff for multithreading
@@ -65,30 +69,37 @@ impl H2O<Host> {
             )?));
 
             wasmtime_wasi_threads::add_to_linker(&mut linker, &store, &module, |h: &mut Host| {
-                h.wasi_threads.as_ref().unwrap()
+                h.wasi_threads
+                    .as_ref()
+                    .context("Failed to get ref of wasi_threads from Host")?
             })?;
         }
 
         // export functions -- version dependent -- has to be done before instantiate
         match &version {
             Some(Version::V0) => {
-                v0::funcs::export_tcp_connect(&mut linker);
-                v0::funcs::export_tcplistener_create(&mut linker);
+                v0::funcs::export_tcp_connect(&mut linker)?;
+                v0::funcs::export_tcplistener_create(&mut linker)?;
             }
             Some(Version::V1) => {
-                v1::funcs::export_tcp_connect(&mut linker);
-                v1::funcs::export_tcplistener_create(&mut linker);
+                v1::funcs::export_tcp_connect(&mut linker)?;
+                v1::funcs::export_tcplistener_create(&mut linker)?;
             }
             _ => {} // add export funcs for other versions here
         }
 
         // export functions -- version independent
-        version_common::funcs::export_config(&mut linker, conf.config_wasm.clone());
+        version_common::funcs::export_config(&mut linker, conf.config_wasm.clone())?;
 
         let instance = linker.instantiate(&mut store, &module)?;
 
         Ok(H2O {
-            version: version.unwrap(),
+            version: match version {
+                Some(v) => v,
+                None => {
+                    return Err(anyhow::anyhow!("Version is None"));
+                }
+            },
 
             engine,
             linker,
@@ -138,8 +149,13 @@ impl H2O<Host> {
 
         // Obtain the directory path and file name from config_wasm
         let full_path = Path::new(&config.config_wasm);
-        let parent_dir = full_path.parent().unwrap(); // Assumes config_wasm has a parent directory
-        let file_name = full_path.file_name().unwrap().to_str().unwrap(); // Assumes file_name is valid UTF-8
+        let parent_dir = full_path
+            .parent()
+            .ok_or_else(|| anyhow::anyhow!("config_wasm does not have a parent directory"))?; // Assumes config_wasm has a parent directory
+        let file_name = full_path
+            .file_name()
+            .and_then(|os_str| os_str.to_str())
+            .ok_or_else(|| anyhow::anyhow!("file_name is not valid UTF-8"))?; // Assumes file_name is valid UTF-8
 
         // Open the parent directory
         let dir = Dir::open_ambient_dir(parent_dir, ambient_authority())?;
@@ -148,7 +164,12 @@ impl H2O<Host> {
 
         let wasi_file = wasmtime_wasi::sync::file::File::from_cap_std(wasi_file);
 
-        let ctx = self.store.data_mut().preview1_ctx.as_mut().unwrap();
+        let ctx = self
+            .store
+            .data_mut()
+            .preview1_ctx
+            .as_mut()
+            .ok_or(anyhow::anyhow!("preview1_ctx in Store is None"))?;
         let config_fd = ctx.push_file(Box::new(wasi_file), FileAccessMode::all())? as i32;
 
         let params = vec![Val::I32(config_fd); config_fn.ty(&self.store).params().len()];

crates/water/src/runtime/listener.rs

@@ -118,7 +118,8 @@ impl WATERListener<Host> {
             .core
             .instance
             .get_func(&mut self.core.store, &conf.entry_fn)
-            .unwrap();
+            .context(format!("Failed to get function {}", &conf.entry_fn))?;
+
         match fnc.call(&mut self.core.store, &[], &mut []) {
             Ok(_) => {}
             Err(e) => {
@@ -153,7 +154,12 @@ impl WATERListener<Host> {
         std::mem::forget(water_writer);
         std::mem::forget(water_reader);
 
-        let ctx = core.store.data_mut().preview1_ctx.as_mut().unwrap();
+        let ctx = core
+            .store
+            .data_mut()
+            .preview1_ctx
+            .as_mut()
+            .ok_or(anyhow::anyhow!("preview1_ctx in Store is None"))?;
         let water_reader_fd = ctx.push_file(Box::new(wasi_water_reader), FileAccessMode::all())?;
         let water_writer_fd = ctx.push_file(Box::new(wasi_water_writer), FileAccessMode::all())?;
 

crates/water/src/runtime/stream.rs

@@ -124,7 +124,6 @@ impl WATERStream<Host> {
         info!("[HOST] WATERStream connecting...");
 
         // TODO: add addr:port sharing with WASM, for now WASM is using config.json's remote_addr:port
-        // let fnc = self.core.instance.get_func(&mut self.core.store, &conf.entry_fn).unwrap();
         let fnc = match self.core.instance.get_func(&mut self.core.store, DIAL_FN) {
             Some(func) => func,
             None => {

crates/water/src/runtime/v0/funcs.rs

@@ -1,10 +1,12 @@
+use anyhow::Ok;
+
 use crate::config::wasm_shared_config::StreamConfig;
 use crate::runtime::*;
 use std::convert::TryInto;
 use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
 
 // TODO: rename this to dial_v1, since it has the ability to let WASM choose ip:port
-pub fn export_tcp_connect(linker: &mut Linker<Host>) {
+pub fn export_tcp_connect(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
     linker
         .func_wrap(
             "env",
@@ -48,22 +50,33 @@ pub fn export_tcp_connect(linker: &mut Linker<Host>) {
                     addr => std::net::TcpStream::connect(addr),
                 }
                 .map(TcpStream::from_std)
-                .context("failed to connect to endpoint")
+                .context(format!(
+                    "Failed to connect to {}:{} in Host exported dial",
+                    host, port
+                ))
                 .unwrap();
 
                 // Connecting Tcp
                 let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();
 
                 // Get the WasiCtx of the caller(WASM), then insert_file into it
-                let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
-                ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
+                let ctx: &mut WasiCtx = caller
+                    .data_mut()
+                    .preview1_ctx
+                    .as_mut()
+                    .context("preview1_ctx in Store is None")
+                    .unwrap();
+                ctx.push_file(socket_file, FileAccessMode::all())
+                    .context("Failed to push file into WASM")
+                    .unwrap() as i32
             },
         )
-        .unwrap();
+        .context("Failed to export Dial function to WASM")?;
+    Ok(())
 }
 
 // TODO: rename this to dial_v1, since it has the ability to let WASM listen on a TcpListener
-pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
+pub fn export_tcplistener_create(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
     linker
         .func_wrap(
             "env",
@@ -89,7 +102,7 @@ pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
                     name: config.name.clone().try_into().unwrap(),
                     port: config.port as u16,
                     addr: config.addr.clone(),
-                }); // TODO: add addr here
+                });
 
                 // Get the pair here addr:port
                 let (addr, port) = match listener_file {
@@ -108,11 +121,19 @@ pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
                 let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();
 
                 // Get the WasiCtx of the caller(WASM), then insert_file into it
-                let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
-                ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
+                let ctx: &mut WasiCtx = caller
+                    .data_mut()
+                    .preview1_ctx
+                    .as_mut()
+                    .context("preview1_ctx in Store is None")
+                    .unwrap();
+                ctx.push_file(socket_file, FileAccessMode::all())
+                    .context("Failed to push file into WASM")
+                    .unwrap() as i32
             },
         )
-        .unwrap();
+        .context("Failed to export TcpListener create function to WASM")?;
+    Ok(())
 }
 
 // Generically link dial functions

crates/water/src/runtime/v1/funcs.rs

@@ -1,10 +1,12 @@
+use anyhow::Ok;
+
 use crate::config::wasm_shared_config::StreamConfig;
 use crate::runtime::*;
 use std::convert::TryInto;
 use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4};
 
 // TODO: rename this to dial_v1, since it has the ability to let WASM choose ip:port
-pub fn export_tcp_connect(linker: &mut Linker<Host>) {
+pub fn export_tcp_connect(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
     linker
         .func_wrap(
             "env",
@@ -48,22 +50,33 @@ pub fn export_tcp_connect(linker: &mut Linker<Host>) {
                     addr => std::net::TcpStream::connect(addr),
                 }
                 .map(TcpStream::from_std)
-                .context("failed to connect to endpoint")
+                .context(format!(
+                    "Failed to connect to {}:{} in Host exported dial",
+                    host, port
+                ))
                 .unwrap();
 
                 // Connecting Tcp
                 let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();
 
                 // Get the WasiCtx of the caller(WASM), then insert_file into it
-                let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
-                ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
+                let ctx: &mut WasiCtx = caller
+                    .data_mut()
+                    .preview1_ctx
+                    .as_mut()
+                    .context("preview1_ctx in Store is None")
+                    .unwrap();
+                ctx.push_file(socket_file, FileAccessMode::all())
+                    .context("Failed to push file into WASM")
+                    .unwrap() as i32
             },
         )
-        .unwrap();
+        .context("Failed to export Dial function to WASM")?;
+    Ok(())
 }
 
 // TODO: rename this to dial_v1, since it has the ability to let WASM listen on a TcpListener
-pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
+pub fn export_tcplistener_create(linker: &mut Linker<Host>) -> Result<(), anyhow::Error> {
     linker
         .func_wrap(
             "env",
@@ -104,14 +117,23 @@ pub fn export_tcplistener_create(linker: &mut Linker<Host>) {
                 // Creating Tcp Listener
                 let tcp = std::net::TcpListener::bind((addr.as_str(), port)).unwrap();
                 let tcp = TcpListener::from_std(tcp);
+                // tcp.set_nonblocking(true);
                 let socket_file: Box<dyn WasiFile> = wasmtime_wasi::net::Socket::from(tcp).into();
 
                 // Get the WasiCtx of the caller(WASM), then insert_file into it
-                let ctx: &mut WasiCtx = caller.data_mut().preview1_ctx.as_mut().unwrap();
-                ctx.push_file(socket_file, FileAccessMode::all()).unwrap() as i32
+                let ctx: &mut WasiCtx = caller
+                    .data_mut()
+                    .preview1_ctx
+                    .as_mut()
+                    .context("preview1_ctx in Store is None")
+                    .unwrap();
+                ctx.push_file(socket_file, FileAccessMode::all())
+                    .context("Failed to push file into WASM")
+                    .unwrap() as i32
             },
         )
-        .unwrap();
+        .context("Failed to export TcpListener create function to WASM")?;
+    Ok(())
 }
 
 // Generically link dial functions

crates/water/src/runtime/version_common/funcs.rs

@@ -1,7 +1,7 @@
 use crate::runtime::*;
 
 // exportint a function for WASM to get CONFIG file
-pub fn export_config(linker: &mut Linker<Host>, config_file: String) {
+pub fn export_config(linker: &mut Linker<Host>, config_file: String) -> Result<(), anyhow::Error> {
     linker
         .func_wrap(
             "env",
@@ -22,5 +22,6 @@ pub fn export_config(linker: &mut Linker<Host>, config_file: String) {
                     .expect("Error with pushing file") as i32
             },
         )
-        .unwrap();
+        .context("Failed to export config function to WASM")?;
+    Ok(())
 }

tests/Cargo.toml

@@ -20,4 +20,4 @@ shadowsocks-service = {version = "1.17.0", features = ["server"]}
 shadowsocks-rust = "1.17.0"
 tokio = { version = "1.24.2", features = ["full", "macros"] }
 futures = "0.3.28"
-tempfile = "3.8.0"
+tempfile = "3.8.0"