wasm_instrument/gas_metering/
mod.rs

1//! This module is used to instrument a Wasm module with the gas metering code.
2//!
3//! The primary public interface is the [`inject`] function which transforms a given
4//! module into one that charges gas for code to be executed. See function documentation for usage
5//! and details.
6
7mod backend;
8
9pub use backend::{host_function, mutable_global, Backend, GasMeter};
10
11#[cfg(test)]
12mod validation;
13
14use alloc::{vec, vec::Vec};
15use core::{cmp::min, mem, num::NonZeroU32};
16use parity_wasm::{
17	builder,
18	elements::{self, IndexMap, Instruction, ValueType},
19};
20
21/// An interface that describes instruction costs.
22pub trait Rules {
23	/// Returns the cost for the passed `instruction`.
24	///
25	/// Returning `None` makes the gas instrumention end with an error. This is meant
26	/// as a way to have a partial rule set where any instruction that is not specifed
27	/// is considered as forbidden.
28	fn instruction_cost(&self, instruction: &Instruction) -> Option<u32>;
29
30	/// Returns the costs for growing the memory using the `memory.grow` instruction.
31	///
32	/// Please note that these costs are in addition to the costs specified by `instruction_cost`
33	/// for the `memory.grow` instruction. Those are meant as dynamic costs which take the
34	/// amount of pages that the memory is grown by into consideration. This is not possible
35	/// using `instruction_cost` because those costs depend on the stack and must be injected as
36	/// code into the function calling `memory.grow`. Therefore returning anything but
37	/// [`MemoryGrowCost::Free`] introduces some overhead to the `memory.grow` instruction.
38	fn memory_grow_cost(&self) -> MemoryGrowCost;
39
40	/// A surcharge cost to calling a function that is added per local of that function.
41	fn call_per_local_cost(&self) -> u32;
42}
43
44/// Dynamic costs for memory growth.
45#[derive(Debug, PartialEq, Eq, Copy, Clone)]
46pub enum MemoryGrowCost {
47	/// Skip per page charge.
48	///
49	/// # Note
50	///
51	/// This makes sense when the amount of pages that a module is allowed to use is limited
52	/// to a rather small number by static validation. In that case it is viable to
53	/// benchmark the costs of `memory.grow` as the worst case (growing to to the maximum
54	/// number of pages).
55	Free,
56	/// Charge the specified amount for each page that the memory is grown by.
57	Linear(NonZeroU32),
58}
59
60impl MemoryGrowCost {
61	/// True iff memory growths code needs to be injected.
62	fn enabled(&self) -> bool {
63		match self {
64			Self::Free => false,
65			Self::Linear(_) => true,
66		}
67	}
68}
69
70/// A type that implements [`Rules`] so that every instruction costs the same.
71///
72/// This is a simplification that is mostly useful for development and testing.
73///
74/// # Note
75///
76/// In a production environment it usually makes no sense to assign every instruction
77/// the same cost. A proper implemention of [`Rules`] should be provided that is probably
78/// created by benchmarking.
79pub struct ConstantCostRules {
80	instruction_cost: u32,
81	memory_grow_cost: u32,
82	call_per_local_cost: u32,
83}
84
85impl ConstantCostRules {
86	/// Create a new [`ConstantCostRules`].
87	///
88	/// Uses `instruction_cost` for every instruction and `memory_grow_cost` to dynamically
89	/// meter the memory growth instruction.
90	pub fn new(instruction_cost: u32, memory_grow_cost: u32, call_per_local_cost: u32) -> Self {
91		Self { instruction_cost, memory_grow_cost, call_per_local_cost }
92	}
93}
94
95impl Default for ConstantCostRules {
96	/// Uses instruction cost of `1` and disables memory growth instrumentation.
97	fn default() -> Self {
98		Self { instruction_cost: 1, memory_grow_cost: 0, call_per_local_cost: 1 }
99	}
100}
101
102impl Rules for ConstantCostRules {
103	fn instruction_cost(&self, _: &Instruction) -> Option<u32> {
104		Some(self.instruction_cost)
105	}
106
107	fn memory_grow_cost(&self) -> MemoryGrowCost {
108		NonZeroU32::new(self.memory_grow_cost).map_or(MemoryGrowCost::Free, MemoryGrowCost::Linear)
109	}
110
111	fn call_per_local_cost(&self) -> u32 {
112		self.call_per_local_cost
113	}
114}
115
116/// Transforms a given module into one that tracks the gas charged during its execution.
117///
118/// The output module uses the `gas` function to track the gas spent. The function could be either
119/// an imported or a local one modifying a mutable global. The argument is the amount of gas
120/// required to continue execution. The execution engine is meant to keep track of the total amount
121/// of gas used and trap or otherwise halt execution of the runtime if the gas usage exceeds some
122/// allowed limit.
123///
124/// The body of each function of the original module is divided into metered blocks, and the calls
125/// to charge gas are inserted at the beginning of every such block of code. A metered block is
126/// defined so that, unless there is a trap, either all of the instructions are executed or none
127/// are. These are similar to basic blocks in a control flow graph, except that in some cases
128/// multiple basic blocks can be merged into a single metered block. This is the case if any path
129/// through the control flow graph containing one basic block also contains another.
130///
131/// Charging gas at the beginning of each metered block ensures that 1) all instructions
132/// executed are already paid for, 2) instructions that will not be executed are not charged for
133/// unless execution traps, and 3) the number of calls to `gas` is minimized. The corollary is
134/// that modules instrumented with this metering code may charge gas for instructions not
135/// executed in the event of a trap.
136///
137/// Additionally, each `memory.grow` instruction found in the module is instrumented to first
138/// make a call to charge gas for the additional pages requested. This cannot be done as part of
139/// the block level gas charges as the gas cost is not static and depends on the stack argument
140/// to `memory.grow`.
141///
142/// The above transformations are performed for every function body defined in the module. This
143/// function also rewrites all function indices references by code, table elements, etc., since
144/// the addition of an imported functions changes the indices of module-defined functions. If
145/// the module has a `NameSection`, added by calling `parse_names`, the indices will also be
146/// updated.
147///
148/// Syncronizing the amount of gas charged with the execution engine can be done in two ways. The
149/// first way is by calling the imported `gas` host function, see [`host_function`] for details. The
150/// second way is by using a local `gas` function together with a mutable global, see
151/// [`mutable_global`] for details.
152///
153/// This routine runs in time linear in the size of the input module.
154///
155/// The function fails if the module contains any operation forbidden by gas rule set, returning
156/// the original module as an `Err`.
157pub fn inject<R: Rules, B: Backend>(
158	module: elements::Module,
159	backend: B,
160	rules: &R,
161) -> Result<elements::Module, elements::Module> {
162	// Prepare module and return the gas function
163	let gas_meter = backend.gas_meter(&module, rules);
164
165	let import_count = module.import_count(elements::ImportCountType::Function) as u32;
166	let functions_space = module.functions_space() as u32;
167	let gas_global_idx = module.globals_space() as u32;
168
169	let mut mbuilder = builder::from_module(module);
170
171	// Calculate the indexes and gas function cost,
172	// for external gas function the cost is counted on the host side
173	let (gas_func_idx, total_func, gas_fn_cost) = match gas_meter {
174		GasMeter::External { module: gas_module, function } => {
175			// Inject the import of the gas function
176			let import_sig = mbuilder
177				.push_signature(builder::signature().with_param(ValueType::I64).build_sig());
178			mbuilder.push_import(
179				builder::import()
180					.module(gas_module)
181					.field(function)
182					.external()
183					.func(import_sig)
184					.build(),
185			);
186
187			(import_count, functions_space + 1, 0)
188		},
189		GasMeter::Internal { global, ref func_instructions, cost } => {
190			// Inject the gas counting global
191			mbuilder.push_global(
192				builder::global()
193					.with_type(ValueType::I64)
194					.mutable()
195					.init_expr(Instruction::I64Const(0))
196					.build(),
197			);
198			// Inject the export entry for the gas counting global
199			let ebuilder = builder::ExportBuilder::new();
200			let global_export = ebuilder
201				.field(global)
202				.with_internal(elements::Internal::Global(gas_global_idx))
203				.build();
204			mbuilder.push_export(global_export);
205
206			let func_idx = functions_space;
207
208			// Build local gas function
209			let gas_func_sig =
210				builder::SignatureBuilder::new().with_param(ValueType::I64).build_sig();
211
212			let function = builder::FunctionBuilder::new()
213				.with_signature(gas_func_sig)
214				.body()
215				.with_instructions(func_instructions.clone())
216				.build()
217				.build();
218
219			// Inject local gas function
220			mbuilder.push_function(function);
221
222			(func_idx, func_idx + 1, cost)
223		},
224	};
225
226	// We need the built the module for making injections to its blocks
227	let mut module = mbuilder.build();
228
229	let mut need_grow_counter = false;
230	let mut error = false;
231
232	// Iterate over module sections and perform needed transformations.
233	// Indexes are needed to be fixed up in `GasMeter::External` case, as it adds an imported
234	// function, which goes to the beginning of the module's functions space.
235	for section in module.sections_mut() {
236		match section {
237			elements::Section::Code(code_section) => {
238				let injection_targets = match gas_meter {
239					GasMeter::External { .. } => code_section.bodies_mut().as_mut_slice(),
240					// Don't inject counters to the local gas function, which is the last one as
241					// it's just added. Cost for its execution is added statically before each
242					// invocation (see `inject_counter()`).
243					GasMeter::Internal { .. } => {
244						let len = code_section.bodies().len();
245						&mut code_section.bodies_mut()[..len - 1]
246					},
247				};
248
249				for func_body in injection_targets {
250					// Increment calling addresses if needed
251					if let GasMeter::External { .. } = gas_meter {
252						for instruction in func_body.code_mut().elements_mut().iter_mut() {
253							if let Instruction::Call(call_index) = instruction {
254								if *call_index >= gas_func_idx {
255									*call_index += 1
256								}
257							}
258						}
259					}
260					let locals_count =
261						func_body.locals().iter().map(|val_type| val_type.count()).sum();
262					if inject_counter(
263						func_body.code_mut(),
264						gas_fn_cost,
265						locals_count,
266						rules,
267						gas_func_idx,
268					)
269					.is_err()
270					{
271						error = true;
272						break
273					}
274					if rules.memory_grow_cost().enabled() &&
275						inject_grow_counter(func_body.code_mut(), total_func) > 0
276					{
277						need_grow_counter = true;
278					}
279				}
280			},
281			elements::Section::Export(export_section) =>
282				if let GasMeter::External { module: _, function: _ } = gas_meter {
283					for export in export_section.entries_mut() {
284						if let elements::Internal::Function(func_index) = export.internal_mut() {
285							if *func_index >= gas_func_idx {
286								*func_index += 1
287							}
288						}
289					}
290				},
291			elements::Section::Element(elements_section) => {
292				// Note that we do not need to check the element type referenced because in the
293				// WebAssembly 1.0 spec, the only allowed element type is funcref.
294				if let GasMeter::External { .. } = gas_meter {
295					for segment in elements_section.entries_mut() {
296						// update all indirect call addresses initial values
297						for func_index in segment.members_mut() {
298							if *func_index >= gas_func_idx {
299								*func_index += 1
300							}
301						}
302					}
303				}
304			},
305			elements::Section::Start(start_idx) =>
306				if let GasMeter::External { .. } = gas_meter {
307					if *start_idx >= gas_func_idx {
308						*start_idx += 1
309					}
310				},
311			elements::Section::Name(s) =>
312				if let GasMeter::External { .. } = gas_meter {
313					for functions in s.functions_mut() {
314						*functions.names_mut() =
315							IndexMap::from_iter(functions.names().iter().map(|(mut idx, name)| {
316								if idx >= gas_func_idx {
317									idx += 1;
318								}
319
320								(idx, name.clone())
321							}));
322					}
323				},
324			_ => {},
325		}
326	}
327
328	if error {
329		return Err(module)
330	}
331
332	if need_grow_counter {
333		Ok(add_grow_counter(module, rules, gas_func_idx))
334	} else {
335		Ok(module)
336	}
337}
338
339/// A control flow block is opened with the `block`, `loop`, and `if` instructions and is closed
340/// with `end`. Each block implicitly defines a new label. The control blocks form a stack during
341/// program execution.
342///
343/// An example of block:
344///
345/// ```ignore
346/// loop
347///   i32.const 1
348///   get_local 0
349///   i32.sub
350///   tee_local 0
351///   br_if 0
352/// end
353/// ```
354///
355/// The start of the block is `i32.const 1`.
356#[derive(Debug)]
357struct ControlBlock {
358	/// The lowest control stack index corresponding to a forward jump targeted by a br, br_if, or
359	/// br_table instruction within this control block. The index must refer to a control block
360	/// that is not a loop, meaning it is a forward jump. Given the way Wasm control flow is
361	/// structured, the lowest index on the stack represents the furthest forward branch target.
362	///
363	/// This value will always be at most the index of the block itself, even if there is no
364	/// explicit br instruction targeting this control block. This does not affect how the value is
365	/// used in the metering algorithm.
366	lowest_forward_br_target: usize,
367
368	/// The active metering block that new instructions contribute a gas cost towards.
369	active_metered_block: MeteredBlock,
370
371	/// Whether the control block is a loop. Loops have the distinguishing feature that branches to
372	/// them jump to the beginning of the block, not the end as with the other control blocks.
373	is_loop: bool,
374}
375
376/// A block of code that metering instructions will be inserted at the beginning of. Metered blocks
377/// are constructed with the property that, in the absence of any traps, either all instructions in
378/// the block are executed or none are.
379#[derive(Debug)]
380struct MeteredBlock {
381	/// Index of the first instruction (aka `Opcode`) in the block.
382	start_pos: usize,
383	/// Sum of costs of all instructions until end of the block.
384	cost: u64,
385}
386
387/// Counter is used to manage state during the gas metering algorithm implemented by
388/// `inject_counter`.
389struct Counter {
390	/// A stack of control blocks. This stack grows when new control blocks are opened with
391	/// `block`, `loop`, and `if` and shrinks when control blocks are closed with `end`. The first
392	/// block on the stack corresponds to the function body, not to any labelled block. Therefore
393	/// the actual Wasm label index associated with each control block is 1 less than its position
394	/// in this stack.
395	stack: Vec<ControlBlock>,
396
397	/// A list of metered blocks that have been finalized, meaning they will no longer change.
398	finalized_blocks: Vec<MeteredBlock>,
399}
400
401impl Counter {
402	fn new() -> Counter {
403		Counter { stack: Vec::new(), finalized_blocks: Vec::new() }
404	}
405
406	/// Open a new control block. The cursor is the position of the first instruction in the block.
407	fn begin_control_block(&mut self, cursor: usize, is_loop: bool) {
408		let index = self.stack.len();
409		self.stack.push(ControlBlock {
410			lowest_forward_br_target: index,
411			active_metered_block: MeteredBlock { start_pos: cursor, cost: 0 },
412			is_loop,
413		})
414	}
415
416	/// Close the last control block. The cursor is the position of the final (pseudo-)instruction
417	/// in the block.
418	fn finalize_control_block(&mut self, cursor: usize) -> Result<(), ()> {
419		// This either finalizes the active metered block or merges its cost into the active
420		// metered block in the previous control block on the stack.
421		self.finalize_metered_block(cursor)?;
422
423		// Pop the control block stack.
424		let closing_control_block = self.stack.pop().ok_or(())?;
425		let closing_control_index = self.stack.len();
426
427		if self.stack.is_empty() {
428			return Ok(())
429		}
430
431		// Update the lowest_forward_br_target for the control block now on top of the stack.
432		{
433			let control_block = self.stack.last_mut().ok_or(())?;
434			control_block.lowest_forward_br_target = min(
435				control_block.lowest_forward_br_target,
436				closing_control_block.lowest_forward_br_target,
437			);
438		}
439
440		// If there may have been a branch to a lower index, then also finalize the active metered
441		// block for the previous control block. Otherwise, finalize it and begin a new one.
442		let may_br_out = closing_control_block.lowest_forward_br_target < closing_control_index;
443		if may_br_out {
444			self.finalize_metered_block(cursor)?;
445		}
446
447		Ok(())
448	}
449
450	/// Finalize the current active metered block.
451	///
452	/// Finalized blocks have final cost which will not change later.
453	fn finalize_metered_block(&mut self, cursor: usize) -> Result<(), ()> {
454		let closing_metered_block = {
455			let control_block = self.stack.last_mut().ok_or(())?;
456			mem::replace(
457				&mut control_block.active_metered_block,
458				MeteredBlock { start_pos: cursor + 1, cost: 0 },
459			)
460		};
461
462		// If the block was opened with a `block`, then its start position will be set to that of
463		// the active metered block in the control block one higher on the stack. This is because
464		// any instructions between a `block` and the first branch are part of the same basic block
465		// as the preceding instruction. In this case, instead of finalizing the block, merge its
466		// cost into the other active metered block to avoid injecting unnecessary instructions.
467		let last_index = self.stack.len() - 1;
468		if last_index > 0 {
469			let prev_control_block = self
470				.stack
471				.get_mut(last_index - 1)
472				.expect("last_index is greater than 0; last_index is stack size - 1; qed");
473			let prev_metered_block = &mut prev_control_block.active_metered_block;
474			if closing_metered_block.start_pos == prev_metered_block.start_pos {
475				prev_metered_block.cost =
476					prev_metered_block.cost.checked_add(closing_metered_block.cost).ok_or(())?;
477				return Ok(())
478			}
479		}
480
481		if closing_metered_block.cost > 0 {
482			self.finalized_blocks.push(closing_metered_block);
483		}
484		Ok(())
485	}
486
487	/// Handle a branch instruction in the program. The cursor is the index of the branch
488	/// instruction in the program. The indices are the stack positions of the target control
489	/// blocks. Recall that the index is 0 for a `return` and relatively indexed from the top of
490	/// the stack by the label of `br`, `br_if`, and `br_table` instructions.
491	fn branch(&mut self, cursor: usize, indices: &[usize]) -> Result<(), ()> {
492		self.finalize_metered_block(cursor)?;
493
494		// Update the lowest_forward_br_target of the current control block.
495		for &index in indices {
496			let target_is_loop = {
497				let target_block = self.stack.get(index).ok_or(())?;
498				target_block.is_loop
499			};
500			if target_is_loop {
501				continue
502			}
503
504			let control_block = self.stack.last_mut().ok_or(())?;
505			control_block.lowest_forward_br_target =
506				min(control_block.lowest_forward_br_target, index);
507		}
508
509		Ok(())
510	}
511
512	/// Returns the stack index of the active control block. Returns None if stack is empty.
513	fn active_control_block_index(&self) -> Option<usize> {
514		self.stack.len().checked_sub(1)
515	}
516
517	/// Get a reference to the currently active metered block.
518	fn active_metered_block(&mut self) -> Result<&mut MeteredBlock, ()> {
519		let top_block = self.stack.last_mut().ok_or(())?;
520		Ok(&mut top_block.active_metered_block)
521	}
522
523	/// Increment the cost of the current block by the specified value.
524	fn increment(&mut self, val: u32) -> Result<(), ()> {
525		let top_block = self.active_metered_block()?;
526		top_block.cost = top_block.cost.checked_add(val.into()).ok_or(())?;
527		Ok(())
528	}
529}
530
531fn inject_grow_counter(instructions: &mut elements::Instructions, grow_counter_func: u32) -> usize {
532	use parity_wasm::elements::Instruction::*;
533	let mut counter = 0;
534	for instruction in instructions.elements_mut() {
535		if let GrowMemory(_) = *instruction {
536			*instruction = Call(grow_counter_func);
537			counter += 1;
538		}
539	}
540	counter
541}
542
543fn add_grow_counter<R: Rules>(
544	module: elements::Module,
545	rules: &R,
546	gas_func: u32,
547) -> elements::Module {
548	use parity_wasm::elements::Instruction::*;
549
550	let cost = match rules.memory_grow_cost() {
551		MemoryGrowCost::Free => return module,
552		MemoryGrowCost::Linear(val) => val.get(),
553	};
554
555	let mut b = builder::from_module(module);
556	b.push_function(
557		builder::function()
558			.signature()
559			.with_param(ValueType::I32)
560			.with_result(ValueType::I32)
561			.build()
562			.body()
563			.with_instructions(elements::Instructions::new(vec![
564				GetLocal(0),
565				GetLocal(0),
566				I64ExtendUI32,
567				I64Const(i64::from(cost)),
568				I64Mul,
569				// todo: there should be strong guarantee that it does not return anything on
570				// stack?
571				Call(gas_func),
572				GrowMemory(0),
573				End,
574			]))
575			.build()
576			.build(),
577	);
578
579	b.build()
580}
581
582fn determine_metered_blocks<R: Rules>(
583	instructions: &elements::Instructions,
584	rules: &R,
585	locals_count: u32,
586) -> Result<Vec<MeteredBlock>, ()> {
587	use parity_wasm::elements::Instruction::*;
588
589	let mut counter = Counter::new();
590
591	// Begin an implicit function (i.e. `func...end`) block.
592	counter.begin_control_block(0, false);
593	// Add locals initialization cost to the function block.
594	let locals_init_cost = rules.call_per_local_cost().checked_mul(locals_count).ok_or(())?;
595	counter.increment(locals_init_cost)?;
596
597	for cursor in 0..instructions.elements().len() {
598		let instruction = &instructions.elements()[cursor];
599		let instruction_cost = rules.instruction_cost(instruction).ok_or(())?;
600		match instruction {
601			Block(_) => {
602				counter.increment(instruction_cost)?;
603
604				// Begin new block. The cost of the following opcodes until `end` or `else` will
605				// be included into this block. The start position is set to that of the previous
606				// active metered block to signal that they should be merged in order to reduce
607				// unnecessary metering instructions.
608				let top_block_start_pos = counter.active_metered_block()?.start_pos;
609				counter.begin_control_block(top_block_start_pos, false);
610			},
611			If(_) => {
612				counter.increment(instruction_cost)?;
613				counter.begin_control_block(cursor + 1, false);
614			},
615			Loop(_) => {
616				counter.increment(instruction_cost)?;
617				counter.begin_control_block(cursor + 1, true);
618			},
619			End => {
620				counter.finalize_control_block(cursor)?;
621			},
622			Else => {
623				counter.finalize_metered_block(cursor)?;
624			},
625			Br(label) | BrIf(label) => {
626				counter.increment(instruction_cost)?;
627
628				// Label is a relative index into the control stack.
629				let active_index = counter.active_control_block_index().ok_or(())?;
630				let target_index = active_index.checked_sub(*label as usize).ok_or(())?;
631				counter.branch(cursor, &[target_index])?;
632			},
633			BrTable(br_table_data) => {
634				counter.increment(instruction_cost)?;
635
636				let active_index = counter.active_control_block_index().ok_or(())?;
637				let target_indices = [br_table_data.default]
638					.iter()
639					.chain(br_table_data.table.iter())
640					.map(|label| active_index.checked_sub(*label as usize))
641					.collect::<Option<Vec<_>>>()
642					.ok_or(())?;
643				counter.branch(cursor, &target_indices)?;
644			},
645			Return => {
646				counter.increment(instruction_cost)?;
647				counter.branch(cursor, &[0])?;
648			},
649			_ => {
650				// An ordinal non control flow instruction increments the cost of the current block.
651				counter.increment(instruction_cost)?;
652			},
653		}
654	}
655
656	counter.finalized_blocks.sort_unstable_by_key(|block| block.start_pos);
657	Ok(counter.finalized_blocks)
658}
659
660fn inject_counter<R: Rules>(
661	instructions: &mut elements::Instructions,
662	gas_function_cost: u64,
663	locals_count: u32,
664	rules: &R,
665	gas_func: u32,
666) -> Result<(), ()> {
667	let blocks = determine_metered_blocks(instructions, rules, locals_count)?;
668	insert_metering_calls(instructions, gas_function_cost, blocks, gas_func)
669}
670
671// Then insert metering calls into a sequence of instructions given the block locations and costs.
672fn insert_metering_calls(
673	instructions: &mut elements::Instructions,
674	gas_function_cost: u64,
675	blocks: Vec<MeteredBlock>,
676	gas_func: u32,
677) -> Result<(), ()> {
678	use parity_wasm::elements::Instruction::*;
679
680	// To do this in linear time, construct a new vector of instructions, copying over old
681	// instructions one by one and injecting new ones as required.
682	let new_instrs_len = instructions.elements().len() + 2 * blocks.len();
683	let original_instrs =
684		mem::replace(instructions.elements_mut(), Vec::with_capacity(new_instrs_len));
685	let new_instrs = instructions.elements_mut();
686
687	let mut block_iter = blocks.into_iter().peekable();
688	for (original_pos, instr) in original_instrs.into_iter().enumerate() {
689		// If there the next block starts at this position, inject metering instructions.
690		let used_block = if let Some(block) = block_iter.peek() {
691			if block.start_pos == original_pos {
692				new_instrs
693					.push(I64Const((block.cost.checked_add(gas_function_cost).ok_or(())?) as i64));
694				new_instrs.push(Call(gas_func));
695				true
696			} else {
697				false
698			}
699		} else {
700			false
701		};
702
703		if used_block {
704			block_iter.next();
705		}
706
707		// Copy over the original instruction.
708		new_instrs.push(instr);
709	}
710
711	if block_iter.next().is_some() {
712		return Err(())
713	}
714
715	Ok(())
716}
717
718#[cfg(test)]
719mod tests {
720	use super::*;
721	use parity_wasm::{builder, elements, elements::Instruction::*, serialize};
722	use pretty_assertions::assert_eq;
723
724	fn get_function_body(
725		module: &elements::Module,
726		index: usize,
727	) -> Option<&[elements::Instruction]> {
728		module
729			.code_section()
730			.and_then(|code_section| code_section.bodies().get(index))
731			.map(|func_body| func_body.code().elements())
732	}
733
734	#[test]
735	fn simple_grow_host_fn() {
736		let module = parse_wat(
737			r#"(module
738			(func (result i32)
739			  global.get 0
740			  memory.grow)
741			(global i32 (i32.const 42))
742			(memory 0 1)
743			)"#,
744		);
745		let backend = host_function::Injector::new("env", "gas");
746		let injected_module =
747			super::inject(module, backend, &ConstantCostRules::new(1, 10_000, 1)).unwrap();
748
749		assert_eq!(
750			get_function_body(&injected_module, 0).unwrap(),
751			&vec![I64Const(2), Call(0), GetGlobal(0), Call(2), End][..]
752		);
753		assert_eq!(
754			get_function_body(&injected_module, 1).unwrap(),
755			&vec![
756				GetLocal(0),
757				GetLocal(0),
758				I64ExtendUI32,
759				I64Const(10000),
760				I64Mul,
761				Call(0),
762				GrowMemory(0),
763				End,
764			][..]
765		);
766
767		let binary = serialize(injected_module).expect("serialization failed");
768		wasmparser::validate(&binary).unwrap();
769	}
770
771	#[test]
772	fn simple_grow_mut_global() {
773		let module = parse_wat(
774			r#"(module
775			(func (result i32)
776			  global.get 0
777			  memory.grow)
778			(global i32 (i32.const 42))
779			(memory 0 1)
780			)"#,
781		);
782		let backend = mutable_global::Injector::new("gas_left");
783		let injected_module =
784			super::inject(module, backend, &ConstantCostRules::new(1, 10_000, 1)).unwrap();
785
786		assert_eq!(
787			get_function_body(&injected_module, 0).unwrap(),
788			&vec![I64Const(13), Call(1), GetGlobal(0), Call(2), End][..]
789		);
790		assert_eq!(
791			get_function_body(&injected_module, 1).unwrap(),
792			&vec![
793				Instruction::GetGlobal(1),
794				Instruction::GetLocal(0),
795				Instruction::I64GeU,
796				Instruction::If(elements::BlockType::NoResult),
797				Instruction::GetGlobal(1),
798				Instruction::GetLocal(0),
799				Instruction::I64Sub,
800				Instruction::SetGlobal(1),
801				Instruction::Else,
802				// sentinel val u64::MAX
803				Instruction::I64Const(-1i64), // non-charged instruction
804				Instruction::SetGlobal(1),    // non-charged instruction
805				Instruction::Unreachable,     // non-charged instruction
806				Instruction::End,
807				Instruction::End,
808			][..]
809		);
810		assert_eq!(
811			get_function_body(&injected_module, 2).unwrap(),
812			&vec![
813				GetLocal(0),
814				GetLocal(0),
815				I64ExtendUI32,
816				I64Const(10000),
817				I64Mul,
818				Call(1),
819				GrowMemory(0),
820				End,
821			][..]
822		);
823
824		let binary = serialize(injected_module).expect("serialization failed");
825		wasmparser::validate(&binary).unwrap();
826	}
827
828	#[test]
829	fn grow_no_gas_no_track_host_fn() {
830		let module = parse_wat(
831			r"(module
832			(func (result i32)
833			  global.get 0
834			  memory.grow)
835			(global i32 (i32.const 42))
836			(memory 0 1)
837			)",
838		);
839		let backend = host_function::Injector::new("env", "gas");
840		let injected_module =
841			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
842
843		assert_eq!(
844			get_function_body(&injected_module, 0).unwrap(),
845			&vec![I64Const(2), Call(0), GetGlobal(0), GrowMemory(0), End][..]
846		);
847
848		assert_eq!(injected_module.functions_space(), 2);
849
850		let binary = serialize(injected_module).expect("serialization failed");
851		wasmparser::validate(&binary).unwrap();
852	}
853
854	#[test]
855	fn grow_no_gas_no_track_mut_global() {
856		let module = parse_wat(
857			r"(module
858			(func (result i32)
859			  global.get 0
860			  memory.grow)
861			(global i32 (i32.const 42))
862			(memory 0 1)
863			)",
864		);
865		let backend = mutable_global::Injector::new("gas_left");
866		let injected_module =
867			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
868
869		assert_eq!(
870			get_function_body(&injected_module, 0).unwrap(),
871			&vec![I64Const(13), Call(1), GetGlobal(0), GrowMemory(0), End][..]
872		);
873
874		assert_eq!(injected_module.functions_space(), 2);
875
876		let binary = serialize(injected_module).expect("serialization failed");
877		wasmparser::validate(&binary).unwrap();
878	}
879
880	#[test]
881	fn call_index_host_fn() {
882		let module = builder::module()
883			.global()
884			.value_type()
885			.i32()
886			.build()
887			.function()
888			.signature()
889			.param()
890			.i32()
891			.build()
892			.body()
893			.build()
894			.build()
895			.function()
896			.signature()
897			.param()
898			.i32()
899			.build()
900			.body()
901			.with_instructions(elements::Instructions::new(vec![
902				Call(0),
903				If(elements::BlockType::NoResult),
904				Call(0),
905				Call(0),
906				Call(0),
907				Else,
908				Call(0),
909				Call(0),
910				End,
911				Call(0),
912				End,
913			]))
914			.build()
915			.build()
916			.build();
917
918		let backend = host_function::Injector::new("env", "gas");
919		let injected_module =
920			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
921
922		assert_eq!(
923			get_function_body(&injected_module, 1).unwrap(),
924			&vec![
925				I64Const(3),
926				Call(0),
927				Call(1),
928				If(elements::BlockType::NoResult),
929				I64Const(3),
930				Call(0),
931				Call(1),
932				Call(1),
933				Call(1),
934				Else,
935				I64Const(2),
936				Call(0),
937				Call(1),
938				Call(1),
939				End,
940				Call(1),
941				End
942			][..]
943		);
944	}
945
946	#[test]
947	fn call_index_mut_global() {
948		let module = builder::module()
949			.global()
950			.value_type()
951			.i32()
952			.build()
953			.function()
954			.signature()
955			.param()
956			.i32()
957			.build()
958			.body()
959			.build()
960			.build()
961			.function()
962			.signature()
963			.param()
964			.i32()
965			.build()
966			.body()
967			.with_instructions(elements::Instructions::new(vec![
968				Call(0),
969				If(elements::BlockType::NoResult),
970				Call(0),
971				Call(0),
972				Call(0),
973				Else,
974				Call(0),
975				Call(0),
976				End,
977				Call(0),
978				End,
979			]))
980			.build()
981			.build()
982			.build();
983
984		let backend = mutable_global::Injector::new("gas_left");
985		let injected_module =
986			super::inject(module, backend, &ConstantCostRules::default()).unwrap();
987
988		assert_eq!(
989			get_function_body(&injected_module, 1).unwrap(),
990			&vec![
991				I64Const(14),
992				Call(2),
993				Call(0),
994				If(elements::BlockType::NoResult),
995				I64Const(14),
996				Call(2),
997				Call(0),
998				Call(0),
999				Call(0),
1000				Else,
1001				I64Const(13),
1002				Call(2),
1003				Call(0),
1004				Call(0),
1005				End,
1006				Call(0),
1007				End
1008			][..]
1009		);
1010	}
1011
1012	fn parse_wat(source: &str) -> elements::Module {
1013		let module_bytes = wat::parse_str(source).unwrap();
1014		elements::deserialize_buffer(module_bytes.as_ref()).unwrap()
1015	}
1016
1017	macro_rules! test_gas_counter_injection {
1018		(names = ($name1:ident, $name2:ident); input = $input:expr; expected = $expected:expr) => {
1019			#[test]
1020			fn $name1() {
1021				let input_module = parse_wat($input);
1022				let expected_module = parse_wat($expected);
1023				let injected_module = super::inject(
1024					input_module,
1025					host_function::Injector::new("env", "gas"),
1026					&ConstantCostRules::default(),
1027				)
1028				.expect("inject_gas_counter call failed");
1029
1030				let actual_func_body = get_function_body(&injected_module, 0)
1031					.expect("injected module must have a function body");
1032				let expected_func_body = get_function_body(&expected_module, 0)
1033					.expect("post-module must have a function body");
1034
1035				assert_eq!(actual_func_body, expected_func_body);
1036			}
1037
1038			#[test]
1039			fn $name2() {
1040				let input_module = parse_wat($input);
1041				let draft_module = parse_wat($expected);
1042				let gas_fun_cost = match mutable_global::Injector::new("gas_left")
1043					.gas_meter(&input_module, &ConstantCostRules::default())
1044				{
1045					GasMeter::Internal { cost, .. } => cost as i64,
1046					_ => 0i64,
1047				};
1048
1049				let injected_module = super::inject(
1050					input_module,
1051					mutable_global::Injector::new("gas_left"),
1052					&ConstantCostRules::default(),
1053				)
1054				.expect("inject_gas_counter call failed");
1055
1056				let actual_func_body = get_function_body(&injected_module, 0)
1057					.expect("injected module must have a function body");
1058				let mut expected_func_body = get_function_body(&draft_module, 0)
1059					.expect("post-module must have a function body")
1060					.to_vec();
1061
1062				// modify expected instructions set for gas_metering::mutable_global
1063				let mut iter = expected_func_body.iter_mut();
1064				while let Some(ins) = iter.next() {
1065					if let I64Const(cost) = ins {
1066						if let Some(ins_next) = iter.next() {
1067							if let Call(0) = ins_next {
1068								*cost += gas_fun_cost;
1069								*ins_next = Call(1);
1070							}
1071						}
1072					}
1073				}
1074
1075				assert_eq!(actual_func_body, &expected_func_body);
1076			}
1077		};
1078	}
1079
1080	test_gas_counter_injection! {
1081		names = (simple_host_fn, simple_mut_global);
1082		input = r#"
1083		(module
1084			(func (result i32)
1085				(get_global 0)))
1086		"#;
1087		expected = r#"
1088		(module
1089			(func (result i32)
1090				(call 0 (i64.const 1))
1091				(get_global 0)))
1092		"#
1093	}
1094
1095	test_gas_counter_injection! {
1096		names = (nested_host_fn, nested_mut_global);
1097		input = r#"
1098		(module
1099			(func (result i32)
1100				(get_global 0)
1101				(block
1102					(get_global 0)
1103					(get_global 0)
1104					(get_global 0))
1105				(get_global 0)))
1106		"#;
1107		expected = r#"
1108		(module
1109			(func (result i32)
1110				(call 0 (i64.const 6))
1111				(get_global 0)
1112				(block
1113					(get_global 0)
1114					(get_global 0)
1115					(get_global 0))
1116				(get_global 0)))
1117		"#
1118	}
1119
1120	test_gas_counter_injection! {
1121		names = (ifelse_host_fn, ifelse_mut_global);
1122		input = r#"
1123		(module
1124			(func (result i32)
1125				(get_global 0)
1126				(if
1127					(then
1128						(get_global 0)
1129						(get_global 0)
1130						(get_global 0))
1131					(else
1132						(get_global 0)
1133						(get_global 0)))
1134				(get_global 0)))
1135		"#;
1136		expected = r#"
1137		(module
1138			(func (result i32)
1139				(call 0 (i64.const 3))
1140				(get_global 0)
1141				(if
1142					(then
1143						(call 0 (i64.const 3))
1144						(get_global 0)
1145						(get_global 0)
1146						(get_global 0))
1147					(else
1148						(call 0 (i64.const 2))
1149						(get_global 0)
1150						(get_global 0)))
1151				(get_global 0)))
1152		"#
1153	}
1154
1155	test_gas_counter_injection! {
1156		names = (branch_innermost_host_fn, branch_innermost_mut_global);
1157		input = r#"
1158		(module
1159			(func (result i32)
1160				(get_global 0)
1161				(block
1162					(get_global 0)
1163					(drop)
1164					(br 0)
1165					(get_global 0)
1166					(drop))
1167				(get_global 0)))
1168		"#;
1169		expected = r#"
1170		(module
1171			(func (result i32)
1172				(call 0 (i64.const 6))
1173				(get_global 0)
1174				(block
1175					(get_global 0)
1176					(drop)
1177					(br 0)
1178					(call 0 (i64.const 2))
1179					(get_global 0)
1180					(drop))
1181				(get_global 0)))
1182		"#
1183	}
1184
1185	test_gas_counter_injection! {
1186		names = (branch_outer_block_host_fn, branch_outer_block_mut_global);
1187		input = r#"
1188		(module
1189			(func (result i32)
1190				(get_global 0)
1191				(block
1192					(get_global 0)
1193					(if
1194						(then
1195							(get_global 0)
1196							(get_global 0)
1197							(drop)
1198							(br_if 1)))
1199					(get_global 0)
1200					(drop))
1201				(get_global 0)))
1202		"#;
1203		expected = r#"
1204		(module
1205			(func (result i32)
1206				(call 0 (i64.const 5))
1207				(get_global 0)
1208				(block
1209					(get_global 0)
1210					(if
1211						(then
1212							(call 0 (i64.const 4))
1213							(get_global 0)
1214							(get_global 0)
1215							(drop)
1216							(br_if 1)))
1217					(call 0 (i64.const 2))
1218					(get_global 0)
1219					(drop))
1220				(get_global 0)))
1221		"#
1222	}
1223
1224	test_gas_counter_injection! {
1225		names = (branch_outer_loop_host_fn, branch_outer_loop_mut_global);
1226		input = r#"
1227		(module
1228			(func (result i32)
1229				(get_global 0)
1230				(loop
1231					(get_global 0)
1232					(if
1233						(then
1234							(get_global 0)
1235							(br_if 0))
1236						(else
1237							(get_global 0)
1238							(get_global 0)
1239							(drop)
1240							(br_if 1)))
1241					(get_global 0)
1242					(drop))
1243				(get_global 0)))
1244		"#;
1245		expected = r#"
1246		(module
1247			(func (result i32)
1248				(call 0 (i64.const 3))
1249				(get_global 0)
1250				(loop
1251					(call 0 (i64.const 4))
1252					(get_global 0)
1253					(if
1254						(then
1255							(call 0 (i64.const 2))
1256							(get_global 0)
1257							(br_if 0))
1258						(else
1259							(call 0 (i64.const 4))
1260							(get_global 0)
1261							(get_global 0)
1262							(drop)
1263							(br_if 1)))
1264					(get_global 0)
1265					(drop))
1266				(get_global 0)))
1267		"#
1268	}
1269
1270	test_gas_counter_injection! {
1271		names = (return_from_func_host_fn, return_from_func_mut_global);
1272		input = r#"
1273		(module
1274			(func (result i32)
1275				(get_global 0)
1276				(if
1277					(then
1278						(return)))
1279				(get_global 0)))
1280		"#;
1281		expected = r#"
1282		(module
1283			(func (result i32)
1284				(call 0 (i64.const 2))
1285				(get_global 0)
1286				(if
1287					(then
1288						(call 0 (i64.const 1))
1289						(return)))
1290				(call 0 (i64.const 1))
1291				(get_global 0)))
1292		"#
1293	}
1294
1295	test_gas_counter_injection! {
1296		names = (branch_from_if_not_else_host_fn, branch_from_if_not_else_mut_global);
1297		input = r#"
1298		(module
1299			(func (result i32)
1300				(get_global 0)
1301				(block
1302					(get_global 0)
1303					(if
1304						(then (br 1))
1305						(else (br 0)))
1306					(get_global 0)
1307					(drop))
1308				(get_global 0)))
1309		"#;
1310		expected = r#"
1311		(module
1312			(func (result i32)
1313				(call 0 (i64.const 5))
1314				(get_global 0)
1315				(block
1316					(get_global 0)
1317					(if
1318						(then
1319							(call 0 (i64.const 1))
1320							(br 1))
1321						(else
1322							(call 0 (i64.const 1))
1323							(br 0)))
1324					(call 0 (i64.const 2))
1325					(get_global 0)
1326					(drop))
1327				(get_global 0)))
1328		"#
1329	}
1330
1331	test_gas_counter_injection! {
1332		names = (empty_loop_host_fn, empty_loop_mut_global);
1333		input = r#"
1334		(module
1335			(func
1336				(loop
1337					(br 0)
1338				)
1339				unreachable
1340			)
1341		)
1342		"#;
1343		expected = r#"
1344		(module
1345			(func
1346				(call 0 (i64.const 2))
1347				(loop
1348					(call 0 (i64.const 1))
1349					(br 0)
1350				)
1351				unreachable
1352			)
1353		)
1354		"#
1355	}
1356}