feat(balance): adding forced balance update by the contract address

integration/balance.test.ts

@@ -60,4 +60,63 @@ describe('Account balance', () => {
     expect(typeof resp.data.balances).toBe('object')
     expect(resp.data.balances).toHaveLength(0)
   })
+
+  it('force update balance for the ERC20 token', async () => {
+    // USDC token contract address on Base
+    const token_contract_address = 'eip155:8453:0x833589fcd6edb6e08f4c7c32d4f71b54bda02913'
+    const endpoint = `/v1/account/${fulfilled_address}/balance`;
+    const queryParams = `?projectId=${projectId}&currency=${currency}&forceUpdate=${token_contract_address}`;
+    const url = `${baseUrl}${endpoint}${queryParams}`;
+    const headers = {
+        'x-sdk-version': sdk_version,
+    };
+    let resp = await httpClient.get(url, { headers });
+    expect(resp.status).toBe(200)
+    expect(typeof resp.data.balances).toBe('object')
+    expect(resp.data.balances.length).toBeGreaterThan(1)
+
+    for (const item of resp.data.balances) {
+      expect(typeof item.name).toBe('string')
+      expect(typeof item.symbol).toBe('string')
+      expect(item.chainId).toEqual(expect.stringMatching(/^(eip155:)?\d+$/))
+      if (item.address !== undefined) {
+        expect(item.address).toEqual(expect.stringMatching(/^(eip155:\d+:0x[0-9a-fA-F]{40})$/))
+      } else {
+        expect(item.address).toBeUndefined()
+      }
+      expect(typeof item.price).toBe('number')
+      expect(typeof item.quantity).toBe('object')
+      expect(typeof item.iconUrl).toBe('string')
+    }
+  })
+
+  it('force update balance for the native token', async () => {
+    // ETH token
+    // We are using `0xe...` as a contract address for native tokens
+    const token_contract_address = 'eip155:1:0xeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee'
+    const endpoint = `/v1/account/${fulfilled_address}/balance`;
+    const queryParams = `?projectId=${projectId}&currency=${currency}&forceUpdate=${token_contract_address}`;
+    const url = `${baseUrl}${endpoint}${queryParams}`;
+    const headers = {
+        'x-sdk-version': sdk_version,
+    };
+    let resp = await httpClient.get(url, { headers });
+    expect(resp.status).toBe(200)
+    expect(typeof resp.data.balances).toBe('object')
+    expect(resp.data.balances.length).toBeGreaterThan(1)
+
+    for (const item of resp.data.balances) {
+      expect(typeof item.name).toBe('string')
+      expect(typeof item.symbol).toBe('string')
+      expect(item.chainId).toEqual(expect.stringMatching(/^(eip155:)?\d+$/))
+      if (item.address !== undefined) {
+        expect(item.address).toEqual(expect.stringMatching(/^(eip155:\d+:0x[0-9a-fA-F]{40})$/))
+      } else {
+        expect(item.address).toBeUndefined()
+      }
+      expect(typeof item.price).toBe('number')
+      expect(typeof item.quantity).toBe('object')
+      expect(typeof item.iconUrl).toBe('string')
+    }
+  })
 })

src/handlers/balance.rs

@@ -1,12 +1,17 @@
 use {
     super::HANDLER_TASK_METRICS,
-    crate::{analytics::BalanceLookupInfo, error::RpcError, state::AppState, utils::network},
+    crate::{
+        analytics::BalanceLookupInfo,
+        error::RpcError,
+        state::AppState,
+        utils::{crypto, network},
+    },
     axum::{
         extract::{ConnectInfo, Path, Query, State},
         response::{IntoResponse, Response},
         Json,
     },
-    ethers::abi::Address,
+    ethers::{abi::Address, types::H160},
     hyper::HeaderMap,
     serde::{Deserialize, Serialize},
     std::{
@@ -16,7 +21,7 @@ use {
         time::{Duration, SystemTime},
     },
     tap::TapFallible,
-    tracing::log::error,
+    tracing::log::{debug, error},
     wc::future::FutureExt,
 };
 
@@ -60,6 +65,8 @@ pub struct BalanceQueryParams {
     pub project_id: String,
     pub currency: BalanceCurrencies,
     pub chain_id: Option<String>,
+    /// Comma separated list of CAIP-10 contract addresses to force update the balance
+    pub force_update: Option<String>,
 }
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Clone)]
@@ -112,7 +119,7 @@ async fn handler_internal(
     Path(address): Path<String>,
 ) -> Result<Response, RpcError> {
     let project_id = query.project_id.clone();
-    address
+    let parsed_address = address
         .parse::<Address>()
         .map_err(|_| RpcError::InvalidAddress)?;
 
@@ -126,7 +133,7 @@ async fn handler_internal(
     }
 
     let start = SystemTime::now();
-    let response = state
+    let mut response = state
         .providers
         .balance_provider
         .get_balance(address.clone(), query.clone().0, state.http_client.clone())
@@ -167,5 +174,76 @@ async fn handler_internal(
         }
     }
 
+    // Check for the cache invalidation for the certain token contract addresses and
+    // update/override balance results for the token from the RPC call
+    if let Some(force_update) = &query.force_update {
+        let rpc_project_id = state
+            .config
+            .server
+            .testing_project_id
+            .as_ref()
+            .ok_or_else(|| {
+                RpcError::InvalidConfiguration(
+                    "Missing testing project id in the configuration for the balance RPC lookups"
+                        .to_string(),
+                )
+            })?;
+        let force_update: Vec<&str> = force_update.split(',').collect();
+        for caip_contract_address in force_update {
+            debug!(
+                "Forcing balance update for the contract address: {}",
+                caip_contract_address
+            );
+            let (namespace, chain_id, contract_address) =
+                crypto::disassemble_caip10(caip_contract_address)
+                    .map_err(|_| RpcError::InvalidAddress)?;
+            let contract_address = contract_address
+                .parse::<Address>()
+                .map_err(|_| RpcError::InvalidAddress)?;
+            let caip2_chain_id = format!("{}:{}", namespace, chain_id);
+            let rpc_balance = crypto::get_erc20_balance(
+                &caip2_chain_id,
+                contract_address,
+                parsed_address,
+                rpc_project_id,
+            )
+            .await?;
+            if let Some(balance) = response
+                .balances
+                .iter_mut()
+                .find(|b| b.address == Some(caip_contract_address.to_string()))
+            {
+                balance.quantity.numeric = crypto::format_token_amount(
+                    rpc_balance,
+                    balance.quantity.decimals.parse::<u32>().unwrap_or(0),
+                );
+                // Recalculating the value with the latest balance
+                balance.value = Some(crypto::convert_token_amount_to_value(
+                    rpc_balance,
+                    balance.price,
+                    balance.quantity.decimals.parse::<u32>().unwrap_or(0),
+                ));
+            }
+            if contract_address == H160::repeat_byte(0xee) {
+                if let Some(balance) = response
+                    .balances
+                    .iter_mut()
+                    .find(|b| b.address.is_none() && b.chain_id == Some(caip2_chain_id.clone()))
+                {
+                    balance.quantity.numeric = crypto::format_token_amount(
+                        rpc_balance,
+                        balance.quantity.decimals.parse::<u32>().unwrap_or(0),
+                    );
+                    // Recalculate the value with the latest balance
+                    balance.value = Some(crypto::convert_token_amount_to_value(
+                        rpc_balance,
+                        balance.price,
+                        balance.quantity.decimals.parse::<u32>().unwrap_or(0),
+                    ));
+                }
+            }
+        }
+    }
+
     Ok(Json(response).into_response())
 }

src/utils/crypto.rs

@@ -1,10 +1,14 @@
 use {
     alloy_primitives::Address,
-    ethers::types::H256,
+    ethers::{
+        prelude::abigen,
+        providers::{Http, Middleware, Provider},
+        types::{H160, H256, U256},
+    },
     once_cell::sync::Lazy,
     regex::Regex,
     relay_rpc::auth::cacao::{signature::eip6492::verify_eip6492, CacaoError},
-    std::str::FromStr,
+    std::{str::FromStr, sync::Arc},
     strum::IntoEnumIterator,
     strum_macros::{Display, EnumIter, EnumString},
     tracing::warn,
@@ -31,6 +35,8 @@ pub enum CryptoUitlsError {
     WrongCaip2Format(String),
     #[error("Wrong CAIP-10 format: {0}")]
     WrongCaip10Format(String),
+    #[error("Provider call error: {0}")]
+    ProviderError(String),
     #[error("Contract call error: {0}")]
     ContractCallError(String),
     #[error("Wrong address format: {0}")]
@@ -104,6 +110,78 @@ pub async fn verify_eip6492_message_signature(
     }
 }
 
+/// Get the balance of the ERC20 token
+#[tracing::instrument]
+pub async fn get_erc20_balance(
+    chain_id: &str,
+    contract: H160,
+    wallet: H160,
+    rpc_project_id: &str,
+) -> Result<U256, CryptoUitlsError> {
+    // Use JSON-RPC call for the balance of the native ERC20 tokens
+    // or call the contract for the custom ERC20 tokens
+    let balance = if contract == H160::repeat_byte(0xee) {
+        get_erc20_jsonrpc_balance(chain_id, wallet, rpc_project_id).await?
+    } else {
+        get_erc20_contract_balance(chain_id, contract, wallet, rpc_project_id).await?
+    };
+
+    Ok(balance)
+}
+
+/// Get the balance of ERC20 token by calling the contract address
+#[tracing::instrument]
+async fn get_erc20_contract_balance(
+    chain_id: &str,
+    contract: H160,
+    wallet: H160,
+    rpc_project_id: &str,
+) -> Result<U256, CryptoUitlsError> {
+    abigen!(
+        ERC20Contract,
+        r#"[
+            function balanceOf(address account) external view returns (uint256)
+        ]"#,
+    );
+
+    let provider = Provider::<Http>::try_from(format!(
+        "https://rpc.walletconnect.com/v1?chainId={}&projectId={}",
+        chain_id, rpc_project_id
+    ))
+    .map_err(|e| CryptoUitlsError::RpcUrlParseError(format!("Failed to parse RPC url: {}", e)))?;
+    let provider = Arc::new(provider);
+
+    let contract = ERC20Contract::new(contract, provider);
+    let balance = contract.balance_of(wallet).call().await.map_err(|e| {
+        CryptoUitlsError::ContractCallError(format!(
+            "Failed to call ERC20 contract for the balance: {}",
+            e
+        ))
+    })?;
+    Ok(balance)
+}
+
+/// Get the balance of ERC20 token using JSON-RPC call
+#[tracing::instrument]
+async fn get_erc20_jsonrpc_balance(
+    chain_id: &str,
+    wallet: H160,
+    rpc_project_id: &str,
+) -> Result<U256, CryptoUitlsError> {
+    let provider = Provider::<Http>::try_from(format!(
+        "https://rpc.walletconnect.com/v1?chainId={}&projectId={}",
+        chain_id, rpc_project_id
+    ))
+    .map_err(|e| CryptoUitlsError::RpcUrlParseError(format!("Failed to parse RPC url: {}", e)))?;
+    let provider = Arc::new(provider);
+
+    let balance = provider
+        .get_balance(wallet, None)
+        .await
+        .map_err(|e| CryptoUitlsError::ProviderError(format!("{}", e)))?;
+    Ok(balance)
+}
+
 /// Convert EVM chain ID to coin type ENSIP-11
 #[tracing::instrument]
 pub fn convert_evm_chain_id_to_coin_type(chain_id: u32) -> u32 {
@@ -289,6 +367,31 @@ pub fn constant_time_eq(a: impl AsRef<[u8]>, b: impl AsRef<[u8]>) -> bool {
     result == 0
 }
 
+/// Format token amount to human readable format according to the token decimals
+pub fn format_token_amount(amount: U256, decimals: u32) -> String {
+    let amount_str = amount.to_string();
+    let decimals_usize = decimals as usize;
+
+    // Handle cases where the total digits are less than or equal to the decimals
+    if amount_str.len() <= decimals_usize {
+        let required_zeros = decimals_usize - amount_str.len();
+        let zeros = "0".repeat(required_zeros);
+        return format!("0.{}{}", zeros, amount_str);
+    }
+
+    // Insert the decimal point at the correct position
+    let (integer_part, decimal_part) = amount_str.split_at(amount_str.len() - decimals_usize);
+    format!("{}.{}", integer_part, decimal_part)
+}
+
+/// Convert token amount to value depending on the token price and decimals
+pub fn convert_token_amount_to_value(balance: U256, price: f64, decimals: u32) -> f64 {
+    let decimals_usize = decimals as usize;
+    let scaling_factor = 10_u64.pow(decimals_usize as u32) as f64;
+    let balance_f64 = balance.as_u64() as f64 / scaling_factor;
+    balance_f64 * price
+}
+
 #[cfg(test)]
 mod tests {
     use {super::*, std::collections::HashMap};
@@ -394,4 +497,31 @@ mod tests {
         let error_result = disassemble_caip10(malformed_caip10);
         assert!(error_result.is_err());
     }
+
+    #[test]
+    fn test_format_token_amount() {
+        // Test case for ethereum 18 decimals
+        let amount_18 = U256::from_dec_str("959694527317077690").unwrap();
+        let decimals_18 = 18;
+        assert_eq!(
+            format_token_amount(amount_18, decimals_18),
+            "0.959694527317077690"
+        );
+
+        // Test case for polygon usdc 6 decimals
+        let amount_6 = U256::from_dec_str("125320550").unwrap();
+        let decimals_6 = 6;
+        assert_eq!(format_token_amount(amount_6, decimals_6), "125.320550");
+    }
+
+    #[test]
+    fn test_convert_token_amount_to_value() {
+        let balance = U256::from_dec_str("959694527317077690").unwrap();
+        let price = 10000.05;
+        let decimals = 18;
+        assert_eq!(
+            convert_token_amount_to_value(balance, price, decimals),
+            0.959_694_527_317_077_7 * price
+        );
+    }
 }