Fix multiple-pack for reactor modules

crates/wasi-vfs-cli/Cargo.toml

@@ -9,5 +9,6 @@ name = "wasi-vfs"
 [dependencies]
 anyhow = "1.0.40"
 structopt = "0.3.21"
+wasm-encoder = "0.38.1"
 wasmparser = "0.106.0"
 wizer = { git = "https://github.com/bytecodealliance/wizer.git", rev = "v3.0.1" }

crates/wasi-vfs-cli/src/lib.rs

@@ -84,10 +84,18 @@ pub fn pack(wasm_bytes: &[u8], map_dirs: Vec<(PathBuf, PathBuf)>) -> Result<Vec<
     for (guest_dir, host_dir) in map_dirs {
         wizer.map_dir(guest_dir, host_dir);
     }
-    if is_wasi_reactor(&wasm_bytes) {
+    // For reactor modules, wasi-vfs needs some special initialization process.
+    // 1st pack: Wizer removes `_initialize` and renames `__wasi_vfs_rt_init` to `_initialize`.
+    //           And adds `__wasi_vfs_rt_init` as a new export duplicated from `_initialize`.
+    // 2nd pack: Wizer removes `_initialize` (which was `__wasi_vfs_rt_init` in the 1st pack)
+    //           and renames `__wasi_vfs_rt_init` to `_initialize`.
+    //           And adds `__wasi_vfs_rt_init` as a new export duplicated from `_initialize`.
+    // 3~n pack: Repeat the 2nd pack.
+    if is_wasi_reactor(wasm_bytes) {
         wizer.func_rename("_initialize", "__wasi_vfs_rt_init");
     }
     let output_bytes = wizer.run(&wasm_bytes)?;
+    let output_bytes = copy_export_entry(&output_bytes, "_initialize", "__wasi_vfs_rt_init")?;
     Ok(output_bytes)
 }
 
@@ -115,5 +123,54 @@ fn is_wasi_reactor(bytes: &[u8]) -> bool {
             _ => continue,
         }
     }
-    false
+    return false;
+}
+
+/// Copy an export entry to another name.
+fn copy_export_entry(bytes: &[u8], source: &str, dest: &str) -> Result<Vec<u8>> {
+    let mut module = wasm_encoder::Module::new();
+
+    let parser = wasmparser::Parser::new(0);
+
+    for payload in parser.parse_all(bytes) {
+        let payload = payload?;
+        match payload {
+            wasmparser::Payload::Version { .. } => continue,
+            wasmparser::Payload::ExportSection(export) => {
+                let mut section = wasm_encoder::ExportSection::new();
+                for entry in export {
+                    let entry = entry?;
+                    section.export(entry.name, translate::export_kind(entry.kind), entry.index);
+                    if entry.name == source {
+                        section.export(dest, translate::export_kind(entry.kind), entry.index);
+                    }
+                }
+                module.section(&section);
+            }
+            wasmparser::Payload::End(_) => continue,
+            _ => {
+                if let Some((id, range)) = payload.as_section() {
+                    let raw = wasm_encoder::RawSection {
+                        id,
+                        data: &bytes[range.start..range.end],
+                    };
+                    module.section(&raw);
+                }
+            }
+        }
+    }
+
+    Ok(module.finish())
+}
+
+mod translate {
+    pub(crate) fn export_kind(x: wasmparser::ExternalKind) -> wasm_encoder::ExportKind {
+        match x {
+            wasmparser::ExternalKind::Func => wasm_encoder::ExportKind::Func,
+            wasmparser::ExternalKind::Table => wasm_encoder::ExportKind::Table,
+            wasmparser::ExternalKind::Memory => wasm_encoder::ExportKind::Memory,
+            wasmparser::ExternalKind::Global => wasm_encoder::ExportKind::Global,
+            wasmparser::ExternalKind::Tag => wasm_encoder::ExportKind::Tag,
+        }
+    }
 }

tests/run-make/pack-twice-reactor/Makefile

@@ -0,0 +1,18 @@
+-include ../tools.mk
+
+check: $(objs)
+	$(CC) $(LDFLAGS) main.c $(LIB_WASI_VFS) -mexec-model=reactor -o $(TMPDIR)/main.wasm
+	$(WASI_VFS_CLI) pack $(TMPDIR)/main.wasm --dir ./mnt0::/mnt0 -o $(TMPDIR)/main.stage0.wasm
+	$(NODE) --experimental-wasi-unstable-preview1 ./check.js $(TMPDIR)/main.stage0.wasm 0
+
+	$(WASI_VFS_CLI) pack $(TMPDIR)/main.stage0.wasm --dir ./mnt1::/mnt1 -o $(TMPDIR)/main.stage1.wasm
+	$(NODE) --experimental-wasi-unstable-preview1 ./check.js $(TMPDIR)/main.stage1.wasm 1
+
+	$(NODE) --experimental-wasi-unstable-preview1 ./check.js \
+	  --dir=./mnt0::/mnt0 --dir=./mnt1::/mnt1 --dir=./mnt1_1::/mnt1 \
+	  $(TMPDIR)/main.wasm 2
+	$(WASI_VFS_CLI) pack $(TMPDIR)/main.stage1.wasm --dir ./mnt1_1::/mnt1 -o $(TMPDIR)/main.stage2.wasm
+	$(NODE) --experimental-wasi-unstable-preview1 ./check.js $(TMPDIR)/main.stage2.wasm 2
+
+clean:
+	rm -rf $(TMPDIR)/*

tests/run-make/pack-twice-reactor/check.js

@@ -0,0 +1,29 @@
+const { WASI } = require("wasi");
+const fs = require("fs");
+const process = require("process");
+const path = require("path");
+
+const args = process.argv.slice(2);
+
+const preopens = {};
+while (args.length > 0 && args[0].startsWith("--dir=")) {
+  // Parse --dir=/foo::/bar
+  const dir = args.shift().substring("--dir=".length);
+  const [host, guest] = dir.split("::");
+  preopens[guest] = host;
+}
+
+const buffer = fs.readFileSync(args[0]);
+const stage = Number(args[1]);
+
+const wasi = new WASI({
+  env: { ...process.env },
+  preopens
+});
+const m = new WebAssembly.Module(buffer);
+const i = new WebAssembly.Instance(m, {
+  wasi_snapshot_preview1: wasi.wasiImport,
+});
+
+wasi.initialize(i);
+i.exports.check(stage);

tests/run-make/pack-twice-reactor/main.c

@@ -0,0 +1,25 @@
+#include "../check.h"
+#include <stdio.h>
+
+#pragma clang diagnostic ignored "-Wunknown-attributes"
+__attribute__((export_name("check")))
+int check(int stage) {
+  if (stage == 0) {
+    check_file_exists("/mnt0/hello.txt");
+  } else if (stage == 1) {
+    check_file_exists("/mnt0/hello.txt");
+    check_file_exists("/mnt1/goodbye.txt");
+  } else if (stage == 2) {
+    check_file_exists("/mnt0/hello.txt");
+    check_file_exists("/mnt1/x.txt");
+    check_file_not_exists("/mnt1/goodbye.txt");
+  } else {
+    fprintf(stderr, "Unknown stage: %d\n", stage);
+    return 1;
+  }
+  return 0;
+}
+
+int main(int argc, char *argv[]) {
+  return 1;
+}