Fix fork error handling (#1571)

Closes #1497 

## Introduced changes

<!-- A brief description of the changes -->

- Fixes incosistencies in fork state behavior when errors should be
raised

TODO: 
- [x] Fix gas calculation (is crashing right now in forking)

## Checklist

<!-- Make sure all of these are complete -->

- [x] Linked relevant issue
- [x] Updated relevant documentation
- [x] Added relevant tests
- [x] Performed self-review of the code
- [x] Added changes to `CHANGELOG.md`

Author: Arcticae

crates/cheatnet/src/forking/state.rs

@@ -13,6 +13,7 @@ use runtime::starknet::context::BlockInfo;
 use sierra_casm::compile;
 use starknet::core::types::{
     BlockId, ContractClass as ContractClassStarknet, FieldElement, MaybePendingBlockWithTxHashes,
+    StarknetError,
 };
 use starknet::providers::jsonrpc::HttpTransport;
 use starknet::providers::{JsonRpcClient, Provider, ProviderError};
@@ -53,6 +54,19 @@ impl ForkStateReader {
     }
 }
 
+#[macro_export]
+macro_rules! other_provider_error {
+    ( $boxed:expr ) => {{
+        let err_str = $boxed.deref().to_string();
+        if err_str.contains("error sending request for url") {
+            return Err(StateReadError(
+                "Unable to reach the node. Check your internet connection and node url".to_string(),
+            ));
+        }
+        Err(StateReadError(format!("JsonRpc provider error: {err_str}")))
+    }};
+}
+
 impl BlockInfoReader for ForkStateReader {
     fn get_block_info(&mut self) -> StateResult<BlockInfo> {
         if let Some(cache_hit) = self.cache.get_block_info() {
@@ -77,30 +91,14 @@ impl BlockInfoReader for ForkStateReader {
             Ok(MaybePendingBlockWithTxHashes::PendingBlock(_)) => {
                 unreachable!("Pending block is not be allowed at the configuration level")
             }
+            Err(ProviderError::Other(boxed)) => other_provider_error!(boxed),
             Err(err) => Err(StateReadError(format!(
-                "Unable to get block with tx hashes from fork, err: {err:?}"
+                "Unable to get block with tx hashes from fork ({err})"
             ))),
         }
     }
 }
 
-#[macro_export]
-macro_rules! other_provider_error {
-    ( $boxed:expr ) => {{
-        let err_str = $boxed.deref().to_string();
-        if err_str.contains("error sending request for url") {
-            return node_connection_error();
-        }
-        Err(StateReadError(format!("JsonRpc provider error: {err_str}")))
-    }};
-}
-
-fn node_connection_error<T>() -> StateResult<T> {
-    Err(StateReadError(
-        "Unable to reach the node. Check your internet connection and node url".to_string(),
-    ))
-}
-
 impl StateReader for ForkStateReader {
     fn get_storage_at(
         &mut self,
@@ -123,8 +121,9 @@ impl StateReader for ForkStateReader {
                 Ok(value_sf)
             }
             Err(ProviderError::Other(boxed)) => other_provider_error!(boxed),
-            Err(_) => Err(StateReadError(format!(
-                "Unable to get storage at address: {contract_address:?} and key: {key:?} from fork"
+            Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => Ok(Default::default()),
+            Err(x) => Err(StateReadError(format!(
+                "Unable to get storage at address: {contract_address:?} and key: {key:?} from fork ({x})"
             ))),
         }
     }
@@ -144,8 +143,11 @@ impl StateReader for ForkStateReader {
                 Ok(nonce)
             }
             Err(ProviderError::Other(boxed)) => other_provider_error!(boxed),
-            Err(_) => Err(StateReadError(format!(
-                "Unable to get nonce at {contract_address:?} from fork"
+            Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => {
+                Ok(Default::default())
+            }
+            Err(x) => Err(StateReadError(format!(
+                "Unable to get nonce at {contract_address:?} from fork ({x})"
             ))),
         }
     }
@@ -165,9 +167,12 @@ impl StateReader for ForkStateReader {
                     .cache_get_class_hash_at(contract_address, class_hash);
                 Ok(class_hash)
             }
+            Err(ProviderError::StarknetError(StarknetError::ContractNotFound)) => {
+                Ok(Default::default())
+            }
             Err(ProviderError::Other(boxed)) => other_provider_error!(boxed),
-            Err(_) => Err(StateReadError(format!(
-                "Unable to get class hash at {contract_address:?} from fork"
+            Err(x) => Err(StateReadError(format!(
+                "Unable to get class hash at {contract_address:?} from fork ({x})"
             ))),
         }
     }
@@ -190,8 +195,13 @@ impl StateReader for ForkStateReader {
 
                         Ok(contract_class)
                     }
+                    Err(ProviderError::StarknetError(StarknetError::ClassHashNotFound)) => {
+                        Err(UndeclaredClassHash(*class_hash))
+                    }
                     Err(ProviderError::Other(boxed)) => other_provider_error!(boxed),
-                    Err(_) => Err(UndeclaredClassHash(*class_hash)),
+                    Err(x) => Err(StateReadError(format!(
+                        "Unable to get compiled class at {class_hash} from fork ({x})"
+                    ))),
                 }
             };
 

crates/cheatnet/src/state.rs

@@ -8,10 +8,7 @@ use blockifier::execution::entry_point::CallEntryPoint;
 use blockifier::state::state_api::State;
 use blockifier::{
     execution::contract_class::ContractClass,
-    state::{
-        errors::StateError,
-        state_api::{StateReader, StateResult},
-    },
+    state::state_api::{StateReader, StateResult},
 };
 use cairo_felt::Felt252;
 use runtime::starknet::context::BlockInfo;
@@ -20,6 +17,7 @@ use runtime::starknet::state::DictStateReader;
 use starknet_api::core::EntryPointSelector;
 
 use crate::constants::build_test_entry_point;
+use blockifier::state::errors::StateError::UndeclaredClassHash;
 use starknet_api::transaction::ContractAddressSalt;
 use starknet_api::{
     core::{ClassHash, CompiledClassHash, ContractAddress, Nonce},
@@ -82,8 +80,8 @@ impl StateReader for ExtendedStateReader {
             .or_else(|_| {
                 self.fork_state_reader
                     .as_mut()
-                    .map_or(Ok(StarkFelt::default()), |reader| {
-                        match_node_response(reader.get_storage_at(contract_address, key))
+                    .map_or(Ok(Default::default()), {
+                        |reader| reader.get_storage_at(contract_address, key)
                     })
             })
     }
@@ -94,8 +92,8 @@ impl StateReader for ExtendedStateReader {
             .or_else(|_| {
                 self.fork_state_reader
                     .as_mut()
-                    .map_or(Ok(Nonce::default()), |reader| {
-                        match_node_response(reader.get_nonce_at(contract_address))
+                    .map_or(Ok(Default::default()), {
+                        |reader| reader.get_nonce_at(contract_address)
                     })
             })
     }
@@ -106,8 +104,8 @@ impl StateReader for ExtendedStateReader {
             .or_else(|_| {
                 self.fork_state_reader
                     .as_mut()
-                    .map_or(Ok(ClassHash::default()), |reader| {
-                        match_node_response(reader.get_class_hash_at(contract_address))
+                    .map_or(Ok(Default::default()), {
+                        |reader| reader.get_class_hash_at(contract_address)
                     })
             })
     }
@@ -119,10 +117,11 @@ impl StateReader for ExtendedStateReader {
         self.dict_state_reader
             .get_compiled_contract_class(class_hash)
             .or_else(|_| {
-                self.fork_state_reader.as_mut().map_or(
-                    Err(StateError::UndeclaredClassHash(*class_hash)),
-                    |reader| reader.get_compiled_contract_class(class_hash),
-                )
+                self.fork_state_reader
+                    .as_mut()
+                    .map_or(Err(UndeclaredClassHash(*class_hash)), |reader| {
+                        reader.get_compiled_contract_class(class_hash)
+                    })
             })
     }
 
@@ -317,16 +316,6 @@ impl TraceData {
     }
 }
 
-fn match_node_response<T: Default>(result: StateResult<T>) -> StateResult<T> {
-    match result {
-        Ok(class_hash) => Ok(class_hash),
-        Err(StateError::StateReadError(msg)) if msg.contains("node") => {
-            Err(StateError::StateReadError(msg))
-        }
-        _ => Ok(Default::default()),
-    }
-}
-
 fn get_cheat_for_contract<T: Clone>(
     global_cheat: &Option<T>,
     contract_cheats: &HashMap<ContractAddress, CheatStatus<T>>,

crates/forge-runner/src/gas.rs

@@ -6,6 +6,7 @@ use blockifier::fee::gas_usage::calculate_tx_gas_usage;
 use blockifier::fee::os_resources::OS_RESOURCES;
 use blockifier::fee::os_usage::get_additional_os_resources;
 use blockifier::state::cached_state::{CachedState, StateChangesCount};
+use blockifier::state::errors::StateError;
 use blockifier::transaction::transaction_types::TransactionType;
 use blockifier::{
     abi::constants, block_context::BlockContext, execution::entry_point::ExecutionResources,
@@ -15,19 +16,16 @@ use cairo_vm::vm::runners::cairo_runner::ExecutionResources as VmExecutionResour
 use cheatnet::runtime_extensions::call_to_blockifier_runtime_extension::rpc::UsedResources;
 use cheatnet::state::ExtendedStateReader;
 
-#[must_use]
 pub fn calculate_used_gas(
     block_context: &BlockContext,
     state: &mut CachedState<ExtendedStateReader>,
     resources: &UsedResources,
-) -> u128 {
+) -> Result<u128, StateError> {
     let total_vm_usage = get_total_vm_usage(&resources.execution_resources);
-    let mut state_changes = state
-        .get_actual_state_changes_for_fee_charge(
-            block_context.fee_token_addresses.eth_fee_token_address,
-            None,
-        )
-        .unwrap();
+    let mut state_changes = state.get_actual_state_changes_for_fee_charge(
+        block_context.fee_token_addresses.eth_fee_token_address,
+        None,
+    )?;
     // compiled_class_hash_updates is used only for keeping track of declares
     // which we don't want to include in gas cost
     state_changes.compiled_class_hash_updates.clear();
@@ -40,8 +38,8 @@ pub fn calculate_used_gas(
 
     let resource_mapping = used_resources_to_resource_mapping(&total_vm_usage, l1_gas_usage);
 
-    calculate_tx_l1_gas_usage(&resource_mapping, block_context)
-        .expect("Calculating gas failed, some resources were not included.")
+    Ok(calculate_tx_l1_gas_usage(&resource_mapping, block_context)
+        .expect("Calculating gas failed, some resources were not included."))
 }
 
 #[must_use]

crates/forge-runner/src/running.rs

@@ -268,7 +268,7 @@ pub fn run_test_case(
     let call_trace_ref = get_call_trace_ref(&mut forge_runtime);
     let execution_resources = get_all_execution_resources(forge_runtime);
 
-    let gas = calculate_used_gas(&block_context, &mut blockifier_state, &execution_resources);
+    let gas = calculate_used_gas(&block_context, &mut blockifier_state, &execution_resources)?;
 
     Ok(RunResultWithInfo {
         run_result,
@@ -305,7 +305,7 @@ fn extract_test_case_summary(
                 Err(err) => bail!(err),
             }
         }
-        // `ForkStateReader.get_block_info`, `get_fork_state_reader` may return an error
+        // `ForkStateReader.get_block_info`, `get_fork_state_reader, `calculate_used_gas` may return an error
         // `available_gas` may be specified with Scarb ~2.4
         Err(error) => Ok(TestCaseSummary::Failed {
             name: case.name.clone(),

crates/forge/src/block_number_map.rs

@@ -4,7 +4,6 @@ use cairo_felt::Felt252;
 use conversions::IntoConv;
 use forge_runner::compiled_runnable::ValidatedForkConfig;
 use num_bigint::BigInt;
-use starknet::core::types::BlockTag::Latest;
 use starknet::core::types::{BlockId, MaybePendingBlockWithTxHashes};
 use starknet::providers::jsonrpc::HttpTransport;
 use starknet::providers::{JsonRpcClient, Provider};
@@ -83,13 +82,11 @@ impl BlockNumberMap {
 async fn get_latest_block_number(url: &Url) -> Result<BlockNumber> {
     let client = JsonRpcClient::new(HttpTransport::new(url.clone()));
 
-    match Handle::current()
-        .spawn(async move { client.get_block_with_tx_hashes(BlockId::Tag(Latest)).await })
+    Handle::current()
+        .spawn(async move { client.block_number().await })
         .await?
-    {
-        Ok(MaybePendingBlockWithTxHashes::Block(block)) => Ok(BlockNumber(block.block_number)),
-        _ => Err(anyhow!("Could not get the latest block number".to_string())),
-    }
+        .map(BlockNumber)
+        .map_err(|x| anyhow!(x.to_string()))
 }
 
 async fn get_block_number_from_hash(url: &Url, block_hash: &Felt252) -> Result<BlockNumber> {