fix(wasix): Prevent recursive merge fallback when mounting packages

Author: theduke

lib/virtual-fs/src/passthru_fs.rs

@@ -2,17 +2,17 @@
 //! needed so that a `Box<dyn VirtualFileSystem>` can be wrapped in an Arc and
 //! shared - some of the interfaces pass around a `Box<dyn VirtualFileSystem>`
 
-use std::path::Path;
+use std::{path::Path, sync::Arc};
 
 use crate::*;
 
 #[derive(Debug)]
 pub struct PassthruFileSystem {
-    fs: Box<dyn FileSystem + Send + Sync + 'static>,
+    fs: Arc<dyn FileSystem + Send + Sync + 'static>,
 }
 
 impl PassthruFileSystem {
-    pub fn new(inner: Box<dyn FileSystem + Send + Sync + 'static>) -> Self {
+    pub fn new(inner: Arc<dyn FileSystem + Send + Sync + 'static>) -> Self {
         Self { fs: inner }
     }
 }
@@ -66,6 +66,8 @@ impl FileSystem for PassthruFileSystem {
 
 #[cfg(test)]
 mod test_builder {
+    use std::sync::Arc;
+
     use tokio::io::{AsyncReadExt, AsyncWriteExt};
 
     use crate::{FileSystem, PassthruFileSystem};
@@ -96,7 +98,7 @@ mod test_builder {
             .unwrap();
         assert_eq!(buf, b"hello");
 
-        let passthru_fs = PassthruFileSystem::new(Box::new(mem_fs.clone()));
+        let passthru_fs = PassthruFileSystem::new(Arc::new(mem_fs.clone()));
         let mut buf = Vec::new();
         passthru_fs
             .new_open_options()

lib/virtual-fs/src/union_fs.rs

@@ -8,7 +8,7 @@ use crate::*;
 
 use std::{path::Path, sync::Arc};
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct MountPoint {
     pub path: PathBuf,
     pub name: String,
@@ -36,6 +36,13 @@ pub struct UnionFileSystem {
     pub mounts: DashMap<PathBuf, MountPoint>,
 }
 
+#[derive(Clone, Copy, Debug)]
+pub enum UnionMergeMode {
+    Replace,
+    Skip,
+    Fail,
+}
+
 impl UnionFileSystem {
     pub fn new() -> Self {
         Self::default()
@@ -58,6 +65,43 @@ impl UnionFileSystem {
                 .to_owned()
         }
     }
+
+    pub fn merge(&self, other: &UnionFileSystem, mode: UnionMergeMode) -> Result<()> {
+        for item in other.mounts.iter() {
+            if self.mounts.contains_key(item.key()) {
+                match mode {
+                    UnionMergeMode::Replace => {
+                        self.mounts.insert(item.key().clone(), item.value().clone());
+                    }
+                    UnionMergeMode::Skip => {
+                        tracing::debug!(
+                            path = %item.key().display(),
+                            "skipping existing mount point while merging two union file systems"
+                        );
+                    }
+                    UnionMergeMode::Fail => {
+                        return Err(FsError::AlreadyExists);
+                    }
+                }
+            } else {
+                self.mounts.insert(item.key().clone(), item.value().clone());
+            }
+        }
+
+        Ok(())
+    }
+
+    // TODO: this was only added for making BinaryPackage::webc_fs clonable,
+    // remove once the TODO in BinaryPackage is resolved
+    pub fn duplicate(&self) -> Self {
+        let mounts = DashMap::new();
+
+        for item in self.mounts.iter() {
+            mounts.insert(item.key().clone(), item.value().clone());
+        }
+
+        Self { mounts }
+    }
 }
 
 impl UnionFileSystem {

lib/wasix/src/bin_factory/binary_package.rs

@@ -3,7 +3,7 @@ use std::{path::Path, sync::Arc};
 use anyhow::Context;
 use once_cell::sync::OnceCell;
 use sha2::Digest;
-use virtual_fs::FileSystem;
+use virtual_fs::UnionFileSystem;
 use wasmer_config::package::{
     PackageHash, PackageId, PackageSource, SuggestedCompilerOptimizations,
 };
@@ -90,7 +90,10 @@ pub struct BinaryPackage {
     /// entrypoint.
     pub entrypoint_cmd: Option<String>,
     pub hash: OnceCell<ModuleHash>,
-    pub webc_fs: Arc<dyn FileSystem + Send + Sync>,
+    // TODO: using a UnionFileSystem here directly is suboptimal, since cloning
+    // it is expensive. Shoudl instead store an immutable map that can easily
+    // be converted into a dashmap.
+    pub webc_fs: Option<Arc<UnionFileSystem>>,
     pub commands: Vec<BinaryPackageCommand>,
     pub uses: Vec<String>,
     pub file_system_memory_footprint: u64,
@@ -264,7 +267,7 @@ impl BinaryPackage {
 mod tests {
     use sha2::Digest;
     use tempfile::TempDir;
-    use virtual_fs::AsyncReadExt;
+    use virtual_fs::{AsyncReadExt, FileSystem as _};
     use wasmer_package::utils::from_disk;
 
     use crate::{
@@ -326,6 +329,8 @@ mod tests {
         // "/public/file.txt" on the guest.
         let mut f = pkg
             .webc_fs
+            .as_ref()
+            .expect("no webc fs")
             .new_open_options()
             .read(true)
             .open("/public/file.txt")

lib/wasix/src/fs/mod.rs

@@ -33,7 +33,10 @@ use futures::{Future, TryStreamExt, future::BoxFuture};
 use serde_derive::{Deserialize, Serialize};
 use tokio::io::AsyncWriteExt;
 use tracing::{debug, trace};
-use virtual_fs::{FileSystem, FsError, OpenOptions, VirtualFile, copy_reference};
+use virtual_fs::{
+    FileSystem, FsError, OpenOptions, UnionFileSystem, VirtualFile, copy_reference,
+    tmp_fs::TmpFileSystem,
+};
 use wasmer_config::package::PackageId;
 use wasmer_wasix_types::{
     types::{__WASI_STDERR_FILENO, __WASI_STDIN_FILENO, __WASI_STDOUT_FILENO},
@@ -362,100 +365,99 @@ impl Default for WasiInodes {
 
 #[derive(Debug, Clone)]
 pub enum WasiFsRoot {
-    Sandbox(Arc<virtual_fs::tmp_fs::TmpFileSystem>),
-    Backing(Arc<Box<dyn FileSystem>>),
-}
-
-impl WasiFsRoot {
-    /// Merge the contents of a filesystem into this one.
-    #[tracing::instrument(level = "debug", skip_all)]
-    pub(crate) async fn merge(
-        &self,
-        other: &Arc<dyn FileSystem + Send + Sync>,
-    ) -> Result<(), virtual_fs::FsError> {
-        match self {
-            WasiFsRoot::Sandbox(fs) => {
-                fs.union(other);
-                Ok(())
-            }
-            WasiFsRoot::Backing(fs) => {
-                merge_filesystems(other, fs).await?;
-                Ok(())
-            }
-        }
-    }
+    Sandbox(TmpFileSystem),
+    Overlay(Arc<virtual_fs::OverlayFileSystem<TmpFileSystem, [UnionFileSystem; 1]>>),
+    Backing(Arc<dyn FileSystem>),
 }
 
 impl FileSystem for WasiFsRoot {
     fn readlink(&self, path: &Path) -> virtual_fs::Result<PathBuf> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.readlink(path),
-            WasiFsRoot::Backing(fs) => fs.readlink(path),
+            Self::Sandbox(fs) => fs.readlink(path),
+            Self::Overlay(overlay) => overlay.readlink(path),
+            Self::Backing(fs) => fs.readlink(path),
         }
     }
 
     fn read_dir(&self, path: &Path) -> virtual_fs::Result<virtual_fs::ReadDir> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.read_dir(path),
-            WasiFsRoot::Backing(fs) => fs.read_dir(path),
+            Self::Sandbox(fs) => fs.read_dir(path),
+            Self::Overlay(overlay) => overlay.read_dir(path),
+            Self::Backing(fs) => fs.read_dir(path),
         }
     }
+
     fn create_dir(&self, path: &Path) -> virtual_fs::Result<()> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.create_dir(path),
-            WasiFsRoot::Backing(fs) => fs.create_dir(path),
+            Self::Sandbox(fs) => fs.create_dir(path),
+            Self::Overlay(overlay) => overlay.create_dir(path),
+            Self::Backing(fs) => fs.create_dir(path),
         }
     }
+
     fn remove_dir(&self, path: &Path) -> virtual_fs::Result<()> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.remove_dir(path),
-            WasiFsRoot::Backing(fs) => fs.remove_dir(path),
+            Self::Sandbox(fs) => fs.remove_dir(path),
+            Self::Overlay(overlay) => overlay.remove_dir(path),
+            Self::Backing(fs) => fs.remove_dir(path),
         }
     }
+
     fn rename<'a>(&'a self, from: &Path, to: &Path) -> BoxFuture<'a, virtual_fs::Result<()>> {
         let from = from.to_owned();
         let to = to.to_owned();
         let this = self.clone();
         Box::pin(async move {
             match this {
-                WasiFsRoot::Sandbox(fs) => fs.rename(&from, &to).await,
-                WasiFsRoot::Backing(fs) => fs.rename(&from, &to).await,
+                Self::Sandbox(fs) => fs.rename(&from, &to).await,
+                Self::Overlay(overlay) => overlay.rename(&from, &to).await,
+                Self::Backing(fs) => fs.rename(&from, &to).await,
             }
         })
     }
+
     fn metadata(&self, path: &Path) -> virtual_fs::Result<virtual_fs::Metadata> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.metadata(path),
-            WasiFsRoot::Backing(fs) => fs.metadata(path),
+            Self::Sandbox(fs) => fs.metadata(path),
+            Self::Overlay(overlay) => overlay.metadata(path),
+            Self::Backing(fs) => fs.metadata(path),
         }
     }
+
     fn symlink_metadata(&self, path: &Path) -> virtual_fs::Result<virtual_fs::Metadata> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.symlink_metadata(path),
-            WasiFsRoot::Backing(fs) => fs.symlink_metadata(path),
+            Self::Sandbox(fs) => fs.symlink_metadata(path),
+            Self::Overlay(overlay) => overlay.symlink_metadata(path),
+            Self::Backing(fs) => fs.symlink_metadata(path),
         }
     }
+
     fn remove_file(&self, path: &Path) -> virtual_fs::Result<()> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.remove_file(path),
-            WasiFsRoot::Backing(fs) => fs.remove_file(path),
+            Self::Sandbox(fs) => fs.remove_file(path),
+            Self::Overlay(overlay) => overlay.remove_file(path),
+            Self::Backing(fs) => fs.remove_file(path),
         }
     }
+
     fn new_open_options(&self) -> OpenOptions<'_> {
         match self {
-            WasiFsRoot::Sandbox(fs) => fs.new_open_options(),
-            WasiFsRoot::Backing(fs) => fs.new_open_options(),
+            Self::Sandbox(fs) => fs.new_open_options(),
+            Self::Overlay(overlay) => overlay.new_open_options(),
+            Self::Backing(fs) => fs.new_open_options(),
         }
     }
+
     fn mount(
         &self,
         name: String,
         path: &Path,
         fs: Box<dyn FileSystem + Send + Sync>,
     ) -> virtual_fs::Result<()> {
         match self {
-            WasiFsRoot::Sandbox(f) => f.mount(name, path, fs),
-            WasiFsRoot::Backing(f) => f.mount(name, path, fs),
+            Self::Sandbox(root) => FileSystem::mount(root, name, path, fs),
+            Self::Overlay(overlay) => FileSystem::mount(overlay.primary(), name, path, fs),
+            Self::Backing(f) => f.mount(name, path, fs),
         }
     }
 }
@@ -638,15 +640,28 @@ impl WasiFs {
         &self,
         binary: &BinaryPackage,
     ) -> Result<(), virtual_fs::FsError> {
-        let needs_to_be_unioned = self.has_unioned.lock().unwrap().insert(binary.id.clone());
+        let Some(webc_fs) = &binary.webc_fs else {
+            return Ok(());
+        };
 
+        let needs_to_be_unioned = self.has_unioned.lock().unwrap().insert(binary.id.clone());
         if !needs_to_be_unioned {
             return Ok(());
         }
 
-        self.root_fs.merge(&binary.webc_fs).await?;
-
-        Ok(())
+        match &self.root_fs {
+            WasiFsRoot::Sandbox(fs) => {
+                // TODO: this can be changed to switch to Self::Overlay instead!
+                let fdyn: Arc<dyn FileSystem + Send + Sync> = webc_fs.clone();
+                fs.union(&fdyn);
+                Ok(())
+            }
+            WasiFsRoot::Overlay(overlay) => {
+                let union = &overlay.secondaries()[0];
+                union.merge(&webc_fs, virtual_fs::UnionMergeMode::Skip)
+            }
+            WasiFsRoot::Backing(backing) => merge_filesystems(webc_fs, backing).await,
+        }
     }
 
     /// Created for the builder API. like `new` but with more information
@@ -2211,14 +2226,14 @@ impl std::fmt::Debug for WasiFs {
 }
 
 /// Returns the default filesystem backing
-pub fn default_fs_backing() -> Box<dyn virtual_fs::FileSystem + Send + Sync> {
+pub fn default_fs_backing() -> Arc<dyn virtual_fs::FileSystem + Send + Sync> {
     cfg_if::cfg_if! {
         if #[cfg(feature = "host-fs")] {
-            Box::new(virtual_fs::host_fs::FileSystem::new(tokio::runtime::Handle::current(), "/").unwrap())
+            Arc::new(virtual_fs::host_fs::FileSystem::new(tokio::runtime::Handle::current(), "/").unwrap())
         } else if #[cfg(not(feature = "host-fs"))] {
-            Box::<virtual_fs::mem_fs::FileSystem>::default()
+            Arc::<virtual_fs::mem_fs::FileSystem>::default()
         } else {
-            Box::<FallbackFileSystem>::default()
+            Arc::<FallbackFileSystem>::default()
         }
     }
 }

lib/wasix/src/runners/wasi.rs

@@ -291,7 +291,12 @@ impl WasiRunner {
                 builder.add_webc(pkg.clone());
                 builder.set_module_hash(pkg.hash());
                 builder.include_packages(pkg.package_ids.clone());
-                Some(Arc::clone(&pkg.webc_fs))
+
+                if let Some(fs) = pkg.webc_fs.as_deref() {
+                    Some(fs.duplicate())
+                } else {
+                    None
+                }
             }
             PackageOrHash::Hash(hash) => {
                 builder.set_module_hash(hash);