fix(aot): execute_metered should error if exit code nonzero (#2336)

Author: jonathanpwang

crates/vm/src/arch/aot/metered_execute.rs

@@ -12,7 +12,7 @@ use crate::{
         },
         execution_mode::{metered::segment_ctx::SegmentationCtx, MeteredCtx, Segment},
         interpreter::{
-            alloc_pre_compute_buf, get_metered_pre_compute_instructions,
+            alloc_pre_compute_buf, check_termination, get_metered_pre_compute_instructions,
             get_metered_pre_compute_max_size, split_pre_compute_buf, PreComputeInstruction,
         },
         AotError, ExecutionError, ExecutorInventory, MeteredExecutor, StaticProgramError, Streams,
@@ -341,16 +341,21 @@ where
         from_state: VmState<F, GuestMemory>,
         ctx: MeteredCtx,
     ) -> Result<(Vec<Segment>, VmState<F, GuestMemory>), ExecutionError> {
-        let vm_exec_state = VmExecState::new(from_state, ctx);
-        let vm_exec_state = self.execute_metered_until_suspend(vm_exec_state)?;
-        // handle execution error
-        match vm_exec_state.exit_code {
-            Ok(_) => Ok((
-                vm_exec_state.ctx.segmentation_ctx.segments,
-                vm_exec_state.vm_state,
-            )),
-            Err(e) => Err(e),
+        let mut exec_state = VmExecState::new(from_state, ctx);
+
+        loop {
+            exec_state = self.execute_metered_until_suspend(exec_state)?;
+            // The execution has terminated.
+            if exec_state.exit_code.is_ok() && exec_state.exit_code.as_ref().unwrap().is_some() {
+                break;
+            }
+            if exec_state.exit_code.is_err() {
+                return Err(exec_state.exit_code.unwrap_err());
+            }
         }
+        check_termination(exec_state.exit_code)?;
+        let VmExecState { vm_state, ctx, .. } = exec_state;
+        Ok((ctx.into_segments(), vm_state))
     }
 
     // TODO: implement execute_metered_until_suspend for AOT if needed

crates/vm/src/arch/interpreter.rs

@@ -836,7 +836,9 @@ fn check_exit_code(exit_code: Result<Option<u32>, ExecutionError>) -> Result<(),
 }
 
 /// Same as [check_exit_code] but errors if program did not terminate.
-fn check_termination(exit_code: Result<Option<u32>, ExecutionError>) -> Result<(), ExecutionError> {
+pub(super) fn check_termination(
+    exit_code: Result<Option<u32>, ExecutionError>,
+) -> Result<(), ExecutionError> {
     let did_terminate = matches!(exit_code.as_ref(), Ok(Some(_)));
     check_exit_code(exit_code)?;
     match did_terminate {