From c90b157a241e4cbc5ca70e945b7f0ceb53de3c3f Mon Sep 17 00:00:00 2001
From: Jan Bujak <jan@parity.io>
Date: Fri, 10 May 2024 13:57:15 +0000
Subject: [PATCH] Handle invalid jumps

---
 crates/polkavm/src/compiler.rs       | 10 +++++++---
 crates/polkavm/src/compiler/amd64.rs |  8 ++++----
 crates/polkavm/src/interpreter.rs    |  9 ++++++---
 crates/polkavm/src/utils.rs          |  5 +++++
 4 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/crates/polkavm/src/compiler.rs b/crates/polkavm/src/compiler.rs
index 10d966ab..6911cc31 100644
--- a/crates/polkavm/src/compiler.rs
+++ b/crates/polkavm/src/compiler.rs
@@ -349,15 +349,19 @@ impl<'a, S> CompilerVisitor<'a, S> where S: Sandbox {
         log::trace!("Compiling {}", self.current_instruction(code_offset));
     }
 
-    fn get_or_forward_declare_label(&mut self, code_offset: u32) -> Label {
+    fn get_or_forward_declare_label(&mut self, code_offset: u32) -> Option<Label> {
         match self.code_offset_to_label.get(code_offset) {
-            Some(label) => label,
+            Some(label) => Some(label),
             None => {
+                if code_offset > self.code_offset_to_label.len() {
+                    return None;
+                }
+
                 let label = self.asm.forward_declare_label();
                 log::trace!("Label: {label} -> {code_offset} (forward declare)");
 
                 self.code_offset_to_label.insert(code_offset, label);
-                label
+                Some(label)
             }
         }
     }
diff --git a/crates/polkavm/src/compiler/amd64.rs b/crates/polkavm/src/compiler/amd64.rs
index 2e0dbd96..d6785b37 100644
--- a/crates/polkavm/src/compiler/amd64.rs
+++ b/crates/polkavm/src/compiler/amd64.rs
@@ -397,7 +397,7 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
     fn branch(&mut self, s1: RawReg, s2: impl Into<RegImm>, target: u32, condition: Condition) {
         let reg_size = self.reg_size();
         let s1 = conv_reg(s1);
-        let label = self.get_or_forward_declare_label(target);
+        let label = self.get_or_forward_declare_label(target).unwrap_or(self.trap_label);
 
         let asm = self.asm.reserve::<U2>();
         let asm = match s2.into() {
@@ -580,7 +580,7 @@ impl<'r, 'a, S> ArchVisitor<'r, 'a, S> where S: Sandbox {
                 self.branch_to_label(Condition::Sign, self.trap_label);
             }
 
-            let target_label = self.get_or_forward_declare_label(export.target_code_offset());
+            let target_label = self.get_or_forward_declare_label(export.target_code_offset()).unwrap_or(self.trap_label);
             self.jump_to_label(target_label);
         }
     }
@@ -1406,13 +1406,13 @@ impl<'r, 'a, S> InstructionVisitor for ArchVisitor<'r, 'a, S> where S: Sandbox {
 
     #[inline(always)]
     fn jump(&mut self, target: u32) -> Self::ReturnTy {
-        let label = self.get_or_forward_declare_label(target);
+        let label = self.get_or_forward_declare_label(target).unwrap_or(self.trap_label);
         self.jump_to_label(label);
     }
 
     #[inline(always)]
     fn load_imm_and_jump(&mut self, ra: RawReg, value: u32, target: u32) -> Self::ReturnTy {
-        let label = self.get_or_forward_declare_label(target);
+        let label = self.get_or_forward_declare_label(target).unwrap_or(self.trap_label);
         let asm = self.asm.reserve::<U2>();
         let asm = asm.push(mov_imm(conv_reg(ra), imm32(value)));
         let asm = jump_to_label(asm, label);
diff --git a/crates/polkavm/src/interpreter.rs b/crates/polkavm/src/interpreter.rs
index 90d3df39..4d54e4b8 100644
--- a/crates/polkavm/src/interpreter.rs
+++ b/crates/polkavm/src/interpreter.rs
@@ -424,7 +424,7 @@ impl InterpretedInstance {
             self.compiled_offset = (compiled_offset.get() - 1) as usize;
         } else {
             self.compiled_offset = self.compiled_instructions.len();
-            self.compile_block::<DEBUG>();
+            self.compile_block::<DEBUG>()?;
         }
 
         if self.gas_remaining.is_some() {
@@ -469,10 +469,11 @@ impl InterpretedInstance {
 
     #[inline(never)]
     #[cold]
-    fn compile_block<const DEBUG: bool>(&mut self) {
+    fn compile_block<const DEBUG: bool>(&mut self) -> Result<(), ExecutionError> {
         let Some(instructions) = self.module.blob().instructions_at(self.instruction_offset) else {
-            return;
+            return Err(ExecutionError::Trap(Default::default()));
         };
+
         if DEBUG {
             log::debug!("Compiling block at {}:", self.instruction_offset);
         }
@@ -496,6 +497,8 @@ impl InterpretedInstance {
 
         self.compiled_offset_for_block
             .insert(self.instruction_offset, NonZeroU32::new((starting_offset + 1) as u32).unwrap());
+
+        Ok(())
     }
 
     fn check_gas(&mut self) -> Result<(), ExecutionError> {
diff --git a/crates/polkavm/src/utils.rs b/crates/polkavm/src/utils.rs
index b3113cf3..4a0662ee 100644
--- a/crates/polkavm/src/utils.rs
+++ b/crates/polkavm/src/utils.rs
@@ -58,6 +58,11 @@ where
         self.inner.get(key as usize).and_then(|value| *value)
     }
 
+    #[inline]
+    pub fn len(&self) -> u32 {
+        self.inner.len() as u32
+    }
+
     #[inline]
     pub fn insert(&mut self, key: u32, value: T) {
         self.inner[key as usize] = Some(value);