diff --git a/gas/src/main.rs b/gas/src/main.rs index b6eba3a..dd50747 100644 --- a/gas/src/main.rs +++ b/gas/src/main.rs @@ -1,38 +1,7 @@ extern crate parity_wasm; +extern crate wasm_utils; use std::env; -use parity_wasm::{builder, elements}; - -pub fn update_call_index(opcodes: &mut elements::Opcodes, inserted_index: u32) { - use parity_wasm::elements::Opcode::*; - for opcode in opcodes.elements_mut().iter_mut() { - match opcode { - &mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => { - update_call_index(block, inserted_index) - }, - &mut Call(ref mut call_index) => { - if *call_index >= inserted_index { *call_index += 1} - }, - _ => { } - } - } -} - -pub fn inject_counter(opcodes: &mut elements::Opcodes, gas_func: u32) { - use parity_wasm::elements::Opcode::*; - for opcode in opcodes.elements_mut().iter_mut() { - match opcode { - &mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => { - inject_counter(block, gas_func) - }, - _ => { } - } - } - - let ops = opcodes.elements_mut().len() as u32; - opcodes.elements_mut().insert(0, I32Const(ops as i32)); - opcodes.elements_mut().insert(1, Call(gas_func)); -} fn main() { @@ -43,69 +12,9 @@ fn main() { } // Loading module - let module = parity_wasm::deserialize_file(&args[1]).unwrap(); + let module = parity_wasm::deserialize_file(&args[1]).expect("Module deserialization to succeed"); - // Injecting gas counting external - let mut mbuilder = builder::from_module(module); - let import_sig = mbuilder.push_signature( - builder::signature() - .param().i32() - .build_sig() - ); + let result = wasm_utils::inject_gas_counter(module); - let mut gas_func = mbuilder.push_import( - builder::import() - .module("env") - .field("gas") - .external().func(import_sig) - .build() - ); - - // back to plain module - let mut module = mbuilder.build(); - - assert!(module.global_section().is_some()); - - // calculate actual function index of the imported definition - // (substract all imports that are NOT functions) - - for import_entry in module.import_section().expect("Builder should have insert the import section").entries() { - match *import_entry.external() { - elements::External::Function(_) => {}, - _ => { gas_func -= 1; } - } - } - - // Updating calling addresses (all calls to function index >= `gas_func` should be incremented) - for section in module.sections_mut() { - match section { - &mut elements::Section::Code(ref mut code_section) => { - for ref mut func_body in code_section.bodies_mut() { - update_call_index(func_body.code_mut(), gas_func); - inject_counter(func_body.code_mut(), gas_func); - } - }, - &mut elements::Section::Export(ref mut export_section) => { - for ref mut export in export_section.entries_mut() { - match export.internal_mut() { - &mut elements::Internal::Function(ref mut func_index) => { - if *func_index >= gas_func { *func_index += 1} - }, - _ => {} - } - } - }, - &mut elements::Section::Element(ref mut elements_section) => { - for ref mut segment in elements_section.entries_mut() { - // update all indirect call addresses initial values - for func_index in segment.members_mut() { - if *func_index >= gas_func { *func_index += 1} - } - } - }, - _ => { } - } - } - - parity_wasm::serialize_to_file(&args[2], module).unwrap(); + parity_wasm::serialize_to_file(&args[2], result).expect("Module serialization to succeed") } diff --git a/runner/build.sh b/runner/build.sh index 6b79c6e..eb08957 100755 --- a/runner/build.sh +++ b/runner/build.sh @@ -21,7 +21,7 @@ then else # c/c++ can be compiled directly by emcc - emcc $file -O3 -s WASM=1 -s SIDE_MODULE=1 -o out/contract.wasm -s + emcc $file -O3 -s WASM=1 -s SIDE_MODULE=1 -o out/contract.wasm # Gas injector cargo run --manifest-path=./../gas/Cargo.toml --release -- ./out/contract.wasm ./out/contract.wasm diff --git a/samples/logger_contract.rs b/samples/logger_contract.rs index 030eb2b..f454bb5 100644 --- a/samples/logger_contract.rs +++ b/samples/logger_contract.rs @@ -6,7 +6,7 @@ use std::slice; -#[link_args = "-s WASM=1 -s NO_EXIT_RUNTIME=1 -s NO_FILESYSTEM=1"] +#[link_args = "-s NO_EXIT_RUNTIME=1 -s NO_FILESYSTEM=1"] extern {} /// Wrapper over storage read/write/size externs diff --git a/src/gas.rs b/src/gas.rs new file mode 100644 index 0000000..e69c288 --- /dev/null +++ b/src/gas.rs @@ -0,0 +1,99 @@ +use parity_wasm::{elements, builder}; + + +pub fn update_call_index(opcodes: &mut elements::Opcodes, inserted_index: u32) { + use parity_wasm::elements::Opcode::*; + for opcode in opcodes.elements_mut().iter_mut() { + match opcode { + &mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => { + update_call_index(block, inserted_index) + }, + &mut Call(ref mut call_index) => { + if *call_index >= inserted_index { *call_index += 1} + }, + _ => { } + } + } +} + +pub fn inject_counter(opcodes: &mut elements::Opcodes, gas_func: u32) { + use parity_wasm::elements::Opcode::*; + for opcode in opcodes.elements_mut().iter_mut() { + match opcode { + &mut Block(_, ref mut block) | &mut If(_, ref mut block) | &mut Loop(_, ref mut block) => { + inject_counter(block, gas_func) + }, + _ => { } + } + } + + let ops = opcodes.elements_mut().len() as u32; + opcodes.elements_mut().insert(0, I32Const(ops as i32)); + opcodes.elements_mut().insert(1, Call(gas_func)); +} + +pub fn inject_gas_counter(module: elements::Module) -> elements::Module { + // Injecting gas counting external + let mut mbuilder = builder::from_module(module); + let import_sig = mbuilder.push_signature( + builder::signature() + .param().i32() + .build_sig() + ); + + let mut gas_func = mbuilder.push_import( + builder::import() + .module("env") + .field("gas") + .external().func(import_sig) + .build() + ); + + // back to plain module + let mut module = mbuilder.build(); + + assert!(module.global_section().is_some()); + + // calculate actual function index of the imported definition + // (substract all imports that are NOT functions) + + for import_entry in module.import_section().expect("Builder should have insert the import section").entries() { + match *import_entry.external() { + elements::External::Function(_) => {}, + _ => { gas_func -= 1; } + } + } + + // Updating calling addresses (all calls to function index >= `gas_func` should be incremented) + for section in module.sections_mut() { + match section { + &mut elements::Section::Code(ref mut code_section) => { + for ref mut func_body in code_section.bodies_mut() { + update_call_index(func_body.code_mut(), gas_func); + inject_counter(func_body.code_mut(), gas_func); + } + }, + &mut elements::Section::Export(ref mut export_section) => { + for ref mut export in export_section.entries_mut() { + match export.internal_mut() { + &mut elements::Internal::Function(ref mut func_index) => { + if *func_index >= gas_func { *func_index += 1} + }, + _ => {} + } + } + }, + &mut elements::Section::Element(ref mut elements_section) => { + for ref mut segment in elements_section.entries_mut() { + // update all indirect call addresses initial values + for func_index in segment.members_mut() { + if *func_index >= gas_func { *func_index += 1} + } + } + }, + _ => { } + } + } + + module +} \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 8d4d9bd..926cc32 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,8 @@ extern crate parity_wasm; mod optimizer; -pub mod symbols; +mod gas; +mod symbols; -pub use optimizer::optimize; \ No newline at end of file +pub use optimizer::optimize; +pub use gas::inject_gas_counter; \ No newline at end of file