add dummyfs

lib/mrf-tool/Cargo.toml

@@ -14,6 +14,7 @@ mrf-manifest = { path = "../mrf-manifest", features = [
     "encode",
     "serialise",
 ] }
+polonius-the-crab = "0.4.1"
 serde_json = "1.0.116"
 wasmparser = "0.206.0"
 

lib/mrf-tool/src/fs.rs

@@ -0,0 +1,99 @@
+use polonius_the_crab::{polonius, polonius_return};
+use std::{
+    collections::HashMap,
+    fs::{self, File},
+    io,
+    path::{Path, PathBuf},
+};
+
+pub trait Filesystem {
+    type File<'a>: io::Write
+    where
+        Self: 'a;
+
+    fn copy(&mut self, src: &Path, dst: &Path) -> io::Result<()>;
+    fn read(&mut self, path: &Path) -> io::Result<Vec<u8>>;
+
+    fn create_or_truncate(&mut self, path: &Path) -> io::Result<Self::File<'_>>;
+    fn open_append(&mut self, path: &Path) -> io::Result<Self::File<'_>>;
+}
+
+pub struct DummyFs {
+    inner: HashMap<PathBuf, Vec<u8>>,
+}
+
+#[inline]
+fn file_not_found() -> io::Error {
+    io::Error::new(io::ErrorKind::NotFound, "file not found")
+}
+
+impl Filesystem for DummyFs {
+    type File<'a> = &'a mut Vec<u8>;
+
+    #[inline]
+    fn copy(&mut self, src: &Path, dst: &Path) -> io::Result<()> {
+        let value = self.inner.get(src).ok_or_else(file_not_found)?;
+        self.inner.insert(dst.to_path_buf(), value.clone());
+
+        Ok(())
+    }
+
+    #[inline]
+    fn read(&mut self, path: &Path) -> io::Result<Vec<u8>> {
+        self.inner.get(path).cloned().ok_or_else(file_not_found)
+    }
+
+    #[inline]
+    fn create_or_truncate(&mut self, path: &Path) -> io::Result<Self::File<'_>> {
+        let mut this = self;
+        // TODO: Remove once we can FINALLY have polonius..
+        polonius!(|this| -> io::Result<&'polonius mut Vec<u8>> {
+            if let Some(value) = this.inner.get_mut(path) {
+                value.clear();
+                polonius_return!(Ok(value));
+            }
+        });
+
+        this.inner.insert(path.to_path_buf(), Vec::new());
+        this.create_or_truncate(path)
+    }
+
+    #[inline]
+    fn open_append(&mut self, path: &Path) -> io::Result<Self::File<'_>> {
+        self.inner.get_mut(path).ok_or_else(file_not_found)
+    }
+}
+
+#[derive(Default)]
+pub struct NativeFs {
+    _priv: (),
+}
+
+impl Filesystem for NativeFs {
+    type File<'a> = File;
+
+    #[inline]
+    fn copy(&mut self, src: &Path, dst: &Path) -> io::Result<()> {
+        fs::copy(src, dst)?;
+        Ok(())
+    }
+
+    #[inline]
+    fn read(&mut self, path: &Path) -> io::Result<Vec<u8>> {
+        fs::read(path)
+    }
+
+    #[inline]
+    fn create_or_truncate(&mut self, path: &Path) -> io::Result<Self::File<'_>> {
+        File::options()
+            .create(true)
+            .truncate(true)
+            .write(true)
+            .open(path)
+    }
+
+    #[inline]
+    fn open_append(&mut self, path: &Path) -> io::Result<Self::File<'_>> {
+        File::options().append(true).open(path)
+    }
+}

lib/mrf-tool/src/lib.rs

@@ -1,14 +1,12 @@
 use self::args::{ManifestSubcommand, ModuleSubcommand, ToolArgs, ToolSubcommand};
 use clap::Parser;
 use color_eyre::{eyre::bail, Result};
-use std::{
-    ffi::OsString,
-    fs::{self, File},
-    io::Write,
-    path::Path,
-};
+use std::{ffi::OsString, io::Write, path::Path};
+
+pub use self::fs::{Filesystem, NativeFs};
 
 mod args;
+mod fs;
 
 pub fn read_manifest<W>(sink: &mut W, module: &[u8]) -> Result<()>
 where
@@ -24,62 +22,64 @@ where
     Ok(())
 }
 
-pub fn remove_manifest(module_path: &Path, output_path: &Path) -> Result<()> {
-    let module = fs::read(module_path)?;
+pub fn remove_manifest<F>(fs: &mut F, module_path: &Path, output_path: &Path) -> Result<()>
+where
+    F: Filesystem,
+{
+    let module = fs.read(module_path)?;
     let Some((_manifest, section_range)) = mrf_manifest::decode(&module)? else {
         bail!("missing manifest in module");
     };
 
-    let mut module_file = File::options()
-        .create(true)
-        .truncate(true)
-        .write(true)
-        .open(output_path)?;
-
+    let mut module_file = fs.create_or_truncate(output_path)?;
     module_file.write_all(&module[..section_range.start])?;
     module_file.write_all(&module[section_range.end..])?;
 
     Ok(())
 }
 
-pub fn write_manifest(manifest: &[u8], module_path: &Path) -> Result<()> {
+pub fn write_manifest<F>(fs: &mut F, manifest: &[u8], module_path: &Path) -> Result<()>
+where
+    F: Filesystem,
+{
     // Parse the manifest and re-encode it in canonical JSON
     let parsed_manifest = serde_json::from_slice(manifest)?;
     let custom_section = mrf_manifest::encode(&parsed_manifest)?;
 
-    let mut file = File::options().append(true).open(module_path)?;
+    let mut file = fs.open_append(module_path)?;
     file.write_all(&custom_section)?;
 
     Ok(())
 }
 
-pub fn handle<I, W>(sink: &mut W, input: I) -> Result<()>
+pub fn handle<F, W, I>(fs: &mut F, sink: &mut W, input: I) -> Result<()>
 where
+    F: Filesystem,
+    W: Write,
     I: IntoIterator,
     <I as IntoIterator>::Item: Into<OsString> + Clone,
-    W: Write,
 {
     let args = ToolArgs::try_parse_from(input)?;
     match args.command {
         ToolSubcommand::Manifest(ManifestSubcommand::Add(args)) => {
-            let manifest = fs::read(args.manifest_path)?;
+            let manifest = fs.read(&args.manifest_path)?;
 
             // Only copy if the paths are distinct
             if args.module_path != args.output {
-                fs::copy(&args.module_path, &args.output)?;
+                fs.copy(&args.module_path, &args.output)?;
             }
 
-            write_manifest(&manifest, &args.output)?;
+            write_manifest(fs, &manifest, &args.output)?;
         }
         ToolSubcommand::Manifest(ManifestSubcommand::Read(args)) => {
-            let data = fs::read(args.module_path)?;
+            let data = fs.read(&args.module_path)?;
             read_manifest(sink, &data)?;
         }
         ToolSubcommand::Manifest(ManifestSubcommand::Remove(args)) => {
-            remove_manifest(&args.module_path, &args.output)?;
+            remove_manifest(fs, &args.module_path, &args.output)?;
         }
         ToolSubcommand::Module(ModuleSubcommand::Validate(args)) => {
-            let data = fs::read(args.module_path)?;
+            let data = fs.read(&args.module_path)?;
             wasmparser::validate(&data)?;
         }
     }

lib/mrf-tool/src/main.rs

@@ -1,9 +1,12 @@
 use color_eyre::Result;
+use mrf_tool::NativeFs;
 use std::{env, io};
 
 fn main() -> Result<()> {
+    let mut fs = NativeFs::default();
     let stdout = io::stdout();
-    if let Err(error) = mrf_tool::handle(&mut stdout.lock(), env::args_os()) {
+
+    if let Err(error) = mrf_tool::handle(&mut fs, &mut stdout.lock(), env::args_os()) {
         if let Some(error) = error.downcast_ref::<clap::Error>() {
             error.exit();
         }