fix(levm): codecopy memory allocation (#1328)

Resolves #1325
lambdaclass · Nov 29, 2024 · 450a1c7 · 450a1c7
1 parent 668780a
commit 450a1c7
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 17 deletions.
diff --git a/crates/vm/levm/src/opcode_handlers/environment.rs b/crates/vm/levm/src/opcode_handlers/environment.rs
@@ -1,6 +1,6 @@
 use crate::{
     call_frame::CallFrame,
-    constants::WORD_SIZE_IN_BYTES_USIZE,
+    constants::{WORD_SIZE, WORD_SIZE_IN_BYTES_USIZE},
     errors::{InternalError, OpcodeSuccess, OutOfGasError, VMError},
     gas_cost,
     vm::{word_to_address, VM},
@@ -210,12 +210,12 @@ impl VM {
         &mut self,
         current_call_frame: &mut CallFrame,
     ) -> Result<OpcodeSuccess, VMError> {
-        let dest_offset: usize = current_call_frame
+        let destination_offset: usize = current_call_frame
             .stack
             .pop()?
             .try_into()
             .map_err(|_| VMError::VeryLargeNumber)?;
-        let offset: usize = current_call_frame
+        let code_offset: usize = current_call_frame
             .stack
             .pop()?
             .try_into()
@@ -226,29 +226,52 @@ impl VM {
             .try_into()
             .map_err(|_| VMError::VeryLargeNumber)?;
 
-        let gas_cost =
-            gas_cost::codecopy(current_call_frame, size, dest_offset).map_err(VMError::OutOfGas)?;
+        let gas_cost = gas_cost::codecopy(current_call_frame, size, destination_offset)
+            .map_err(VMError::OutOfGas)?;
 
         self.increase_consumed_gas(current_call_frame, gas_cost)?;
 
         if size == 0 {
             return Ok(OpcodeSuccess::Continue);
         }
 
-        let bytecode_len = current_call_frame.bytecode.len();
-        let code = if offset < bytecode_len {
-            current_call_frame.bytecode.slice(
-                offset
-                    ..(offset.checked_add(size).ok_or(VMError::Internal(
+        let new_memory_size = (destination_offset
+            .checked_add(size)
+            .ok_or(VMError::Internal(
+                InternalError::ArithmeticOperationOverflow,
+            ))?)
+        .checked_next_multiple_of(WORD_SIZE)
+        .ok_or(VMError::Internal(
+            InternalError::ArithmeticOperationOverflow,
+        ))?;
+        let current_memory_size = current_call_frame.memory.data.len();
+
+        if current_memory_size < new_memory_size {
+            current_call_frame
+                .memory
+                .data
+                .try_reserve(new_memory_size)
+                .map_err(|_err| VMError::MemorySizeOverflow)?;
+            current_call_frame.memory.data.resize(new_memory_size, 0);
+        }
+
+        for i in 0..size {
+            if let Some(memory_byte) =
+                current_call_frame
+                    .memory
+                    .data
+                    .get_mut(destination_offset.checked_add(i).ok_or(VMError::Internal(
                         InternalError::ArithmeticOperationOverflow,
                     ))?)
-                    .min(bytecode_len),
-            )
-        } else {
-            vec![0u8; size].into()
-        };
-
-        current_call_frame.memory.store_bytes(dest_offset, &code)?;
+            {
+                *memory_byte = *current_call_frame
+                    .bytecode
+                    .get(code_offset.checked_add(i).ok_or(VMError::Internal(
+                        InternalError::ArithmeticOperationOverflow,
+                    ))?)
+                    .unwrap_or(&0u8);
+            }
+        }
 
         Ok(OpcodeSuccess::Continue)
     }

diff --git a/crates/vm/levm/tests/edge_case_tests.rs b/crates/vm/levm/tests/edge_case_tests.rs
@@ -262,3 +262,15 @@ fn test_non_compliance_log() {
         &U256::zero()
     );
 }
+
+#[test]
+fn test_non_compliance_codecopy_memory_resize() {
+    let mut vm =
+        new_vm_with_bytecode(Bytes::copy_from_slice(&[97, 56, 57, 0x5f, 0x5f, 57, 89])).unwrap();
+    let mut current_call_frame = vm.call_frames.pop().unwrap();
+    vm.execute(&mut current_call_frame).unwrap();
+    assert_eq!(
+        current_call_frame.stack.stack.first().unwrap(),
+        &U256::from(14400)
+    );
+}