Add init-plugin subcommand (#798)

crates/cli/src/commands.rs

@@ -48,6 +48,9 @@ pub enum Command {
     /// Emits the plugin binary that is required to run dynamically
     /// linked WebAssembly modules.
     EmitProvider(EmitProviderCommandOpts),
+    /// Initializes a plugin binary.
+    #[command(arg_required_else_help = true)]
+    InitPlugin(InitPluginCommandOpts),
 }
 
 #[derive(Debug, Parser)]
@@ -111,6 +114,16 @@ pub struct EmitProviderCommandOpts {
     pub out: Option<PathBuf>,
 }
 
+#[derive(Debug, Parser)]
+pub struct InitPluginCommandOpts {
+    #[arg(value_name = "PLUGIN", required = true)]
+    /// Path to the plugin to initialize.
+    pub plugin: PathBuf,
+    #[arg(short, long = "out")]
+    /// Output path for the initialized plugin binary (default is stdout).
+    pub out: Option<PathBuf>,
+}
+
 impl<T> ValueParserFactory for GroupOption<T>
 where
     T: GroupDescriptor,

crates/cli/src/main.rs

@@ -15,7 +15,7 @@ use codegen::{CodeGenBuilder, CodeGenType};
 use commands::CodegenOptionGroup;
 use js::JS;
 use js_config::JsConfig;
-use plugins::Plugin;
+use plugins::{Plugin, UninitializedPlugin};
 use std::fs;
 use std::fs::File;
 use std::io::Write;
@@ -80,6 +80,19 @@ fn main() -> Result<()> {
             fs::write(&opts.output, wasm)?;
             Ok(())
         }
+        Command::InitPlugin(opts) => {
+            let plugin_bytes = fs::read(&opts.plugin)?;
+
+            let uninitialized_plugin = UninitializedPlugin::new(&plugin_bytes)?;
+            let initialized_plugin_bytes = uninitialized_plugin.initialize()?;
+
+            let mut out: Box<dyn Write> = match opts.out.as_ref() {
+                Some(path) => Box::new(File::create(path)?),
+                None => Box::new(std::io::stdout()),
+            };
+            out.write_all(&initialized_plugin_bytes)?;
+            Ok(())
+        }
     }
 }
 

crates/cli/src/plugins.rs

@@ -1,12 +1,15 @@
 use crate::bytecode;
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, bail, Result};
 use serde::Deserialize;
 use std::{
+    fs,
     io::{Read, Seek},
     str,
 };
+use walrus::{ExportItem, ValType};
 use wasi_common::{pipe::WritePipe, sync::WasiCtxBuilder};
 use wasmtime::{AsContextMut, Engine, Linker};
+use wizer::Wizer;
 
 const PLUGIN_MODULE: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/plugin.wasm"));
 
@@ -116,8 +119,171 @@ impl Plugin {
     }
 }
 
+/// A validated but uninitialized plugin.
+pub(super) struct UninitializedPlugin<'a> {
+    bytes: &'a [u8],
+}
+
+impl<'a> UninitializedPlugin<'a> {
+    /// Creates a validated but uninitialized plugin.
+    pub fn new(bytes: &'a [u8]) -> Result<Self> {
+        Self::validate(bytes)?;
+        Ok(Self { bytes })
+    }
+
+    fn validate(plugin_bytes: &'a [u8]) -> Result<()> {
+        let mut errors = vec![];
+
+        let module = walrus::Module::from_buffer(plugin_bytes)?;
+
+        if let Err(err) = Self::validate_exported_func(&module, "initialize_runtime", &[], &[]) {
+            errors.push(err);
+        }
+        if let Err(err) = Self::validate_exported_func(
+            &module,
+            "compile_src",
+            &[ValType::I32, ValType::I32],
+            &[ValType::I32],
+        ) {
+            errors.push(err);
+        }
+        if let Err(err) = Self::validate_exported_func(
+            &module,
+            "invoke",
+            &[ValType::I32, ValType::I32, ValType::I32, ValType::I32],
+            &[],
+        ) {
+            errors.push(err);
+        }
+
+        let has_memory = module
+            .exports
+            .iter()
+            .any(|export| export.name == "memory" && matches!(export.item, ExportItem::Memory(_)));
+        if !has_memory {
+            errors.push("missing exported memory named `memory`".to_string());
+        }
+
+        let has_import_namespace = module
+            .customs
+            .iter()
+            .any(|(_, section)| section.name() == "import_namespace");
+        if !has_import_namespace {
+            errors.push("missing custom section named `import_namespace`".to_string());
+        }
+
+        if !errors.is_empty() {
+            bail!("Problems with module: {}", errors.join(", "))
+        }
+        Ok(())
+    }
+
+    /// Initializes the plugin.
+    pub fn initialize(&self) -> Result<Vec<u8>> {
+        let initialized_plugin = Wizer::new()
+            .allow_wasi(true)?
+            .init_func("initialize_runtime")
+            .keep_init_func(true)
+            .wasm_bulk_memory(true)
+            .run(self.bytes)?;
+
+        let tempdir = tempfile::tempdir()?;
+        let in_tempfile_path = tempdir.path().join("in_temp.wasm");
+        let out_tempfile_path = tempdir.path().join("out_temp.wasm");
+        fs::write(&in_tempfile_path, initialized_plugin)?;
+        wasm_opt::OptimizationOptions::new_opt_level_3() // Aggressively optimize for speed.
+            .shrink_level(wasm_opt::ShrinkLevel::Level0) // Don't optimize for size at the expense of performance.
+            .debug_info(false)
+            .run(&in_tempfile_path, &out_tempfile_path)?;
+        Ok(fs::read(out_tempfile_path)?)
+    }
+
+    fn validate_exported_func(
+        module: &walrus::Module,
+        name: &str,
+        expected_params: &[ValType],
+        expected_results: &[ValType],
+    ) -> Result<(), String> {
+        let func_id = module
+            .exports
+            .get_func(name)
+            .map_err(|_| format!("missing export for function named `{name}`"))?;
+        let function = module.funcs.get(func_id);
+        let ty_id = function.ty();
+        let ty = module.types.get(ty_id);
+        let params = ty.params();
+        let has_correct_params = params == expected_params;
+        let results = ty.results();
+        let has_correct_results = results == expected_results;
+        if !has_correct_params || !has_correct_results {
+            return Err(format!("type for function `{name}` is incorrect"));
+        }
+
+        Ok(())
+    }
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 struct ConfigSchema {
     supported_properties: Vec<JsConfigProperty>,
 }
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use walrus::{FunctionBuilder, ModuleConfig, ValType};
+
+    use crate::plugins::UninitializedPlugin;
+
+    #[test]
+    fn test_validate_plugin_with_everything_missing() -> Result<()> {
+        let mut empty_module = walrus::Module::with_config(ModuleConfig::default());
+        let plugin_bytes = empty_module.emit_wasm();
+        let error = UninitializedPlugin::new(&plugin_bytes).err().unwrap();
+        assert_eq!(
+            error.to_string(),
+            "Problems with module: missing export for function named \
+            `initialize_runtime`, missing export for function named \
+            `compile_src`, missing export for function named `invoke`, \
+            missing exported memory named `memory`, missing custom section \
+            named `import_namespace`"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_validate_plugin_with_wrong_params_for_initialize_runtime() -> Result<()> {
+        let mut module = walrus::Module::with_config(ModuleConfig::default());
+        let initialize_runtime = FunctionBuilder::new(&mut module.types, &[ValType::I32], &[])
+            .finish(vec![], &mut module.funcs);
+        module.exports.add("initialize_runtime", initialize_runtime);
+
+        let plugin_bytes = module.emit_wasm();
+        let error = UninitializedPlugin::new(&plugin_bytes).err().unwrap();
+        let expected_part_of_error =
+            "Problems with module: type for function `initialize_runtime` is incorrect,";
+        if !error.to_string().contains(expected_part_of_error) {
+            panic!("Expected error to contain '{expected_part_of_error}' but it did not. Full error is: '{error}'");
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn test_validate_plugin_with_wrong_results_for_initialize_runtime() -> Result<()> {
+        let mut module = walrus::Module::with_config(ModuleConfig::default());
+        let mut initialize_runtime = FunctionBuilder::new(&mut module.types, &[], &[ValType::I32]);
+        initialize_runtime.func_body().i32_const(0);
+        let initialize_runtime = initialize_runtime.finish(vec![], &mut module.funcs);
+        module.exports.add("initialize_runtime", initialize_runtime);
+
+        let plugin_bytes = module.emit_wasm();
+        let error = UninitializedPlugin::new(&plugin_bytes).err().unwrap();
+        let expected_part_of_error =
+            "Problems with module: type for function `initialize_runtime` is incorrect,";
+        if !error.to_string().contains(expected_part_of_error) {
+            panic!("Expected error to contain '{expected_part_of_error}' but it did not. Full error is: '{error}'");
+        }
+        Ok(())
+    }
+}

crates/cli/tests/integration_test.rs

@@ -1,6 +1,8 @@
-use anyhow::Result;
+use anyhow::{bail, Result};
 use javy_runner::{Builder, Runner, RunnerError};
-use std::str;
+use std::{path::PathBuf, process::Command, str};
+use wasi_common::sync::WasiCtxBuilder;
+use wasmtime::{AsContextMut, Engine, Linker, Module, Store};
 
 use javy_test_macros::javy_cli_test;
 
@@ -287,6 +289,61 @@ fn test_exported_default_fn(builder: &mut Builder) -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_init_plugin() -> Result<()> {
+    // This test works by trying to call the `compile_src` function on the
+    // default plugin. The unwizened version should fail because the
+    // underlying Javy runtime has not been initialized yet. Using `init-plugin` on
+    // the unwizened plugin should initialize the runtime so calling
+    // `compile-src` on this module should succeed.
+    let engine = Engine::default();
+    let mut linker = Linker::new(&engine);
+    wasi_common::sync::add_to_linker(&mut linker, |s| s)?;
+    let wasi = WasiCtxBuilder::new().build();
+    let mut store = Store::new(&engine, wasi);
+
+    let uninitialized_plugin = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("..")
+        .join("..")
+        .join(
+            std::path::Path::new("target")
+                .join("wasm32-wasip1")
+                .join("release")
+                .join("plugin.wasm"),
+        );
+
+    // Check that plugin is in fact uninitialized at this point.
+    let module = Module::from_file(&engine, &uninitialized_plugin)?;
+    let instance = linker.instantiate(store.as_context_mut(), &module)?;
+    let result = instance
+        .get_typed_func::<(i32, i32), i32>(store.as_context_mut(), "compile_src")?
+        .call(store.as_context_mut(), (0, 0));
+    // This should fail because the runtime is uninitialized.
+    assert!(result.is_err());
+
+    // Initialize the plugin.
+    let output = Command::new(env!("CARGO_BIN_EXE_javy"))
+        .arg("init-plugin")
+        .arg(uninitialized_plugin.to_str().unwrap())
+        .output()?;
+    if !output.status.success() {
+        bail!(
+            "Running init-command failed with output {}",
+            str::from_utf8(&output.stderr)?,
+        );
+    }
+    let initialized_plugin = output.stdout;
+
+    // Check the plugin is initialized and runs.
+    let module = Module::new(&engine, &initialized_plugin)?;
+    let instance = linker.instantiate(store.as_context_mut(), &module)?;
+    // This should succeed because the runtime is initialized.
+    instance
+        .get_typed_func::<(i32, i32), i32>(store.as_context_mut(), "compile_src")?
+        .call(store.as_context_mut(), (0, 0))?;
+    Ok(())
+}
+
 fn run_with_u8s(r: &mut Runner, stdin: u8) -> (u8, String, u64) {
     let (output, logs, fuel_consumed) = run(r, &stdin.to_le_bytes());
     assert_eq!(1, output.len());