CLI tool update

crates/water/src/config/mod.rs

@@ -1,5 +1,7 @@
 pub mod wasm_shared_config;
 
+/// WATER configuration
+#[derive(Clone)]
 pub struct WATERConfig {
     pub filepath: String,
     pub entry_fn: String,
@@ -27,24 +29,25 @@ impl WATERConfig {
     }
 }
 
+/// WATER client type
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub enum WaterBinType {
-    Unknown,
-    Wrap,
     Dial,
     Listen,
     Relay,
     Runner,
+    Wrap,
+    Unknown,
 }
 
 impl From<u32> for WaterBinType {
     fn from(num: u32) -> Self {
         match num {
             0 => WaterBinType::Dial,
             1 => WaterBinType::Listen,
-            2 => WaterBinType::Runner,
-            3 => WaterBinType::Wrap,
-            4 => WaterBinType::Relay,
+            2 => WaterBinType::Relay,
+            3 => WaterBinType::Runner,
+            4 => WaterBinType::Wrap,
             _ => WaterBinType::Unknown,
         }
     }

crates/water/src/globals.rs

@@ -1,7 +1,7 @@
 #![allow(dead_code)]
 
 pub const WASM_PATH: &str = "./proxy.wasm";
-pub const CONFIG_WASM_PATH: &str = "./conf.json";
+pub const CONFIG_WASM_PATH: &str = "./config.json";
 
 pub const MAIN: &str = "main";
 pub const VERSION_FN: &str = "_water_version";

crates/water/src/runtime/client.rs

@@ -81,6 +81,32 @@ impl WATERClient {
         })
     }
 
+    pub fn keep_listen(&mut self) -> Result<Self, anyhow::Error> {
+        info!("[HOST] WATERClient keep listening...",);
+
+        let water = match &mut self.stream {
+            WATERClientType::Listener(ref mut listener) => WATERClientType::Listener(Box::new(
+                v0::listener::WATERListener::migrate_listener(&self.config, listener.get_core())?,
+            )
+                as Box<dyn WATERListenerTrait>),
+            WATERClientType::Relay(ref mut relay) => WATERClientType::Relay(Box::new(
+                v0::relay::WATERRelay::migrate_listener(&self.config, relay.get_core())?,
+            )
+                as Box<dyn WATERRelayTrait>),
+            _ => {
+                return Err(anyhow::anyhow!(
+                    "[HOST] This client is neither a Listener nor a Relay"
+                ));
+            }
+        };
+
+        Ok(WATERClient {
+            config: self.config.clone(),
+            debug: self.debug,
+            stream: water,
+        })
+    }
+
     pub fn set_debug(&mut self, debug: bool) {
         self.debug = debug;
     }
@@ -163,16 +189,13 @@ impl WATERClient {
 
         match &mut self.stream {
             WATERClientType::Dialer(dialer) => dialer.run_entry_fn(&self.config),
-            WATERClientType::Listener(listener) => {
-                // TODO: clone listener here, since we are doing one WATM instance / accept
-                listener.run_entry_fn(&self.config)
-            }
+            WATERClientType::Listener(listener) => listener.run_entry_fn(&self.config),
             WATERClientType::Relay(relay) => relay.run_entry_fn(&self.config),
             _ => Err(anyhow::anyhow!("This client is not a Runner")),
         }
     }
 
-    // this will run the extry_fn(WATM) in the current thread -- replace Host when running
+    // this will run the entry_fn(WATM) in the current thread -- replace Host when running
     pub fn execute(&mut self) -> Result<(), anyhow::Error> {
         info!("[HOST] WATERClient Executing ...");
 

crates/water/src/runtime/core.rs

@@ -23,16 +23,20 @@ impl H2O<Host> {
     pub fn init(conf: &WATERConfig) -> Result<Self, anyhow::Error> {
         info!("[HOST] WATERCore H2O initing...");
 
-        let mut wasm_config = wasmtime::Config::new();
-        wasm_config.wasm_threads(true);
+        let wasm_config = wasmtime::Config::new();
+
+        #[cfg(feature = "multithread")]
+        {
+            wasm_config.wasm_threads(true);
+        }
 
         let engine = Engine::new(&wasm_config)?;
-        let mut linker: Linker<Host> = Linker::new(&engine);
+        let linker: Linker<Host> = Linker::new(&engine);
 
         let module = Module::from_file(&engine, &conf.filepath)?;
 
         let host = Host::default();
-        let mut store = Store::new(&engine, host);
+        let store = Store::new(&engine, host);
 
         let mut error_occured = None;
 
@@ -64,6 +68,17 @@ impl H2O<Host> {
             return Err(anyhow::Error::msg("WATM module version not found"));
         }
 
+        Self::create_core(conf, linker, store, module, engine, version)
+    }
+
+    pub fn create_core(
+        conf: &WATERConfig,
+        mut linker: Linker<Host>,
+        mut store: Store<Host>,
+        module: Module,
+        engine: Engine,
+        version: Option<Version>,
+    ) -> Result<Self, anyhow::Error> {
         store.data_mut().preview1_ctx = Some(WasiCtxBuilder::new().inherit_stdio().build());
 
         if store.data().preview1_ctx.is_none() {
@@ -93,22 +108,9 @@ impl H2O<Host> {
         match &version {
             Some(Version::V0(ref config)) => match config {
                 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))?;
-
-                    // // if client_type is Listen, then create a listener with the same config
-                    // if conf.client_type == WaterBinType::Listen {
-                    //     match v0_conf.lock() {
-                    //         Ok(mut v0_conf) => {
-                    //             v0_conf.create_listener()?;
-                    //         }
-                    //         Err(e) => {
-                    //             return Err(anyhow::anyhow!("Failed to lock v0_conf: {}", e))?;
-                    //         }
-                    //     }
-                    // }
                 }
                 None => {
                     return Err(anyhow::anyhow!(
@@ -148,6 +150,57 @@ impl H2O<Host> {
         })
     }
 
+    // This function is for migrating the v0 core for listener and relay
+    // to handle every new connection is creating a new separate core (as v0 spec)
+    pub fn v0_migrate_core(conf: &WATERConfig, core: &H2O<Host>) -> Result<Self, anyhow::Error> {
+        info!("[HOST] WATERCore H2O v0_migrating...");
+
+        // reseting the listener accepted_fd or the relay's accepted_fd & dial_fd
+        // when migrating from existed listener / relay
+        let version = match &core.version {
+            Version::V0(v0conf) => {
+                match v0conf {
+                    Some(og_v0_conf) => match og_v0_conf.lock() {
+                        Ok(og_v0_conf) => {
+                            let mut new_v0_conf_inner = og_v0_conf.clone();
+                            // reset the new cloned v0conf
+                            new_v0_conf_inner.reset_listener_or_relay();
+
+                            Version::V0(Some(Arc::new(Mutex::new(new_v0_conf_inner))))
+                        }
+                        Err(e) => {
+                            return Err(anyhow::anyhow!("Failed to lock v0_conf: {}", e))?;
+                        }
+                    },
+                    None => {
+                        return Err(anyhow::anyhow!("v0_conf is None"))?;
+                    }
+                }
+            }
+            _ => {
+                return Err(anyhow::anyhow!("This is not a V0 core"))?;
+            }
+        };
+
+        // NOTE: Some of the followings can reuse the existing core, leave to later explore
+        let wasm_config = wasmtime::Config::new();
+
+        #[cfg(feature = "multithread")]
+        {
+            wasm_config.wasm_threads(true);
+        }
+
+        let engine = Engine::new(&wasm_config)?;
+        let linker: Linker<Host> = Linker::new(&engine);
+
+        let module = Module::from_file(&engine, &conf.filepath)?;
+
+        let host = Host::default();
+        let store = Store::new(&engine, host);
+
+        Self::create_core(conf, linker, store, module, engine, Some(version))
+    }
+
     pub fn _prepare(&mut self, conf: &WATERConfig) -> Result<(), anyhow::Error> {
         self._init(conf.debug)?;
         self._process_config(conf)?; // This is for now needed only by v1_preview

crates/water/src/runtime/transport.rs

@@ -190,30 +190,4 @@ pub trait WATERTransportTrait: Send {
 
         Ok(handle)
     }
-
-    // fn read(&mut self, _buf: &mut Vec<u8>) -> Result<i64, anyhow::Error> {
-    //     Err(anyhow::anyhow!("Method not supported"))
-    // }
-
-    // fn write(&mut self, _buf: &[u8]) -> Result<(), anyhow::Error> {
-    //     Err(anyhow::anyhow!("Method not supported"))
-    // }
-
-    // // v0 only
-    // fn cancel_with(&mut self, _conf: &WATERConfig) -> Result<(), anyhow::Error> {
-    //     Err(anyhow::anyhow!("Method not supported"))
-    // }
-
-    // // v0 only
-    // fn cancel(&mut self, _conf: &WATERConfig) -> Result<(), anyhow::Error> {
-    //     Err(anyhow::anyhow!("Method not supported"))
-    // }
-
-    // // v0 only
-    // fn run_entry_fn(
-    //     &mut self,
-    //     _conf: &WATERConfig,
-    // ) -> Result<std::thread::JoinHandle<Result<(), anyhow::Error>>, anyhow::Error> {
-    //     Err(anyhow::anyhow!("Method not supported"))
-    // }
 }

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

@@ -51,8 +51,8 @@ impl Config {
 pub enum V0CRole {
     Unknown,
     Dialer(i32),
-    Listener(i32),
-    Relay(i32, i32), // listener_fd, dialer_fd
+    Listener(i32, i32),   // listener_fd, accepted_fd
+    Relay(i32, i32, i32), // listener_fd, accepted_fd, dialer_fd
 }
 
 // V0 specific configurations
@@ -93,8 +93,12 @@ impl V0Config {
         info!("[HOST] WATERCore V0 connecting to {}", addr);
 
         match &mut self.conn {
-            V0CRole::Relay(_lis, ref mut conn_fd) => {
+            V0CRole::Relay(_, _, ref mut conn_fd) => {
                 // now relay has been built, need to dial
+                if *conn_fd != -1 {
+                    return Err(anyhow::Error::msg("Relay already connected"));
+                }
+
                 let conn = std::net::TcpStream::connect(addr)?;
                 *conn_fd = conn.as_raw_fd();
                 Ok(conn)
@@ -116,27 +120,40 @@ impl V0Config {
         let listener = std::net::TcpListener::bind(addr)?;
 
         if is_relay {
-            self.conn = V0CRole::Relay(listener.into_raw_fd(), 0);
+            self.conn = V0CRole::Relay(listener.into_raw_fd(), -1, -1);
         } else {
-            self.conn = V0CRole::Listener(listener.into_raw_fd());
+            self.conn = V0CRole::Listener(listener.into_raw_fd(), -1);
         }
         Ok(())
     }
 
     pub fn accept(&mut self) -> Result<std::net::TcpStream, anyhow::Error> {
         info!("[HOST] WATERCore V0 accept with conn {:?} ...", self.conn);
 
-        match &self.conn {
-            V0CRole::Listener(listener) => {
-                let listener = unsafe { std::net::TcpListener::from_raw_fd(*listener) };
+        match self.conn {
+            V0CRole::Listener(ref mut listener_fd, ref mut accepted_fd) => {
+                if *accepted_fd != -1 {
+                    return Err(anyhow::Error::msg("Listener already accepted"));
+                }
+
+                let listener = unsafe { std::net::TcpListener::from_raw_fd(*listener_fd) };
+
                 let (stream, _) = listener.accept()?;
-                self.conn = V0CRole::Listener(listener.into_raw_fd()); // makde sure it is not closed after scope
+
+                *listener_fd = listener.into_raw_fd(); // makde sure the listener is not closed after scope
+                *accepted_fd = stream.as_raw_fd();
+
                 Ok(stream)
             }
-            V0CRole::Relay(listener, _) => {
-                let listener = unsafe { std::net::TcpListener::from_raw_fd(*listener) };
+            V0CRole::Relay(ref mut listener_fd, ref mut accepted_fd, _) => {
+                if *accepted_fd != -1 {
+                    return Err(anyhow::Error::msg("Relay already accepted"));
+                }
+
+                let listener = unsafe { std::net::TcpListener::from_raw_fd(*listener_fd) };
                 let (stream, _) = listener.accept()?;
-                self.conn = V0CRole::Relay(listener.into_raw_fd(), 0); // makde sure it is not closed after scope
+                *listener_fd = listener.into_raw_fd(); // makde sure the listener is not closed after scope
+                *accepted_fd = stream.as_raw_fd();
                 Ok(stream)
             }
             _ => Err(anyhow::Error::msg("not a listener")),
@@ -146,22 +163,56 @@ impl V0Config {
     pub fn defer(&mut self) {
         info!("[HOST] WATERCore V0 defer with conn {:?} ...", self.conn);
 
-        match &self.conn {
-            V0CRole::Listener(_listener) => {
-                // TODO: Listener shouldn't be deferred, but the stream it connected to should be
-                // let listener = unsafe { std::net::TcpListener::from_raw_fd(*listener) };
-                // drop(listener);
+        match self.conn {
+            V0CRole::Listener(_, ref mut accepted_fd) => {
+                // The accepted stream should be defered, not the listener
+                let accepted_conn = unsafe { std::net::TcpStream::from_raw_fd(*accepted_fd) };
+                drop(accepted_conn);
+                *accepted_fd = -1; // set it back to default
             }
-            V0CRole::Dialer(conn) => {
-                let conn = unsafe { std::net::TcpStream::from_raw_fd(*conn) };
+            V0CRole::Dialer(conn_fd) => {
+                let conn = unsafe { std::net::TcpStream::from_raw_fd(conn_fd) };
                 drop(conn);
             }
-            V0CRole::Relay(_listener, conn) => {
-                // Listener shouldn't be deferred, like the above reason
-                // let listener = unsafe { std::net::TcpListener::from_raw_fd(*listener) };
-                // drop(listener);
-                let conn = unsafe { std::net::TcpStream::from_raw_fd(*conn) };
+            V0CRole::Relay(_, ref mut accepted_fd, ref mut conn_fd) => {
+                let accepted_conn = unsafe { std::net::TcpStream::from_raw_fd(*accepted_fd) };
+                drop(accepted_conn);
+                *accepted_fd = -1; // set it back to default
+
+                let conn = unsafe { std::net::TcpStream::from_raw_fd(*conn_fd) };
                 drop(conn);
+                *conn_fd = -1; // set it back to default
+            }
+            _ => {}
+        }
+    }
+
+    pub fn reset_listener_or_relay(&mut self) {
+        info!(
+            "[HOST] WATERCore v0 reset lisener / relay with conn {:?} ...",
+            self.conn
+        );
+
+        match self.conn {
+            V0CRole::Listener(_, ref mut accepted_fd) => {
+                if *accepted_fd != -1 {
+                    let accepted_conn = unsafe { std::net::TcpStream::from_raw_fd(*accepted_fd) };
+                    drop(accepted_conn);
+                    *accepted_fd = -1; // set it back to default
+                }
+            }
+            V0CRole::Relay(_, ref mut accepted_fd, ref mut conn_fd) => {
+                if *accepted_fd != -1 {
+                    let accepted_conn = unsafe { std::net::TcpStream::from_raw_fd(*accepted_fd) };
+                    drop(accepted_conn);
+                    *accepted_fd = -1; // set it back to default
+                }
+
+                if *conn_fd != -1 {
+                    let conn = unsafe { std::net::TcpStream::from_raw_fd(*conn_fd) };
+                    drop(conn);
+                    *conn_fd = -1; // set it back to default
+                }
             }
             _ => {}
         }