fix(metering): Prevent execution past limit

Inject metering checks ahead of executing the next code block, instead
of after. This ensures that execution will always return with an error
as soon as it is determined that the execution limit will be exceeded,
without ever actually exceeding the limit.

If execution returns with ExecutionLimitExceededError, get_points_used
will return some number of points greater than the limit, since the
INTERNAL_FIELD_USED is incremented and saved prior to the limit check.

Fix wasmerio#999

lib/middleware-common/src/metering.rs

@@ -22,12 +22,82 @@ static INTERNAL_FIELD_LIMIT: InternalField = InternalField::allocate();
 /// the same function calls so we can say that the metering is deterministic.
 ///
 pub struct Metering {
-    current_block: u64,
+    cost_operator_idx: usize,
+    current_block_cost: u64,
 }
 
 impl Metering {
     pub fn new() -> Metering {
-        Metering { current_block: 0 }
+        Metering {
+            cost_operator_idx: 0,
+            current_block_cost: 0,
+        }
+    }
+
+    /// inject_metering injects a series of opcodes that adds the cost of the next code block to
+    /// INTERNAL_FIELD_USED and then checks if it has exceeded INTERNAL_FIELD_LIMIT. Since the cost
+    /// of the next code block is not known at the point of injection, Operator::Unreachable is
+    /// used in place of Operator::I64Const{COST}, and the position of this Operator in the sink is
+    /// saved, so that it can later be replaced with the correct cost, once it is know later on in
+    /// parsing Events.
+    fn inject_metering<'a, 'b: 'a>(&mut self, sink: &mut EventSink<'a, 'b>) {
+        // PUSH USED
+        sink.push(Event::Internal(InternalEvent::GetInternal(
+            INTERNAL_FIELD_USED.index() as _,
+        )));
+
+        // placeholder for PUSH COST
+        self.cost_operator_idx = sink.buffer.len();
+        sink.push(Event::WasmOwned(Operator::Unreachable));
+
+        // USED + COST
+        sink.push(Event::WasmOwned(Operator::I64Add));
+
+        // SAVE USED
+        sink.push(Event::Internal(InternalEvent::SetInternal(
+            INTERNAL_FIELD_USED.index() as _,
+        )));
+
+        // PUSH USED
+        sink.push(Event::Internal(InternalEvent::GetInternal(
+            INTERNAL_FIELD_USED.index() as _,
+        )));
+
+        // PUSH LIMIT
+        sink.push(Event::Internal(InternalEvent::GetInternal(
+            INTERNAL_FIELD_LIMIT.index() as _,
+        )));
+
+        // IF USED > LIMIT
+        sink.push(Event::WasmOwned(Operator::I64GtU));
+        sink.push(Event::WasmOwned(Operator::If {
+            ty: WpTypeOrFuncType::Type(WpType::EmptyBlockType),
+        }));
+
+        //          TRAP! EXECUTION LIMIT EXCEEDED
+        sink.push(Event::Internal(InternalEvent::Breakpoint(Box::new(|_| {
+            Err(Box::new(ExecutionLimitExceededError))
+        }))));
+
+        // ENDIF
+        sink.push(Event::WasmOwned(Operator::End));
+    }
+
+    /// set_cost_op_reset_cost replaces the previous placeholder Operator::Unreachable with the
+    /// correct Operator::I64Const{current_block_cost} and resets current_block_cost to 0.
+    fn set_cost_op_reset_cost<'a, 'b: 'a>(&mut self, sink: &mut EventSink<'a, 'b>) {
+        sink.buffer[self.cost_operator_idx] = Event::WasmOwned(Operator::I64Const {
+            value: self.current_block_cost as i64,
+        });
+        self.current_block_cost = 0;
+    }
+
+    /// increment_cost adds 1 to the current_block_cost.
+    ///
+    /// Later this may be replaced with a cost map for assigning custom unique cost values to
+    /// specific Operators.
+    fn increment_cost<'a, 'b: 'a>(&mut self, _op: &Event<'a, 'b>) {
+        self.current_block_cost += 1;
     }
 }
 
@@ -42,65 +112,32 @@ impl FunctionMiddleware for Metering {
         _module_info: &ModuleInfo,
         sink: &mut EventSink<'a, 'b>,
     ) -> Result<(), Self::Error> {
+        self.increment_cost(&op);
+
         match op {
             Event::Internal(InternalEvent::FunctionBegin(_)) => {
-                self.current_block = 0;
+                self.inject_metering(sink);
             }
-            Event::Wasm(&ref op) | Event::WasmOwned(ref op) => {
-                self.current_block += 1;
-                match *op {
-                    Operator::Loop { .. }
-                    | Operator::Block { .. }
-                    | Operator::End
-                    | Operator::If { .. }
-                    | Operator::Else
-                    | Operator::Unreachable
-                    | Operator::Br { .. }
-                    | Operator::BrTable { .. }
-                    | Operator::BrIf { .. }
-                    | Operator::Call { .. }
-                    | Operator::CallIndirect { .. }
-                    | Operator::Return => {
-                        sink.push(Event::Internal(InternalEvent::GetInternal(
-                            INTERNAL_FIELD_USED.index() as _,
-                        )));
-                        sink.push(Event::WasmOwned(Operator::I64Const {
-                            value: self.current_block as i64,
-                        }));
-                        sink.push(Event::WasmOwned(Operator::I64Add));
-                        sink.push(Event::Internal(InternalEvent::SetInternal(
-                            INTERNAL_FIELD_USED.index() as _,
-                        )));
-                        self.current_block = 0;
-                    }
-                    _ => {}
-                }
-                match *op {
-                    Operator::Br { .. }
-                    | Operator::BrTable { .. }
-                    | Operator::BrIf { .. }
-                    | Operator::Call { .. }
-                    | Operator::CallIndirect { .. } => {
-                        sink.push(Event::Internal(InternalEvent::GetInternal(
-                            INTERNAL_FIELD_USED.index() as _,
-                        )));
-                        sink.push(Event::Internal(InternalEvent::GetInternal(
-                            INTERNAL_FIELD_LIMIT.index() as _,
-                        )));
-                        sink.push(Event::WasmOwned(Operator::I64GeU));
-                        sink.push(Event::WasmOwned(Operator::If {
-                            ty: WpTypeOrFuncType::Type(WpType::EmptyBlockType),
-                        }));
-                        sink.push(Event::Internal(InternalEvent::Breakpoint(Box::new(|_| {
-                            Err(Box::new(ExecutionLimitExceededError))
-                        }))));
-                        sink.push(Event::WasmOwned(Operator::End));
-                    }
-                    _ => {}
+            Event::Wasm(&ref op) | Event::WasmOwned(ref op) => match *op {
+                Operator::Loop { .. }
+                | Operator::End
+                | Operator::If { .. }
+                | Operator::Else
+                | Operator::Br { .. }
+                | Operator::BrTable { .. }
+                | Operator::BrIf { .. }
+                | Operator::Return => {
+                    self.set_cost_op_reset_cost(sink);
+                    self.inject_metering(sink);
                 }
+                _ => {}
+            },
+            Event::Internal(InternalEvent::FunctionEnd) => {
+                self.set_cost_op_reset_cost(sink);
             }
             _ => {}
         }
+
         sink.push(op);
         Ok(())
     }