use std::vec::Vec;
use parity_wasm::{elements, builder};
use rules;
pub fn update_call_index(instructions: &mut elements::Instructions, inserted_index: u32) {
use parity_wasm::elements::Instruction::*;
for instruction in instructions.elements_mut().iter_mut() {
if let &mut Call(ref mut call_index) = instruction {
if *call_index >= inserted_index { *call_index += 1}
}
}
}
#[derive(Debug)]
struct BlockEntry {
start_pos: usize,
cost: u32,
}
struct Counter {
blocks: Vec<BlockEntry>,
stack: Vec<usize>,
}
impl Counter {
fn new() -> Counter {
Counter {
stack: Vec::new(),
blocks: Vec::new(),
}
}
fn begin(&mut self, cursor: usize) {
let block_idx = self.blocks.len();
self.blocks.push(BlockEntry {
start_pos: cursor,
cost: 1,
});
self.stack.push(block_idx);
}
fn finalize(&mut self) -> Result<(), ()> {
self.stack.pop().ok_or_else(|| ())?;
Ok(())
}
fn increment(&mut self, val: u32) -> Result<(), ()> {
let stack_top = self.stack.last_mut().ok_or_else(|| ())?;
let top_block = self.blocks.get_mut(*stack_top).ok_or_else(|| ())?;
top_block.cost = top_block.cost.checked_add(val).ok_or_else(|| ())?;
Ok(())
}
}
fn inject_grow_counter(instructions: &mut elements::Instructions, grow_counter_func: u32) -> usize {
use parity_wasm::elements::Instruction::*;
let mut counter = 0;
for instruction in instructions.elements_mut() {
if let GrowMemory(_) = *instruction {
*instruction = 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::Instruction::*;
let mut b = builder::from_module(module);
b.push_function(
builder::function()
.signature().params().i32().build().with_return_type(Some(elements::ValueType::I32)).build()
.body()
.with_instructions(elements::Instructions::new(vec![
GetLocal(0),
GetLocal(0),
I32Const(rules.grow_cost() as i32),
I32Mul,
Call(gas_func),
GrowMemory(0),
End,
]))
.build()
.build()
);
b.build()
}
pub fn inject_counter(
instructions: &mut elements::Instructions,
rules: &rules::Set,
gas_func: u32,
) -> Result<(), ()> {
use parity_wasm::elements::Instruction::*;
let mut counter = Counter::new();
counter.begin(0);
for cursor in 0..instructions.elements().len() {
let instruction = &instructions.elements()[cursor];
match *instruction {
Block(_) | If(_) | Loop(_) => {
let instruction_cost = rules.process(instruction)?;
counter.increment(instruction_cost)?;
counter.begin(cursor + 1);
}
End => {
counter.finalize()?;
},
Else => {
counter.finalize()?;
counter.begin(cursor + 1);
}
_ => {
let instruction_cost = rules.process(instruction)?;
counter.increment(instruction_cost)?;
}
}
}
let mut cumulative_offset = 0;
for block in counter.blocks {
let effective_pos = block.start_pos + cumulative_offset;
instructions.elements_mut().insert(effective_pos, I32Const(block.cost as i32));
instructions.elements_mut().insert(effective_pos+1, Call(gas_func));
cumulative_offset += 2;
}
Ok(())
}
pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set)
-> Result<elements::Module, elements::Module>
{
let mut mbuilder = builder::from_module(module);
let import_sig = mbuilder.push_signature(
builder::signature()
.param().i32()
.build_sig()
);
mbuilder.push_import(
builder::import()
.module("env")
.field("gas")
.external().func(import_sig)
.build()
);
let mut module = mbuilder.build();
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;
let mut error = false;
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);
if let Err(_) = inject_counter(func_body.code_mut(), rules, gas_func) {
error = true;
break;
}
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) => {
for ref mut export in export_section.entries_mut() {
if let &mut elements::Internal::Function(ref mut func_index) = export.internal_mut() {
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() {
for func_index in segment.members_mut() {
if *func_index >= gas_func { *func_index += 1}
}
}
},
&mut elements::Section::Start(ref mut start_idx) => {
if *start_idx >= gas_func { *start_idx += 1}
},
_ => { }
}
}
if error { return Err(module); }
if need_grow_counter { Ok(add_grow_counter(module, rules, gas_func)) } else { Ok(module) }
}
#[cfg(test)]
mod tests {
extern crate wabt;
use parity_wasm::{serialize, builder, elements};
use super::*;
use rules;
#[test]
fn simple_grow() {
use parity_wasm::elements::Instruction::*;
let module = builder::module()
.global()
.value_type().i32()
.build()
.function()
.signature().param().i32().build()
.body()
.with_instructions(elements::Instructions::new(
vec![
GetGlobal(0),
GrowMemory(0),
End
]
))
.build()
.build()
.build();
let injected_module = inject_gas_counter(module, &rules::Set::default().with_grow_cost(10000)).unwrap();
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),
GetLocal(0),
I32Const(10000),
I32Mul,
Call(0),
GrowMemory(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]
fn grow_no_gas_no_track() {
use parity_wasm::elements::Instruction::*;
let module = builder::module()
.global()
.value_type().i32()
.build()
.function()
.signature().param().i32().build()
.body()
.with_instructions(elements::Instructions::new(
vec![
GetGlobal(0),
GrowMemory(0),
End
]
))
.build()
.build()
.build();
let injected_module = inject_gas_counter(module, &rules::Set::default()).unwrap();
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);
let binary = serialize(injected_module).expect("serialization failed");
self::wabt::wasm2wat(&binary).unwrap();
}
#[test]
fn simple() {
use parity_wasm::elements::Instruction::*;
let module = builder::module()
.global()
.value_type().i32()
.build()
.function()
.signature().param().i32().build()
.body()
.with_instructions(elements::Instructions::new(
vec![
GetGlobal(0),
End
]
))
.build()
.build()
.build();
let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
assert_eq!(
&vec![
I32Const(2),
Call(0),
GetGlobal(0),
End
][..],
injected_module
.code_section().expect("function section should exist").bodies()[0]
.code().elements()
);
}
#[test]
fn nested() {
use parity_wasm::elements::Instruction::*;
let module = builder::module()
.global()
.value_type().i32()
.build()
.function()
.signature().param().i32().build()
.body()
.with_instructions(elements::Instructions::new(
vec![
GetGlobal(0),
Block(elements::BlockType::NoResult),
GetGlobal(0),
GetGlobal(0),
GetGlobal(0),
End,
GetGlobal(0),
End
]
))
.build()
.build()
.build();
let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
assert_eq!(
&vec![
I32Const(4),
Call(0),
GetGlobal(0),
Block(elements::BlockType::NoResult),
I32Const(4),
Call(0),
GetGlobal(0),
GetGlobal(0),
GetGlobal(0),
End,
GetGlobal(0),
End
][..],
injected_module
.code_section().expect("function section should exist").bodies()[0]
.code().elements()
);
}
#[test]
fn ifelse() {
use parity_wasm::elements::Instruction::*;
let module = builder::module()
.global()
.value_type().i32()
.build()
.function()
.signature().param().i32().build()
.body()
.with_instructions(elements::Instructions::new(
vec![
GetGlobal(0),
If(elements::BlockType::NoResult),
GetGlobal(0),
GetGlobal(0),
GetGlobal(0),
Else,
GetGlobal(0),
GetGlobal(0),
End,
GetGlobal(0),
End
]
))
.build()
.build()
.build();
let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
assert_eq!(
&vec![
I32Const(4),
Call(0),
GetGlobal(0),
If(elements::BlockType::NoResult),
I32Const(4),
Call(0),
GetGlobal(0),
GetGlobal(0),
GetGlobal(0),
Else,
I32Const(3),
Call(0),
GetGlobal(0),
GetGlobal(0),
End,
GetGlobal(0),
End
][..],
injected_module
.code_section().expect("function section should exist").bodies()[0]
.code().elements()
);
}
#[test]
fn call_index() {
use parity_wasm::elements::Instruction::*;
let module = builder::module()
.global()
.value_type().i32()
.build()
.function()
.signature().param().i32().build()
.body().build()
.build()
.function()
.signature().param().i32().build()
.body()
.with_instructions(elements::Instructions::new(
vec![
Call(0),
If(elements::BlockType::NoResult),
Call(0),
Call(0),
Call(0),
Else,
Call(0),
Call(0),
End,
Call(0),
End
]
))
.build()
.build()
.build();
let injected_module = inject_gas_counter(module, &Default::default()).unwrap();
assert_eq!(
&vec![
I32Const(4),
Call(0),
Call(1),
If(elements::BlockType::NoResult),
I32Const(4),
Call(0),
Call(1),
Call(1),
Call(1),
Else,
I32Const(3),
Call(0),
Call(1),
Call(1),
End,
Call(1),
End
][..],
injected_module
.code_section().expect("function section should exist").bodies()[1]
.code().elements()
);
}
#[test]
fn forbidden() {
use parity_wasm::elements::Instruction::*;
let module = builder::module()
.global()
.value_type().i32()
.build()
.function()
.signature().param().i32().build()
.body()
.with_instructions(elements::Instructions::new(
vec![
F32Const(555555),
End
]
))
.build()
.build()
.build();
let rules = rules::Set::default().with_forbidden_floats();
if let Err(_) = inject_gas_counter(module, &rules) { }
else { panic!("Should be error because of the forbidden operation")}
}
}