feat(EOF): Add CALLF/JUMPF stack checks (#1417)

* feat(EOF): Add CALLF stack check

* fix tests

* jump stack check

crates/interpreter/src/instructions/control.rs

@@ -87,13 +87,24 @@ pub fn callf<H: Host + ?Sized>(interpreter: &mut Interpreter, _host: &mut H) {
     gas!(interpreter, gas::LOW);
 
     let idx = unsafe { read_u16(interpreter.instruction_pointer) } as usize;
-    // TODO Check stack with EOF types.
 
-    if interpreter.function_stack.return_stack_len() == 1024 {
+    if interpreter.function_stack.return_stack_len() >= 1024 {
         interpreter.instruction_result = InstructionResult::EOFFunctionStackOverflow;
         return;
     }
 
+    // get target types
+    let Some(types) = interpreter.eof().unwrap().body.types_section.get(idx) else {
+        panic!("Invalid EOF in execution, expecting correct intermediate in callf")
+    };
+
+    // Check max stack height for target code section.
+    // safe to subtract as max_stack_height is always more than inputs.
+    if interpreter.stack.len() + (types.max_stack_size - types.inputs as u16) as usize > 1024 {
+        interpreter.instruction_result = InstructionResult::StackOverflow;
+        return;
+    }
+
     // push current idx and PC to the callf stack.
     // PC is incremented by 2 to point to the next instruction after callf.
     interpreter
@@ -120,7 +131,17 @@ pub fn jumpf<H: Host + ?Sized>(interpreter: &mut Interpreter, _host: &mut H) {
 
     let idx = unsafe { read_u16(interpreter.instruction_pointer) } as usize;
 
-    // TODO(EOF) do types stack checks
+    // get target types
+    let Some(types) = interpreter.eof().unwrap().body.types_section.get(idx) else {
+        panic!("Invalid EOF in execution, expecting correct intermediate in jumpf")
+    };
+
+    // Check max stack height for target code section.
+    // safe to subtract as max_stack_height is always more than inputs.
+    if interpreter.stack.len() + (types.max_stack_size - types.inputs as u16) as usize > 1024 {
+        interpreter.instruction_result = InstructionResult::StackOverflow;
+        return;
+    }
 
     interpreter.function_stack.set_current_code_idx(idx);
     interpreter.load_eof_code(idx, 0)
@@ -183,7 +204,7 @@ pub fn unknown<H: Host + ?Sized>(interpreter: &mut Interpreter, _host: &mut H) {
 
 #[cfg(test)]
 mod test {
-    use revm_primitives::{bytes, Bytecode, Eof, PragueSpec};
+    use revm_primitives::{bytes, eof::TypesSection, Bytecode, Eof, PragueSpec};
 
     use super::*;
     use crate::{
@@ -281,27 +302,39 @@ mod test {
         Eof::decode(bytes).unwrap()
     }
 
-    #[test]
-    fn callf_retf_jumpf() {
-        let table = make_instruction_table::<_, PragueSpec>();
-        let mut host = DummyHost::default();
+    fn eof_setup(bytes1: Bytes, bytes2: Bytes) -> Interpreter {
+        eof_setup_with_types(bytes1, bytes2, TypesSection::default())
+    }
+
+    /// Two code section and types section is for last code.
+    fn eof_setup_with_types(bytes1: Bytes, bytes2: Bytes, types: TypesSection) -> Interpreter {
         let mut eof = dummy_eof();
 
         eof.body.code_section.clear();
+        eof.body.types_section.clear();
         eof.header.code_sizes.clear();
 
-        let bytes1 = Bytes::from([CALLF, 0x00, 0x01, JUMPF, 0x00, 0x01]);
         eof.header.code_sizes.push(bytes1.len() as u16);
         eof.body.code_section.push(bytes1.clone());
-        let bytes2 = Bytes::from([STOP, RETF]);
+        eof.body.types_section.push(TypesSection::new(0, 0, 11));
+
         eof.header.code_sizes.push(bytes2.len() as u16);
         eof.body.code_section.push(bytes2.clone());
+        eof.body.types_section.push(types);
 
         let mut interp = Interpreter::new_bytecode(Bytecode::Eof(eof));
         interp.gas = Gas::new(10000);
+        interp
+    }
 
-        assert_eq!(interp.function_stack.current_code_idx, 0);
-        assert!(interp.function_stack.return_stack.is_empty());
+    #[test]
+    fn callf_retf_stop() {
+        let table = make_instruction_table::<_, PragueSpec>();
+        let mut host = DummyHost::default();
+
+        let bytes1 = Bytes::from([CALLF, 0x00, 0x01, STOP]);
+        let bytes2 = Bytes::from([RETF]);
+        let mut interp = eof_setup(bytes1, bytes2.clone());
 
         // CALLF
         interp.step(&table, &mut host);
@@ -313,19 +346,94 @@ mod test {
         );
         assert_eq!(interp.instruction_pointer, bytes2.as_ptr());
 
-        // STOP
-        interp.step(&table, &mut host);
         // RETF
         interp.step(&table, &mut host);
 
         assert_eq!(interp.function_stack.current_code_idx, 0);
         assert_eq!(interp.function_stack.return_stack, Vec::new());
         assert_eq!(interp.program_counter(), 3);
 
+        // STOP
+        interp.step(&table, &mut host);
+        assert_eq!(interp.instruction_result, InstructionResult::Stop);
+    }
+
+    #[test]
+    fn callf_stop() {
+        let table = make_instruction_table::<_, PragueSpec>();
+        let mut host = DummyHost::default();
+
+        let bytes1 = Bytes::from([CALLF, 0x00, 0x01]);
+        let bytes2 = Bytes::from([STOP]);
+        let mut interp = eof_setup(bytes1, bytes2.clone());
+
+        // CALLF
+        interp.step(&table, &mut host);
+
+        assert_eq!(interp.function_stack.current_code_idx, 1);
+        assert_eq!(
+            interp.function_stack.return_stack[0],
+            FunctionReturnFrame::new(0, 3)
+        );
+        assert_eq!(interp.instruction_pointer, bytes2.as_ptr());
+
+        // STOP
+        interp.step(&table, &mut host);
+        assert_eq!(interp.instruction_result, InstructionResult::Stop);
+    }
+
+    #[test]
+    fn callf_stack_overflow() {
+        let table = make_instruction_table::<_, PragueSpec>();
+        let mut host = DummyHost::default();
+
+        let bytes1 = Bytes::from([CALLF, 0x00, 0x01]);
+        let bytes2 = Bytes::from([STOP]);
+        let mut interp =
+            eof_setup_with_types(bytes1, bytes2.clone(), TypesSection::new(0, 0, 1025));
+
+        // CALLF
+        interp.step(&table, &mut host);
+
+        // stack overflow
+        assert_eq!(interp.instruction_result, InstructionResult::StackOverflow);
+    }
+
+    #[test]
+    fn jumpf_stop() {
+        let table = make_instruction_table::<_, PragueSpec>();
+        let mut host = DummyHost::default();
+
+        let bytes1 = Bytes::from([JUMPF, 0x00, 0x01]);
+        let bytes2 = Bytes::from([STOP]);
+        let mut interp = eof_setup(bytes1, bytes2.clone());
+
         // JUMPF
         interp.step(&table, &mut host);
+
         assert_eq!(interp.function_stack.current_code_idx, 1);
-        assert_eq!(interp.function_stack.return_stack, Vec::new());
+        assert!(interp.function_stack.return_stack.is_empty());
         assert_eq!(interp.instruction_pointer, bytes2.as_ptr());
+
+        // STOP
+        interp.step(&table, &mut host);
+        assert_eq!(interp.instruction_result, InstructionResult::Stop);
+    }
+
+    #[test]
+    fn jumpf_stack_overflow() {
+        let table = make_instruction_table::<_, PragueSpec>();
+        let mut host = DummyHost::default();
+
+        let bytes1 = Bytes::from([JUMPF, 0x00, 0x01]);
+        let bytes2 = Bytes::from([STOP]);
+        let mut interp =
+            eof_setup_with_types(bytes1, bytes2.clone(), TypesSection::new(0, 0, 1025));
+
+        // JUMPF
+        interp.step(&table, &mut host);
+
+        // stack overflow
+        assert_eq!(interp.instruction_result, InstructionResult::StackOverflow);
     }
 }

crates/interpreter/src/interpreter.rs

@@ -137,7 +137,7 @@ impl Interpreter {
     pub(crate) fn load_eof_code(&mut self, idx: usize, pc: usize) {
         // SAFETY: eof flag is true only if bytecode is Eof.
         let Bytecode::Eof(eof) = &self.contract.bytecode else {
-            panic!("Expected EOF bytecode")
+            panic!("Expected EOF code section")
         };
         let Some(code) = eof.body.code(idx) else {
             panic!("Code not found")

crates/primitives/src/bytecode/eof/types_section.rs

@@ -20,6 +20,14 @@ pub struct TypesSection {
 }
 
 impl TypesSection {
+    /// Returns new `TypesSection` with the given inputs, outputs, and max_stack_size.
+    pub fn new(inputs: u8, outputs: u8, max_stack_size: u16) -> Self {
+        Self {
+            inputs,
+            outputs,
+            max_stack_size,
+        }
+    }
     /// Calculates the difference between the number of input and output stack elements.
     #[inline]
     pub const fn io_diff(&self) -> i32 {