Get replacement instructions from custom section

cranelift/wasm/src/environ/spec.rs

@@ -21,7 +21,9 @@ use cranelift_entity::PrimaryMap;
 use cranelift_frontend::FunctionBuilder;
 use std::boxed::Box;
 use std::string::ToString;
-use wasmparser::{FuncValidator, FunctionBody, Operator, ValidatorResources, WasmFeatures};
+use wasmparser::{
+    FuncValidator, FunctionBody, InstReplacement, Operator, ValidatorResources, WasmFeatures,
+};
 
 /// The value of a WebAssembly global variable.
 #[derive(Clone, Copy)]
@@ -551,6 +553,12 @@ pub trait FuncEnvironment: TargetEnvironment {
     fn is_x86(&self) -> bool {
         false
     }
+
+    /// Get the instruction replacements for this environment.
+    fn inst_replacement(&self, offset: usize) -> Option<InstReplacement> {
+        drop(offset);
+        None
+    }
 }
 
 /// An object satisfying the `ModuleEnvironment` trait can be passed as argument to the

cranelift/wasm/src/func_translator.rs

@@ -4,17 +4,21 @@
 //! function to Cranelift IR guided by a `FuncEnvironment` which provides information about the
 //! WebAssembly module and the runtime environment.
 
-use crate::code_translator::{bitcast_wasm_returns, translate_operator};
-use crate::environ::FuncEnvironment;
-use crate::state::FuncTranslationState;
-use crate::translation_utils::get_vmctx_value_label;
-use crate::WasmResult;
 use core::convert::TryInto;
+use log::debug;
+
+use wasmparser::{self, BinaryReader, FunctionBody, FuncValidator, WasmModuleResources};
+
 use cranelift_codegen::entity::EntityRef;
 use cranelift_codegen::ir::{self, Block, InstBuilder, ValueLabel};
 use cranelift_codegen::timing;
 use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext, Variable};
-use wasmparser::{self, BinaryReader, FuncValidator, FunctionBody, WasmModuleResources};
+
+use crate::code_translator::{bitcast_wasm_returns, translate_operator};
+use crate::environ::FuncEnvironment;
+use crate::state::FuncTranslationState;
+use crate::translation_utils::get_vmctx_value_label;
+use crate::WasmResult;
 
 /// WebAssembly to Cranelift IR function translator.
 ///
@@ -203,12 +207,12 @@ fn declare_locals<FE: FuncEnvironment + ?Sized>(
             builder.ins().vconst(ir::types::I8X16, constant_handle)
         }
         Ref(wasmparser::RefType {
-            nullable: true,
-            heap_type,
-        }) => environ.translate_ref_null(builder.cursor(), heap_type.try_into()?)?,
+                nullable: true,
+                heap_type,
+            }) => environ.translate_ref_null(builder.cursor(), heap_type.try_into()?)?,
         Ref(wasmparser::RefType {
-            nullable: false, ..
-        }) => unreachable!(),
+                nullable: false, ..
+            }) => unreachable!(),
     };
 
     let ty = builder.func.dfg.value_type(zeroval);
@@ -237,14 +241,34 @@ fn parse_function_body<FE: FuncEnvironment + ?Sized>(
     debug_assert_eq!(state.control_stack.len(), 1, "State not initialized");
 
     environ.before_translate_function(builder, state)?;
+
     while !reader.eof() {
         let pos = reader.original_position();
         builder.set_srcloc(cur_srcloc(&reader));
-        let op = reader.read_operator()?;
-        validator.op(pos, &op)?;
-        environ.before_translate_operator(&op, builder, state)?;
-        translate_operator(validator, &op, builder, state, environ)?;
-        environ.after_translate_operator(&op, builder, state)?;
+
+        // Code duplication is necessary here because in the first branch, op lives only as long
+        // as &replacement.inst
+        if let Some(replacement) = environ.inst_replacement(pos) {
+            reader.skip(|r| {
+                let skipped_bytes = r.read_bytes(replacement.skip_bytes as usize)?;
+                debug!("Replacing {} bytes ({:x?}) with {:x?}", skipped_bytes.len(), skipped_bytes, &replacement.inst);
+                Ok(())
+            })?;
+            let mut reader = BinaryReader::new(&replacement.inst);
+            let op = reader.read_operator()?;
+
+            validator.op(pos, &op)?;
+            environ.before_translate_operator(&op, builder, state)?;
+            translate_operator(validator, &op, builder, state, environ)?;
+            environ.after_translate_operator(&op, builder, state)?;
+        } else {
+            let op = reader.read_operator()?;
+
+            validator.op(pos, &op)?;
+            environ.before_translate_operator(&op, builder, state)?;
+            translate_operator(validator, &op, builder, state, environ)?;
+            environ.after_translate_operator(&op, builder, state)?;
+        }
     }
     environ.after_translate_function(builder, state)?;
     let pos = reader.original_position();
@@ -278,16 +302,19 @@ fn cur_srcloc(reader: &BinaryReader) -> ir::SourceLoc {
 
 #[cfg(test)]
 mod tests {
-    use super::FuncTranslator;
-    use crate::environ::DummyEnvironment;
-    use cranelift_codegen::ir::types::I32;
-    use cranelift_codegen::{ir, isa, settings, Context};
     use log::debug;
     use target_lexicon::PointerWidth;
     use wasmparser::{
-        FuncValidator, FunctionBody, Parser, ValidPayload, Validator, ValidatorResources,
+        FunctionBody, FuncValidator, Parser, Validator, ValidatorResources, ValidPayload,
     };
 
+    use cranelift_codegen::{Context, ir, isa, settings};
+    use cranelift_codegen::ir::types::I32;
+
+    use crate::environ::DummyEnvironment;
+
+    use super::FuncTranslator;
+
     #[test]
     fn small1() {
         // Implicit return.
@@ -321,8 +348,8 @@ mod tests {
 
         let (body, mut validator) = extract_func(&wasm);
         trans
-            .translate_body(&mut validator, body, &mut ctx.func, &mut runtime.func_env())
-            .unwrap();
+        .translate_body(&mut validator, body, &mut ctx.func, &mut runtime.func_env())
+        .unwrap();
         debug!("{}", ctx.func.display());
         ctx.verify(&flags).unwrap();
     }
@@ -360,8 +387,8 @@ mod tests {
 
         let (body, mut validator) = extract_func(&wasm);
         trans
-            .translate_body(&mut validator, body, &mut ctx.func, &mut runtime.func_env())
-            .unwrap();
+        .translate_body(&mut validator, body, &mut ctx.func, &mut runtime.func_env())
+        .unwrap();
         debug!("{}", ctx.func.display());
         ctx.verify(&flags).unwrap();
     }
@@ -403,8 +430,8 @@ mod tests {
 
         let (body, mut validator) = extract_func(&wasm);
         trans
-            .translate_body(&mut validator, body, &mut ctx.func, &mut runtime.func_env())
-            .unwrap();
+        .translate_body(&mut validator, body, &mut ctx.func, &mut runtime.func_env())
+        .unwrap();
         debug!("{}", ctx.func.display());
         ctx.verify(&flags).unwrap();
     }

crates/cranelift/src/func_environ.rs

@@ -14,7 +14,7 @@ use cranelift_wasm::{
 };
 use std::convert::TryFrom;
 use std::mem;
-use wasmparser::Operator;
+use wasmparser::{InstReplacement, Operator};
 use wasmtime_environ::{
     BuiltinFunctionIndex, MemoryPlan, MemoryStyle, Module, ModuleTranslation, ModuleTypes, PtrSize,
     TableStyle, Tunables, VMOffsets, WASM_PAGE_SIZE,
@@ -2165,4 +2165,8 @@ impl<'module_environment> cranelift_wasm::FuncEnvironment for FuncEnvironment<'m
     fn is_x86(&self) -> bool {
         self.isa.triple().architecture == target_lexicon::Architecture::X86_64
     }
+
+    fn inst_replacement(&self, offset: usize) -> Option<InstReplacement> {
+        self.translation.instr_replacements.get(&(offset - self.translation.code_offset)).cloned()
+    }
 }

crates/environ/src/component/translate.rs

@@ -337,16 +337,18 @@ impl<'a, 'data> Translator<'a, 'data> {
         // translation process. When doing this only a `Translation` is created
         // which is a simple representation of a component.
         let mut remaining = component;
+        let mut offset = 0;
         loop {
             let payload = match self.parser.parse(remaining, true)? {
                 Chunk::Parsed { payload, consumed } => {
                     remaining = &remaining[consumed..];
+                    offset += consumed;
                     payload
                 }
                 Chunk::NeedMoreData(_) => unreachable!(),
             };
 
-            match self.translate_payload(payload, component)? {
+            match self.translate_payload(payload, component, offset)? {
                 Action::KeepGoing => {}
                 Action::Skip(n) => remaining = &remaining[n..],
                 Action::Done => break,
@@ -379,6 +381,7 @@ impl<'a, 'data> Translator<'a, 'data> {
         &mut self,
         payload: Payload<'data>,
         component: &'data [u8],
+        offset: usize,
     ) -> Result<Action> {
         match payload {
             Payload::Version {
@@ -694,7 +697,11 @@ impl<'a, 'data> Translator<'a, 'data> {
             //
             // FIXME(WebAssembly/component-model#14): probably want to specify
             // and parse a `name` section here.
-            Payload::CustomSection { .. } => {}
+            Payload::CustomSection { 0: reader } => {
+                if reader.name() == "pcsections.mem-safety" {
+                    self.validator.memory_safety_section(&reader, offset)?;
+                }
+            }
 
             // Anything else is either not reachable since we never enable the
             // feature in Wasmtime or we do enable it and it's a bug we don't

crates/environ/src/module_environ.rs

@@ -15,8 +15,8 @@ use std::path::PathBuf;
 use std::sync::Arc;
 use wasmparser::{
     types::Types, CustomSectionReader, DataKind, ElementItems, ElementKind, Encoding, ExternalKind,
-    FuncToValidate, FunctionBody, NameSectionReader, Naming, Operator, Parser, Payload, Type,
-    TypeRef, Validator, ValidatorResources,
+    FuncToValidate, FunctionBody, InstReplacement, NameSectionReader, Naming, Operator, Parser,
+    Payload, Type, TypeRef, Validator, ValidatorResources,
 };
 
 /// Object containing the standalone environment information.
@@ -94,6 +94,13 @@ pub struct ModuleTranslation<'data> {
     /// The type information of the current module made available at the end of the
     /// validation process.
     types: Option<Types>,
+
+    /// At which offset in the binary wasm file the code starts. This is used for instruction
+    /// replacements and is quite hacky.
+    pub code_offset: usize,
+
+    /// Replacement for instructions
+    pub instr_replacements: HashMap<usize, InstReplacement>,
 }
 
 impl<'data> ModuleTranslation<'data> {
@@ -502,6 +509,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
             Payload::CodeSectionStart { count, range, .. } => {
                 self.validator.code_section_start(count, &range)?;
                 let cnt = usize::try_from(count).unwrap();
+                self.result.code_offset = range.start;
                 self.result.function_body_inputs.reserve_exact(cnt);
                 self.result.debuginfo.wasm_file.code_section_offset = range.start as u64;
             }
@@ -648,6 +656,11 @@ and for re-adding support for interface types you can see this issue:
                 ))
             }
 
+            Payload::CustomSection(s) if s.name() == "pcsections.mem-safety" => {
+                self.result
+                    .instr_replacements
+                    .extend(self.validator.memory_safety_section(&s, s.range().start)?);
+            }
             Payload::CustomSection(s) => {
                 self.register_dwarf_section(&s);
             }

tests/mem_safety/stack.wast

@@ -16,14 +16,18 @@
         (local.get $ptr)
         (i64.const 0)
         (i64.const 16)
-        (segment.stack_free)
+        drop
+        drop
+        drop
+        ;;(segment.stack_free)
         (return)
   )
   (func (export "bar") (result i32)
         (local $ptr i64)
         (i64.const 0)
         (i64.const 16)
-        (segment.stack_new)
+        drop
+        ;;(segment.stack_new)
         (local.tee $ptr)
 
         ;; this should generate a mte fault
@@ -37,7 +41,11 @@
         (local.get $ptr)
         (i64.const 0)
         (i64.const 16)
-        (segment.stack_free)
+        ;; 0x1a, len 3 -> segment.stack_free
+        drop
+        drop
+        drop
+        ;;(segment.stack_free)
         (return)
   )
 )