Skip to content

Commit

Permalink
[Rosetta] Add staking_contract distribute implementation (#6810)
Browse files Browse the repository at this point in the history 
* fix conflicts

* lint

---------

Co-authored-by: Yukai <[email protected]>
  • Loading branch information
@xiaoxiaff @yukaitu-cb
xiaoxiaff and yukaitu-cb authored Mar 8, 2023
1 parent af4e858 commit 4831cf6
Show file tree
Hide file tree
Showing 6 changed files with 321 additions and 3 deletions.
39 changes: 39 additions & 0 deletions crates/aptos-rosetta/src/client.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -387,6 +387,45 @@ impl RosettaClient {
.await
}

pub async fn distribute_staking_rewards(
&self,
network_identifier: &NetworkIdentifier,
private_key: &Ed25519PrivateKey,
operator: AccountAddress,
staker: AccountAddress,
expiry_time_secs: u64,
sequence_number: Option<u64>,
max_gas: Option<u64>,
gas_unit_price: Option<u64>,
) -> anyhow::Result<TransactionIdentifier> {
let sender = self
.get_account_address(network_identifier.clone(), private_key)
.await?;
let mut keys = HashMap::new();
keys.insert(sender, private_key);

let operations = vec![Operation::distribute_staking_rewards(
0,
None,
sender,
AccountIdentifier::base_account(operator),
AccountIdentifier::base_account(staker),
)];

self.submit_operations(
sender,
network_identifier.clone(),
&keys,
operations,
expiry_time_secs,
sequence_number,
max_gas,
gas_unit_price,
false,
)
.await
}

pub async fn create_stake_pool(
&self,
network_identifier: &NetworkIdentifier,
Expand Down
46 changes: 46 additions & 0 deletions crates/aptos-rosetta/src/construction.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -550,6 +550,11 @@ async fn construction_parse(
(AccountAddress::ONE, STAKING_CONTRACT_MODULE, UNLOCK_STAKE_FUNCTION) => {
parse_unlock_stake_operation(sender, &type_args, &args)?
},
(
AccountAddress::ONE,
STAKING_CONTRACT_MODULE,
DISTRIBUTE_STAKING_REWARDS_FUNCTION,
) => parse_distribute_staking_rewards_operation(sender, &type_args, &args)?,
_ => {
return Err(ApiError::TransactionParseError(Some(format!(
"Unsupported entry function type {:x}::{}::{}",
Expand Down Expand Up @@ -850,6 +855,30 @@ pub fn parse_unlock_stake_operation(
)])
}

pub fn parse_distribute_staking_rewards_operation(
sender: AccountAddress,
type_args: &[TypeTag],
args: &[Vec<u8>],
) -> ApiResult<Vec<Operation>> {
if !type_args.is_empty() {
return Err(ApiError::TransactionParseError(Some(format!(
"Distribute should not have type arguments: {:?}",
type_args
))));
}

let operator: AccountAddress = parse_function_arg("distribute_staking_rewards", args, 0)?;
let staker: AccountAddress = parse_function_arg("distribute_staking_rewards", args, 1)?;

Ok(vec![Operation::distribute_staking_rewards(
0,
None,
sender,
AccountIdentifier::base_account(operator),
AccountIdentifier::base_account(staker),
)])
}

/// Construction payloads command (OFFLINE)
///
/// Constructs payloads for given known operations
Expand Down Expand Up @@ -966,6 +995,23 @@ async fn construction_payloads(
))));
}
},
InternalOperation::DistributeStakingRewards(inner) => {
if let InternalOperation::DistributeStakingRewards(ref metadata_op) =
metadata.internal_operation
{
if inner.operator != metadata_op.operator || inner.staker != metadata_op.staker {
return Err(ApiError::InvalidInput(Some(format!(
"Distribute staking rewards operation doesn't match metadata {:?} vs {:?}",
inner, metadata.internal_operation
))));
}
} else {
return Err(ApiError::InvalidInput(Some(format!(
"Distribute staking rewards operation doesn't match metadata {:?} vs {:?}",
inner, metadata.internal_operation
))));
}
},
}

// Encode operation
Expand Down
5 changes: 5 additions & 0 deletions crates/aptos-rosetta/src/types/misc.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,13 +94,15 @@ pub enum OperationType {
InitializeStakePool,
ResetLockup,
UnlockStake,
DistributeStakingRewards,
// Fee must always be last for ordering
Fee,
}

impl OperationType {
const CREATE_ACCOUNT: &'static str = "create_account";
const DEPOSIT: &'static str = "deposit";
const DISTRIBUTE_STAKING_REWARDS: &'static str = "distribute_staking_rewards";
const FEE: &'static str = "fee";
const INITIALIZE_STAKE_POOL: &'static str = "initialize_stake_pool";
const RESET_LOCKUP: &'static str = "reset_lockup";
Expand All @@ -123,6 +125,7 @@ impl OperationType {
InitializeStakePool,
ResetLockup,
UnlockStake,
DistributeStakingRewards,
]
}
}
Expand All @@ -142,6 +145,7 @@ impl FromStr for OperationType {
Self::INITIALIZE_STAKE_POOL => Ok(OperationType::InitializeStakePool),
Self::RESET_LOCKUP => Ok(OperationType::ResetLockup),
Self::UNLOCK_STAKE => Ok(OperationType::UnlockStake),
Self::DISTRIBUTE_STAKING_REWARDS => Ok(OperationType::DistributeStakingRewards),
_ => Err(ApiError::DeserializationFailed(Some(format!(
"Invalid OperationType: {}",
s
Expand All @@ -163,6 +167,7 @@ impl Display for OperationType {
InitializeStakePool => Self::INITIALIZE_STAKE_POOL,
ResetLockup => Self::RESET_LOCKUP,
UnlockStake => Self::UNLOCK_STAKE,
DistributeStakingRewards => Self::DISTRIBUTE_STAKING_REWARDS,
Fee => Self::FEE,
})
}
Expand Down
1 change: 1 addition & 0 deletions crates/aptos-rosetta/src/types/move_types.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ pub const SWITCH_OPERATOR_WITH_SAME_COMMISSION_FUNCTION: &str =
"switch_operator_with_same_commission";
pub const UPDATE_VOTER_FUNCTION: &str = "update_voter";
pub const UNLOCK_STAKE_FUNCTION: &str = "unlock_stake";
pub const DISTRIBUTE_STAKING_REWARDS_FUNCTION: &str = "distribute";

pub const DECIMALS_FIELD: &str = "decimal";
pub const DEPOSIT_EVENTS_FIELD: &str = "deposit_events";
Expand Down
113 changes: 111 additions & 2 deletions crates/aptos-rosetta/src/types/objects.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,9 @@
use crate::{
common::{is_native_coin, native_coin, native_coin_tag},
construction::{
parse_create_stake_pool_operation, parse_reset_lockup_operation,
parse_set_operator_operation, parse_set_voter_operation, parse_unlock_stake_operation,
parse_create_stake_pool_operation, parse_distribute_staking_rewards_operation,
parse_reset_lockup_operation, parse_set_operator_operation, parse_set_voter_operation,
parse_unlock_stake_operation,
},
error::ApiResult,
types::{
Expand Down Expand Up @@ -387,6 +388,25 @@ impl Operation {
Some(OperationMetadata::unlock_stake(operator, amount)),
)
}

pub fn distribute_staking_rewards(
operation_index: u64,
status: Option<OperationStatusType>,
account: AccountAddress,
operator: AccountIdentifier,
staker: AccountIdentifier,
) -> Operation {
Operation::new(
OperationType::DistributeStakingRewards,
operation_index,
status,
AccountIdentifier::base_account(account),
None,
Some(OperationMetadata::distribute_staking_rewards(
operator, staker,
)),
)
}
}

impl std::cmp::PartialOrd for Operation {
Expand Down Expand Up @@ -437,6 +457,8 @@ pub struct OperationMetadata {
pub commission_percentage: Option<U64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub amount: Option<U64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub staker: Option<AccountIdentifier>,
}

impl OperationMetadata {
Expand Down Expand Up @@ -497,6 +519,17 @@ impl OperationMetadata {
..Default::default()
}
}

pub fn distribute_staking_rewards(
operator: AccountIdentifier,
staker: AccountIdentifier,
) -> Self {
OperationMetadata {
operator: Some(operator),
staker: Some(staker),
..Default::default()
}
}
}

/// Public key used for the rosetta implementation. All private keys will never be handled
Expand Down Expand Up @@ -807,6 +840,19 @@ fn parse_failed_operations_from_txn_payload(
warn!("Failed to parse unlock stake {:?}", inner);
}
},
(AccountAddress::ONE, STAKING_CONTRACT_MODULE, DISTRIBUTE_STAKING_REWARDS_FUNCTION) => {
if let Ok(mut ops) = parse_distribute_staking_rewards_operation(
sender,
inner.ty_args(),
inner.args(),
) {
if let Some(operation) = ops.get_mut(0) {
operation.status = Some(OperationStatusType::Failure.to_string());
}
} else {
warn!("Failed to parse distribute staking rewards {:?}", inner);
}
},
_ => {
// If we don't recognize the transaction payload, then we can't parse operations
},
Expand Down Expand Up @@ -1298,6 +1344,34 @@ async fn parse_staking_contract_resource_changes(
}
operations.push(operation);
}

// Handle distribute events, there are no events on the stake pool
let distribute_staking_rewards_events =
filter_events(events, store.distribute_events.key(), |event_key, event| {
if let Ok(event) = bcs::from_bytes::<DistributeEvent>(event.event_data()) {
Some(event)
} else {
// If we can't parse the withdraw event, then there's nothing
warn!(
"Failed to parse distribute event! Skipping for {}:{}",
event_key.get_creator_address(),
event_key.get_creation_number()
);
None
}
});

// For every distribute events, add staking reward operation
for event in distribute_staking_rewards_events {
operations.push(Operation::staking_reward(
operation_index,
Some(OperationStatusType::Success),
AccountIdentifier::base_account(event.recipient),
native_coin(),
event.amount,
));
operation_index += 1;
}
}

Ok(operations)
Expand Down Expand Up @@ -1400,6 +1474,7 @@ pub enum InternalOperation {
InitializeStakePool(InitializeStakePool),
ResetLockup(ResetLockup),
UnlockStake(UnlockStake),
DistributeStakingRewards(DistributeStakingRewards),
}

impl InternalOperation {
Expand Down Expand Up @@ -1544,6 +1619,25 @@ impl InternalOperation {
}));
}
},
Ok(OperationType::DistributeStakingRewards) => {
if let (
Some(OperationMetadata {
operator: Some(operator),
staker: Some(staker),
..
}),
Some(account),
) = (&operation.metadata, &operation.account)
{
return Ok(Self::DistributeStakingRewards(
DistributeStakingRewards {
sender: account.account_address()?,
operator: operator.account_address()?,
staker: staker.account_address()?,
},
));
}
},
_ => {},
}
}
Expand Down Expand Up @@ -1572,6 +1666,7 @@ impl InternalOperation {
Self::InitializeStakePool(inner) => inner.owner,
Self::ResetLockup(inner) => inner.owner,
Self::UnlockStake(inner) => inner.owner,
Self::DistributeStakingRewards(inner) => inner.sender,
}
}

Expand Down Expand Up @@ -1639,6 +1734,13 @@ impl InternalOperation {
),
unlock_stake.owner,
),
InternalOperation::DistributeStakingRewards(distribute_staking_rewards) => (
aptos_stdlib::staking_contract_distribute(
distribute_staking_rewards.operator,
distribute_staking_rewards.staker,
),
distribute_staking_rewards.sender,
),
})
}
}
Expand Down Expand Up @@ -1804,3 +1906,10 @@ pub struct UnlockStake {
pub operator: AccountAddress,
pub amount: u64,
}

#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
pub struct DistributeStakingRewards {
pub sender: AccountAddress,
pub operator: AccountAddress,
pub staker: AccountAddress,
}
Loading

0 comments on commit 4831cf6

Please sign in to comment.