Add optional extra-checks to Wasmi executor and make checking more …

…consistent (#1217)

* use unreachable_unchecked utiltiy in FuncEntity

* use unreachable_unchecked in get_entity codegen

* use unreachable_unchecked in branch executors

* use unreachable_unchecked in call executors

* use unreachable_unchecked in copy executors

* use unreachable_unchecked in load executors

* use unreachable_unchecked in memory executors

* use unreachable_unchecked in return executors

* use unreachable_unchecked in table executors

* use unreachable_unchecked in select executors

* use unreachable_unchecked in store executors

* use unreachable_unchecked for invalid_instruction_word

* add extra-checks crate feature

crates/wasmi/Cargo.toml

@@ -57,6 +57,18 @@ std = [
 # An example of such an environment is `wasm32-unknown-unknown`.
 no-hash-maps = ["wasmi_collections/no-hash-maps"]
 
+# Enables extra checks performed during Wasmi bytecode execution.
+#
+# These checks are unnecessary as long as Wasmi translation works as intended.
+# If Wasmi translation invariants are broken due to bugs, these checks prevent
+# Wasmi execution to exhibit undefined behavior (UB) in certain cases.
+#
+# Expected execution overhead is upt to 20%, if enabled.
+#
+# - Enable if your focus is on safety.
+# - Disable if your focus is on execution speed.
+extra-checks = []
+
 [[bench]]
 name = "benches"
 harness = false

crates/wasmi/src/engine/code_map.rs

@@ -487,7 +487,9 @@ impl FuncEntity {
                 // Safety: we just asserted that `self` must be an uncompiled function
                 //         since otherwise we would have returned `None` above.
                 //         Since this is a performance critical path we need to leave out this check.
-                unsafe { core::hint::unreachable_unchecked() }
+                unsafe {
+                    unreachable_unchecked!("expected uncompiled function but found: {self:?}")
+                }
             }
         }
     }

crates/wasmi/src/engine/executor/instrs.rs

@@ -7,6 +7,7 @@ use crate::{
         bytecode::{index, BlockFuel, Const16, Instruction, Reg},
         code_map::CodeMap,
         executor::stack::{CallFrame, FrameRegisters, ValueStack},
+        utils::unreachable_unchecked,
         DedupFuncType,
         EngineFunc,
     },
@@ -1418,9 +1419,14 @@ macro_rules! get_entity {
                 unsafe { self.cache.$name(index) }
                     .unwrap_or_else(|| {
                         const ENTITY_NAME: &'static str = ::core::stringify!($id_ty);
-                        ::core::unreachable!(
-                            "missing {ENTITY_NAME} at index {index:?} for the currently used instance",
-                        )
+                        // Safety: within the Wasmi executor it is assumed that store entity
+                        //         indices within the Wasmi bytecode are always valid for the
+                        //         store. This is an invariant of the Wasmi translation.
+                        unsafe {
+                            unreachable_unchecked!(
+                                "missing {ENTITY_NAME} at index {index:?} for the currently used instance",
+                            )
+                        }
                     })
             }
         )*
@@ -1727,7 +1733,13 @@ impl<'engine> Executor<'engine> {
     /// This includes [`Instruction`] variants such as [`Instruction::TableIndex`]
     /// that primarily carry parameters for actually executable [`Instruction`].
     fn invalid_instruction_word(&mut self) -> Result<(), Error> {
-        self.execute_trap(TrapCode::UnreachableCodeReached)
+        // Safety: Wasmi translation guarantees that branches are never taken to instruction parameters directly.
+        unsafe {
+            unreachable_unchecked!(
+                "expected instruction but found instruction parameter: {:?}",
+                *self.ip.get()
+            )
+        }
     }
 
     /// Executes a Wasm `unreachable` instruction.

crates/wasmi/src/engine/executor/instrs/branch.rs

@@ -1,14 +1,17 @@
 use super::Executor;
 use crate::{
     core::UntypedVal,
-    engine::bytecode::{
-        BranchOffset,
-        BranchOffset16,
-        Comparator,
-        ComparatorAndOffset,
-        Const16,
-        Instruction,
-        Reg,
+    engine::{
+        bytecode::{
+            BranchOffset,
+            BranchOffset16,
+            Comparator,
+            ComparatorAndOffset,
+            Const16,
+            Instruction,
+            Reg,
+        },
+        utils::unreachable_unchecked,
     },
 };
 use core::cmp;
@@ -58,7 +61,14 @@ impl<'engine> Executor<'engine> {
             Instruction::Const32 { value } => UntypedVal::from(u32::from(value)),
             Instruction::I64Const32 { value } => UntypedVal::from(i64::from(value)),
             Instruction::F64Const32 { value } => UntypedVal::from(f64::from(value)),
-            _ => unreachable!(),
+            unexpected => {
+                // Safety: one of the above instruction parameters is guaranteed to exist by the Wasmi translation.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected instruction parameter for `Instruction::BranchTable1` but found: {unexpected:?}"
+                    )
+                }
+            }
         };
         self.ip.add(offset);
         if let Instruction::BranchTableTarget { results, offset } = *self.ip.get() {
@@ -72,8 +82,16 @@ impl<'engine> Executor<'engine> {
     pub fn execute_branch_table_2(&mut self, index: Reg, len_targets: u32) {
         let offset = self.fetch_branch_table_offset(index, len_targets);
         self.ip.add(1);
-        let Instruction::Register2 { regs } = *self.ip.get() else {
-            unreachable!()
+        let regs = match *self.ip.get() {
+            Instruction::Register2 { regs } => regs,
+            unexpected => {
+                // Safety: Wasmi translation guarantees that `Instruction::Register2` follows.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected `Instruction::Register2` but found: {unexpected:?}"
+                    )
+                }
+            }
         };
         self.ip.add(offset);
         if let Instruction::BranchTableTarget { results, offset } = *self.ip.get() {
@@ -91,8 +109,16 @@ impl<'engine> Executor<'engine> {
     pub fn execute_branch_table_3(&mut self, index: Reg, len_targets: u32) {
         let offset = self.fetch_branch_table_offset(index, len_targets);
         self.ip.add(1);
-        let Instruction::Register3 { regs } = *self.ip.get() else {
-            unreachable!()
+        let regs = match *self.ip.get() {
+            Instruction::Register3 { regs } => regs,
+            unexpected => {
+                // Safety: Wasmi translation guarantees that `Instruction::Register3` follows.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected `Instruction::Register3` but found: {unexpected:?}"
+                    )
+                }
+            }
         };
         self.ip.add(offset);
         if let Instruction::BranchTableTarget { results, offset } = *self.ip.get() {
@@ -110,8 +136,16 @@ impl<'engine> Executor<'engine> {
     pub fn execute_branch_table_span(&mut self, index: Reg, len_targets: u32) {
         let offset = self.fetch_branch_table_offset(index, len_targets);
         self.ip.add(1);
-        let Instruction::RegisterSpan { span: values } = *self.ip.get() else {
-            unreachable!()
+        let values = match *self.ip.get() {
+            Instruction::RegisterSpan { span } => span,
+            unexpected => {
+                // Safety: Wasmi translation guarantees that `Instruction::RegisterSpan` follows.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected `Instruction::RegisterSpan` but found: {unexpected:?}"
+                    )
+                }
+            }
         };
         let len = values.len();
         let values = values.span();
@@ -154,7 +188,12 @@ impl<'engine> Executor<'engine> {
                 // will point to `Instruction::Return` which does the job for us.
                 // This has some technical advantages for us.
             }
-            _ => unreachable!(),
+            unexpected => {
+                // Safety: Wasmi translator guarantees that one of the above `Instruction` variants exists.
+                unsafe {
+                    unreachable_unchecked!("expected target for `Instruction::BranchTableMany` but found: {unexpected:?}")
+                }
+            }
         }
     }
 

crates/wasmi/src/engine/executor/instrs/call.rs

@@ -5,6 +5,7 @@ use crate::{
         bytecode::{index, Instruction, Reg, RegSpan},
         code_map::CompiledFuncRef,
         executor::stack::{CallFrame, FrameParams, ValueStack},
+        utils::unreachable_unchecked,
         EngineFunc,
         FuncParams,
     },
@@ -184,9 +185,14 @@ impl<'engine> Executor<'engine> {
                 let index = u32::from(self.get_register(index));
                 (index, table)
             }
-            unexpected => unreachable!(
-                "expected `Instruction::CallIndirectParams[Imm16]` but found {unexpected:?}"
-            ),
+            unexpected => {
+                // Safety: Wasmi translation guarantees that correct instruction parameter follows.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected `Instruction::CallIndirectParams` but found {unexpected:?}"
+                    )
+                }
+            }
         }
     }
 
@@ -208,9 +214,14 @@ impl<'engine> Executor<'engine> {
                 let index = u32::from(index);
                 (index, table)
             }
-            unexpected => unreachable!(
-                "expected `Instruction::CallIndirectParams[Imm16]` but found {unexpected:?}"
-            ),
+            unexpected => {
+                // Safety: Wasmi translation guarantees that correct instruction parameter follows.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected `Instruction::CallIndirectParamsImm16` but found {unexpected:?}"
+                    )
+                }
+            }
         }
     }
 
@@ -260,9 +271,12 @@ impl<'engine> Executor<'engine> {
                 self.copy_regs(uninit_params, regs);
             }
             unexpected => {
-                unreachable!(
-                    "unexpected Instruction found while copying call parameters: {unexpected:?}"
-                )
+                // Safety: Wasmi translation guarantees that register list finalizer exists.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected register-list finalizer but found: {unexpected:?}"
+                    )
+                }
             }
         }
     }

crates/wasmi/src/engine/executor/instrs/copy.rs

@@ -1,7 +1,10 @@
 use super::{Executor, InstructionPtr};
 use crate::{
     core::UntypedVal,
-    engine::bytecode::{AnyConst32, Const32, FixedRegSpan, Instruction, Reg, RegSpan},
+    engine::{
+        bytecode::{AnyConst32, Const32, FixedRegSpan, Instruction, Reg, RegSpan},
+        utils::unreachable_unchecked,
+    },
 };
 use core::slice;
 use smallvec::SmallVec;
@@ -132,9 +135,14 @@ impl<'engine> Executor<'engine> {
             Instruction::Register { reg } => slice::from_ref(reg),
             Instruction::Register2 { regs } => regs,
             Instruction::Register3 { regs } => regs,
-            unexpected => unreachable!(
-                "unexpected Instruction found while copying many values: {unexpected:?}"
-            ),
+            unexpected => {
+                // Safety: Wasmi translator guarantees that register-list finalizer exists.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected register-list finalizer but found: {unexpected:?}"
+                    )
+                }
+            }
         };
         tmp.extend(values.iter().map(|value| self.get_register(*value)));
         for (result, value) in results.iter_sized(tmp.len()).zip(tmp) {
@@ -175,7 +183,14 @@ impl<'engine> Executor<'engine> {
             Instruction::Register { reg } => slice::from_ref(reg),
             Instruction::Register2 { regs } => regs,
             Instruction::Register3 { regs } => regs,
-            unexpected => unreachable!("unexpected Instruction found while copying many non-overlapping values: {unexpected:?}"),
+            unexpected => {
+                // Safety: Wasmi translator guarantees that register-list finalizer exists.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected register-list finalizer but found: {unexpected:?}"
+                    )
+                }
+            }
         };
         copy_values(values);
         ip

crates/wasmi/src/engine/executor/instrs/load.rs

@@ -4,6 +4,7 @@ use crate::{
     engine::{
         bytecode::{Const16, Reg},
         executor::instr_ptr::InstructionPtr,
+        utils::unreachable_unchecked,
     },
     ir::{index::Memory, Instruction},
     store::StoreInner,
@@ -22,7 +23,12 @@ impl<'engine> Executor<'engine> {
         match *addr.get() {
             Instruction::RegisterAndImm32 { reg, imm } => (reg, u32::from(imm)),
             instr => {
-                unreachable!("expected an `Instruction::RegisterAndImm32` but found: {instr:?}")
+                // Safety: Wasmi translation guarantees that `Instruction::RegisterAndImm32` exists.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected an `Instruction::RegisterAndImm32` but found: {instr:?}"
+                    )
+                }
             }
         }
     }

crates/wasmi/src/engine/executor/instrs/memory.rs

@@ -1,7 +1,10 @@
 use super::{Executor, InstructionPtr};
 use crate::{
     core::TrapCode,
-    engine::bytecode::{index::Data, Const16, Instruction, Reg},
+    engine::{
+        bytecode::{index::Data, Const16, Instruction, Reg},
+        utils::unreachable_unchecked,
+    },
     error::EntityGrowError,
     ir::index::Memory,
     store::{ResourceLimiterRef, StoreInner},
@@ -17,7 +20,12 @@ impl<'engine> Executor<'engine> {
         match *addr.get() {
             Instruction::MemoryIndex { index } => index,
             unexpected => {
-                unreachable!("expected `Instruction::MemoryIndex` but found: {unexpected:?}")
+                // Safety: Wasmi translation guarantees that [`Instruction::MemoryIndex`] exists.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected `Instruction::MemoryIndex` but found: {unexpected:?}"
+                    )
+                }
             }
         }
     }
@@ -29,7 +37,12 @@ impl<'engine> Executor<'engine> {
         match *addr.get() {
             Instruction::DataIndex { index } => index,
             unexpected => {
-                unreachable!("expected `Instruction::DataIndex` but found: {unexpected:?}")
+                // Safety: Wasmi translation guarantees that [`Instruction::DataIndex`] exists.
+                unsafe {
+                    unreachable_unchecked!(
+                        "expected `Instruction::DataIndex` but found: {unexpected:?}"
+                    )
+                }
             }
         }
     }

crates/wasmi/src/engine/executor/instrs/return_.rs

@@ -4,6 +4,7 @@ use crate::{
     engine::{
         bytecode::{AnyConst32, BoundedRegSpan, Const32, Instruction, Reg, RegSpan},
         executor::stack::FrameRegisters,
+        utils::unreachable_unchecked,
     },
     store::StoreInner,
 };
@@ -236,7 +237,14 @@ impl<'engine> Executor<'engine> {
             Instruction::Register { reg } => slice::from_ref(reg),
             Instruction::Register2 { regs } => regs,
             Instruction::Register3 { regs } => regs,
-            unexpected => unreachable!("unexpected `Instruction` found while executing `Instruction::ReturnMany`: {unexpected:?}"),
+            unexpected => {
+                // Safety: Wasmi translation guarantees that a register-list finalizer exists.
+                unsafe {
+                    unreachable_unchecked!(
+                        "unexpected register-list finalizer but found: {unexpected:?}"
+                    )
+                }
+            }
         };
         copy_results(values);
     }