funds-manager: Add endpoint to fetch token balances for hot wallets

funds-manager/funds-manager-api/src/lib.rs

@@ -30,9 +30,6 @@ pub const GET_FEE_WALLETS_ROUTE: &str = "get-fee-wallets";
 /// The route to withdraw a fee balance
 pub const WITHDRAW_FEE_BALANCE_ROUTE: &str = "withdraw-fee-balance";
 
-/// The route to create a new hot wallet
-pub const CREATE_HOT_WALLET_ROUTE: &str = "create-hot-wallet";
-
 // -------------
 // | Api Types |
 // -------------
@@ -102,3 +99,28 @@ pub struct CreateHotWalletResponse {
     /// The address of the hot wallet
     pub address: String,
 }
+
+/// The response containing hot wallet balances
+#[derive(Debug, Serialize, Deserialize)]
+pub struct HotWalletBalancesResponse {
+    /// The list of hot wallets with their balances
+    pub wallets: Vec<WalletWithBalances>,
+}
+
+/// A hot wallet with its balances
+#[derive(Debug, Serialize, Deserialize)]
+pub struct WalletWithBalances {
+    /// The address of the hot wallet
+    pub address: String,
+    /// The balances of various tokens
+    pub balances: Vec<TokenBalance>,
+}
+
+/// A balance for a specific token
+#[derive(Debug, Serialize, Deserialize)]
+pub struct TokenBalance {
+    /// The mint address of the token
+    pub mint: String,
+    /// The balance amount
+    pub amount: u128,
+}

funds-manager/funds-manager-server/src/custody_client/hot_wallets.rs

@@ -3,14 +3,19 @@
 //! We store funds in hot wallets to prevent excessive in/out-flow from
 //! Fireblocks
 
+use std::sync::Arc;
+
 use ethers::{
+    providers::{Http, Provider},
     signers::{LocalWallet, Signer},
+    types::Address,
     utils::hex::ToHexExt,
 };
+use funds_manager_api::{TokenBalance, WalletWithBalances};
 use rand::thread_rng;
 use tracing::info;
 
-use super::CustodyClient;
+use super::{CustodyClient, ERC20};
 use crate::{error::FundsManagerError, helpers::create_secrets_manager_entry_with_description};
 
 impl CustodyClient {
@@ -40,4 +45,50 @@ impl CustodyClient {
         info!("Created hot wallet with address: {} for vault: {}", address, vault);
         Ok(address)
     }
+
+    /// Get balances for all hot wallets
+    pub async fn get_hot_wallet_balances(
+        &self,
+        mints: &[String],
+    ) -> Result<Vec<WalletWithBalances>, FundsManagerError> {
+        let hot_wallets = self.get_all_hot_wallets().await?;
+        let provider = Arc::new(self.get_rpc_provider()?);
+
+        let mut hot_wallet_balances = Vec::new();
+        for wallet in hot_wallets.iter().map(|w| w.address.clone()) {
+            // Fetch token balances for the wallet
+            let mut balances = Vec::new();
+            for mint in mints.iter() {
+                let balance = self.get_token_balance(&wallet, mint, provider.clone()).await?;
+                balances.push(TokenBalance { mint: mint.clone(), amount: balance });
+            }
+
+            hot_wallet_balances.push(WalletWithBalances { address: wallet, balances });
+        }
+
+        Ok(hot_wallet_balances)
+    }
+
+    /// Fetch the token balance at the given address for a wallet
+    async fn get_token_balance(
+        &self,
+        wallet_address: &str,
+        token_address: &str,
+        provider: Arc<Provider<Http>>,
+    ) -> Result<u128, FundsManagerError> {
+        let wallet_address: Address = wallet_address.parse().map_err(|_| {
+            FundsManagerError::parse(format!("Invalid wallet address: {wallet_address}"))
+        })?;
+        let token_address: Address = token_address.parse().map_err(|_| {
+            FundsManagerError::parse(format!("Invalid token address: {token_address}"))
+        })?;
+
+        let token = ERC20::new(token_address, provider);
+        token
+            .balance_of(wallet_address)
+            .call()
+            .await
+            .map(|balance| balance.as_u128())
+            .map_err(FundsManagerError::arbitrum)
+    }
 }

funds-manager/funds-manager-server/src/custody_client/mod.rs

@@ -26,6 +26,7 @@ use crate::error::FundsManagerError;
 abigen!(
     ERC20,
     r#"[
+        function balanceOf(address owner) external view returns (uint256)
         function symbol() external view returns (string memory)
     ]"#
 );
@@ -91,15 +92,20 @@ impl CustodyClient {
             .map_err(FundsManagerError::fireblocks)
     }
 
+    /// Get a JSON RPC provider for the given RPC url
+    pub fn get_rpc_provider(&self) -> Result<Provider<Http>, FundsManagerError> {
+        Provider::<Http>::try_from(&self.arbitrum_rpc_url)
+            .map_err(err_str!(FundsManagerError::Arbitrum))
+    }
+
     /// Get the symbol for an ERC20 token at the given address
     pub(self) async fn get_erc20_token_symbol(
         &self,
         token_address: &str,
     ) -> Result<String, FundsManagerError> {
         let addr =
             Address::from_str(token_address).map_err(err_str!(FundsManagerError::Arbitrum))?;
-        let provider = Provider::<Http>::try_from(&self.arbitrum_rpc_url)
-            .map_err(err_str!(FundsManagerError::Arbitrum))?;
+        let provider = self.get_rpc_provider()?;
         let client = Arc::new(provider);
         let erc20 = ERC20::new(addr, client);
 

funds-manager/funds-manager-server/src/custody_client/queries.rs

@@ -9,6 +9,17 @@ use crate::error::FundsManagerError;
 use crate::CustodyClient;
 
 impl CustodyClient {
+    /// Get all hot wallets
+    pub async fn get_all_hot_wallets(&self) -> Result<Vec<HotWallet>, FundsManagerError> {
+        let mut conn = self.get_db_conn().await?;
+        let wallets = hot_wallets::table
+            .load::<HotWallet>(&mut conn)
+            .await
+            .map_err(err_str!(FundsManagerError::Db))?;
+
+        Ok(wallets)
+    }
+
     /// Insert a new hot wallet into the database
     pub async fn insert_hot_wallet(
         &self,

funds-manager/funds-manager-server/src/handlers.rs

@@ -6,14 +6,15 @@ use crate::Server;
 use bytes::Bytes;
 use funds_manager_api::{
     CreateHotWalletRequest, CreateHotWalletResponse, DepositAddressResponse, FeeWalletsResponse,
-    WithdrawFeeBalanceRequest, WithdrawFundsRequest, WithdrawGasRequest,
+    HotWalletBalancesResponse, WithdrawFeeBalanceRequest, WithdrawFundsRequest, WithdrawGasRequest,
 };
+use itertools::Itertools;
 use std::collections::HashMap;
 use std::sync::Arc;
 use warp::reply::Json;
 
-/// The "mint" query param
-pub const MINT_QUERY_PARAM: &str = "mint";
+/// The "mints" query param
+pub const MINTS_QUERY_PARAM: &str = "mints";
 /// The asset used for gas (ETH)
 pub const GAS_ASSET_NAME: &str = "ETH";
 /// The maximum amount of gas that can be withdrawn at a given time
@@ -67,8 +68,8 @@ pub(crate) async fn get_deposit_address_handler(
     query_params: HashMap<String, String>,
     server: Arc<Server>,
 ) -> Result<Json, warp::Rejection> {
-    let mint = query_params.get(MINT_QUERY_PARAM).ok_or_else(|| {
-        warp::reject::custom(ApiError::BadRequest("Missing 'mint' query parameter".to_string()))
+    let mint = query_params.get(MINTS_QUERY_PARAM).ok_or_else(|| {
+        warp::reject::custom(ApiError::BadRequest("Missing 'mints' query parameter".to_string()))
     })?;
 
     let address = server
@@ -144,3 +145,24 @@ pub(crate) async fn create_hot_wallet_handler(
     let resp = CreateHotWalletResponse { address };
     Ok(warp::reply::json(&resp))
 }
+
+/// Handler for getting hot wallet balances
+pub(crate) async fn get_hot_wallet_balances_handler(
+    _body: Bytes, // unused
+    query_params: HashMap<String, String>,
+    server: Arc<Server>,
+) -> Result<Json, warp::Rejection> {
+    let mints = query_params
+        .get(MINTS_QUERY_PARAM)
+        .map(|s| s.split(',').map(String::from).collect_vec())
+        .unwrap_or_default();
+
+    let wallets = server
+        .custody_client
+        .get_hot_wallet_balances(&mints)
+        .await
+        .map_err(|e| warp::reject::custom(ApiError::InternalError(e.to_string())))?;
+
+    let resp = HotWalletBalancesResponse { wallets };
+    Ok(warp::reply::json(&resp))
+}

funds-manager/funds-manager-server/src/main.rs

@@ -21,14 +21,14 @@ use error::FundsManagerError;
 use ethers::signers::LocalWallet;
 use fee_indexer::Indexer;
 use funds_manager_api::{
-    CreateHotWalletRequest, WithdrawFeeBalanceRequest, WithdrawGasRequest, CREATE_HOT_WALLET_ROUTE,
+    CreateHotWalletRequest, WithdrawFeeBalanceRequest, WithdrawGasRequest,
     GET_DEPOSIT_ADDRESS_ROUTE, GET_FEE_WALLETS_ROUTE, INDEX_FEES_ROUTE, PING_ROUTE,
     REDEEM_FEES_ROUTE, WITHDRAW_CUSTODY_ROUTE, WITHDRAW_FEE_BALANCE_ROUTE, WITHDRAW_GAS_ROUTE,
 };
 use handlers::{
     create_hot_wallet_handler, get_deposit_address_handler, get_fee_wallets_handler,
-    index_fees_handler, quoter_withdraw_handler, redeem_fees_handler, withdraw_fee_balance_handler,
-    withdraw_gas_handler,
+    get_hot_wallet_balances_handler, index_fees_handler, quoter_withdraw_handler,
+    redeem_fees_handler, withdraw_fee_balance_handler, withdraw_gas_handler,
 };
 use middleware::{identity, with_hmac_auth, with_json_body};
 use relayer_client::RelayerClient;
@@ -342,13 +342,19 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     let create_hot_wallet = warp::post()
         .and(warp::path("hot-wallets"))
-        .and(warp::path(CREATE_HOT_WALLET_ROUTE))
         .and(with_hmac_auth(server.clone()))
         .map(with_json_body::<CreateHotWalletRequest>)
         .and_then(identity)
         .and(with_server(server.clone()))
         .and_then(create_hot_wallet_handler);
 
+    let get_hot_wallet_balances = warp::get()
+        .and(warp::path("hot-wallets"))
+        .and(with_hmac_auth(server.clone()))
+        .and(warp::query::<HashMap<String, String>>())
+        .and(with_server(server.clone()))
+        .and_then(get_hot_wallet_balances_handler);
+
     let routes = ping
         .or(index_fees)
         .or(redeem_fees)
@@ -358,6 +364,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
         .or(get_balances)
         .or(withdraw_fee_balance)
         .or(create_hot_wallet)
+        .or(get_hot_wallet_balances)
         .recover(handle_rejection);
     warp::serve(routes).run(([0, 0, 0, 0], cli.port)).await;