From 24dc9e3f15d3b93de9d5bb1c7d9d61c7db1d4180 Mon Sep 17 00:00:00 2001 From: NikVolf Date: Mon, 5 Feb 2018 00:10:53 +0300 Subject: [PATCH 1/2] add grow memory injecting --- src/gas.rs | 159 ++++++++++++++++++++++++++++++++++++++++++++++++--- src/rules.rs | 12 +++- 2 files changed, 161 insertions(+), 10 deletions(-) diff --git a/src/gas.rs b/src/gas.rs index 21c018e..031e3b9 100644 --- a/src/gas.rs +++ b/src/gas.rs @@ -20,7 +20,52 @@ enum InjectAction { IncrementSpawn, } -pub fn inject_counter(opcodes: &mut elements::Opcodes, rules: &rules::Set, gas_func: u32) { +fn inject_grow_counter(opcodes: &mut elements::Opcodes, grow_counter_func: u32) -> usize { + use parity_wasm::elements::Opcode::*; + let mut counter = 0; + for opcode in opcodes.elements_mut() { + match *opcode { + GrowMemory(_) => { + *opcode = Call(grow_counter_func); + counter += 1; + }, + _ => {} + } + } + counter +} + +fn add_grow_counter(module: elements::Module, rules: &rules::Set, gas_func: u32) -> elements::Module { + use parity_wasm::elements::Opcode::*; + + let mut b = builder::from_module(module); + b.push_function( + builder::function() + .signature().params().i32().i32().build().with_return_type(Some(elements::ValueType::I32)).build() + .body() + .with_opcodes(elements::Opcodes::new(vec![ + GetLocal(0), + I32Const(rules.grow_cost() as i32), + I32Mul, + TeeLocal(1), + // todo: there should be strong guarantee that it does not return anything on stack? + Call(gas_func), + GetLocal(1), + GrowMemory(0), + GetLocal(0), + ])) + .build() + .build() + ); + + b.build() +} + +pub fn inject_counter( + opcodes: &mut elements::Opcodes, + rules: &rules::Set, + gas_func: u32, +) { use parity_wasm::elements::Opcode::*; let mut stack: Vec<(usize, usize)> = Vec::new(); @@ -91,7 +136,7 @@ pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set) -> eleme .build_sig() ); - let mut gas_func = mbuilder.push_import( + mbuilder.push_import( builder::import() .module("env") .field("gas") @@ -105,12 +150,9 @@ pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set) -> eleme // 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; } - } - } + let gas_func = module.import_count(elements::ImportCountType::Function) as u32 - 1; + let total_func = module.functions_space() as u32; + let mut need_grow_counter = false; // Updating calling addresses (all calls to function index >= `gas_func` should be incremented) for section in module.sections_mut() { @@ -119,6 +161,11 @@ pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set) -> eleme 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(), rules, gas_func); + if rules.grow_cost() > 0 { + if inject_grow_counter(func_body.code_mut(), total_func) > 0 { + need_grow_counter = true; + } + } } }, &mut elements::Section::Export(ref mut export_section) => { @@ -143,7 +190,7 @@ pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set) -> eleme } } - module + if need_grow_counter { add_grow_counter(module, rules, gas_func) } else { module } } #[cfg(test)] @@ -151,6 +198,100 @@ mod tests { use parity_wasm::{builder, elements}; use super::*; + use rules; + + #[test] + fn simple_grow() { + use parity_wasm::elements::Opcode::*; + + let module = builder::module() + .global() + .value_type().i32() + .build() + .function() + .signature().param().i32().build() + .body() + .with_opcodes(elements::Opcodes::new( + vec![ + GetGlobal(0), + GrowMemory(0), + End + ] + )) + .build() + .build() + .build(); + + let injected_module = inject_gas_counter(module, &rules::Set::default().with_grow_cost(10000)); + + assert_eq!( + &vec![ + I32Const(3), + Call(0), + GetGlobal(0), + Call(2), + End + ][..], + injected_module + .code_section().expect("function section should exist").bodies()[0] + .code().elements() + ); + assert_eq!( + &vec![ + GetLocal(0), + I32Const(10000), + I32Mul, + TeeLocal(1), + Call(0), + GetLocal(1), + GrowMemory(0), + GetLocal(0), + ][..], + injected_module + .code_section().expect("function section should exist").bodies()[1] + .code().elements() + ); + } + + #[test] + fn grow_no_gas_no_track() { + use parity_wasm::elements::Opcode::*; + + let module = builder::module() + .global() + .value_type().i32() + .build() + .function() + .signature().param().i32().build() + .body() + .with_opcodes(elements::Opcodes::new( + vec![ + GetGlobal(0), + GrowMemory(0), + End + ] + )) + .build() + .build() + .build(); + + let injected_module = inject_gas_counter(module, &rules::Set::default()); + + assert_eq!( + &vec![ + I32Const(3), + Call(0), + GetGlobal(0), + GrowMemory(0), + End + ][..], + injected_module + .code_section().expect("function section should exist").bodies()[0] + .code().elements() + ); + + assert_eq!(injected_module.functions_space(), 2); + } #[test] fn simple() { diff --git a/src/rules.rs b/src/rules.rs index 4f68889..a4d739a 100644 --- a/src/rules.rs +++ b/src/rules.rs @@ -253,14 +253,24 @@ impl InstructionType { #[derive(Debug, Default)] pub struct Set { entries: HashMap, + grow: u32, } impl Set { pub fn new(entries: HashMap) -> Self { - Set { entries: entries } + Set { entries: entries, grow: 0, } } pub fn process(&self, opcode: &elements::Opcode) -> u32 { self.entries.get(&InstructionType::op(opcode)).map(|x| *x).unwrap_or(1) } + + pub fn grow_cost(&self) -> u32 { + self.grow + } + + pub fn with_grow_cost(mut self, val: u32) -> Self { + self.grow = val; + self + } } \ No newline at end of file From d220107dab7cad039ba68cf8d937fa1e2de0959e Mon Sep 17 00:00:00 2001 From: NikVolf Date: Mon, 5 Feb 2018 00:40:28 +0300 Subject: [PATCH 2/2] add validation via wabt-rs and fix --- Cargo.toml | 1 + src/gas.rs | 12 +++++++++++- 2 files changed, 12 insertions(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index 9ed330c..dbb2323 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ byteorder = "1" [dev-dependencies] tempdir = "0.3" +wabt = "0.1" [lib] diff --git a/src/gas.rs b/src/gas.rs index 031e3b9..5f3e9bc 100644 --- a/src/gas.rs +++ b/src/gas.rs @@ -53,6 +53,7 @@ fn add_grow_counter(module: elements::Module, rules: &rules::Set, gas_func: u32) GetLocal(1), GrowMemory(0), GetLocal(0), + End, ])) .build() .build() @@ -196,7 +197,9 @@ pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set) -> eleme #[cfg(test)] mod tests { - use parity_wasm::{builder, elements}; + extern crate wabt; + + use parity_wasm::{serialize, builder, elements}; use super::*; use rules; @@ -246,11 +249,15 @@ mod tests { GetLocal(1), GrowMemory(0), GetLocal(0), + End, ][..], injected_module .code_section().expect("function section should exist").bodies()[1] .code().elements() ); + + let binary = serialize(injected_module).expect("serialization failed"); + self::wabt::wasm2wat(&binary).unwrap(); } #[test] @@ -291,6 +298,9 @@ mod tests { ); assert_eq!(injected_module.functions_space(), 2); + + let binary = serialize(injected_module).expect("serialization failed"); + self::wabt::wasm2wat(&binary).unwrap(); } #[test]