feat(wasm-gen): customize logic on reaching stack limit (#3448)

Cargo.toml

@@ -183,13 +183,13 @@ tempfile = "3.8.1"
 #
 # fork of `parity-wasm` with sign-ext enabled by default.
 #
-# https://github.com/gear-tech/parity-wasm/tree/v0.45.0-sign-ext
-# gear-wasm = "0.45.0"
+# https://github.com/gear-tech/parity-wasm/tree/v0.45.1-sign-ext
+# gear-wasm = "0.45.1"
 #
 # fork of `wasm-instrument`
 #
-# https://github.com/gear-tech/wasm-instrument/tree/v0.2.1-sign-ext
-gwasm-instrument = { version = "0.2.1", default-features = false }
+# https://github.com/gear-tech/wasm-instrument/tree/v0.2.3-sign-ext
+gwasm-instrument = { version = "0.2.3", default-features = false }
 
 # Internal deps
 authorship = { package = "gear-authorship", path = "node/authorship" }

pallets/gear/src/schedule.rs

@@ -56,7 +56,7 @@ pub const INSTR_BENCHMARK_BATCH_SIZE: u32 = 500;
 // To avoid potential stack overflow problems we have a panic in sandbox in case,
 // execution is ended with stack overflow error. So, process queue execution will be
 // stopped and we will be able to investigate the problem and decrease this constant if needed.
-pub const STACK_HEIGHT_LIMIT: u32 = 18_369;
+pub const STACK_HEIGHT_LIMIT: u32 = 36_743;
 
 /// Definition of the cost schedule and other parameterization for the wasm vm.
 ///
@@ -739,8 +739,7 @@ impl<T: Config> Default for Schedule<T> {
 impl Default for Limits {
     fn default() -> Self {
         Self {
-            // TODO #3435. Disabled stack height is a temp solution.
-            stack_height: cfg!(not(feature = "fuzz")).then_some(STACK_HEIGHT_LIMIT),
+            stack_height: Some(STACK_HEIGHT_LIMIT),
             globals: 256,
             locals: 1024,
             parameters: 128,

pallets/gear/src/tests.rs

@@ -8252,6 +8252,11 @@ fn gas_spent_vs_balance() {
 
 #[test]
 fn gas_spent_precalculated() {
+    use gear_wasm_instrument::parity_wasm::{
+        self,
+        elements::{Instruction, Module},
+    };
+
     // After instrumentation will be:
     // (export "handle" (func $handle_export))
     // (func $add
@@ -8292,102 +8297,25 @@ fn gas_spent_precalculated() {
         )
     )"#;
 
-    // After instrumentation will be:
-    // (export "init" (func $init_export))
-    // (func $init)
-    // (func $init_export
-    //      <-- call gas_charge -->
-    //      <-- stack limit check and increase -->
-    //      call $init
-    //      <-- stack limit decrease -->
-    // )
-    let wat_empty_init = r#"
-    (module
-        (import "env" "memory" (memory 1))
-        (export "init" (func $init))
-        (func $init)
-    )"#;
-
-    // After instrumentation will be:
-    // (export "init" (func $init_export))
-    // (func $f1)
-    // (func $init
-    //      <-- call gas_charge -->
-    //      <-- stack limit check and increase -->
-    //      call $f1
-    //      <-- stack limit decrease -->
-    // )
-    // (func $init_export
-    //      <-- call gas_charge -->
-    //      <-- stack limit check and increase -->
-    //      call $init
-    //      <-- stack limit decrease -->
-    // )
-    let wat_two_stack_limits = r#"
-    (module
-        (import "env" "memory" (memory 1))
-        (export "init" (func $init))
-        (func $f1)
-        (func $init
-            (call $f1)
-        )
-    )"#;
-
-    // After instrumentation will be:
-    // (export "init" (func $init_export))
-    // (func $init
-    //      <-- call gas_charge -->
-    //      <-- stack limit check and increase -->
-    //      i32.const 1
-    //      local.set $1
-    //      <-- stack limit decrease -->
-    // )
-    // (func $init_export
-    //      <-- call gas_charge -->
-    //      <-- stack limit check and increase -->
-    //      call $init
-    //      <-- stack limit decrease -->
-    // )
-    let wat_two_gas_charge = r#"
-    (module
-        (import "env" "memory" (memory 1))
-        (export "init" (func $init))
-        (func $init
-            (local $1 i32)
-            i32.const 1
-            local.set $1
-        )
-    )"#;
-
     init_logger();
     new_test_ext().execute_with(|| {
         let pid = upload_program_default(USER_1, ProgramCodeKind::Custom(wat))
             .expect("submit result was asserted");
-        let empty_init_pid =
-            upload_program_default(USER_3, ProgramCodeKind::Custom(wat_empty_init))
-                .expect("submit result was asserted");
-        let init_two_gas_charge_pid =
-            upload_program_default(USER_3, ProgramCodeKind::Custom(wat_two_gas_charge))
-                .expect("submit result was asserted");
-        let init_two_stack_limits_pid =
-            upload_program_default(USER_3, ProgramCodeKind::Custom(wat_two_stack_limits))
-                .expect("submit result was asserted");
 
         run_to_block(2, None);
 
-        let get_program_code_len = |pid| {
+        let get_program_code = |pid| {
             let code_id = CodeId::from_origin(
                 ProgramStorageOf::<Test>::get_program(pid)
                     .and_then(|program| common::ActiveProgram::try_from(program).ok())
                     .expect("program must exist")
                     .code_hash,
             );
-            <Test as Config>::CodeStorage::get_code(code_id)
-                .unwrap()
-                .code()
-                .len() as u64
+            <Test as Config>::CodeStorage::get_code(code_id).unwrap()
         };
 
+        let get_program_code_len = |pid| get_program_code(pid).code().len() as u64;
+
         let get_gas_charged_for_code = |pid| {
             let schedule = <Test as Config>::Schedule::get();
             let per_byte_cost = schedule.db_read_per_byte.ref_time();
@@ -8398,53 +8326,52 @@ fn gas_spent_precalculated() {
                 + module_instantiation_per_byte * code_len
         };
 
-        let calc_gas_spent_for_init = |wat| {
-            Gear::calculate_gas_info(
-                USER_1.into_origin(),
-                HandleKind::Init(ProgramCodeKind::Custom(wat).to_bytes()),
-                EMPTY_PAYLOAD.to_vec(),
-                0,
-                true,
-                true,
-            )
-            .unwrap()
-            .min_limit
-        };
+        let instrumented_code = get_program_code(pid);
+        let module = parity_wasm::deserialize_buffer::<Module>(instrumented_code.code())
+            .expect("invalid wasm bytes");
 
-        let gas_two_gas_charge = calc_gas_spent_for_init(wat_two_gas_charge);
-        let gas_two_stack_limits = calc_gas_spent_for_init(wat_two_stack_limits);
-        let gas_empty_init = calc_gas_spent_for_init(wat_empty_init);
+        let (handle_export_func_body, gas_charge_func_body) = module
+            .code_section()
+            .and_then(|section| match section.bodies() {
+                [.., handle_export, gas_charge] => Some((handle_export, gas_charge)),
+                _ => None,
+            })
+            .expect("failed to locate `handle_export()` and `gas_charge()` functions");
+
+        let gas_charge_call_cost = gas_charge_func_body
+            .code()
+            .elements()
+            .iter()
+            .find_map(|instruction| match instruction {
+                Instruction::I64Const(cost) => Some(*cost as u64),
+                _ => None,
+            })
+            .expect("failed to get cost of `gas_charge()` function");
+
+        let handle_export_instructions = handle_export_func_body.code().elements();
+        assert!(matches!(
+            handle_export_instructions,
+            [
+                Instruction::I32Const(_), //stack check limit cost
+                Instruction::Call(_),     //call to `gas_charge()`
+                ..
+            ]
+        ));
 
         macro_rules! cost {
             ($name:ident) => {
                 <Test as Config>::Schedule::get().instruction_weights.$name as u64
             };
         }
 
-        // `wat_empty_init` has 1 gas_charge call and
-        // `wat_two_gas_charge` has 2 gas_charge calls, so we can calculate
-        // gas_charge function call cost as difference between them,
-        // taking in account difference in other aspects.
-        let gas_charge_call_cost = (gas_two_gas_charge - gas_empty_init)
-            // Take in account difference in executed instructions
-            - cost!(i64const)
-            - cost!(local_set)
-            // Take in account difference in gas depended on code len
-            - (get_gas_charged_for_code(init_two_gas_charge_pid)
-                - get_gas_charged_for_code(empty_init_pid));
-
-        // `wat_empty_init` has 1 stack limit check and
-        // `wat_two_stack_limits` has 2 stack limit checks, so we can calculate
-        // stack limit check cost as difference between them,
-        // taking in account difference in other aspects.
-        let stack_check_limit_cost = (gas_two_stack_limits - gas_empty_init)
-            // Take in account difference in executed instructions
-            - cost!(call)
-            // Take in account additional gas_charge call
-            - gas_charge_call_cost
-            // Take in account difference in gas depended on code len
-            - (get_gas_charged_for_code(init_two_stack_limits_pid)
-                - get_gas_charged_for_code(empty_init_pid));
+        let stack_check_limit_cost = handle_export_instructions
+            .iter()
+            .find_map(|instruction| match instruction {
+                Instruction::I32Const(cost) => Some(*cost as u64),
+                _ => None,
+            })
+            .expect("failed to get stack check limit cost")
+            - cost!(call);
 
         let gas_spent_expected = {
             let execution_cost = cost!(call) * 2

utils/calc-stack-height/src/main.rs

@@ -60,7 +60,7 @@ fn main() -> anyhow::Result<()> {
 
     let mut stack_height = 0;
 
-    loop {
+    while low <= high {
         let mid = (low + high) / 2;
 
         let code = Code::try_new(
@@ -76,8 +76,6 @@ fn main() -> anyhow::Result<()> {
         let init = instance.exports.get_function("init")?;
         let err = init.call(&[]).unwrap_err();
 
-        let stop = low == high;
-
         match err.to_trap() {
             Some(TrapCode::UnreachableCodeReached) => {
                 low = mid + 1;
@@ -93,10 +91,6 @@ fn main() -> anyhow::Result<()> {
             }
             code => panic!("unexpected trap code: {:?}", code),
         }
-
-        if stop {
-            break;
-        }
     }
 
     println!(

utils/wasm-gen/src/generator.rs

@@ -147,6 +147,8 @@ impl<'a, 'b> GearWasmGenerator<'a, 'b> {
             module
         };
 
+        let module = utils::inject_stack_limiter(module);
+
         Ok(if config.remove_recursions {
             log::trace!("Removing recursions");
             utils::remove_recursion(module)

utils/wasm-gen/src/utils.rs

@@ -26,6 +26,7 @@ use gear_wasm_instrument::{
         },
     },
     syscalls::SysCallName,
+    wasm_instrument::{self, InjectionConfig},
 };
 use gsys::HashWithValue;
 use std::{
@@ -222,6 +223,37 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
+pub fn inject_stack_limiter(module: Module) -> Module {
+    wasm_instrument::inject_stack_limiter_with_config(
+        module,
+        InjectionConfig {
+            stack_limit: 30_003,
+            injection_fn: |signature| {
+                let results = signature.results();
+                let mut body = Vec::with_capacity(results.len() + 1);
+
+                for result in results {
+                    let instruction = match result {
+                        ValueType::I32 => Instruction::I32Const(u32::MAX as i32),
+                        ValueType::I64 => Instruction::I64Const(u64::MAX as i64),
+                        ValueType::F32 | ValueType::F64 => {
+                            unreachable!("f32/64 types are not supported")
+                        }
+                    };
+
+                    body.push(instruction);
+                }
+
+                body.push(Instruction::Return);
+
+                body
+            },
+            stack_height_export_name: None,
+        },
+    )
+    .expect("Failed to inject stack height limits")
+}
+
 /// Injects a critical gas limit to a given wasm module.
 ///
 /// Code before injection gas limiter: