fix: add JsonStorageKey for eth_getProof and eth_getStorageAt (#2067)

crates/primitives/src/serde_helper/mod.rs

@@ -1,5 +1,8 @@
 //! Various serde utilities
 
+mod storage_key;
+pub use storage_key::*;
+
 mod jsonu256;
 pub use jsonu256::*;
 

crates/primitives/src/serde_helper/storage_key.rs

@@ -0,0 +1,38 @@
+use crate::{H256, U256};
+use serde::{Deserialize, Serialize};
+use std::fmt::Write;
+
+/// A storage key type that can be serialized to and from a hex string up to 40 characters. Used
+/// for `eth_getStorageAt` and `eth_getProof` RPCs.
+#[derive(Clone, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
+#[serde(from = "U256", into = "String")]
+pub struct JsonStorageKey(pub H256);
+
+impl From<U256> for JsonStorageKey {
+    fn from(value: U256) -> Self {
+        // SAFETY: Address (H256) and U256 have the same number of bytes
+        JsonStorageKey(H256::from(value.to_be_bytes()))
+    }
+}
+
+impl From<JsonStorageKey> for String {
+    fn from(value: JsonStorageKey) -> Self {
+        // SAFETY: Address (H256) and U256 have the same number of bytes
+        let uint = U256::from_be_bytes(value.0 .0);
+
+        // serialize byte by byte
+        //
+        // this is mainly so we can return an output that hive testing expects, because the
+        // `eth_getProof` implementation in geth simply mirrors the input
+        //
+        // see the use of `hexKey` in the `eth_getProof` response:
+        // <https://github.com/ethereum/go-ethereum/blob/00a73fbcce3250b87fc4160f3deddc44390848f4/internal/ethapi/api.go#L658-L690>
+        let bytes = uint.to_be_bytes_trimmed_vec();
+        let mut hex = String::with_capacity(2 + bytes.len() * 2);
+        hex.push_str("0x");
+        for byte in bytes {
+            write!(hex, "{:02x}", byte).unwrap();
+        }
+        hex
+    }
+}

crates/rpc/rpc-api/src/eth.rs

@@ -1,6 +1,7 @@
 use jsonrpsee::{core::RpcResult as Result, proc_macros::rpc};
 use reth_primitives::{
-    AccessListWithGasUsed, Address, BlockId, BlockNumberOrTag, Bytes, H256, H64, U256, U64,
+    serde_helper::JsonStorageKey, AccessListWithGasUsed, Address, BlockId, BlockNumberOrTag, Bytes,
+    H256, H64, U256, U64,
 };
 use reth_rpc_types::{
     state::StateOverride, CallRequest, EIP1186AccountProofResponse, FeeHistory, Index, RichBlock,
@@ -116,7 +117,7 @@ pub trait EthApi {
     async fn storage_at(
         &self,
         address: Address,
-        index: U256,
+        index: JsonStorageKey,
         block_number: Option<BlockId>,
     ) -> Result<H256>;
 
@@ -244,7 +245,7 @@ pub trait EthApi {
     async fn get_proof(
         &self,
         address: Address,
-        keys: Vec<H256>,
+        keys: Vec<JsonStorageKey>,
         block_number: Option<BlockId>,
     ) -> Result<EIP1186AccountProofResponse>;
 }

crates/rpc/rpc-builder/tests/it/http.rs

@@ -68,7 +68,7 @@ where
     EthApiClient::fee_history(client, 0.into(), block_number.into(), None).await.unwrap();
     EthApiClient::balance(client, address, None).await.unwrap();
     EthApiClient::transaction_count(client, address, None).await.unwrap();
-    EthApiClient::storage_at(client, address, U256::default(), None).await.unwrap();
+    EthApiClient::storage_at(client, address, U256::default().into(), None).await.unwrap();
     EthApiClient::block_by_hash(client, hash, false).await.unwrap();
     EthApiClient::block_by_number(client, block_number, false).await.unwrap();
     EthApiClient::block_transaction_count_by_number(client, block_number).await.unwrap();

crates/rpc/rpc-types/src/eth/account.rs

@@ -1,5 +1,5 @@
 #![allow(missing_docs)]
-use reth_primitives::{Address, Bytes, H256, H512, U256, U64};
+use reth_primitives::{serde_helper::JsonStorageKey, Address, Bytes, H256, H512, U256, U64};
 use serde::{Deserialize, Serialize};
 
 /// Account information.
@@ -14,7 +14,7 @@ pub struct AccountInfo {
 #[serde(rename_all = "camelCase")]
 pub struct StorageProof {
     /// Storage key.
-    pub key: U256,
+    pub key: JsonStorageKey,
     /// Value that the key holds
     pub value: U256,
     /// proof for the pair

crates/rpc/rpc/src/eth/api/server.rs

@@ -11,7 +11,8 @@ use crate::{
 };
 use jsonrpsee::core::RpcResult as Result;
 use reth_primitives::{
-    AccessListWithGasUsed, Address, BlockId, BlockNumberOrTag, Bytes, Header, H256, H64, U256, U64,
+    serde_helper::JsonStorageKey, AccessListWithGasUsed, Address, BlockId, BlockNumberOrTag, Bytes,
+    Header, H256, H64, U256, U64,
 };
 use reth_provider::{BlockProvider, EvmEnvProvider, HeaderProvider, StateProviderFactory};
 use reth_rpc_api::EthApiServer;
@@ -176,7 +177,7 @@ where
     async fn storage_at(
         &self,
         address: Address,
-        index: U256,
+        index: JsonStorageKey,
         block_number: Option<BlockId>,
     ) -> Result<H256> {
         trace!(target: "rpc::eth", ?address, ?block_number, "Serving eth_getStorageAt");
@@ -421,7 +422,7 @@ where
     async fn get_proof(
         &self,
         address: Address,
-        keys: Vec<H256>,
+        keys: Vec<JsonStorageKey>,
         block_number: Option<BlockId>,
     ) -> Result<EIP1186AccountProofResponse> {
         trace!(target: "rpc::eth", ?address, ?keys, ?block_number, "Serving eth_getProof");

crates/rpc/rpc/src/eth/api/state.rs

@@ -4,7 +4,10 @@ use crate::{
     eth::error::{EthApiError, EthResult},
     EthApi,
 };
-use reth_primitives::{Address, BlockId, BlockNumberOrTag, Bytes, H256, KECCAK_EMPTY, U256};
+use reth_primitives::{
+    serde_helper::JsonStorageKey, Address, BlockId, BlockNumberOrTag, Bytes, H256, KECCAK_EMPTY,
+    U256,
+};
 use reth_provider::{
     AccountProvider, BlockProvider, EvmEnvProvider, StateProvider, StateProviderFactory,
 };
@@ -39,19 +42,18 @@ where
     pub(crate) fn storage_at(
         &self,
         address: Address,
-        index: U256,
+        index: JsonStorageKey,
         block_id: Option<BlockId>,
     ) -> EthResult<H256> {
         let state = self.state_at_block_id_or_latest(block_id)?;
-        let storage_key = H256(index.to_be_bytes());
-        let value = state.storage(address, storage_key)?.unwrap_or_default();
+        let value = state.storage(address, index.0)?.unwrap_or_default();
         Ok(H256(value.to_be_bytes()))
     }
 
     pub(crate) fn get_proof(
         &self,
         address: Address,
-        keys: Vec<H256>,
+        keys: Vec<JsonStorageKey>,
         block_id: Option<BlockId>,
     ) -> EthResult<EIP1186AccountProofResponse> {
         let chain_info = self.client().chain_info()?;
@@ -74,14 +76,15 @@ where
 
         let state = self.state_at_block_id(block_id)?;
 
-        let (account_proof, storage_hash, stg_proofs) = state.proof(address, &keys)?;
+        let hash_keys = keys.iter().map(|key| key.0).collect::<Vec<_>>();
+        let (account_proof, storage_hash, stg_proofs) = state.proof(address, &hash_keys)?;
 
         let storage_proof = keys
             .into_iter()
             .zip(stg_proofs)
             .map(|(key, proof)| {
-                state.storage(address, key).map(|op| StorageProof {
-                    key: U256::from_be_bytes(*key.as_fixed_bytes()),
+                state.storage(address, key.0).map(|op| StorageProof {
+                    key,
                     value: op.unwrap_or_default(),
                     proof,
                 })
@@ -97,7 +100,7 @@ where
             ..Default::default()
         };
 
-        if let Some(account) = state.basic_account(address)? {
+        if let Some(account) = state.basic_account(proof.address)? {
             proof.balance = account.balance;
             proof.nonce = account.nonce.into();
             proof.code_hash = account.get_bytecode_hash();