alloy_contract/
storage_slot.rs1use alloy_network::{Network, TransactionBuilder};
2use alloy_primitives::{Address, Bytes, B256, U256};
3use alloy_provider::Provider;
4use alloy_rpc_types_eth::state::{AccountOverride, StateOverridesBuilder};
5use alloy_sol_types::{sol, SolCall, SolValue};
6use alloy_transport::TransportError;
7
8#[derive(Debug, Clone)]
38pub struct StorageSlotFinder<P, N>
39where
40 N: Network,
41{
42 provider: P,
43 contract: Address,
44 calldata: Bytes,
45 expected_value: U256,
46 base_request: N::TransactionRequest,
47}
48
49impl<P, N> StorageSlotFinder<P, N>
50where
51 P: Provider<N>,
52 N: Network,
53{
54 pub fn new(provider: P, contract: Address, calldata: Bytes, expected_value: U256) -> Self {
65 Self {
66 provider,
67 contract,
68 calldata,
69 expected_value,
70 base_request: N::TransactionRequest::default(),
71 }
72 }
73
74 pub fn balance_of(provider: P, token_address: Address, user: Address) -> Self {
86 sol! {
87 contract IERC20 {
88 function balanceOf(address target) external view returns (uint256);
89 }
90 }
91 let calldata = IERC20::balanceOfCall { target: user }.abi_encode().into();
92 Self::new(provider, token_address, calldata, U256::from(1337))
93 }
94
95 pub const fn with_expected_value(mut self, value: U256) -> Self {
97 self.expected_value = value;
98 self
99 }
100
101 pub fn with_request(mut self, base_request: N::TransactionRequest) -> Self {
106 self.base_request = base_request;
107 self
108 }
109
110 pub async fn find_slot(self) -> Result<Option<B256>, TransportError> {
131 let Self { provider, contract, calldata, expected_value, base_request } = self;
132
133 let tx = base_request.with_to(contract).with_input(calldata);
134
135 let access_list_result = provider.create_access_list(&tx).await?;
137 let access_list = access_list_result.access_list;
138 for item in access_list.0 {
141 if item.address != contract {
142 continue;
143 };
144 for slot in &item.storage_keys {
145 let account_override = AccountOverride::default().with_state_diff(std::iter::once(
146 (*slot, B256::from(expected_value.to_be_bytes())),
147 ));
148
149 let state_override =
150 StateOverridesBuilder::default().append(contract, account_override).build();
151
152 let Ok(result) = provider.call(tx.clone()).overrides(state_override).await else {
153 continue;
155 };
156
157 let Ok(result_value) = U256::abi_decode(&result) else {
158 continue;
160 };
161
162 if result_value == expected_value {
163 return Ok(Some(*slot));
164 }
165 }
166 }
167 Ok(None)
168 }
169}
170
171#[cfg(test)]
172mod tests {
173 use crate::StorageSlotFinder;
174 use alloy_network::TransactionBuilder;
175 use alloy_primitives::{address, ruint::uint, Address, B256, U256};
176 use alloy_provider::{ext::AnvilApi, Provider, ProviderBuilder};
177 use alloy_rpc_types_eth::TransactionRequest;
178 use alloy_sol_types::sol;
179 const FORK_URL: &str = "https://reth-ethereum.ithaca.xyz/rpc";
180 use alloy_sol_types::SolCall;
181
182 async fn test_erc20_token_set_balance(token: Address) {
183 let provider = ProviderBuilder::new().connect_anvil_with_config(|a| a.fork(FORK_URL));
184 let user = address!("0xd8dA6BF26964aF9D7eEd9e03E53415D37aA96045");
185 let amount = U256::from(500u64);
186 let finder = StorageSlotFinder::balance_of(provider.clone(), token, user);
187 let storage_slot = U256::from_be_bytes(finder.find_slot().await.unwrap().unwrap().0);
188
189 provider
190 .anvil_set_storage_at(token, storage_slot, B256::from(amount.to_be_bytes()))
191 .await
192 .unwrap();
193
194 sol! {
195 function balanceOf(address owner) view returns (uint256);
196 }
197
198 let balance_of_call = balanceOfCall::new((user,));
199 let input = balanceOfCall::abi_encode(&balance_of_call);
200
201 let result = provider
202 .call(TransactionRequest::default().with_to(token).with_input(input))
203 .await
204 .unwrap();
205 let balance = balanceOfCall::abi_decode_returns(&result).unwrap();
206
207 assert_eq!(balance, amount);
208 }
209
210 #[tokio::test]
211 async fn test_erc20_dai_set_balance() {
212 let dai = address!("0x6B175474E89094C44Da98b954EedeAC495271d0F");
213 test_erc20_token_set_balance(dai).await
214 }
215
216 #[tokio::test]
217 async fn test_erc20_usdc_set_balance() {
218 let usdc = address!("0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48");
219 test_erc20_token_set_balance(usdc).await
220 }
221
222 #[tokio::test]
223 async fn test_erc20_tether_set_balance() {
224 let tether = address!("0xdAC17F958D2ee523a2206206994597C13D831ec7");
225 test_erc20_token_set_balance(tether).await
226 }
227 #[tokio::test]
228 async fn test_erc20_token_polygon() {
229 let provider =
230 ProviderBuilder::new().connect_http("https://polygon-rpc.com".parse().unwrap());
231 let usdt = address!("0xc2132D05D31c914a87C6611C10748AEb04B58e8F"); let user = address!("0x0aD71c9106455801eAe0e11D5A1Dd5232537E662");
233 let finder = StorageSlotFinder::balance_of(provider.clone(), usdt, user)
234 .with_request(TransactionRequest::default().gas_limit(100000));
235 let storage_slot = U256::from_be_bytes(finder.find_slot().await.unwrap().unwrap().0);
236 assert_eq!(
237 storage_slot,
238 uint!(
239 38414845661641411266428303013962925072609060211040678298987263275302781786590_U256
240 )
241 );
242 }
243}