diff --git a/lib/middlewares/src/metering.rs b/lib/middlewares/src/metering.rs index 864d6650ce4..18bc7997146 100644 --- a/lib/middlewares/src/metering.rs +++ b/lib/middlewares/src/metering.rs @@ -368,15 +368,39 @@ mod tests { fn bytecode() -> Vec { wat2wasm( - br#" - (module + br#"(module (type $add_t (func (param i32) (result i32))) (func $add_one_f (type $add_t) (param $value i32) (result i32) local.get $value i32.const 1 i32.add) - (export "add_one" (func $add_one_f))) - "#, + (func $short_loop_f + (local $x f64) (local $j i32) + (local.set $x (f64.const 5.5)) + + (loop $named_loop + ;; $j++ + local.get $j + i32.const 1 + i32.add + local.set $j + + ;; if $j < 5, one more time + local.get $j + i32.const 5 + i32.lt_s + br_if $named_loop + ) + ) + (func $infi_loop_f + (loop $infi_loop_start + br $infi_loop_start + ) + ) + (export "add_one" (func $add_one_f)) + (export "short_loop" (func $short_loop_f)) + (export "infi_loop" (func $infi_loop_f)) + )"#, ) .unwrap() .into() @@ -486,4 +510,68 @@ mod tests { MeteringPoints::Remaining(4) ); } + + #[test] + fn metering_works_for_loops() { + const INITIAL_POINTS: u64 = 10_000; + + fn cost(operator: &Operator) -> u64 { + match operator { + Operator::Loop { .. } => 1000, + Operator::Br { .. } | Operator::BrIf { .. } => 10, + Operator::F64Const { .. } => 7, + _ => 0, + } + } + + // Short loop + + let metering = Arc::new(Metering::new(INITIAL_POINTS, cost)); + let mut compiler_config = Cranelift::default(); + compiler_config.push_middleware(metering); + let mut store = Store::new(EngineBuilder::new(compiler_config)); + let module = Module::new(&store, bytecode()).unwrap(); + + let instance = Instance::new(&mut store, &module, &imports! {}).unwrap(); + let short_loop: TypedFunction<(), ()> = instance + .exports + .get_function("short_loop") + .unwrap() + .typed(&store) + .unwrap(); + short_loop.call(&mut store).unwrap(); + + let points_used: u64 = match get_remaining_points(&mut store, &instance) { + MeteringPoints::Exhausted => panic!("Unexpected exhausted"), + MeteringPoints::Remaining(remaining) => INITIAL_POINTS - remaining, + }; + + assert_eq!( + points_used, + 7 /* pre-loop instructions */ + + 1000 /* loop instruction */ + 50 /* five conditional breaks */ + ); + + // Infinite loop + + let metering = Arc::new(Metering::new(INITIAL_POINTS, cost)); + let mut compiler_config = Cranelift::default(); + compiler_config.push_middleware(metering); + let mut store = Store::new(EngineBuilder::new(compiler_config)); + let module = Module::new(&store, bytecode()).unwrap(); + + let instance = Instance::new(&mut store, &module, &imports! {}).unwrap(); + let infi_loop: TypedFunction<(), ()> = instance + .exports + .get_function("infi_loop") + .unwrap() + .typed(&store) + .unwrap(); + infi_loop.call(&mut store).unwrap_err(); // exhausted leads to runtime error + + assert_eq!( + get_remaining_points(&mut store, &instance), + MeteringPoints::Exhausted + ); + } }