Ensure WasiEnv cleanup in CLI

lib/cli/src/commands/run.rs

@@ -167,17 +167,17 @@ impl RunWithPathBuf {
         })
     }
 
-    fn inner_module_run(&self, mut store: Store, instance: Instance) -> Result<()> {
+    fn inner_module_run(&self, store: &mut Store, instance: Instance) -> Result<()> {
         // If this module exports an _initialize function, run that first.
         if let Ok(initialize) = instance.exports.get_function("_initialize") {
             initialize
-                .call(&mut store, &[])
+                .call(store, &[])
                 .with_context(|| "failed to run _initialize function")?;
         }
 
         // Do we want to invoke a function?
         if let Some(ref invoke) = self.invoke {
-            let result = self.invoke_function(&mut store, &instance, invoke, &self.args)?;
+            let result = self.invoke_function(store, &instance, invoke, &self.args)?;
             println!(
                 "{}",
                 result
@@ -188,7 +188,7 @@ impl RunWithPathBuf {
             );
         } else {
             let start: Function = self.try_find_function(&instance, "_start", &[])?;
-            let result = start.call(&mut store, &[]);
+            let result = start.call(store, &[]);
             #[cfg(feature = "wasi")]
             self.wasi.handle_result(result)?;
             #[cfg(not(feature = "wasi"))]
@@ -299,16 +299,19 @@ impl RunWithPathBuf {
                                 .map(|f| f.to_string_lossy().to_string())
                         })
                         .unwrap_or_default();
-                    let (_ctx, instance) = self
+                    let (ctx, instance) = self
                         .wasi
                         .instantiate(&mut store, &module, program_name, self.args.clone())
                         .with_context(|| "failed to instantiate WASI module")?;
-                    self.inner_module_run(store, instance)
+                    let res = self.inner_module_run(&mut store, instance);
+
+                    ctx.cleanup(&mut store, None);
+                    res
                 }
                 // not WASI
                 _ => {
                     let instance = Instance::new(&mut store, &module, &imports! {})?;
-                    self.inner_module_run(store, instance)
+                    self.inner_module_run(&mut store, instance)
                 }
             }
         };

lib/cli/src/commands/run/wasi.rs

@@ -4,13 +4,13 @@ use std::collections::HashMap;
 use std::path::PathBuf;
 use std::sync::Arc;
 use std::{collections::BTreeSet, path::Path};
-use wasmer::{AsStoreMut, FunctionEnv, Instance, Module, RuntimeError, Value};
+use wasmer::{AsStoreMut, Instance, Module, RuntimeError, Value};
 use wasmer_vfs::FileSystem;
 use wasmer_vfs::{DeviceFile, PassthruFileSystem, RootFileSystemBuilder};
 use wasmer_wasi::types::__WASI_STDIN_FILENO;
 use wasmer_wasi::{
     default_fs_backing, get_wasi_versions, PluggableRuntimeImplementation, WasiEnv, WasiError,
-    WasiVersion,
+    WasiFunctionEnv, WasiVersion,
 };
 
 use clap::Parser;
@@ -108,7 +108,7 @@ impl Wasi {
         module: &Module,
         program_name: String,
         args: Vec<String>,
-    ) -> Result<(FunctionEnv<WasiEnv>, Instance)> {
+    ) -> Result<(WasiFunctionEnv, Instance)> {
         let args = args.iter().cloned().map(|arg| arg.into_bytes());
 
         let map_commands = self
@@ -182,7 +182,7 @@ impl Wasi {
         }
 
         let (instance, wasi_env) = builder.instantiate(module.clone(), store)?;
-        Ok((wasi_env.env, instance))
+        Ok((wasi_env, instance))
     }
 
     /// Helper function for handling the result of a Wasi _start function.

lib/vfs/src/builder.rs

@@ -72,31 +72,31 @@ impl RootFileSystemBuilder {
         if self.add_wasmer_command {
             let _ = tmp
                 .new_open_options_ext()
-                .insert_custom_file(PathBuf::from("/bin/wasmer"), Box::new(NullFile::default()));
+                .insert_device_file(PathBuf::from("/bin/wasmer"), Box::new(NullFile::default()));
         }
         if self.default_dev_files {
             let _ = tmp
                 .new_open_options_ext()
-                .insert_custom_file(PathBuf::from("/dev/null"), Box::new(NullFile::default()));
+                .insert_device_file(PathBuf::from("/dev/null"), Box::new(NullFile::default()));
             let _ = tmp
                 .new_open_options_ext()
-                .insert_custom_file(PathBuf::from("/dev/zero"), Box::new(ZeroFile::default()));
+                .insert_device_file(PathBuf::from("/dev/zero"), Box::new(ZeroFile::default()));
-            let _ = tmp.new_open_options_ext().insert_custom_file(
+            let _ = tmp.new_open_options_ext().insert_device_file(
                 PathBuf::from("/dev/stdin"),
                 self.stdin
                     .unwrap_or_else(|| Box::new(DeviceFile::new(DeviceFile::STDIN))),
             );
-            let _ = tmp.new_open_options_ext().insert_custom_file(
+            let _ = tmp.new_open_options_ext().insert_device_file(
                 PathBuf::from("/dev/stdout"),
                 self.stdout
                     .unwrap_or_else(|| Box::new(DeviceFile::new(DeviceFile::STDOUT))),
             );
-            let _ = tmp.new_open_options_ext().insert_custom_file(
+            let _ = tmp.new_open_options_ext().insert_device_file(
                 PathBuf::from("/dev/stderr"),
                 self.stderr
                     .unwrap_or_else(|| Box::new(DeviceFile::new(DeviceFile::STDERR))),
             );
-            let _ = tmp.new_open_options_ext().insert_custom_file(
+            let _ = tmp.new_open_options_ext().insert_device_file(
                 PathBuf::from("/dev/tty"),
                 self.tty.unwrap_or_else(|| Box::new(NullFile::default())),
             );

lib/vfs/src/mem_fs/file_opener.rs

@@ -202,7 +202,7 @@ impl FileSystem {
 
     /// Inserts a arc file into the file system that references another file
     /// in another file system (does not copy the real data)
-    pub fn insert_custom_file(
+    pub fn insert_device_file(
         &self,
         path: PathBuf,
         file: Box<dyn crate::VirtualFile + Send + Sync>,
@@ -214,51 +214,47 @@ impl FileSystem {
         let inode_of_parent = match inode_of_parent {
             InodeResolution::Found(a) => a,
             InodeResolution::Redirect(..) => {
+                // TODO: should remove the inode again!
                 return Err(FsError::InvalidInput);
             }
         };
 
-        match maybe_inode_of_file {
+        if let Some(_inode_of_file) = maybe_inode_of_file {
-            // The file already exists, then it can not be inserted.
+            // TODO: restore previous inode?
-            Some(_inode_of_file) => return Err(FsError::AlreadyExists),
+            return Err(FsError::AlreadyExists);
+        }
+        // Write lock.
+        let mut fs_lock = self.inner.write().map_err(|_| FsError::Lock)?;
 
-            // The file doesn't already exist; it's OK to create it if
+        // Creating the file in the storage.
-            None => {
+        let inode_of_file = fs_lock.storage.vacant_entry().key();
-                // Write lock.
+        let real_inode_of_file = fs_lock.storage.insert(Node::CustomFile(CustomFileNode {
-                let mut fs_lock = self.inner.write().map_err(|_| FsError::Lock)?;
+            inode: inode_of_file,
-
+            name: name_of_file,
-                // Creating the file in the storage.
+            file: Mutex::new(file),
-                let inode_of_file = fs_lock.storage.vacant_entry().key();
+            metadata: {
-                let real_inode_of_file = fs_lock.storage.insert(Node::CustomFile(CustomFileNode {
+                let time = time();
-                    inode: inode_of_file,
+                Metadata {
-                    name: name_of_file,
+                    ft: FileType {
-                    file: Mutex::new(file),
+                        file: true,
-                    metadata: {
+                        ..Default::default()
-                        let time = time();
-                        Metadata {
-                            ft: FileType {
-                                file: true,
-                                ..Default::default()
-                            },
-                            accessed: time,
-                            created: time,
-                            modified: time,
-                            len: 0,
-                        }
                     },
-                }));
+                    accessed: time,
+                    created: time,
+                    modified: time,
+                    len: 0,
+                }
+            },
+        }));
 
-                assert_eq!(
+        assert_eq!(
-                    inode_of_file, real_inode_of_file,
+            inode_of_file, real_inode_of_file,
-                    "new file inode should have been correctly calculated",
+            "new file inode should have been correctly calculated",
-                );
+        );
 
-                // Adding the new directory to its parent.
+        // Adding the new directory to its parent.
-                fs_lock.add_child_to_node(inode_of_parent, inode_of_file)?;
+        fs_lock.add_child_to_node(inode_of_parent, inode_of_file)?;
 
-                inode_of_file
-            }
-        };
         Ok(())
     }
 

lib/vfs/src/special_file.rs

@@ -31,6 +31,7 @@ impl AsyncSeek for DeviceFile {
     fn start_seek(self: Pin<&mut Self>, _position: SeekFrom) -> io::Result<()> {
         Ok(())
     }
+
     fn poll_complete(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
         Poll::Ready(Ok(0))
     }
@@ -44,12 +45,15 @@ impl AsyncWrite for DeviceFile {
     ) -> Poll<io::Result<usize>> {
         Poll::Ready(Ok(buf.len()))
     }
+
     fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
         Poll::Ready(Ok(()))
     }
+
     fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
         Poll::Ready(Ok(()))
     }
+
     fn poll_write_vectored(
         self: Pin<&mut Self>,
         _cx: &mut Context<'_>,
@@ -57,6 +61,7 @@ impl AsyncWrite for DeviceFile {
     ) -> Poll<io::Result<usize>> {
         Poll::Ready(Ok(bufs.len()))
     }
+
     fn is_write_vectored(&self) -> bool {
         false
     }

lib/wasi/src/fs/mod.rs

@@ -371,40 +371,22 @@ impl WasiFs {
     pub async fn close_all(&self) {
         // TODO: this should close all uniquely owned files instead of just flushing.
 
-        let mut map = self.fd_map.write().unwrap();
+        let to_close = {
-
+            if let Ok(map) = self.fd_map.read() {
-        for (_, item) in map.iter() {
+                map.keys().copied().collect::<Vec<_>>()
-            if let Ok(mut item) = item.inode.inner.kind.write() {
+            } else {
-                match &mut *item {
+                Vec::new()
-                    Kind::File { ref mut handle, .. } => {
-                        if let Some(file_arc) = handle.take() {
-                            match Arc::try_unwrap(file_arc) {
-                                Ok(file_lock) => {
-                                    if let Ok(mut file) = file_lock.into_inner() {
-                                        file.flush().await.ok();
-                                        file.shutdown().await.ok();
-                                    }
-                                }
-                                Err(file_arc) => {
-                                    if let Ok(mut file) = file_arc.write() {
-                                        file.flush().await.ok();
-                                    }
-                                    *handle = Some(file_arc);
-                                }
-                            }
-                        }
-                    }
-                    Kind::Socket { socket: _ } => {}
-                    Kind::Pipe { pipe } => {
-                        pipe.flush().await.ok();
-                        pipe.close();
-                    }
-                    _ => {}
-                }
             }
+        };
+
+        for fd in to_close {
+            self.flush(fd).await.ok();
+            self.close_fd(fd).ok();
         }
 
-        map.clear();
+        if let Ok(mut map) = self.fd_map.write() {
+            map.clear();
+        }
     }
 
     /// Will conditionally union the binary file system with this one