update: keep the library code safe

refraction-networking · Nov 1, 2023 · f2f1d72 · f2f1d72
1 parent 177cbc0
commit f2f1d72
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 14 deletions.
diff --git a/crates/water/src/runtime/core.rs b/crates/water/src/runtime/core.rs
@@ -94,12 +94,19 @@ impl H2O<Host> {
 
         // export functions -- version dependent -- has to be done before instantiate
         match &version {
-            Some(Version::V0(v0_conf)) => {
-                let v0_conf = Arc::new(Mutex::new(v0_conf.clone()));
-                v0::funcs::export_tcp_connect(&mut linker, Arc::clone(&v0_conf))?;
-                v0::funcs::export_accept(&mut linker, Arc::clone(&v0_conf))?;
-                v0::funcs::export_defer(&mut linker, Arc::clone(&v0_conf))?;
-            }
+            Some(Version::V0(ref conf)) => match conf {
+                Some(v0_conf) => {
+                    let v0_conf = Arc::new(Mutex::new(v0_conf.clone()));
+                    v0::funcs::export_tcp_connect(&mut linker, Arc::clone(&v0_conf))?;
+                    v0::funcs::export_accept(&mut linker, Arc::clone(&v0_conf))?;
+                    v0::funcs::export_defer(&mut linker, Arc::clone(&v0_conf))?;
+                }
+                None => {
+                    return Err(anyhow::anyhow!(
+                        "v0_conf wasn't initialized / setup correctly"
+                    ))?;
+                }
+            },
             Some(Version::V1) => {
                 v1::funcs::export_tcp_connect(&mut linker)?;
                 v1::funcs::export_tcplistener_create(&mut linker)?;

diff --git a/crates/water/src/runtime/v0/stream.rs b/crates/water/src/runtime/v0/stream.rs
@@ -179,7 +179,9 @@ impl WATERStreamTrait for WATERStream<Host> {
 
         let store = Arc::clone(&self.core.store);
         let entry_fn = {
-            let mut store = store.lock().unwrap();
+            let mut store = store
+                .lock()
+                .map_err(|e| anyhow::Error::msg(format!("Failed to lock store: {}", e)))?;
             match self
                 .core
                 .instance
@@ -196,7 +198,9 @@ impl WATERStreamTrait for WATERStream<Host> {
         };
 
         let handle = std::thread::spawn(move || {
-            let mut store = store.lock().unwrap();
+            let mut store = store
+                .lock()
+                .map_err(|e| anyhow::Error::msg(format!("Failed to lock store: {}", e)))?;
             let mut res = vec![Val::I32(0); entry_fn.ty(&mut *store).results().len()];
             match entry_fn.call(&mut *store, &[], &mut res) {
                 Ok(_) => Ok(()),

diff --git a/crates/water/src/runtime/version.rs b/crates/water/src/runtime/version.rs
@@ -6,7 +6,7 @@ use crate::runtime::*;
 
 pub enum Version {
     Unknown,
-    V0(V0Config),
+    V0(Option<V0Config>),
     V1,
     V2,
 }
@@ -30,7 +30,7 @@ impl Version {
                     wasm_config.remote_port,
                 )?;
                 // v0_conf.conn = V0CRole::Dialer(std::net::TcpStream::connect(format!("{}:{}", wasm_config.remote_address, wasm_config.remote_port))?);
-                Version::V0(v0_conf)
+                Version::V0(Some(v0_conf))
             }
             WaterBinType::Listen => {
                 let v0_conf = V0Config::init(
@@ -39,7 +39,7 @@ impl Version {
                     wasm_config.local_port,
                 )?;
                 // v0_conf.conn = V0CRole::Listener(std::net::TcpListener::bind(format!("{}:{}", wasm_config.local_address, wasm_config.local_port))?);
-                Version::V0(v0_conf)
+                Version::V0(Some(v0_conf))
             }
             WaterBinType::Unknown => {
                 Version::Unknown // WATER is setting up?
@@ -68,9 +68,7 @@ impl FromStr for Version {
 
     fn from_str(s: &str) -> Result<Version, ()> {
         match s {
-            "_water_v0" => Ok(Version::V0(
-                V0Config::init("".into(), "".into(), 0).unwrap(),
-            )),
+            "_water_v0" => Ok(Version::V0(None)),
             "_water_v1" => Ok(Version::V1),
             "_water_v2" => Ok(Version::V2),
             _ => Err(()),