From 409ced002abed816f9aef2b7b0c903e4148c6d0f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Alexander=20Thei=C3=9Fen?= Date: Tue, 4 Aug 2020 12:53:53 +0200 Subject: [PATCH] Allow specifying the module of the imported 'gas' function (#140) * No need for mem::replace when doing a simple assignment * Allow specifying the module of the imported 'gas' function This allows users to place the imported function inside a custom module instead of the generic 'env' module. --- cli/gas/main.rs | 2 +- src/gas/mod.rs | 20 ++++++++++++-------- src/optimizer.rs | 7 +++---- tests/diff.rs | 3 ++- 4 files changed, 18 insertions(+), 14 deletions(-) diff --git a/cli/gas/main.rs b/cli/gas/main.rs index 0e9139f..dc185dc 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, &Default::default() + module, &Default::default(), "env" ).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 5869a0d..d162291 100644 --- a/src/gas/mod.rs +++ b/src/gas/mod.rs @@ -389,7 +389,7 @@ fn insert_metering_calls( /// Transforms a given module into one that charges gas for code to be executed by proxy of an /// imported gas metering function. /// -/// The output module imports a function "gas" from the module "env" with type signature +/// The output module imports a function "gas" from the specified module with type signature /// [i32] -> []. The argument is the amount of gas required to continue execution. The external /// function is meant to keep track of the total amount of gas used and trap or otherwise halt /// execution of the runtime if the gas usage exceeds some allowed limit. @@ -420,7 +420,7 @@ fn insert_metering_calls( /// /// The function fails if the module contains any operation forbidden by gas rule set, returning /// the original module as an Err. -pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set) +pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set, gas_module_name: &str) -> Result { // Injecting gas counting external @@ -433,7 +433,7 @@ pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set) mbuilder.push_import( builder::import() - .module("env") + .module(gas_module_name) .field("gas") .external().func(import_sig) .build() @@ -534,7 +534,11 @@ mod tests { .build() .build(); - let injected_module = inject_gas_counter(module, &rules::Set::default().with_grow_cost(10000)).unwrap(); + let injected_module = inject_gas_counter( + module, + &rules::Set::default().with_grow_cost(10000), + "env", + ).unwrap(); assert_eq!( get_function_body(&injected_module, 0).unwrap(), @@ -583,7 +587,7 @@ mod tests { .build() .build(); - let injected_module = inject_gas_counter(module, &rules::Set::default()).unwrap(); + let injected_module = inject_gas_counter(module, &rules::Set::default(), "env").unwrap(); assert_eq!( get_function_body(&injected_module, 0).unwrap(), @@ -634,7 +638,7 @@ mod tests { .build() .build(); - let injected_module = inject_gas_counter(module, &Default::default()).unwrap(); + let injected_module = inject_gas_counter(module, &Default::default(), "env").unwrap(); assert_eq!( get_function_body(&injected_module, 1).unwrap(), @@ -681,7 +685,7 @@ mod tests { let rules = rules::Set::default().with_forbidden_floats(); - if let Err(_) = inject_gas_counter(module, &rules) { } + if let Err(_) = inject_gas_counter(module, &rules, "env") { } else { panic!("Should be error because of the forbidden operation")} } @@ -702,7 +706,7 @@ mod tests { let input_module = parse_wat($input); let expected_module = parse_wat($expected); - let injected_module = inject_gas_counter(input_module, &Default::default()) + let injected_module = inject_gas_counter(input_module, &Default::default(), "env") .expect("inject_gas_counter call failed"); let actual_func_body = get_function_body(&injected_module, 0) diff --git a/src/optimizer.rs b/src/optimizer.rs index 8e2388f..f3e2e9b 100644 --- a/src/optimizer.rs +++ b/src/optimizer.rs @@ -29,7 +29,7 @@ pub fn optimize( let module_temp = module_temp .parse_names() .unwrap_or_else(|(_err, module)| module); - mem::replace(module, module_temp); + *module = module_temp; // Algo starts from the top, listing all items that should stay let mut stay = Set::new(); @@ -309,7 +309,7 @@ pub fn optimize( let totalle = eliminated_funcs.iter().take_while(|i| (**i as u32) < index).count() as u32; (index - totalle, value) }).collect(); - mem::replace(func_name.names_mut(), updated_map); + *func_name.names_mut() = updated_map; } if let Some(ref mut local_name) = name_section.locals_mut() { @@ -321,8 +321,7 @@ pub fn optimize( let totalle = eliminated_funcs.iter().take_while(|i| (**i as u32) < index).count() as u32; (index - totalle, value) }).collect(); - - mem::replace(local_name.local_names_mut(), updated_map); + *local_name.local_names_mut() = updated_map; } } _ => { } diff --git a/tests/diff.rs b/tests/diff.rs index 9026fbb..9756474 100644 --- a/tests/diff.rs +++ b/tests/diff.rs @@ -110,7 +110,8 @@ 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).expect("Failed to instrument with gas metering"); + let instrumented = utils::inject_gas_counter(module, &rules, "env") + .expect("Failed to instrument with gas metering"); elements::serialize(instrumented).expect("Failed to serialize") }); }