Patch: use checked sum for locals counter (#40)

follow-up patch to #38

Co-authored-by: Alexander Theißen <alex.theissen@me.com>

src/gas_metering/backend.rs

@@ -112,8 +112,8 @@ pub mod mutable_global {
 			];
 
 			// calculate gas used for the gas charging func execution itself
-			let mut gas_fn_cost = func_instructions.iter().fold(0, |cost, instruction| {
-				cost + (rules.instruction_cost(instruction).unwrap_or(0) as u64)
+			let mut gas_fn_cost = func_instructions.iter().fold(0, |cost: u64, instruction| {
+				cost.saturating_add(rules.instruction_cost(instruction).unwrap_or(u32::MAX).into())
 			});
 			// don't charge for the instructions used to fail when out of gas
 			let fail_cost = vec![
@@ -122,10 +122,11 @@ pub mod mutable_global {
 				Instruction::Unreachable,               // non-charged instruction
 			]
 			.iter()
-			.fold(0, |cost, instruction| {
-				cost + (rules.instruction_cost(instruction).unwrap_or(0) as u64)
+			.fold(0, |cost: u64, instruction| {
+				cost.saturating_add(rules.instruction_cost(instruction).unwrap_or(u32::MAX).into())
 			});
 
+			// the fail costs are a subset of the overall costs and hence this never underflows
 			gas_fn_cost -= fail_cost;
 
 			GasMeter::Internal {

src/gas_metering/mod.rs

@@ -166,7 +166,7 @@ pub fn inject<R: Rules, B: Backend>(
 	let functions_space = module.functions_space() as u32;
 	let gas_global_idx = module.globals_space() as u32;
 
-	let mut mbuilder = builder::from_module(module);
+	let mut mbuilder = builder::from_module(module.clone());
 
 	// Calculate the indexes and gas function cost,
 	// for external gas function the cost is counted on the host side
@@ -224,15 +224,14 @@ pub fn inject<R: Rules, B: Backend>(
 	};
 
 	// We need the built the module for making injections to its blocks
-	let mut module = mbuilder.build();
+	let mut resulting_module = mbuilder.build();
 
 	let mut need_grow_counter = false;
-	let mut error = false;
-
+	let mut result = Ok(());
 	// Iterate over module sections and perform needed transformations.
 	// Indexes are needed to be fixed up in `GasMeter::External` case, as it adds an imported
 	// function, which goes to the beginning of the module's functions space.
-	for section in module.sections_mut() {
+	'outer: for section in resulting_module.sections_mut() {
 		match section {
 			elements::Section::Code(code_section) => {
 				let injection_targets = match gas_meter {
@@ -257,19 +256,22 @@ pub fn inject<R: Rules, B: Backend>(
 							}
 						}
 					}
-					let locals_count =
-						func_body.locals().iter().map(|val_type| val_type.count()).sum();
-					if inject_counter(
-						func_body.code_mut(),
-						gas_fn_cost,
-						locals_count,
-						rules,
-						gas_func_idx,
-					)
-					.is_err()
-					{
-						error = true;
-						break
+					result = func_body
+						.locals()
+						.iter()
+						.try_fold(0u32, |count, val_type| count.checked_add(val_type.count()))
+						.ok_or(())
+						.and_then(|locals_count| {
+							inject_counter(
+								func_body.code_mut(),
+								gas_fn_cost,
+								locals_count,
+								rules,
+								gas_func_idx,
+							)
+						});
+					if result.is_err() {
+						break 'outer
 					}
 					if rules.memory_grow_cost().enabled() &&
 						inject_grow_counter(func_body.code_mut(), total_func) > 0
@@ -325,14 +327,12 @@ pub fn inject<R: Rules, B: Backend>(
 		}
 	}
 
-	if error {
-		return Err(module)
-	}
+	result.map_err(|_| module)?;
 
 	if need_grow_counter {
-		Ok(add_grow_counter(module, rules, gas_func_idx))
+		Ok(add_grow_counter(resulting_module, rules, gas_func_idx))
 	} else {
-		Ok(module)
+		Ok(resulting_module)
 	}
 }
 

src/gas_metering/validation.rs

@@ -135,8 +135,12 @@ fn build_control_flow_graph(
 	let mut stack = vec![ControlFrame::new(entry_node_id, terminal_node_id, false)];
 	let mut metered_blocks_iter = blocks.iter().peekable();
 
-	let locals_count = body.locals().iter().fold(0, |count, val_type| count + val_type.count());
-	let locals_init_cost = (rules.call_per_local_cost()).checked_mul(locals_count).ok_or(())?;
+	let locals_count = body
+		.locals()
+		.iter()
+		.try_fold(0u32, |count, val_type| count.checked_add(val_type.count()))
+		.ok_or(())?;
+	let locals_init_cost = rules.call_per_local_cost().checked_mul(locals_count).ok_or(())?;
 
 	for (cursor, instruction) in body.code().elements().iter().enumerate() {
 		let active_node_id = stack
@@ -350,8 +354,7 @@ mod tests {
 
 			for func_body in module.code_section().iter().flat_map(|section| section.bodies()) {
 				let rules = ConstantCostRules::default();
-				let locals_count =
-					func_body.locals().iter().fold(0, |count, val_type| count + val_type.count());
+				let locals_count = func_body.locals().iter().map(|val_type| val_type.count()).sum();
 
 				let metered_blocks =
 					determine_metered_blocks(func_body.code(), &rules, locals_count).unwrap();