diff --git a/cli/gas/main.rs b/cli/gas/main.rs index 6ec45ee..53a875d 100644 --- a/cli/gas/main.rs +++ b/cli/gas/main.rs @@ -17,7 +17,7 @@ fn main() { let module = parity_wasm::deserialize_file(&args[1]).expect("Module deserialization to succeed"); let result = utils::inject_gas_counter( - module, &utils::rules::Set::default(), "env" + module, &utils::rules::Set::default(), "env", utils::ProcessNames::No, ).expect("Failed to inject gas. Some forbidden opcodes?"); parity_wasm::serialize_to_file(&args[2], result).expect("Module serialization to succeed") diff --git a/src/gas/mod.rs b/src/gas/mod.rs index b34f6ac..6a3f523 100644 --- a/src/gas/mod.rs +++ b/src/gas/mod.rs @@ -9,6 +9,7 @@ mod validation; use crate::std::cmp::min; use crate::std::mem; +use crate::std::string::String; use crate::std::vec::Vec; use parity_wasm::{elements, elements::ValueType, builder}; @@ -396,6 +397,12 @@ fn insert_metering_calls( Ok(()) } +/// If process name section while injecting gas counter +pub enum ProcessNames { + No, + UnsafeYes, +} + /// Transforms a given module into one that charges gas for code to be executed by proxy of an /// imported gas metering function. /// @@ -434,6 +441,7 @@ pub fn inject_gas_counter( module: elements::Module, rules: &R, gas_module_name: &str, + process_names: ProcessNames, ) -> Result { @@ -501,6 +509,25 @@ pub fn inject_gas_counter( elements::Section::Start(start_idx) => { if *start_idx >= gas_func { *start_idx += 1} }, + elements::Section::Name(section) => { + match process_names { + ProcessNames::No => {}, + ProcessNames::UnsafeYes => { + if let Some(funcs) = section.functions_mut() { + let mut new_names = elements::IndexMap::::default(); + for (idx, func) in funcs.names().iter() { + if idx >= gas_func { + new_names.insert(idx + 1, String::from(func)); + } else { + new_names.insert(idx, String::from(func)); + } + } + + *funcs.names_mut() = new_names; + } + } + } + } _ => { } } } @@ -552,6 +579,7 @@ mod tests { module, &rules::Set::default().with_grow_cost(10000), "env", + ProcessNames::No, ).unwrap(); assert_eq!( @@ -601,7 +629,7 @@ mod tests { .build() .build(); - let injected_module = inject_gas_counter(module, &rules::Set::default(), "env").unwrap(); + let injected_module = inject_gas_counter(module, &rules::Set::default(), "env", ProcessNames::No).unwrap(); assert_eq!( get_function_body(&injected_module, 0).unwrap(), @@ -652,7 +680,7 @@ mod tests { .build() .build(); - let injected_module = inject_gas_counter(module, &rules::Set::default(), "env").unwrap(); + let injected_module = inject_gas_counter(module, &rules::Set::default(), "env", ProcessNames::No).unwrap(); assert_eq!( get_function_body(&injected_module, 1).unwrap(), @@ -700,7 +728,7 @@ mod tests { let rules = rules::Set::default().with_forbidden_floats(); - if inject_gas_counter(module, &rules, "env").is_ok() { + if inject_gas_counter(module, &rules, "env", ProcessNames::No).is_ok() { panic!("Should be error because of the forbidden operation") } } @@ -721,7 +749,7 @@ mod tests { let input_module = parse_wat($input); let expected_module = parse_wat($expected); - let injected_module = inject_gas_counter(input_module, &rules::Set::default(), "env") + let injected_module = inject_gas_counter(input_module, &rules::Set::default(), "env", ProcessNames::No) .expect("inject_gas_counter call failed"); let actual_func_body = get_function_body(&injected_module, 0) diff --git a/src/lib.rs b/src/lib.rs index 08bd80c..852021e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -34,7 +34,7 @@ pub use build::{build, Error as BuildError, SourceTarget}; pub use ext::{ externalize, externalize_mem, shrink_unknown_stack, underscore_funcs, ununderscore_funcs, }; -pub use gas::inject_gas_counter; +pub use gas::{inject_gas_counter, ProcessNames}; pub use optimizer::{optimize, Error as OptimizerError}; pub use pack::{pack_instance, Error as PackingError}; pub use runtime_type::inject_runtime_type; diff --git a/tests/diff.rs b/tests/diff.rs index 9756474..c8cffaf 100644 --- a/tests/diff.rs +++ b/tests/diff.rs @@ -110,7 +110,7 @@ mod gas { let rules = utils::rules::Set::default(); let module = elements::deserialize_buffer(input).expect("Failed to deserialize"); - let instrumented = utils::inject_gas_counter(module, &rules, "env") + let instrumented = utils::inject_gas_counter(module, &rules, "env", utils::ProcessNames::No) .expect("Failed to instrument with gas metering"); elements::serialize(instrumented).expect("Failed to serialize") });