Skip to content

Commit

Permalink
[rosetta] Fix staking balances
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@gregnazario
gregnazario committed Nov 9, 2024
1 parent e174527 commit 67b6773
Showing 1 changed file with 150 additions and 101 deletions.
251 changes: 150 additions & 101 deletions crates/aptos-rosetta/src/account.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,15 @@ 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::{
ident_str,
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
Expand Down Expand Up @@ -105,43 +106,62 @@ async fn get_balances(

let mut balances = vec![];
let mut lockup_expiration: u64 = 0;
let mut total_requested_balance: Option<u64> = 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<u64> {
// Retrieve sequence number
let sequence_number = match rest_client
.get_account_resource_at_version_bcs(owner_address, "0x1::account::Account", version)
Expand Down Expand Up @@ -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<Amount>, u64, Option<Vec<AccountAddress>>)> {
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<Amount>, 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<Currency>,
) -> ApiResult<Vec<Amount>> {
let mut balances = vec![];

// Retrieve the fungible asset balances and the coin balances
for currency in currencies_to_lookup.iter() {
Expand Down Expand Up @@ -330,10 +384,5 @@ async fn get_balances(
}
}

Ok((
sequence_number,
maybe_operators,
balances,
lockup_expiration,
))
Ok(balances)
}

0 comments on commit 67b6773

Please sign in to comment.