From 4564b8ee08992c1449db5c040f8799f4ba9a0175 Mon Sep 17 00:00:00 2001 From: StackOverflowExcept1on <109800286+StackOverflowExcept1on@users.noreply.github.com> Date: Sun, 19 Nov 2023 10:42:04 +0300 Subject: [PATCH] feat(wasm-gen): add limits for recursions and heavy loops (#3391) --- utils/wasm-gen/src/config.rs | 13 +- utils/wasm-gen/src/config/generator.rs | 10 ++ utils/wasm-gen/src/generator.rs | 8 + utils/wasm-gen/src/tests.rs | 34 ++++ utils/wasm-gen/src/utils.rs | 225 ++++++++++++++++++++++++- 5 files changed, 285 insertions(+), 5 deletions(-) diff --git a/utils/wasm-gen/src/config.rs b/utils/wasm-gen/src/config.rs index 3b797ed5d15..b715baf62b8 100644 --- a/utils/wasm-gen/src/config.rs +++ b/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 { pub existing_addresses: Option>, /// 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, /// Injection type for each syscall. pub injection_types: SysCallsInjectionTypes, /// Config of gear wasm call entry-points (exports). @@ -158,6 +166,7 @@ impl Default for StandardGearWasmConfigsBundle { 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> ConfigsBundle for StandardGearWasmConfigsBundle { log_info, existing_addresses, remove_recursion, + critical_gas_limit, injection_types, entry_points_set, initial_pages, @@ -199,6 +209,7 @@ impl> ConfigsBundle for StandardGearWasmConfigsBundle { 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) diff --git a/utils/wasm-gen/src/config/generator.rs b/utils/wasm-gen/src/config/generator.rs index 2c1da46750f..2b47e0058f7 100644 --- a/utils/wasm-gen/src/config/generator.rs +++ b/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) -> 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, } /// Memory pages config used by [`crate::MemoryGenerator`]. diff --git a/utils/wasm-gen/src/generator.rs b/utils/wasm-gen/src/generator.rs index 5277a5fb22e..dd6642f6708 100644 --- a/utils/wasm-gen/src/generator.rs +++ b/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) diff --git a/utils/wasm-gen/src/tests.rs b/utils/wasm-gen/src/tests.rs index 1d4ff88eccf..03926826690 100644 --- a/utils/wasm-gen/src/tests.rs +++ b/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::(&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#" diff --git a/utils/wasm-gen/src/utils.rs b/utils/wasm-gen/src/utils.rs index c89440967a6..0dc4b1e510c 100644 --- a/utils/wasm-gen/src/utils.rs +++ b/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 . -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( 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::::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::() 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 { let address_data_size = mem::size_of::(); let address_data_slice = unsafe {