Preserve locals before entering control flow during compilation (#945)

* add some regression tests

* rename tests

* simplify expression test

* add commented out ideal and correct codegen

* move checks into sync_local_refs

* implement naive local preservation

This implements naive local preservation in order to see if this fixes the problems with ffmpeg.wasm. This is not efficient during compilation, introduces new attack surfaces at compilation and produces inefficient code.

After asserting that this indeed fixes the issues without major performance loss the following needs to happen:

1. We need a new data structure for locals on the compilation stack to efficiently cope with the fact that all locals on the stack need to be preserved frequently, not just a single local.
2. We need to find a way to reduce dynamic allocations in FuncTranslator::preserve_locals.
3. We need to merge copies whenever possible. Since preservation register indices are contiguous we should be able to encode fused copy instructions instead of many single copy instructions.

* adjust new tests to reflect changes

Seems like the fixes applied to the tests as expected.

* new formatting for conditional preserve tests

* add if based local preservation tests

* apply rustfmt

* apply clippy suggestion

* do local preservation after input copy for loops

* adjust existing translation tests for new codegen

* use local register instead of local index

* remove local preservation before loops

While coming up with test cases I couldn't find a single one that would have invalid behavior before these changes. This is because with only a loop control flow structure there is no way to skip certain parts without using block or if control flow structures within the loop body. Someone please proof me wrong. Getting rid of as many local preservations is good for performance if not needed.

* introduce TranslationBuffers type

This is to unify reusable utility buffers of the Wasmi translator.

* reuse preserved locals buffer

* add clear before buffer usage

This is probably not needed but an additional safety guard

* make encode_copies return Option<Instr>

This is so that the caller can use its encoded root instruction index for notification of preservation in a future commit.

* improve debug_assert in notify_preserved_register

Relax its allowed copy instructions since this is going to be needed with fused copies of local preservation instructions.

* decrease the first preservation slot index by 1

This fixes a bug with RegisterSpanIter going out of bounds when trying to create a register span that starts with the initial preservation register before its index is being defragmented which caused an i16 integer overflow.

Note: we could have also altered the implementation of RegisterSpanIter, however, this would have decreased its efficiency eventually which was less desirable.

* fuse copies for local preservation upon entering blocks or ifs

* make use of slice-group-by no_std support

* consider copy-span encodings for local preservation

* add more local preservation tests with more locals

* shift method up

* count locals on the compilation stack

This removes a potential attack vector with a malicous Wasm file having lots of values on the compilation stack and a series of blocks or ifs that each have to preserve all locals on the compilation stack.
In order to further improve the attack vector resistance we also need to preserve locals on the compilation stack in reverse order of their appearance, e.g. as the compilation stack would have popped from.
This commit also contains some needed cleanups.

* avoid heap memory allocation in preserve_all_locals_inplace

Since this method is only called for small heaps we can get away by using an ArrayVec and linear search instead of using a BTreeMap.

* reverse order of local preservation

This softens certain non-critical attack vectors on the small stack mode implementation.

* add local preservation tests with params

* make LocalRefs more resistant to malicious attackers

* fix internal doc links

* fix bug in new LocalRefs impl

* add conditional preserve with 30 stacked locals

crates/wasmi/Cargo.toml

@@ -25,6 +25,8 @@ smallvec = { version = "1.13.1", features = ["union"] }
 multi-stash = { version = "0.2.0" }
 num-traits = { version = "0.2", default-features = false }
 num-derive = "0.4"
+slice-group-by = { version = "0.3", default-features = false }
+arrayvec = { version = "0.7.4", default-features = false }
 
 [dev-dependencies]
 wat = "1"
@@ -35,7 +37,15 @@ criterion = { version = "0.5", default-features = false }
 
 [features]
 default = ["std"]
-std = ["wasmi_core/std", "wasmi_arena/std", "wasmparser/std", "spin/std", "num-traits/std"]
+std = [
+    "wasmi_core/std",
+    "wasmi_arena/std",
+    "wasmparser/std",
+    "spin/std",
+    "num-traits/std",
+    "slice-group-by/std",
+    "arrayvec/std",
+]
 
 [[bench]]
 name = "benches"

crates/wasmi/src/engine/translator/instr_encoder.rs

@@ -385,7 +385,7 @@ impl InstrEncoder {
         mut results: RegisterSpanIter,
         values: &[TypedProvider],
         fuel_info: FuelInfo,
-    ) -> Result<(), Error> {
+    ) -> Result<Option<Instr>, Error> {
         assert_eq!(results.len(), values.len());
         if let Some((TypedProvider::Register(value), rest)) = values.split_first() {
             if results.span().head() == *value {
@@ -400,24 +400,20 @@ impl InstrEncoder {
         match values {
             [] => {
                 // The copy sequence is empty, nothing to encode in this case.
-                Ok(())
-            }
-            [v0] => {
-                self.encode_copy(stack, result, *v0, fuel_info)?;
-                Ok(())
+                Ok(None)
             }
+            [v0] => self.encode_copy(stack, result, *v0, fuel_info),
             [v0, v1] => {
                 if TypedProvider::Register(result.next()) == *v1 {
                     // Case: the second of the 2 copies is a no-op which we can avoid
                     // Note: we already asserted that the first copy is not a no-op
-                    self.encode_copy(stack, result, *v0, fuel_info)?;
-                    return Ok(());
+                    return self.encode_copy(stack, result, *v0, fuel_info);
                 }
                 let reg0 = Self::provider2reg(stack, v0)?;
                 let reg1 = Self::provider2reg(stack, v1)?;
                 self.bump_fuel_consumption(fuel_info, FuelCosts::base)?;
-                self.push_instr(Instruction::copy2(results.span(), reg0, reg1))?;
-                Ok(())
+                let instr = self.push_instr(Instruction::copy2(results.span(), reg0, reg1))?;
+                Ok(Some(instr))
             }
             [v0, v1, rest @ ..] => {
                 debug_assert!(!rest.is_empty());
@@ -437,22 +433,22 @@ impl InstrEncoder {
                         true => Instruction::copy_span,
                         false => Instruction::copy_span_non_overlapping,
                     };
-                    self.push_instr(make_instr(
+                    let instr = self.push_instr(make_instr(
                         results.span(),
                         values.span(),
                         values.len_as_u16(),
                     ))?;
-                    return Ok(());
+                    return Ok(Some(instr));
                 }
                 let make_instr = match Self::has_overlapping_copies(results, values) {
                     true => Instruction::copy_many,
                     false => Instruction::copy_many_non_overlapping,
                 };
                 let reg0 = Self::provider2reg(stack, v0)?;
                 let reg1 = Self::provider2reg(stack, v1)?;
-                self.push_instr(make_instr(results.span(), reg0, reg1))?;
+                let instr = self.push_instr(make_instr(results.span(), reg0, reg1))?;
                 self.encode_register_list(stack, rest)?;
-                Ok(())
+                Ok(Some(instr))
             }
         }
     }
@@ -807,10 +803,20 @@ impl InstrEncoder {
     ///
     /// This will ignore any preservation notifications after the first one.
     pub fn notify_preserved_register(&mut self, preserve_instr: Instr) {
-        debug_assert!(
-            matches!(self.instrs.get(preserve_instr), Instruction::Copy { .. }),
-            "a preserve instruction is always a register copy instruction"
-        );
+        {
+            let preserved = self.instrs.get(preserve_instr);
+            debug_assert!(
+                matches!(
+                    preserved,
+                    Instruction::Copy { .. }
+                        | Instruction::Copy2 { .. }
+                        | Instruction::CopySpanNonOverlapping { .. }
+                        | Instruction::CopyManyNonOverlapping { .. }
+                ),
+                "a preserve instruction is always a register copy instruction but found: {:?}",
+                preserved,
+            );
+        }
         if self.notified_preservation.is_none() {
             self.notified_preservation = Some(preserve_instr);
         }

crates/wasmi/src/engine/translator/mod.rs

@@ -62,6 +62,7 @@ use crate::{
     FuncType,
 };
 use core::fmt;
+use slice_group_by::GroupBy as _;
 use std::vec::Vec;
 use wasmi_core::{TrapCode, UntypedValue, ValueType};
 use wasmparser::{
@@ -82,10 +83,41 @@ pub struct FuncTranslatorAllocations {
     instr_encoder: InstrEncoder,
     /// The control stack.
     control_stack: ControlStack,
-    /// Buffer to store providers when popped from the [`ValueStack`] in bulk.
-    buffer: Vec<TypedProvider>,
-    /// Buffer to temporarily store `br_table` target depths.
+    /// Some reusable buffers for translation purposes.
+    buffer: TranslationBuffers,
+}
+
+/// Reusable allocations for utility buffers.
+#[derive(Debug, Default)]
+pub struct TranslationBuffers {
+    /// Buffer to temporarily hold a bunch of [`TypedProvider`] when bulk-popped from the [`ValueStack`].
+    providers: Vec<TypedProvider>,
+    /// Buffer to temporarily hold `br_table` target depths.
     br_table_targets: Vec<u32>,
+    /// Buffer to temporarily hold a bunch of preserved [`Register`] locals.
+    preserved: Vec<PreservedLocal>,
+}
+
+/// A pair of local [`Register`] and its preserved [`Register`].
+#[derive(Debug, Copy, Clone)]
+pub struct PreservedLocal {
+    local: Register,
+    preserved: Register,
+}
+
+impl PreservedLocal {
+    /// Creates a new [`PreservedLocal`].
+    pub fn new(local: Register, preserved: Register) -> Self {
+        Self { local, preserved }
+    }
+}
+
+impl TranslationBuffers {
+    fn reset(&mut self) {
+        self.providers.clear();
+        self.br_table_targets.clear();
+        self.preserved.clear();
+    }
 }
 
 impl FuncTranslatorAllocations {
@@ -94,8 +126,7 @@ impl FuncTranslatorAllocations {
         self.stack.reset();
         self.instr_encoder.reset();
         self.control_stack.reset();
-        self.buffer.clear();
-        self.br_table_targets.clear();
+        self.buffer.reset();
     }
 }
 
@@ -739,6 +770,61 @@ impl FuncTranslator {
         self.push_fueled_instr(instr, FuelCosts::base)
     }
 
+    /// Preserve all locals that are currently on the emulated stack.
+    ///
+    /// # Note
+    ///
+    /// This is required for correctness upon entering the compilation
+    /// of a Wasm control flow structure such as a Wasm `block`, `if` or `loop`.
+    /// Locals on the stack might be manipulated conditionally witihn the
+    /// control flow structure and therefore need to be preserved before
+    /// this might happen.
+    /// For efficiency reasons all locals are preserved independent of their
+    /// actual use in the entered control flow structure since the analysis
+    /// of their uses would be too costly.
+    fn preserve_locals(&mut self) -> Result<(), Error> {
+        let fuel_info = self.fuel_info();
+        let preserved = &mut self.alloc.buffer.preserved;
+        preserved.clear();
+        self.alloc.stack.preserve_all_locals(|preserved_local| {
+            preserved.push(preserved_local);
+            Ok(())
+        })?;
+        preserved.reverse();
+        let copy_groups = preserved.linear_group_by(|a, b| {
+            // Note: we group copies into groups with continuous result register indices
+            //       because this is what allows us to fuse single `Copy` instructions into
+            //       more efficient `Copy2` or `CopyManyNonOverlapping` instructions.
+            //
+            // At the time of this writing the author was not sure if all result registers
+            // of all preserved locals are always continuous so this can be understood as
+            // a safety guard.
+            (b.preserved.to_i16() - a.preserved.to_i16()) == 1
+        });
+        for copy_group in copy_groups {
+            let len = copy_group.len();
+            let results = RegisterSpan::new(copy_group[0].preserved).iter(len);
+            let providers = &mut self.alloc.buffer.providers;
+            providers.clear();
+            providers.extend(
+                copy_group
+                    .iter()
+                    .map(|p| p.local)
+                    .map(TypedProvider::Register),
+            );
+            let instr = self.alloc.instr_encoder.encode_copies(
+                &mut self.alloc.stack,
+                results,
+                &providers[..],
+                fuel_info,
+            )?;
+            if let Some(instr) = instr {
+                self.alloc.instr_encoder.notify_preserved_register(instr)
+            }
+        }
+        Ok(())
+    }
+
     /// Convenience function to copy the parameters when branching to a control frame.
     fn translate_copy_branch_params(
         &mut self,
@@ -749,12 +835,12 @@ impl FuncTranslator {
             return Ok(());
         }
         let fuel_info = self.fuel_info();
-        let params = &mut self.alloc.buffer;
+        let params = &mut self.alloc.buffer.providers;
         self.alloc.stack.pop_n(branch_params.len(), params);
         self.alloc.instr_encoder.encode_copies(
             &mut self.alloc.stack,
             branch_params,
-            &self.alloc.buffer[..],
+            &self.alloc.buffer.providers[..],
             fuel_info,
         )?;
         Ok(())
@@ -1103,7 +1189,7 @@ impl FuncTranslator {
         len_block_params: usize,
         len_branch_params: usize,
     ) -> Result<RegisterSpan, Error> {
-        let params = &mut self.alloc.buffer;
+        let params = &mut self.alloc.buffer.providers;
         // Pop the block parameters off the stack.
         self.alloc.stack.pop_n(len_block_params, params);
         // Peek the branch parameter registers which are going to be returned.
@@ -2444,7 +2530,7 @@ impl FuncTranslator {
     fn translate_return_with(&mut self, fuel_info: FuelInfo) -> Result<(), Error> {
         let func_type = self.func_type();
         let results = func_type.results();
-        let values = &mut self.alloc.buffer;
+        let values = &mut self.alloc.buffer.providers;
         self.alloc.stack.pop_n(results.len(), values);
         self.alloc
             .instr_encoder
@@ -2458,7 +2544,7 @@ impl FuncTranslator {
         bail_unreachable!(self);
         let len_results = self.func_type().results().len();
         let fuel_info = self.fuel_info();
-        let values = &mut self.alloc.buffer;
+        let values = &mut self.alloc.buffer.providers;
         self.alloc.stack.peek_n(len_results, values);
         self.alloc.instr_encoder.encode_return_nez(
             &mut self.alloc.stack,