From c5e3b2e719961b14509fbf38cf8e686015fc6531 Mon Sep 17 00:00:00 2001
From: Elias Tazartes <66871571+Eikix@users.noreply.github.com>
Date: Fri, 23 Feb 2024 11:18:41 +0100
Subject: [PATCH] fix storage at type (#784)

* fix storage at type

* fix b256

* Update tests/eth_provider.rs

Co-authored-by: greged93 <82421016+greged93@users.noreply.github.com>

* Update src/eth_provider/provider.rs

Co-authored-by: greged93 <82421016+greged93@users.noreply.github.com>

* fix type

---------

Co-authored-by: greged93 <82421016+greged93@users.noreply.github.com>
---
 src/eth_provider/provider.rs   | 8 +++++---
 src/eth_rpc/api/eth_api.rs     | 2 +-
 src/eth_rpc/servers/eth_rpc.rs | 2 +-
 tests/eth_provider.rs          | 6 +++---
 4 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/src/eth_provider/provider.rs b/src/eth_provider/provider.rs
index ff39fd323..24da8c77a 100644
--- a/src/eth_provider/provider.rs
+++ b/src/eth_provider/provider.rs
@@ -99,7 +99,7 @@ pub trait EthereumProvider {
     /// Returns the balance of an address in native eth.
     async fn balance(&self, address: Address, block_id: Option<BlockId>) -> EthProviderResult<U256>;
     /// Returns the storage of an address at a certain index.
-    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> EthProviderResult<U256>;
+    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> EthProviderResult<B256>;
     /// Returns the nonce for the address at the given block.
     async fn transaction_count(&self, address: Address, block_id: Option<BlockId>) -> EthProviderResult<U256>;
     /// Returns the code for the address at the given block.
@@ -272,7 +272,7 @@ where
         Ok(low + (high << 128))
     }
 
-    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> EthProviderResult<U256> {
+    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> EthProviderResult<B256> {
         let eth_block_id = EthBlockId::new(block_id.unwrap_or(BlockId::Number(BlockNumberOrTag::Latest)));
         let starknet_block_id: StarknetBlockId = eth_block_id.try_into()?;
 
@@ -286,7 +286,9 @@ where
 
         let low: U256 = into_via_wrapper!(storage.low);
         let high: U256 = into_via_wrapper!(storage.high);
-        Ok(low + (high << 128))
+        let storage: U256 = low + (high << 128);
+
+        Ok(storage.into())
     }
 
     async fn transaction_count(&self, address: Address, block_id: Option<BlockId>) -> EthProviderResult<U256> {
diff --git a/src/eth_rpc/api/eth_api.rs b/src/eth_rpc/api/eth_api.rs
index 6c3e72f0e..14610a59f 100644
--- a/src/eth_rpc/api/eth_api.rs
+++ b/src/eth_rpc/api/eth_api.rs
@@ -90,7 +90,7 @@ pub trait EthApi {
 
     /// Returns the value from a storage position at a given address
     #[method(name = "getStorageAt")]
-    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> Result<U256>;
+    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> Result<B256>;
 
     /// Returns the number of transactions sent from an address at given block number.
     #[method(name = "getTransactionCount")]
diff --git a/src/eth_rpc/servers/eth_rpc.rs b/src/eth_rpc/servers/eth_rpc.rs
index 3f47cb2c4..012b1b70a 100644
--- a/src/eth_rpc/servers/eth_rpc.rs
+++ b/src/eth_rpc/servers/eth_rpc.rs
@@ -135,7 +135,7 @@ where
     }
 
     #[tracing::instrument(skip_all, ret, err, fields(address = %address, index = ?index, block_id = ?block_id))]
-    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> Result<U256> {
+    async fn storage_at(&self, address: Address, index: U256, block_id: Option<BlockId>) -> Result<B256> {
         Ok(self.eth_provider.storage_at(address, index, block_id).await?)
     }
 
diff --git a/tests/eth_provider.rs b/tests/eth_provider.rs
index e1931485a..3d9f2c68c 100644
--- a/tests/eth_provider.rs
+++ b/tests/eth_provider.rs
@@ -1,6 +1,6 @@
 #![cfg(feature = "testing")]
 use std::cmp::min;
-use std::str::FromStr as _;
+use std::str::FromStr;
 
 use kakarot_rpc::eth_provider::provider::EthereumProvider;
 use kakarot_rpc::models::felt::Felt252Wrapper;
@@ -13,7 +13,7 @@ use reth_rpc_types::request::TransactionInput;
 use reth_rpc_types::TransactionRequest;
 use rstest::*;
 
-use reth_primitives::{Address, BlockNumberOrTag, Bytes, U256, U64};
+use reth_primitives::{Address, BlockNumberOrTag, Bytes, B256, U256, U64};
 
 #[rstest]
 #[awt]
@@ -134,7 +134,7 @@ async fn test_storage_at(#[future] counter: (Katana, KakarotEvmContract), _setup
 
     // Then
     let count = eth_provider.storage_at(counter_address, U256::from(0), None).await.unwrap();
-    assert_eq!(U256::from(1), count);
+    assert_eq!(B256::left_padding_from(&[0x1]), count);
 }
 
 #[rstest]