Merge #1839

1839: Metering middleware. r=nlewycky a=nlewycky

Based on wasmerio/wasmer-reborn#129 . Original description:

- [x] Necessary APIs for modifying module info.
- [x] Global transformation.
- [x] Per-function transformation.
- [x] Testing and examples.


Co-authored-by: losfair <zhy20000919@hotmail.com>
Co-authored-by: Syrus <me@syrusakbary.com>
Co-authored-by: Nick Lewycky <nick@wasmer.io>
Co-authored-by: nlewycky <nick@wasmer.io>

Cargo.toml

@@ -23,6 +23,7 @@ wasmer-wasi = { version = "1.0.0-alpha5", path = "lib/wasi", optional = true }
 wasmer-wast = { version = "1.0.0-alpha5", path = "tests/lib/wast", optional = true }
 wasmer-cache = { version = "1.0.0-alpha5", path = "lib/cache", optional = true }
 wasmer-types = { version = "1.0.0-alpha5", path = "lib/wasmer-types" }
+wasmer-middlewares = { version = "1.0.0-alpha5", path = "lib/middlewares", optional = true }
 cfg-if = "1.0"
 
 [workspace]
@@ -79,6 +80,7 @@ default = [
     "cache",
     "wasi",
     # "emscripten",
+    "middlewares",
 ]
 engine = []
 jit = [
@@ -117,6 +119,7 @@ llvm = [
     "wasmer-compiler-llvm",
     "compiler",
 ]
+middlewares = ["wasmer-middlewares"]
 
 # Testing features
 test-singlepass = [

lib/api/src/lib.rs

@@ -71,8 +71,7 @@ pub use crate::utils::is_wasm;
 pub use target_lexicon::{Architecture, CallingConvention, OperatingSystem, Triple, HOST};
 #[cfg(feature = "compiler")]
 pub use wasmer_compiler::{
-    wasmparser, CompilerConfig, FunctionMiddleware, FunctionMiddlewareGenerator,
-    MiddlewareReaderState,
+    wasmparser, CompilerConfig, FunctionMiddleware, MiddlewareReaderState, ModuleMiddleware,
 };
 pub use wasmer_compiler::{CpuFeature, Features, Target};
 pub use wasmer_engine::{

lib/compiler-cranelift/src/compiler.rs

@@ -19,11 +19,13 @@ use cranelift_codegen::{binemit, Context};
 #[cfg(feature = "unwind")]
 use gimli::write::{Address, EhFrame, FrameTable};
 use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
+use std::sync::Arc;
 use wasmer_compiler::CompileError;
 use wasmer_compiler::{CallingConvention, ModuleTranslationState, Target};
 use wasmer_compiler::{
     Compilation, CompileModuleInfo, CompiledFunction, CompiledFunctionFrameInfo,
-    CompiledFunctionUnwindInfo, Compiler, Dwarf, FunctionBody, FunctionBodyData, SectionIndex,
+    CompiledFunctionUnwindInfo, Compiler, Dwarf, FunctionBody, FunctionBodyData,
+    ModuleMiddlewareChain, SectionIndex,
 };
 use wasmer_types::entity::{EntityRef, PrimaryMap};
 use wasmer_types::{FunctionIndex, LocalFunctionIndex, SignatureIndex};
@@ -54,14 +56,17 @@ impl Compiler for CraneliftCompiler {
     fn compile_module(
         &self,
         target: &Target,
-        compile_info: &CompileModuleInfo,
+        compile_info: &mut CompileModuleInfo,
         module_translation_state: &ModuleTranslationState,
         function_body_inputs: PrimaryMap<LocalFunctionIndex, FunctionBodyData<'_>>,
     ) -> Result<Compilation, CompileError> {
         let isa = self.config().isa(target);
         let frontend_config = isa.frontend_config();
         let memory_styles = &compile_info.memory_styles;
         let table_styles = &compile_info.table_styles;
+        let mut module = (*compile_info.module).clone();
+        self.config.middlewares.apply_on_module_info(&mut module);
+        compile_info.module = Arc::new(module);
         let module = &compile_info.module;
         let signatures = module
             .signatures
@@ -77,7 +82,7 @@ impl Compiler for CraneliftCompiler {
             // FDEs will cause some issues in Linux.
             None
         } else {
-            use std::sync::{Arc, Mutex};
+            use std::sync::Mutex;
             match target.triple().default_calling_convention() {
                 Ok(CallingConvention::SystemV) => {
                     match isa.create_systemv_cie() {
@@ -125,7 +130,7 @@ impl Compiler for CraneliftCompiler {
                 )?;
 
                 let mut code_buf: Vec<u8> = Vec::new();
-                let mut reloc_sink = RelocSink::new(module, func_index);
+                let mut reloc_sink = RelocSink::new(&module, func_index);
                 let mut trap_sink = TrapSink::new();
                 let mut stackmap_sink = binemit::NullStackMapSink {};
                 context

lib/compiler-cranelift/src/config.rs

@@ -3,7 +3,7 @@ use cranelift_codegen::isa::{lookup, TargetIsa};
 use cranelift_codegen::settings::{self, Configurable};
 use std::sync::Arc;
 use wasmer_compiler::{
-    Architecture, Compiler, CompilerConfig, CpuFeature, FunctionMiddlewareGenerator, Target,
+    Architecture, Compiler, CompilerConfig, CpuFeature, ModuleMiddleware, Target,
 };
 
 // Runtime Environment
@@ -35,7 +35,7 @@ pub struct Cranelift {
     enable_pic: bool,
     opt_level: OptLevel,
     /// The middleware chain.
-    pub(crate) middlewares: Vec<Arc<dyn FunctionMiddlewareGenerator>>,
+    pub(crate) middlewares: Vec<Arc<dyn ModuleMiddleware>>,
 }
 
 impl Cranelift {
@@ -199,7 +199,7 @@ impl CompilerConfig for Cranelift {
     }
 
     /// Pushes a middleware onto the back of the middleware chain.
-    fn push_middleware(&mut self, middleware: Arc<dyn FunctionMiddlewareGenerator>) {
+    fn push_middleware(&mut self, middleware: Arc<dyn ModuleMiddleware>) {
         self.middlewares.push(middleware);
     }
 }

lib/compiler-cranelift/src/translator/func_translator.rs

@@ -19,7 +19,7 @@ use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
 use tracing::info;
 use wasmer_compiler::wasmparser;
 use wasmer_compiler::{
-    to_wasm_error, wasm_unsupported, GenerateMiddlewareChain, MiddlewareBinaryReader,
+    to_wasm_error, wasm_unsupported, MiddlewareBinaryReader, ModuleMiddlewareChain,
     ModuleTranslationState, WasmResult,
 };
 use wasmer_types::LocalFunctionIndex;
@@ -75,7 +75,7 @@ impl FuncTranslator {
         reader.set_middleware_chain(
             config
                 .middlewares
-                .generate_middleware_chain(local_function_index),
+                .generate_function_middleware_chain(local_function_index),
         );
         self.translate_from_reader(module_translation_state, reader, func, environ)
     }

lib/compiler-llvm/src/compiler.rs

@@ -8,15 +8,16 @@ use inkwell::module::{Linkage, Module};
 use inkwell::targets::FileType;
 use inkwell::DLLStorageClass;
 use rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
+use std::sync::Arc;
 use wasmer_compiler::{
     Compilation, CompileError, CompileModuleInfo, Compiler, CustomSection, CustomSectionProtection,
-    Dwarf, FunctionBodyData, ModuleTranslationState, RelocationTarget, SectionBody, SectionIndex,
-    Symbol, SymbolRegistry, Target,
+    Dwarf, FunctionBodyData, ModuleMiddlewareChain, ModuleTranslationState, RelocationTarget,
+    SectionBody, SectionIndex, Symbol, SymbolRegistry, Target,
 };
 use wasmer_types::entity::{EntityRef, PrimaryMap};
 use wasmer_types::{FunctionIndex, LocalFunctionIndex, SignatureIndex};
 
-//use std::sync::{Arc, Mutex};
+//use std::sync::Mutex;
 
 /// A compiler that compiles a WebAssembly module with LLVM, translating the Wasm to LLVM IR,
 /// optimizing it and then translating to assembly.
@@ -210,17 +211,21 @@ impl Compiler for LLVMCompiler {
     fn experimental_native_compile_module<'data, 'module>(
         &self,
         target: &Target,
-        module: &'module CompileModuleInfo,
+        compile_info: &'module mut CompileModuleInfo,
         module_translation: &ModuleTranslationState,
         // The list of function bodies
         function_body_inputs: &PrimaryMap<LocalFunctionIndex, FunctionBodyData<'data>>,
         symbol_registry: &dyn SymbolRegistry,
         // The metadata to inject into the wasmer_metadata section of the object file.
         wasmer_metadata: &[u8],
     ) -> Option<Result<Vec<u8>, CompileError>> {
+        let mut module = (*compile_info.module).clone();
+        self.config.middlewares.apply_on_module_info(&mut module);
+        compile_info.module = Arc::new(module);
+
         Some(self.compile_native_object(
             target,
-            module,
+            compile_info,
             module_translation,
             function_body_inputs,
             symbol_registry,
@@ -233,13 +238,17 @@ impl Compiler for LLVMCompiler {
     fn compile_module<'data, 'module>(
         &self,
         target: &Target,
-        compile_info: &'module CompileModuleInfo,
+        compile_info: &'module mut CompileModuleInfo,
         module_translation: &ModuleTranslationState,
         function_body_inputs: PrimaryMap<LocalFunctionIndex, FunctionBodyData<'data>>,
     ) -> Result<Compilation, CompileError> {
         //let data = Arc::new(Mutex::new(0));
         let memory_styles = &compile_info.memory_styles;
         let table_styles = &compile_info.table_styles;
+
+        let mut module = (*compile_info.module).clone();
+        self.config.middlewares.apply_on_module_info(&mut module);
+        compile_info.module = Arc::new(module);
         let module = &compile_info.module;
 
         // TODO: merge constants in sections.
@@ -260,7 +269,7 @@ impl Compiler for LLVMCompiler {
                     // TODO: remove (to serialize)
                     //let _data = data.lock().unwrap();
                     func_translator.translate(
-                        &module,
+                        module,
                         module_translation,
                         i,
                         input,

lib/compiler-llvm/src/config.rs

@@ -8,7 +8,7 @@ use itertools::Itertools;
 use std::fmt::Debug;
 use std::sync::Arc;
 use target_lexicon::Architecture;
-use wasmer_compiler::{Compiler, CompilerConfig, FunctionMiddlewareGenerator, Target, Triple};
+use wasmer_compiler::{Compiler, CompilerConfig, ModuleMiddleware, Target, Triple};
 use wasmer_types::{FunctionType, LocalFunctionIndex};
 
 /// The InkWell ModuleInfo type
@@ -45,7 +45,7 @@ pub struct LLVM {
     is_pic: bool,
     pub(crate) callbacks: Option<Arc<dyn LLVMCallbacks>>,
     /// The middleware chain.
-    pub(crate) middlewares: Vec<Arc<dyn FunctionMiddlewareGenerator>>,
+    pub(crate) middlewares: Vec<Arc<dyn ModuleMiddleware>>,
 }
 
 impl LLVM {
@@ -212,7 +212,7 @@ impl CompilerConfig for LLVM {
     }
 
     /// Pushes a middleware onto the back of the middleware chain.
-    fn push_middleware(&mut self, middleware: Arc<dyn FunctionMiddlewareGenerator>) {
+    fn push_middleware(&mut self, middleware: Arc<dyn ModuleMiddleware>) {
         self.middlewares.push(middleware);
     }
 }

lib/compiler-llvm/src/translator/code.rs

@@ -26,8 +26,8 @@ use crate::config::{CompiledKind, LLVM};
 use crate::object_file::{load_object_file, CompiledFunction};
 use wasmer_compiler::wasmparser::{MemoryImmediate, Operator};
 use wasmer_compiler::{
-    to_wasm_error, wptype_to_type, CompileError, FunctionBodyData, GenerateMiddlewareChain,
-    MiddlewareBinaryReader, ModuleTranslationState, RelocationTarget, Symbol, SymbolRegistry,
+    to_wasm_error, wptype_to_type, CompileError, FunctionBodyData, MiddlewareBinaryReader,
+    ModuleMiddlewareChain, ModuleTranslationState, RelocationTarget, Symbol, SymbolRegistry,
 };
 use wasmer_types::entity::PrimaryMap;
 use wasmer_types::{
@@ -150,7 +150,7 @@ impl FuncTranslator {
         reader.set_middleware_chain(
             config
                 .middlewares
-                .generate_middleware_chain(*local_func_index),
+                .generate_function_middleware_chain(*local_func_index),
         );
 
         let mut params = vec![];

lib/compiler-singlepass/src/compiler.rs

@@ -13,7 +13,7 @@ use wasmer_compiler::wasmparser::BinaryReaderError;
 use wasmer_compiler::TrapInformation;
 use wasmer_compiler::{Compilation, CompileError, CompiledFunction, Compiler, SectionIndex};
 use wasmer_compiler::{
-    CompileModuleInfo, CompilerConfig, GenerateMiddlewareChain, MiddlewareBinaryReader,
+    CompileModuleInfo, CompilerConfig, MiddlewareBinaryReader, ModuleMiddlewareChain,
     ModuleTranslationState, Target,
 };
 use wasmer_compiler::{FunctionBody, FunctionBodyData};
@@ -47,16 +47,19 @@ impl Compiler for SinglepassCompiler {
     fn compile_module(
         &self,
         _target: &Target,
-        compile_info: &CompileModuleInfo,
+        compile_info: &mut CompileModuleInfo,
         _module_translation: &ModuleTranslationState,
         function_body_inputs: PrimaryMap<LocalFunctionIndex, FunctionBodyData<'_>>,
     ) -> Result<Compilation, CompileError> {
         if compile_info.features.multi_value {
             return Err(CompileError::UnsupportedFeature("multivalue".to_string()));
         }
-        let vmoffsets = VMOffsets::new(8, &compile_info.module);
         let memory_styles = &compile_info.memory_styles;
         let table_styles = &compile_info.table_styles;
+        let mut module = (*compile_info.module).clone();
+        self.config.middlewares.apply_on_module_info(&mut module);
+        compile_info.module = Arc::new(module);
+        let vmoffsets = VMOffsets::new(8, &compile_info.module);
         let module = &compile_info.module;
         let import_trampolines: PrimaryMap<SectionIndex, _> = (0..module.num_imported_functions)
             .map(FunctionIndex::new)
@@ -73,7 +76,10 @@ impl Compiler for SinglepassCompiler {
             .collect::<Vec<(LocalFunctionIndex, &FunctionBodyData<'_>)>>()
             .par_iter()
             .map(|(i, input)| {
-                let middleware_chain = self.config.middlewares.generate_middleware_chain(*i);
+                let middleware_chain = self
+                    .config
+                    .middlewares
+                    .generate_function_middleware_chain(*i);
                 let mut reader =
                     MiddlewareBinaryReader::new_with_offset(input.data, input.module_offset);
                 reader.set_middleware_chain(middleware_chain);

lib/compiler-singlepass/src/config.rs

@@ -3,15 +3,15 @@
 
 use crate::compiler::SinglepassCompiler;
 use std::sync::Arc;
-use wasmer_compiler::{Compiler, CompilerConfig, CpuFeature, FunctionMiddlewareGenerator, Target};
+use wasmer_compiler::{Compiler, CompilerConfig, CpuFeature, ModuleMiddleware, Target};
 use wasmer_types::Features;
 
 #[derive(Debug, Clone)]
 pub struct Singlepass {
     pub(crate) enable_nan_canonicalization: bool,
     pub(crate) enable_stack_check: bool,
     /// The middleware chain.
-    pub(crate) middlewares: Vec<Arc<dyn FunctionMiddlewareGenerator>>,
+    pub(crate) middlewares: Vec<Arc<dyn ModuleMiddleware>>,
 }
 
 impl Singlepass {
@@ -66,7 +66,7 @@ impl CompilerConfig for Singlepass {
     }
 
     /// Pushes a middleware onto the back of the middleware chain.
-    fn push_middleware(&mut self, middleware: Arc<dyn FunctionMiddlewareGenerator>) {
+    fn push_middleware(&mut self, middleware: Arc<dyn ModuleMiddleware>) {
         self.middlewares.push(middleware);
     }
 }

lib/compiler/README.md

@@ -34,7 +34,7 @@ pub trait Compiler {
     fn compile_module<'data, 'module>(
         &self,
         target: &Target,
-        compile_info: &'module CompileModuleInfo,
+        compile_info: &'module mut CompileModuleInfo,
         module_translation: &ModuleTranslationState,
         // The list of function bodies
         function_body_inputs: PrimaryMap<LocalFunctionIndex, FunctionBodyData<'data>>,

lib/compiler/src/compiler.rs

@@ -7,7 +7,7 @@ use crate::lib::std::boxed::Box;
 use crate::lib::std::sync::Arc;
 use crate::module::CompileModuleInfo;
 use crate::target::Target;
-use crate::translator::FunctionMiddlewareGenerator;
+use crate::translator::ModuleMiddleware;
 use crate::FunctionBodyData;
 use crate::ModuleTranslationState;
 use crate::SectionIndex;
@@ -45,7 +45,7 @@ pub trait CompilerConfig {
     }
 
     /// Pushes a middleware onto the back of the middleware chain.
-    fn push_middleware(&mut self, middleware: Arc<dyn FunctionMiddlewareGenerator>);
+    fn push_middleware(&mut self, middleware: Arc<dyn ModuleMiddleware>);
 }
 
 /// An implementation of a Compiler from parsed WebAssembly module to Compiled native code.
@@ -84,7 +84,7 @@ pub trait Compiler {
     fn compile_module<'data, 'module>(
         &self,
         target: &Target,
-        module: &'module CompileModuleInfo,
+        module: &'module mut CompileModuleInfo,
         module_translation: &ModuleTranslationState,
         // The list of function bodies
         function_body_inputs: PrimaryMap<LocalFunctionIndex, FunctionBodyData<'data>>,
@@ -96,7 +96,7 @@ pub trait Compiler {
     fn experimental_native_compile_module<'data, 'module>(
         &self,
         _target: &Target,
-        _module: &'module CompileModuleInfo,
+        _module: &'module mut CompileModuleInfo,
         _module_translation: &ModuleTranslationState,
         // The list of function bodies
         _function_body_inputs: &PrimaryMap<LocalFunctionIndex, FunctionBodyData<'data>>,

lib/compiler/src/lib.rs

@@ -86,8 +86,8 @@ pub use crate::target::{
 #[cfg(feature = "translator")]
 pub use crate::translator::{
     to_wasm_error, translate_module, wptype_to_type, FunctionBodyData, FunctionMiddleware,
-    FunctionMiddlewareGenerator, GenerateMiddlewareChain, MiddlewareBinaryReader,
-    MiddlewareReaderState, ModuleEnvironment, ModuleInfoTranslation, ModuleTranslationState,
+    MiddlewareBinaryReader, MiddlewareReaderState, ModuleEnvironment, ModuleInfoTranslation,
+    ModuleMiddleware, ModuleMiddlewareChain, ModuleTranslationState,
 };
 pub use crate::trap::TrapInformation;
 pub use crate::unwind::CompiledFunctionUnwindInfo;