Make the rules passed to gas metering injection generic

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(), "env"
+		module, &utils::rules::Set::default(), "env"
 	).expect("Failed to inject gas. Some forbidden opcodes?");
 
 	parity_wasm::serialize_to_file(&args[2], result).expect("Module serialization to succeed")

src/gas/mod.rs

@@ -12,7 +12,7 @@ use std::mem;
 use std::vec::Vec;
 
 use parity_wasm::{elements, builder};
-use rules;
+use rules::Rules;
 
 pub fn update_call_index(instructions: &mut elements::Instructions, inserted_index: u32) {
 	use parity_wasm::elements::Instruction::*;
@@ -234,8 +234,18 @@ fn inject_grow_counter(instructions: &mut elements::Instructions, grow_counter_f
 	counter
 }
 
-fn add_grow_counter(module: elements::Module, rules: &rules::Set, gas_func: u32) -> elements::Module {
+fn add_grow_counter<R: Rules>(
+	module: elements::Module,
+	rules: &R,
+	gas_func: u32
+) -> elements::Module {
 	use parity_wasm::elements::Instruction::*;
+	use crate::rules::MemoryGrowCost;
+
+	let cost = match rules.memory_grow_cost() {
+		None => return module,
+		Some(MemoryGrowCost::Linear(val)) => val.get(),
+	};
 
 	let mut b = builder::from_module(module);
 	b.push_function(
@@ -245,7 +255,7 @@ fn add_grow_counter(module: elements::Module, rules: &rules::Set, gas_func: u32)
 				.with_instructions(elements::Instructions::new(vec![
 					GetLocal(0),
 					GetLocal(0),
-					I32Const(rules.grow_cost() as i32),
+					I32Const(cost as i32),
 					I32Mul,
 					// todo: there should be strong guarantee that it does not return anything on stack?
 					Call(gas_func),
@@ -259,9 +269,9 @@ fn add_grow_counter(module: elements::Module, rules: &rules::Set, gas_func: u32)
 	b.build()
 }
 
-pub(crate) fn determine_metered_blocks(
+pub(crate) fn determine_metered_blocks<R: Rules>(
 	instructions: &elements::Instructions,
-	rules: &rules::Set,
+	rules: &R,
 ) -> Result<Vec<MeteredBlock>, ()> {
 	use parity_wasm::elements::Instruction::*;
 
@@ -272,7 +282,7 @@ pub(crate) fn determine_metered_blocks(
 
 	for cursor in 0..instructions.elements().len() {
 		let instruction = &instructions.elements()[cursor];
-		let instruction_cost = rules.process(instruction)?;
+		let instruction_cost = rules.instruction_cost(instruction).ok_or(())?;
 		match instruction {
 			Block(_) => {
 				counter.increment(instruction_cost)?;
@@ -333,9 +343,9 @@ pub(crate) fn determine_metered_blocks(
 	Ok(counter.finalized_blocks)
 }
 
-pub fn inject_counter(
+pub fn inject_counter<R: Rules>(
 	instructions: &mut elements::Instructions,
-	rules: &rules::Set,
+	rules: &R,
 	gas_func: u32,
 ) -> Result<(), ()> {
 	let blocks = determine_metered_blocks(instructions, rules)?;
@@ -420,7 +430,11 @@ 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, gas_module_name: &str)
+pub fn inject_gas_counter<R: Rules>(
+	module: elements::Module,
+	rules: &R,
+	gas_module_name: &str,
+)
 	-> Result<elements::Module, elements::Module>
 {
 	// Injecting gas counting external
@@ -460,7 +474,7 @@ pub fn inject_gas_counter(module: elements::Module, rules: &rules::Set, gas_modu
 						error = true;
 						break;
 					}
-					if rules.grow_cost() > 0 && inject_grow_counter(func_body.code_mut(), total_func) > 0 {
+					if rules.memory_grow_cost().is_some() && inject_grow_counter(func_body.code_mut(), total_func) > 0 {
 						need_grow_counter = true;
 					}
 				}
@@ -636,7 +650,7 @@ mod tests {
 				.build()
 			.build();
 
-		let injected_module = inject_gas_counter(module, &Default::default(), "env").unwrap();
+		let injected_module = inject_gas_counter(module, &rules::Set::default(), "env").unwrap();
 
 		assert_eq!(
 			get_function_body(&injected_module, 1).unwrap(),
@@ -705,7 +719,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(), "env")
+				let injected_module = inject_gas_counter(input_module, &rules::Set::default(), "env")
 					.expect("inject_gas_counter call failed");
 
 				let actual_func_body = get_function_body(&injected_module, 0)

src/gas/validation.rs

@@ -10,6 +10,7 @@
 
 use super::MeteredBlock;
 use rules::Set as RuleSet;
+use rules::Rules;
 use parity_wasm::elements::{FuncBody, Instruction};
 
 use std::collections::HashMap;
@@ -167,7 +168,7 @@ fn build_control_flow_graph(
 			graph.increment_charged_cost(active_node_id, next_metered_block.cost);
 		}
 
-		let instruction_cost = rules.process(instruction)?;
+		let instruction_cost = rules.instruction_cost(instruction).ok_or(())?;
 		match instruction {
 			Instruction::Block(_) => {
 				graph.increment_actual_cost(active_node_id, instruction_cost);

src/rules.rs

@@ -3,10 +3,39 @@ use std::collections::{HashMap as Map};
 #[cfg(not(features = "std"))]
 use std::collections::{BTreeMap as Map};
 
+use std::num::NonZeroU32;
+
 use parity_wasm::elements;
 
 pub struct UnknownInstruction;
 
+/// An interface that describes instruction costs.
+pub trait Rules {
+	/// Returns the cost for the passed `instruction`.
+	///
+	/// Returning `None` makes the gas instrumention end with an error. This is meant
+	/// as a way to have a partial rule set where any instruction that is not specifed
+	/// is considered as forbidden.
+	fn instruction_cost(&self, instruction: &elements::Instruction) -> Option<u32>;
+
+	/// Returns the costs for growing the memory using the `memory.grow` instruction.
+	///
+	/// Please note that these costs are in addition to the costs specified by `instruction_cost`
+	/// for the `memory.grow` instruction. Specifying `None` leads to no additional charge.
+	/// Those are meant as dynamic costs which take the amount of pages that the memory is
+	/// grown by into consideration. This is not possible using `instruction_cost` because
+	/// those costs depend on the stack and must be injected as code into the function calling
+	/// `memory.grow`. Therefore returning `Some` comes with a performance cost.
+	fn memory_grow_cost(&self) -> Option<MemoryGrowCost>;
+}
+
+/// Dynamic costs for memory growth.
+#[derive(Debug, PartialEq, Eq, Copy, Clone)]
+pub enum MemoryGrowCost {
+	/// Charge the specified amount for each page that the memory is grown by.
+	Linear(NonZeroU32),
+}
+
 #[derive(Debug, PartialEq, Eq, Copy, Clone)]
 pub enum Metering {
 	Regular,
@@ -288,14 +317,6 @@ impl Set {
 		Set { regular, entries, grow: 0 }
 	}
 
-	pub fn process(&self, instruction: &elements::Instruction) -> Result<u32, ()>  {
-		match self.entries.get(&InstructionType::op(instruction)).cloned() {
-			None | Some(Metering::Regular) => Ok(self.regular),
-			Some(Metering::Forbidden) => Err(()),
-			Some(Metering::Fixed(val)) => Ok(val),
-		}
-	}
-
 	pub fn grow_cost(&self) -> u32 {
 		self.grow
 	}
@@ -313,3 +334,21 @@ impl Set {
 		self
 	}
 }
+
+impl Rules for Set {
+	fn instruction_cost(&self, instruction: &elements::Instruction) -> Option<u32> {
+		match self.entries.get(&InstructionType::op(instruction)) {
+			None | Some(Metering::Regular) => Some(self.regular),
+			Some(Metering::Fixed(val)) => Some(*val),
+			Some(Metering::Forbidden) => None,
+		}
+	}
+
+	fn memory_grow_cost(&self) -> Option<MemoryGrowCost> {
+		if let Some(val) = NonZeroU32::new(self.grow) {
+			Some(MemoryGrowCost::Linear(val))
+		} else {
+			None
+		}
+	}
+}