feat(wasm-gen): add limits for recursions and heavy loops (#3391)

utils/wasm-gen/src/config.rs

@@ -62,7 +62,8 @@
 //! let wasm_gen_config = GearWasmGeneratorConfig {
 //!     memory_config: memory_pages_config,
 //!     entry_points_config: entry_points_set,
-//!     remove_recursions: true,
+//!     remove_recursions: false,
+//!     critical_gas_limit: Some(1_000_000),
 //!     syscalls_config,
 //! };
 //! ```
@@ -140,6 +141,13 @@ pub struct StandardGearWasmConfigsBundle<T = [u8; 32]> {
     pub existing_addresses: Option<NonEmpty<T>>,
     /// Flag which signals whether recursions must be removed.
     pub remove_recursion: bool,
+    /// If the limit is set to `Some(_)`, programs will try to stop execution
+    /// after reaching a critical gas limit, which can be useful to exit from
+    /// heavy loops and recursions that waste all gas.
+    ///
+    /// The `gr_gas_available` syscall is called at the beginning of each
+    /// function and for each control instruction (blocks, loops, conditions).
+    pub critical_gas_limit: Option<u64>,
     /// Injection type for each syscall.
     pub injection_types: SysCallsInjectionTypes,
     /// Config of gear wasm call entry-points (exports).
@@ -158,6 +166,7 @@ impl<T> Default for StandardGearWasmConfigsBundle<T> {
             log_info: Some("StandardGearWasmConfigsBundle".into()),
             existing_addresses: None,
             remove_recursion: false,
+            critical_gas_limit: Some(1_000_000),
             injection_types: SysCallsInjectionTypes::all_once(),
             entry_points_set: Default::default(),
             initial_pages: DEFAULT_INITIAL_SIZE,
@@ -173,6 +182,7 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
             log_info,
             existing_addresses,
             remove_recursion,
+            critical_gas_limit,
             injection_types,
             entry_points_set,
             initial_pages,
@@ -199,6 +209,7 @@ impl<T: Into<Hash>> ConfigsBundle for StandardGearWasmConfigsBundle<T> {
             upper_limit: None,
         };
         let gear_wasm_generator_config = GearWasmGeneratorConfigBuilder::new()
+            .with_critical_gas_limit(critical_gas_limit)
             .with_recursions_removed(remove_recursion)
             .with_syscalls_config(syscalls_config_builder.build())
             .with_entry_points_config(entry_points_set)

utils/wasm-gen/src/config/generator.rs

@@ -60,6 +60,13 @@ impl GearWasmGeneratorConfigBuilder {
         self
     }
 
+    /// Defines whether programs should have a critical gas limit.
+    pub fn with_critical_gas_limit(mut self, critical_gas_limit: Option<u64>) -> Self {
+        self.0.critical_gas_limit = critical_gas_limit;
+
+        self
+    }
+
     /// Build the gear wasm generator.
     pub fn build(self) -> GearWasmGeneratorConfig {
         self.0
@@ -81,6 +88,9 @@ pub struct GearWasmGeneratorConfig {
     /// Flag, signalizing whether recursions
     /// should be removed from resulting module.
     pub remove_recursions: bool,
+    /// The critical gas limit after which the program
+    /// will attempt to terminate successfully.
+    pub critical_gas_limit: Option<u64>,
 }
 
 /// Memory pages config used by [`crate::MemoryGenerator`].

utils/wasm-gen/src/generator.rs

@@ -139,6 +139,14 @@ impl<'a, 'b> GearWasmGenerator<'a, 'b> {
             .into_wasm_module()
             .into_inner();
 
+        let module = if let Some(critical_gas_limit) = config.critical_gas_limit {
+            log::trace!("Injecting critical gas limit");
+            utils::inject_critical_gas_limit(module, critical_gas_limit)
+        } else {
+            log::trace!("Critical gas limit is not set");
+            module
+        };
+
         Ok(if config.remove_recursions {
             log::trace!("Removing recursions");
             utils::remove_recursion(module)

utils/wasm-gen/src/tests.rs

@@ -47,6 +47,40 @@ use std::{mem, num::NonZeroUsize};
 
 const UNSTRUCTURED_SIZE: usize = 1_000_000;
 
+#[test]
+fn inject_critical_gas_limit_works() {
+    let wat1 = r#"
+    (module
+        (memory $memory0 (import "env" "memory") 16)
+        (export "handle" (func $handle))
+        (func $handle
+            call $f
+            drop
+        )
+        (func $f (result i64)
+            call $f
+        )
+        (func $g
+            (loop $my_loop
+                br $my_loop
+            )
+        )
+    )"#;
+
+    let wasm_bytes = wat::parse_str(wat1).expect("invalid wat");
+    let module =
+        parity_wasm::deserialize_buffer::<Module>(&wasm_bytes).expect("invalid wasm bytes");
+    let module_with_critical_gas_limit = utils::inject_critical_gas_limit(module, 1_000_000);
+
+    let wasm_bytes = module_with_critical_gas_limit
+        .into_bytes()
+        .expect("invalid pw module");
+    assert!(wasmparser::validate(&wasm_bytes).is_ok());
+
+    let wat = wasmprinter::print_bytes(&wasm_bytes).expect("failed printing bytes");
+    println!("wat = {wat}");
+}
+
 #[test]
 fn remove_trivial_recursions() {
     let wat1 = r#"

utils/wasm-gen/src/utils.rs

@@ -16,9 +16,16 @@
 // You should have received a copy of the GNU General Public License
 // along with this program. If not, see <https://www.gnu.org/licenses/>.
 
-use gear_wasm_instrument::parity_wasm::{
-    builder,
-    elements::{self, FuncBody, ImportCountType, Instruction, Module, Type, ValueType},
+use crate::wasm::PageCount as WasmPageCount;
+use gear_wasm_instrument::{
+    parity_wasm::{
+        builder,
+        elements::{
+            BlockType, External, FuncBody, ImportCountType, Instruction, Instructions, Internal,
+            Module, Section, Type, ValueType,
+        },
+    },
+    syscalls::SysCallName,
 };
 use gsys::HashWithValue;
 use std::{
@@ -95,7 +102,7 @@ pub fn remove_recursion(module: Module) -> Module {
                     .with_results(signature.results().to_vec())
                     .build()
                     .body()
-                    .with_instructions(elements::Instructions::new(body))
+                    .with_instructions(Instructions::new(body))
                     .build()
                     .build(),
             )
@@ -215,6 +222,216 @@ fn find_recursion_impl<Callback>(
     path.pop();
 }
 
+/// Injects a critical gas limit to a given wasm module.
+///
+/// Code before injection gas limiter:
+/// ```ignore
+/// fn func() {
+///     func();
+///     loop { }
+/// }
+/// ```
+///
+/// Code after injection gas limiter:
+/// ```ignore
+/// use gcore::exec;
+///
+/// const CRITICAL_GAS_LIMIT: u64 = 1_000_000;
+///
+/// fn func() {
+///     // exit from recursions
+///     if exec::gas_available() <= CRITICAL_GAS_LIMIT {
+///         return;
+///     }
+///     func();
+///     loop {
+///         // exit from heavy loops
+///         if exec::gas_available() <= CRITICAL_GAS_LIMIT {
+///             return;
+///         }
+///     }
+/// }
+/// ```
+pub fn inject_critical_gas_limit(module: Module, critical_gas_limit: u64) -> Module {
+    // get initial memory size of program
+    let Some(mem_size) = module
+        .import_section()
+        .and_then(|section| {
+            section
+                .entries()
+                .iter()
+                .find_map(|entry| match entry.external() {
+                    External::Memory(mem_ty) => Some(mem_ty.limits().initial()),
+                    _ => None,
+                })
+        })
+        .map(Into::<WasmPageCount>::into)
+        .map(|page_count| page_count.memory_size())
+    else {
+        return module;
+    };
+
+    // store available gas pointer on the last memory page
+    let gas_ptr = mem_size - mem::size_of::<u64>() as u32;
+
+    // add gr_gas_available import if needed
+    let maybe_gr_gas_available_index = module.import_section().and_then(|section| {
+        section
+            .entries()
+            .iter()
+            .filter(|entry| matches!(entry.external(), External::Function(_)))
+            .enumerate()
+            .find_map(|(i, entry)| {
+                (entry.module() == "env" && entry.field() == SysCallName::GasAvailable.to_str())
+                    .then_some(i as u32)
+            })
+    });
+    // sections should only be rewritten if the module did not previously have gr_gas_available import
+    let rewrite_sections = maybe_gr_gas_available_index.is_none();
+
+    let (gr_gas_available_index, mut module) = match maybe_gr_gas_available_index {
+        Some(gr_gas_available_index) => (gr_gas_available_index, module),
+        None => {
+            let mut mbuilder = builder::from_module(module);
+
+            // fn gr_gas_available(gas: *mut u64);
+            let import_sig = mbuilder
+                .push_signature(builder::signature().with_param(ValueType::I32).build_sig());
+
+            mbuilder.push_import(
+                builder::import()
+                    .module("env")
+                    .field(SysCallName::GasAvailable.to_str())
+                    .external()
+                    .func(import_sig)
+                    .build(),
+            );
+
+            // back to plain module
+            let module = mbuilder.build();
+
+            let import_count = module.import_count(ImportCountType::Function);
+            let gr_gas_available_index = import_count as u32 - 1;
+
+            (gr_gas_available_index, module)
+        }
+    };
+
+    let (Some(type_section), Some(function_section)) =
+        (module.type_section(), module.function_section())
+    else {
+        return module;
+    };
+
+    let types = type_section.types().to_vec();
+    let signature_entries = function_section.entries().to_vec();
+
+    let Some(code_section) = module.code_section_mut() else {
+        return module;
+    };
+
+    for (index, func_body) in code_section.bodies_mut().iter_mut().enumerate() {
+        let signature_index = &signature_entries[index];
+        let signature = &types[signature_index.type_ref() as usize];
+        let Type::Function(signature) = signature;
+        let results = signature.results();
+
+        // create the body of the gas limiter:
+        let mut body = Vec::with_capacity(results.len() + 9);
+        body.extend_from_slice(&[
+            // gr_gas_available(gas_ptr)
+            Instruction::I32Const(gas_ptr as i32),
+            Instruction::Call(gr_gas_available_index),
+            // gas_available = *gas_ptr
+            Instruction::I32Const(gas_ptr as i32),
+            Instruction::I64Load(3, 0),
+            Instruction::I64Const(critical_gas_limit as i64),
+            // if gas_available <= critical_gas_limit { return result; }
+            Instruction::I64LeU,
+            Instruction::If(BlockType::NoResult),
+        ]);
+
+        // exit the current function with dummy results
+        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.extend_from_slice(&[Instruction::Return, Instruction::End]);
+
+        let instructions = func_body.code_mut().elements_mut();
+
+        let original_instructions =
+            mem::replace(instructions, Vec::with_capacity(instructions.len()));
+        let new_instructions = instructions;
+
+        // insert gas limiter at the beginning of each function to limit recursions
+        new_instructions.extend_from_slice(&body);
+
+        // also insert gas limiter at the beginning of each block, loop and condition
+        // to limit control instructions
+        for instruction in original_instructions {
+            match instruction {
+                Instruction::Block(_) | Instruction::Loop(_) | Instruction::If(_) => {
+                    new_instructions.push(instruction);
+                    new_instructions.extend_from_slice(&body);
+                }
+                Instruction::Call(call_index)
+                    if rewrite_sections && call_index >= gr_gas_available_index =>
+                {
+                    // fix function indexes if import gr_gas_available was inserted
+                    new_instructions.push(Instruction::Call(call_index + 1));
+                }
+                _ => {
+                    new_instructions.push(instruction);
+                }
+            }
+        }
+    }
+
+    // fix other sections if import gr_gas_available was inserted
+    if rewrite_sections {
+        let sections = module.sections_mut();
+        sections.retain(|section| !matches!(section, Section::Custom(_)));
+
+        for section in sections {
+            match section {
+                Section::Export(export_section) => {
+                    for export in export_section.entries_mut() {
+                        if let Internal::Function(func_index) = export.internal_mut() {
+                            if *func_index >= gr_gas_available_index {
+                                *func_index += 1
+                            }
+                        }
+                    }
+                }
+                Section::Element(elements_section) => {
+                    for segment in elements_section.entries_mut() {
+                        for func_index in segment.members_mut() {
+                            if *func_index >= gr_gas_available_index {
+                                *func_index += 1
+                            }
+                        }
+                    }
+                }
+                Section::Start(start_idx) => {
+                    if *start_idx >= gr_gas_available_index {
+                        *start_idx += 1;
+                    }
+                }
+                _ => {}
+            }
+        }
+    }
+
+    module
+}
+
 pub(crate) fn hash_with_value_to_vec(hash_with_value: &HashWithValue) -> Vec<u8> {
     let address_data_size = mem::size_of::<HashWithValue>();
     let address_data_slice = unsafe {