From d69cd07e324a95974bfdb83787f55c4952fd59f0 Mon Sep 17 00:00:00 2001 From: Greg Nazario Date: Fri, 8 Nov 2024 20:14:25 -0500 Subject: [PATCH] [rosetta] Fix staking balances the last refactor made it get staking balances and the account's balance this should fix it so it never cross contaminates between staking, delegated staking, and normal balances. It's much more straightforward and should be easier to understand. --- crates/aptos-rosetta/src/account.rs | 251 +++++++++++++++++----------- 1 file changed, 150 insertions(+), 101 deletions(-) diff --git a/crates/aptos-rosetta/src/account.rs b/crates/aptos-rosetta/src/account.rs index 6e24fd04e2dde..ac5608e907297 100644 --- a/crates/aptos-rosetta/src/account.rs +++ b/crates/aptos-rosetta/src/account.rs @@ -18,6 +18,7 @@ use aptos_logger::{debug, trace, warn}; use aptos_rest_client::{ aptos_api_types::{AptosError, AptosErrorCode, ViewFunction}, error::{AptosErrorResponse, RestError}, + Client, }; use aptos_types::{account_address::AccountAddress, account_config::AccountResource}; use move_core_types::{ @@ -25,7 +26,7 @@ use move_core_types::{ language_storage::{ModuleId, StructTag, TypeTag}, parser::parse_type_tag, }; -use std::str::FromStr; +use std::{collections::HashSet, str::FromStr}; use warp::Filter; /// Account routes e.g. balance @@ -105,43 +106,62 @@ async fn get_balances( let mut balances = vec![]; let mut lockup_expiration: u64 = 0; - let mut total_requested_balance: Option = None; + let mut maybe_operators = None; - // Lookup the delegation pool, if it's provided in the account information - if pool_address.is_some() { - match get_delegation_stake_balances( - rest_client.as_ref(), - &account, - owner_address, - pool_address.unwrap(), - version, - ) - .await - { - Ok(Some(balance_result)) => { - if let Some(balance) = balance_result.balance { - total_requested_balance = Some( - total_requested_balance.unwrap_or_default() - + u64::from_str(&balance.value).unwrap_or_default(), - ); - } - lockup_expiration = balance_result.lockup_expiration; - if let Some(balance) = total_requested_balance { - balances.push(Amount { - value: balance.to_string(), - currency: native_coin(), - }) - } - }, - result => { - warn!( - "Failed to retrieve requested balance for delegator_address: {}, pool_address: {}: {:?}", - owner_address, pool_address.unwrap(), result - ) - }, + // Handle the things that must always happen + + // Retrieve the sequence number + let sequence_number = get_sequence_number(&rest_client, owner_address, version).await?; + + // Filter currencies to lookup + let currencies_to_lookup = if let Some(currencies) = maybe_filter_currencies { + currencies.into_iter().collect() + } else { + server_context.currencies.clone() + }; + + // Regular account, FA and Coin + if account.is_base_account() { + balances = + get_base_balances(&rest_client, owner_address, version, currencies_to_lookup).await?; + } else if pool_address.is_some() { + // Lookup the delegation pool, if it's provided in the account information + // Filter appropriately, must have native coin + if currencies_to_lookup.contains(&native_coin()) { + (balances, lockup_expiration) = get_delegation_info( + &rest_client, + &account, + owner_address, + pool_address.unwrap(), + version, + ) + .await?; + } + } else { + // Retrieve staking information (if it applies) + // Only non-pool addresses, and non base accounts + // + // These are special cases around querying the stake amounts + // Filter appropriately, must have native coin + if currencies_to_lookup.contains(&native_coin()) { + (balances, lockup_expiration, maybe_operators) = + get_staking_info(&rest_client, &account, owner_address, version).await?; } } + Ok(( + sequence_number, + maybe_operators, + balances, + lockup_expiration, + )) +} + +async fn get_sequence_number( + rest_client: &Client, + owner_address: AccountAddress, + version: u64, +) -> ApiResult { // Retrieve sequence number let sequence_number = match rest_client .get_account_resource_at_version_bcs(owner_address, "0x1::account::Account", version) @@ -178,77 +198,111 @@ async fn get_balances( }, }; - // Retrieve staking information (if it applies) - // Only non-pool addresses, and non base accounts + Ok(sequence_number) +} + +async fn get_staking_info( + rest_client: &Client, + account: &AccountIdentifier, + owner_address: AccountAddress, + version: u64, +) -> ApiResult<(Vec, u64, Option>)> { + let mut balances = vec![]; + let mut lockup_expiration: u64 = 0; let mut maybe_operators = None; - if !account.is_base_account() && pool_address.is_none() { - if let Ok(response) = rest_client - .get_account_resource_at_version_bcs( - owner_address, - "0x1::staking_contract::Store", - version, - ) - .await - { - let store: Store = response.into_inner(); - maybe_operators = Some(vec![]); - for (operator, contract) in store.staking_contracts { - // Keep track of operators - maybe_operators.as_mut().unwrap().push(operator); - match get_stake_balances( - rest_client.as_ref(), + let mut total_balance = 0; + let mut has_staking = false; + + if let Ok(response) = rest_client + .get_account_resource_at_version_bcs(owner_address, "0x1::staking_contract::Store", version) + .await + { + let store: Store = response.into_inner(); + maybe_operators = Some(vec![]); + for (operator, contract) in store.staking_contracts { + // Keep track of operators + maybe_operators.as_mut().unwrap().push(operator); + match get_stake_balances(rest_client, account, contract.pool_address, version).await { + Ok(Some(balance_result)) => { + if let Some(balance) = balance_result.balance { + has_staking = true; + total_balance += u64::from_str(&balance.value).unwrap_or_default(); + } + // TODO: This seems like it only works if there's only one staking contract (hopefully it stays that way) + lockup_expiration = balance_result.lockup_expiration; + }, + result => { + warn!( + "Failed to retrieve requested balance for account: {}, address: {}: {:?}", + owner_address, contract.pool_address, result + ) + }, + } + } + if has_staking { + balances.push(Amount { + value: total_balance.to_string(), + currency: native_coin(), + }) + } + + /* TODO: Right now operator stake is not supported + else if account.is_operator_stake() { + // For operator stake, filter on operator address + let operator_address = account.operator_address()?; + if let Some(contract) = store.staking_contracts.get(&operator_address) { + balances.push(get_total_stake( + rest_client, &account, contract.pool_address, version, - ) - .await - { - Ok(Some(balance_result)) => { - if let Some(balance) = balance_result.balance { - total_requested_balance = Some( - total_requested_balance.unwrap_or_default() - + u64::from_str(&balance.value).unwrap_or_default(), - ); - } - lockup_expiration = balance_result.lockup_expiration; - }, - result => { - warn!( - "Failed to retrieve requested balance for account: {}, address: {}: {:?}", - owner_address, contract.pool_address, result - ) - }, - } + ).await?); } - if let Some(balance) = total_requested_balance { + }*/ + } + + Ok((balances, lockup_expiration, maybe_operators)) +} + +async fn get_delegation_info( + rest_client: &Client, + account: &AccountIdentifier, + owner_address: AccountAddress, + pool_address: AccountAddress, + version: u64, +) -> ApiResult<(Vec, u64)> { + let mut balances = vec![]; + let mut lockup_expiration: u64 = 0; + + match get_delegation_stake_balances(rest_client, account, owner_address, pool_address, version) + .await + { + Ok(Some(balance_result)) => { + if let Some(balance) = balance_result.balance { balances.push(Amount { - value: balance.to_string(), + value: balance.value, currency: native_coin(), - }) + }); } - - /* TODO: Right now operator stake is not supported - else if account.is_operator_stake() { - // For operator stake, filter on operator address - let operator_address = account.operator_address()?; - if let Some(contract) = store.staking_contracts.get(&operator_address) { - balances.push(get_total_stake( - rest_client, - &account, - contract.pool_address, - version, - ).await?); - } - }*/ - } + lockup_expiration = balance_result.lockup_expiration; + }, + result => { + warn!( + "Failed to retrieve requested balance for delegator_address: {}, pool_address: {}: {:?}", + owner_address, pool_address, result + ) + }, } + Ok((balances, lockup_expiration)) +} - // Filter currencies to lookup - let currencies_to_lookup = if let Some(currencies) = maybe_filter_currencies { - currencies.into_iter().collect() - } else { - server_context.currencies.clone() - }; +async fn get_base_balances( + rest_client: &Client, + owner_address: AccountAddress, + version: u64, + currencies_to_lookup: HashSet, +) -> ApiResult> { + let mut balances = vec![]; // Retrieve the fungible asset balances and the coin balances for currency in currencies_to_lookup.iter() { @@ -330,10 +384,5 @@ async fn get_balances( } } - Ok(( - sequence_number, - maybe_operators, - balances, - lockup_expiration, - )) + Ok(balances) }