Handle storage read/write syscalls in cheatable runtime

Towards #3331

commit-id:984675b6

Author: ddoktorski

crates/cheatnet/src/runtime_extensions/call_to_blockifier_runtime_extension/execution/cheated_syscalls.rs

@@ -8,7 +8,8 @@ use blockifier::execution::syscalls::hint_processor::{
 use blockifier::execution::syscalls::syscall_base::SyscallResult;
 use blockifier::execution::syscalls::vm_syscall_utils::{
     CallContractRequest, DeployRequest, DeployResponse, EmptyRequest, GetBlockHashRequest,
-    GetBlockHashResponse, GetExecutionInfoResponse, LibraryCallRequest, SyscallResponse,
+    GetBlockHashResponse, GetExecutionInfoResponse, LibraryCallRequest, StorageReadRequest,
+    StorageReadResponse, StorageWriteRequest, StorageWriteResponse, SyscallResponse,
     SyscallSelector, WriteResponseResult,
 };
 use blockifier::execution::{call_info::CallInfo, entry_point::ConstructorContext};
@@ -28,6 +29,8 @@ use blockifier::{
 };
 use cairo_vm::types::relocatable::Relocatable;
 use cairo_vm::vm::vm_core::VirtualMachine;
+use conversions::string::TryFromHexStr;
+use runtime::starknet::constants::TEST_ADDRESS;
 use starknet_api::core::calculate_contract_address;
 use starknet_api::{
     contract_class::EntryPointType,
@@ -240,6 +243,76 @@ pub fn get_block_hash_syscall(
     Ok(GetBlockHashResponse { block_hash })
 }
 
+#[allow(clippy::needless_pass_by_value)]
+pub fn storage_read(
+    request: StorageReadRequest,
+    _vm: &mut VirtualMachine,
+    syscall_handler: &mut SyscallHintProcessor<'_>,
+    cheatnet_state: &mut CheatnetState,
+    _remaining_gas: &mut u64,
+) -> SyscallResult<StorageReadResponse> {
+    let original_storage_address = syscall_handler.base.call.storage_address;
+    maybe_modify_storage_address(syscall_handler, cheatnet_state);
+
+    let value = syscall_handler
+        .base
+        .storage_read(request.address)
+        .inspect_err(|_| {
+            // Restore state on error before bubbling up
+            syscall_handler.base.call.storage_address = original_storage_address;
+        })?;
+
+    // Restore the original storage_address
+    syscall_handler.base.call.storage_address = original_storage_address;
+
+    Ok(StorageReadResponse { value })
+}
+
+#[allow(clippy::needless_pass_by_value)]
+pub fn storage_write(
+    request: StorageWriteRequest,
+    _vm: &mut VirtualMachine,
+    syscall_handler: &mut SyscallHintProcessor<'_>,
+    cheatnet_state: &mut CheatnetState,
+    _remaining_gas: &mut u64,
+) -> SyscallResult<StorageWriteResponse> {
+    let original_storage_address = syscall_handler.base.call.storage_address;
+    maybe_modify_storage_address(syscall_handler, cheatnet_state);
+
+    syscall_handler
+        .base
+        .storage_write(request.address, request.value)
+        .inspect_err(|_| {
+            // Restore state on error before bubbling up
+            syscall_handler.base.call.storage_address = original_storage_address;
+        })?;
+
+    // Restore the original storage_address
+    syscall_handler.base.call.storage_address = original_storage_address;
+
+    Ok(StorageWriteResponse {})
+}
+
+// This logic is used to modify the storage address to enable using `contract_state_for_testing`
+// inside `interact_with_state` closure cheatcode.
+fn maybe_modify_storage_address(
+    syscall_handler: &mut SyscallHintProcessor<'_>,
+    cheatnet_state: &mut CheatnetState,
+) {
+    let contract_address = syscall_handler.storage_address();
+
+    if contract_address
+        != TryFromHexStr::try_from_hex_str(TEST_ADDRESS).expect("Failed to parse TEST_ADDRESS")
+    {
+        return;
+    }
+
+    let cheated_data = cheatnet_state.get_cheated_data(contract_address);
+    if let Some(actual_address) = cheated_data.contract_address {
+        syscall_handler.base.call.storage_address = actual_address;
+    }
+}
+
 #[derive(Debug)]
 // crates/blockifier/src/execution/syscalls/mod.rs:127 (SingleSegmentResponse)
 // It is created here because fields in the original structure are private

crates/cheatnet/src/runtime_extensions/cheatable_starknet_runtime_extension.rs

@@ -87,6 +87,22 @@ impl<'a> ExtensionLogic for CheatableStarknetRuntimeExtension<'a> {
                     SyscallSelector::GetBlockHash,
                 )
                 .map(|()| SyscallHandlingResult::Handled),
+            SyscallSelector::StorageRead => self
+                .execute_syscall(
+                    syscall_handler,
+                    vm,
+                    cheated_syscalls::storage_read,
+                    SyscallSelector::StorageRead,
+                )
+                .map(|()| SyscallHandlingResult::Handled),
+            SyscallSelector::StorageWrite => self
+                .execute_syscall(
+                    syscall_handler,
+                    vm,
+                    cheated_syscalls::storage_write,
+                    SyscallSelector::StorageWrite,
+                )
+                .map(|()| SyscallHandlingResult::Handled),
             _ => Ok(SyscallHandlingResult::Forwarded),
         }
     }
@@ -132,6 +148,8 @@ fn get_syscall_cost(
         SyscallSelector::Deploy => gas_costs.syscalls.deploy,
         SyscallSelector::GetExecutionInfo => gas_costs.syscalls.get_execution_info,
         SyscallSelector::GetBlockHash => gas_costs.syscalls.get_block_hash,
+        SyscallSelector::StorageRead => gas_costs.syscalls.storage_read,
+        SyscallSelector::StorageWrite => gas_costs.syscalls.storage_write,
         _ => unreachable!("Syscall has no associated cost"),
     }
 }