ethers_erc20balances/
lib.rs

1use async_trait::async_trait;
2use auto_impl::auto_impl;
3use ethers::prelude::{
4    abigen, Address, JsonRpcClient, Lazy, Middleware, Provider, ProviderError, Signer,
5    SignerMiddleware, U256,
6};
7use ethers::types::BlockId;
8use std::collections::HashMap;
9use std::str::FromStr;
10use std::sync::Arc;
11
12pub static CONTRACTS: Lazy<HashMap<U256, Address>> = Lazy::new(|| {
13    HashMap::from([
14        (
15            1.into(),
16            Address::from_str("0xb1f8e55c7f64d203c1400b9d8555d050f94adf39")
17                .expect("Failed to parse address"),
18        ), // mainnet
19        (
20            3.into(),
21            Address::from_str("0x8D9708f3F514206486D7E988533f770a16d074a7")
22                .expect("Failed to parse address"),
23        ), // ropsten
24        (
25            4.into(),
26            Address::from_str("0x3183B673f4816C94BeF53958BaF93C671B7F8Cf2")
27                .expect("Failed to parse address"),
28        ), // rinkeby
29        (
30            69.into(),
31            Address::from_str("0x55ABBa8d669D60A10c104CC493ec5ef389EC92bb")
32                .expect("Failed to parse address"),
33        ), // kovan
34        (
35            5.into(),
36            Address::from_str("0x9788C4E93f9002a7ad8e72633b11E8d1ecd51f9b")
37                .expect("Failed to parse address"),
38        ), // goerli
39        (
40            56.into(),
41            Address::from_str("0x2352c63A83f9Fd126af8676146721Fa00924d7e4")
42                .expect("Failed to parse address"),
43        ), // binance smart chain mainnet
44        (
45            97.into(),
46            Address::from_str("0x2352c63A83f9Fd126af8676146721Fa00924d7e4")
47                .expect("Failed to parse address"),
48        ), // binance smart chain testnet
49        (
50            137.into(),
51            Address::from_str("0x2352c63A83f9Fd126af8676146721Fa00924d7e4")
52                .expect("Failed to parse address"),
53        ), // polygon
54        (
55            80001.into(),
56            Address::from_str("0x2352c63A83f9Fd126af8676146721Fa00924d7e4")
57                .expect("Failed to parse address"),
58        ), // mumbai
59        (
60            10.into(),
61            Address::from_str("0xB1c568e9C3E6bdaf755A60c7418C269eb11524FC")
62                .expect("Failed to parse address"),
63        ), // optimism
64        (
65            69.into(),
66            Address::from_str("0xB1c568e9C3E6bdaf755A60c7418C269eb11524FC")
67                .expect("Failed to parse address"),
68        ), // optimism kovan
69        (
70            42161.into(),
71            Address::from_str("0x151E24A486D7258dd7C33Fb67E4bB01919B7B32c")
72                .expect("Failed to parse address"),
73        ), // arbitrum
74        (
75            43114.into(),
76            Address::from_str("0xD023D153a0DFa485130ECFdE2FAA7e612EF94818")
77                .expect("Failed to parse address"),
78        ), // avax
79        (
80            250.into(),
81            Address::from_str("0x07f697424ABe762bB808c109860c04eA488ff92B")
82                .expect("Failed to parse address"),
83        ), // fantom
84        (
85            25.into(),
86            Address::from_str("0x56a4420cb0ef5b0d14ce1bbe380992fa31d6a907")
87                .expect("Failed to parse address"),
88        ), // cronos
89        (
90            66.into(),
91            Address::from_str("0x25B3584f4799F788c0189dd6496b0AA02cBA4605")
92                .expect("Failed to parse address"),
93        ), // okt
94        (
95            1666600000.into(),
96            Address::from_str("0x549b6A5A3027F9B73A23Db4bb95701bAcb9b9573")
97                .expect("Failed to parse address"),
98        ), // harmony
99        (
100            17000.into(),
101            Address::from_str("0x437DF28584e878948aE2417E86e15690cCf822F4")
102                .expect("Failed to parse address"),
103        ), // holesky
104    ])
105});
106
107abigen!(BalanceChecker, "abi/BalanceChecker.abi.json",);
108
109#[async_trait]
110#[auto_impl(&, Arc, Box)]
111pub trait Erc20BalancesMiddleware {
112    type Error;
113
114    async fn get_erc20_balances(
115        &self,
116        address: Vec<Address>,
117        token_addresses: Vec<Address>,
118    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error>;
119    async fn get_erc20_balances_with_chain_id(
120        &self,
121        address: Vec<Address>,
122        token_addresses: Vec<Address>,
123        chain_id: U256,
124    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error>;
125    async fn get_erc20_balances_at_block(
126        &self,
127        address: Vec<Address>,
128        token_addresses: Vec<Address>,
129        block: BlockId,
130    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error>;
131    async fn get_erc20_balances_at_block_with_chain_id(
132        &self,
133        address: Vec<Address>,
134        token_addresses: Vec<Address>,
135        chain_id: U256,
136        block: BlockId,
137    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error>;
138}
139
140#[async_trait]
141impl<P> Erc20BalancesMiddleware for Provider<P>
142where
143    P: JsonRpcClient,
144    Self: Clone,
145{
146    type Error = ProviderError;
147
148    async fn get_erc20_balances(
149        &self,
150        address: Vec<Address>,
151        token_addresses: Vec<Address>,
152    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
153        let chain_id = self.get_chainid().await?;
154        self.get_erc20_balances_with_chain_id(address, token_addresses, chain_id)
155            .await
156    }
157
158    async fn get_erc20_balances_with_chain_id(
159        &self,
160        address: Vec<Address>,
161        token_addresses: Vec<Address>,
162        chain_id: U256,
163    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
164        let contract_address = CONTRACTS
165            .get(&chain_id)
166            .ok_or_else(|| {
167                ProviderError::CustomError(format!("Chain id {} not supported", chain_id))
168            })?
169            .clone();
170        let contract = BalanceChecker::new(contract_address, Arc::new(self.clone()));
171        let balances = contract.balances(address.clone(), token_addresses.clone()).call().await.map_err(|e|
172            ProviderError::CustomError(format!("Failed to get balances for addresses {:?} and tokens {:?} on chain id {}; {:?}", address, token_addresses, chain_id, e))
173        )?;
174        Ok(reformat(address, token_addresses, balances))
175    }
176
177    async fn get_erc20_balances_at_block(
178        &self,
179        address: Vec<Address>,
180        token_addresses: Vec<Address>,
181        block: BlockId,
182    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
183        let chain_id = self.get_chainid().await?;
184        self.get_erc20_balances_at_block_with_chain_id(address, token_addresses, chain_id, block)
185            .await
186    }
187
188    async fn get_erc20_balances_at_block_with_chain_id(
189        &self,
190        address: Vec<Address>,
191        token_addresses: Vec<Address>,
192        chain_id: U256,
193        block: BlockId,
194    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
195        let contract_address = CONTRACTS
196            .get(&chain_id)
197            .ok_or_else(|| {
198                ProviderError::CustomError(format!("Chain id {} not supported", chain_id))
199            })?
200            .clone();
201
202        let contract = BalanceChecker::new(contract_address, Arc::new(self.clone()));
203        let balances = contract.balances(address.clone(), token_addresses.clone()).block(block).call().await.map_err(|e|
204            ProviderError::CustomError(format!("Failed to get balances for addresses {:?} and tokens {:?} on chain id {}; {:?}", address, token_addresses, chain_id, e))
205        )?;
206
207        Ok(reformat(address, token_addresses, balances))
208    }
209}
210
211#[async_trait]
212impl<M, S> Erc20BalancesMiddleware for SignerMiddleware<M, S>
213where
214    M: Middleware,
215    S: Signer,
216    Self: Clone,
217{
218    type Error = ProviderError;
219
220    async fn get_erc20_balances(
221        &self,
222        address: Vec<Address>,
223        token_addresses: Vec<Address>,
224    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
225        let chain_id = self
226            .get_chainid()
227            .await
228            .map_err(|e| ProviderError::CustomError(format!("Failed to get chain id; {:?}", e)))?;
229        Self::get_erc20_balances_with_chain_id(self, address, token_addresses, chain_id).await
230    }
231
232    async fn get_erc20_balances_with_chain_id(
233        &self,
234        address: Vec<Address>,
235        token_addresses: Vec<Address>,
236        chain_id: U256,
237    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
238        let contract_address = CONTRACTS
239            .get(&chain_id)
240            .ok_or_else(|| {
241                ProviderError::CustomError(format!("Chain id {} not supported", chain_id))
242            })?
243            .clone();
244        let contract = BalanceChecker::new(contract_address, Arc::new(self.clone()));
245        let balances = contract.balances(address.clone(), token_addresses.clone()).call().await.map_err(|e|
246            ProviderError::CustomError(format!("Failed to get balances for addresses {:?} and tokens {:?} on chain id {}; {:?}", address, token_addresses, chain_id, e))
247        )?;
248        Ok(reformat(address, token_addresses, balances))
249    }
250
251    async fn get_erc20_balances_at_block(
252        &self,
253        address: Vec<Address>,
254        token_addresses: Vec<Address>,
255        block: BlockId,
256    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
257        let chain_id = self
258            .get_chainid()
259            .await
260            .map_err(|e| ProviderError::CustomError(format!("Failed to get chain id; {:?}", e)))?;
261        self.get_erc20_balances_at_block_with_chain_id(address, token_addresses, chain_id, block)
262            .await
263    }
264
265    async fn get_erc20_balances_at_block_with_chain_id(
266        &self,
267        address: Vec<Address>,
268        token_addresses: Vec<Address>,
269        chain_id: U256,
270        block: BlockId,
271    ) -> Result<HashMap<Address, HashMap<Address, U256>>, Self::Error> {
272        let contract_address = CONTRACTS
273            .get(&chain_id)
274            .ok_or_else(|| {
275                ProviderError::CustomError(format!("Chain id {} not supported", chain_id))
276            })?
277            .clone();
278
279        let contract = BalanceChecker::new(contract_address, Arc::new(self.clone()));
280        let balances = contract.balances(address.clone(), token_addresses.clone()).block(block).call().await.map_err(|e|
281            ProviderError::CustomError(format!("Failed to get balances for addresses {:?} and tokens {:?} on chain id {}; {:?}", address, token_addresses, chain_id, e))
282        )?;
283
284        Ok(reformat(address, token_addresses, balances))
285    }
286}
287
288fn reformat(
289    addresses: Vec<Address>,
290    token_addresses: Vec<Address>,
291    balances: Vec<U256>,
292) -> HashMap<Address, HashMap<Address, U256>> {
293    balances
294        .chunks(token_addresses.len())
295        .enumerate()
296        .map(|(i, balances)| {
297            let address = addresses[i];
298            let balances = token_addresses
299                .iter()
300                .zip(balances.iter())
301                .map(|(token_address, balance)| (token_address.clone(), balance.clone()))
302                .collect();
303            (address, balances)
304        })
305        .collect()
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311
312    #[test]
313    fn test_reformat() {
314        // Create sample addresses
315        let addresses = vec![
316            Address::from_str("0xb1f8e55c7f64d203c1400b9d8555d050f94adf39").unwrap(),
317            Address::from_str("0x8D9708f3F514206486D7E988533f770a16d074a7").unwrap(),
318        ];
319
320        // Create sample token addresses
321        let token_addresses = vec![
322            Address::from_str("0x3183B673f4816C94BeF53958BaF93C671B7F8Cf2").unwrap(),
323            Address::from_str("0x55ABBa8d669D60A10c104CC493ec5ef389EC92bb").unwrap(),
324        ];
325
326        // Create sample balances
327        let balances = vec![
328            U256::from(10), // balance for address[0] and token[0]
329            U256::from(7),  // balance for address[0] and token[1]
330            U256::from(5),  // balance for address[1] and token[0]
331            U256::from(7),  // balance for address[1] and token[1]
332        ];
333
334        // Get result from reformat function
335        let result = reformat(addresses.clone(), token_addresses.clone(), balances);
336
337        // Create expected result
338        let mut expected = HashMap::new();
339        let mut address_0_balances = HashMap::new();
340        address_0_balances.insert(token_addresses[0], U256::from(10));
341
342        let mut address_1_balances = HashMap::new();
343        address_1_balances.insert(token_addresses[0], U256::from(5));
344        address_1_balances.insert(token_addresses[1], U256::from(7));
345
346        expected.insert(addresses[0], address_0_balances);
347        expected.insert(addresses[1], address_1_balances);
348
349        assert_eq!(result, expected);
350    }
351}
ethers_erc20balances/lib.rs

ethers_erc20balances/
lib.rs
