From 40b68830dabf299735e38cb7a75ae60b7c0c8084 Mon Sep 17 00:00:00 2001 From: Jeffrey Charles Date: Fri, 1 Nov 2024 10:32:53 -0400 Subject: [PATCH] Add init-plugin subcommand (#798) --- crates/cli/src/commands.rs | 13 +++ crates/cli/src/main.rs | 15 ++- crates/cli/src/plugins.rs | 168 ++++++++++++++++++++++++++- crates/cli/tests/integration_test.rs | 61 +++++++++- 4 files changed, 253 insertions(+), 4 deletions(-) diff --git a/crates/cli/src/commands.rs b/crates/cli/src/commands.rs index 848d7736..33500603 100644 --- a/crates/cli/src/commands.rs +++ b/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, } +#[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, +} + impl ValueParserFactory for GroupOption where T: GroupDescriptor, diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 1e0e0160..733ae85f 100644 --- a/crates/cli/src/main.rs +++ b/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 = 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(()) + } } } diff --git a/crates/cli/src/plugins.rs b/crates/cli/src/plugins.rs index 57e9aea9..a54306ac 100644 --- a/crates/cli/src/plugins.rs +++ b/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::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> { + 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, } + +#[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(()) + } +} diff --git a/crates/cli/tests/integration_test.rs b/crates/cli/tests/integration_test.rs index 79f6cd25..698f060f 100644 --- a/crates/cli/tests/integration_test.rs +++ b/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());