Allows operators to set beneficiaries

aptos-move/framework/aptos-framework/doc/vesting.md

@@ -89,6 +89,7 @@ withdrawable, admin can call admin_withdraw to withdraw all funds to the vesting
 -  [Function `reset_beneficiary`](#0x1_vesting_reset_beneficiary)
 -  [Function `set_management_role`](#0x1_vesting_set_management_role)
 -  [Function `set_beneficiary_resetter`](#0x1_vesting_set_beneficiary_resetter)
+-  [Function `set_beneficiary_for_operator`](#0x1_vesting_set_beneficiary_for_operator)
 -  [Function `get_role_holder`](#0x1_vesting_get_role_holder)
 -  [Function `get_vesting_account_signer`](#0x1_vesting_get_vesting_account_signer)
 -  [Function `get_vesting_account_signer_internal`](#0x1_vesting_get_vesting_account_signer_internal)
@@ -133,6 +134,7 @@ withdrawable, admin can call admin_withdraw to withdraw all funds to the vesting
     -  [Function `reset_beneficiary`](#@Specification_1_reset_beneficiary)
     -  [Function `set_management_role`](#@Specification_1_set_management_role)
     -  [Function `set_beneficiary_resetter`](#@Specification_1_set_beneficiary_resetter)
+    -  [Function `set_beneficiary_for_operator`](#@Specification_1_set_beneficiary_for_operator)
     -  [Function `get_role_holder`](#@Specification_1_get_role_holder)
     -  [Function `get_vesting_account_signer`](#@Specification_1_get_vesting_account_signer)
     -  [Function `get_vesting_account_signer_internal`](#@Specification_1_get_vesting_account_signer_internal)
@@ -2460,6 +2462,34 @@ account.
 
 
 
+</details>
+
+<a name="0x1_vesting_set_beneficiary_for_operator"></a>
+
+## Function `set_beneficiary_for_operator`
+
+Set the beneficiary for the operator.
+
+
+<pre><code><b>public</b> entry <b>fun</b> <a href="vesting.md#0x1_vesting_set_beneficiary_for_operator">set_beneficiary_for_operator</a>(operator: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>, new_beneficiary: <b>address</b>)
+</code></pre>
+
+
+
+<details>
+<summary>Implementation</summary>
+
+
+<pre><code><b>public</b> entry <b>fun</b> <a href="vesting.md#0x1_vesting_set_beneficiary_for_operator">set_beneficiary_for_operator</a>(
+    operator: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>,
+    new_beneficiary: <b>address</b>,
+) {
+    <a href="staking_contract.md#0x1_staking_contract_set_beneficiary_for_operator">staking_contract::set_beneficiary_for_operator</a>(operator, new_beneficiary);
+}
+</code></pre>
+
+
+
 </details>
 
 <a name="0x1_vesting_get_role_holder"></a>
@@ -3465,6 +3495,22 @@ This address should be deterministic for the same admin and vesting contract cre
 
 
 
+<a name="@Specification_1_set_beneficiary_for_operator"></a>
+
+### Function `set_beneficiary_for_operator`
+
+
+<pre><code><b>public</b> entry <b>fun</b> <a href="vesting.md#0x1_vesting_set_beneficiary_for_operator">set_beneficiary_for_operator</a>(operator: &<a href="../../aptos-stdlib/../move-stdlib/doc/signer.md#0x1_signer">signer</a>, new_beneficiary: <b>address</b>)
+</code></pre>
+
+
+
+
+<pre><code><b>pragma</b> verify = <b>false</b>;
+</code></pre>
+
+
+
 <a name="@Specification_1_get_role_holder"></a>
 
 ### Function `get_role_holder`

aptos-move/framework/aptos-framework/sources/staking_contract.spec.move

@@ -243,6 +243,16 @@ spec aptos_framework::staking_contract {
         aborts_if simple_map::spec_contains_key(staking_contracts, new_operator);
     }
 
+    spec set_beneficiary_for_operator(operator: &signer, new_beneficiary: address) {
+        // TODO: temporary mockup
+        pragma verify = false;
+    }
+
+    spec beneficiary_for_operator(operator: address): address {
+        // TODO: temporary mockup
+        pragma verify = false;
+    }
+
     /// Staking_contract exists the stacker/operator pair.
     spec distribute(staker: address, operator: address) {
         // TODO: Call `distribute_internal` and could not verify `update_distribution_pool`.

aptos-move/framework/aptos-framework/sources/vesting.move

@@ -905,6 +905,14 @@ module aptos_framework::vesting {
         set_management_role(admin, contract_address, utf8(ROLE_BENEFICIARY_RESETTER), beneficiary_resetter);
     }
 
+    /// Set the beneficiary for the operator.
+    public entry fun set_beneficiary_for_operator(
+        operator: &signer,
+        new_beneficiary: address,
+    ) {
+        staking_contract::set_beneficiary_for_operator(operator, new_beneficiary);
+    }
+
     public fun get_role_holder(contract_address: address, role: String): address acquires VestingAccountManagement {
         assert!(exists<VestingAccountManagement>(contract_address), error::not_found(EVESTING_ACCOUNT_HAS_NO_ROLES));
         let roles = &borrow_global<VestingAccountManagement>(contract_address).roles;
@@ -1011,6 +1019,12 @@ module aptos_framework::vesting {
     #[test_only]
     const VALIDATOR_STATUS_INACTIVE: u64 = 4;
 
+    #[test_only]
+    const MODULE_EVENT: u64 = 26;
+
+    #[test_only]
+    const OPERATOR_BENEFICIARY: u64 = 38;
+
     #[test_only]
     public fun setup(aptos_framework: &signer, accounts: &vector<address>) {
         use aptos_framework::aptos_account::create_account;
@@ -1023,6 +1037,8 @@ module aptos_framework::vesting {
                 create_account(addr);
             };
         });
+
+        std::features::change_feature_flags(aptos_framework, vector[MODULE_EVENT, OPERATOR_BENEFICIARY], vector[]);
     }
 
     #[test_only]
@@ -1534,6 +1550,85 @@ module aptos_framework::vesting {
         assert!(coin::balance<AptosCoin>(operator_address) == expected_commission, 1);
     }
 
+    #[test(aptos_framework = @0x1, admin = @0x123, shareholder = @0x234, operator1 = @0x345, beneficiary = @0x456, operator2 = @0x567)]
+    public entry fun test_set_beneficiary_for_operator(
+        aptos_framework: &signer,
+        admin: &signer,
+        shareholder: &signer,
+        operator1: &signer,
+        beneficiary: &signer,
+        operator2: &signer,
+    ) acquires AdminStore, VestingContract {
+        let admin_address = signer::address_of(admin);
+        let operator_address1 = signer::address_of(operator1);
+        let operator_address2 = signer::address_of(operator2);
+        let shareholder_address = signer::address_of(shareholder);
+        let beneficiary_address = signer::address_of(beneficiary);
+        setup(aptos_framework, &vector[admin_address, shareholder_address, operator_address1, beneficiary_address]);
+        let contract_address = setup_vesting_contract(
+            admin, &vector[shareholder_address], &vector[GRANT_AMOUNT], admin_address, 0);
+        assert!(operator_commission_percentage(contract_address) == 0, 0);
+        let stake_pool_address = stake_pool_address(contract_address);
+        // 10% commission will be paid to the operator.
+        update_operator(admin, contract_address, operator_address1, 10);
+        assert!(staking_contract::beneficiary_for_operator(operator_address1) == operator_address1, 0);
+        set_beneficiary_for_operator(operator1, beneficiary_address);
+        assert!(staking_contract::beneficiary_for_operator(operator_address1) == beneficiary_address, 0);
+
+        // Operator needs to join the validator set for the stake pool to earn rewards.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator1, stake_pool_address, true);
+        stake::assert_stake_pool(stake_pool_address, GRANT_AMOUNT, 0, 0, 0);
+        assert!(get_accumulated_rewards(contract_address) == 0, 0);
+        assert!(remaining_grant(contract_address) == GRANT_AMOUNT, 0);
+
+        // Stake pool earns some rewards.
+        stake::end_epoch();
+        let (_, accumulated_rewards, _) = staking_contract::staking_contract_amounts(contract_address,
+            operator_address1
+        );
+        // Commission is calculated using the previous commission percentage which is 10%.
+        let expected_commission = accumulated_rewards / 10;
+
+        // Request commission.
+        staking_contract::request_commission(operator1, contract_address, operator_address1);
+        // Unlocks the commission.
+        stake::fast_forward_to_unlock(stake_pool_address);
+        expected_commission = with_rewards(expected_commission);
+
+        // Distribute the commission to the operator.
+        distribute(contract_address);
+
+        // Assert that the beneficiary receives the expected commission.
+        assert!(coin::balance<AptosCoin>(operator_address1) == 0, 1);
+        assert!(coin::balance<AptosCoin>(beneficiary_address) == expected_commission, 1);
+        let old_beneficiay_balance = coin::balance<AptosCoin>(beneficiary_address);
+
+        // switch operator to operator2. The rewards should go to operator2 not to the beneficiay of operator1.
+        update_operator(admin, contract_address, operator_address2, 10);
+
+        stake::end_epoch();
+        let (_, accumulated_rewards, _) = staking_contract::staking_contract_amounts(contract_address,
+            operator_address2
+        );
+
+        let expected_commission = accumulated_rewards / 10;
+
+        // Request commission.
+        staking_contract::request_commission(operator2, contract_address, operator_address2);
+        // Unlocks the commission.
+        stake::fast_forward_to_unlock(stake_pool_address);
+        expected_commission = with_rewards(expected_commission);
+
+        // Distribute the commission to the operator.
+        distribute(contract_address);
+
+        // Assert that the rewards go to operator2, and the balance of the operator1's beneficiay remains the same.
+        assert!(coin::balance<AptosCoin>(operator_address2) >= expected_commission, 1);
+        assert!(coin::balance<AptosCoin>(beneficiary_address) == old_beneficiay_balance, 1);
+
+    }
+
     #[test(aptos_framework = @0x1, admin = @0x123, shareholder = @0x234)]
     #[expected_failure(abort_code = 0x30008, location = Self)]
     public entry fun test_cannot_unlock_rewards_after_contract_is_terminated(

aptos-move/framework/aptos-framework/sources/vesting.spec.move

@@ -365,6 +365,14 @@ spec aptos_framework::vesting {
         include SetManagementRoleAbortsIf;
     }
 
+    spec set_beneficiary_for_operator(
+        operator: &signer,
+        new_beneficiary: address,
+    ) {
+        // TODO: temporary mockup
+        pragma verify = false;
+    }
+
     spec get_role_holder(contract_address: address, role: String): address {
         aborts_if !exists<VestingAccountManagement>(contract_address);
         let roles = global<VestingAccountManagement>(contract_address).roles;

aptos-move/framework/aptos-framework/tests/delegation_pool_integration_tests.move

@@ -32,6 +32,9 @@ module aptos_framework::delegation_pool_integration_tests {
     #[test_only]
     const DELEGATION_POOLS: u64 = 11;
 
+    #[test_only]
+    const MODULE_EVENT: u64 = 26;
+
     #[test_only]
     public fun initialize_for_test(aptos_framework: &signer) {
         initialize_for_test_custom(
@@ -76,7 +79,7 @@ module aptos_framework::delegation_pool_integration_tests {
             voting_power_increase_limit
         );
         reconfiguration::initialize_for_test(aptos_framework);
-        features::change_feature_flags(aptos_framework, vector[DELEGATION_POOLS], vector[]);
+        features::change_feature_flags(aptos_framework, vector[DELEGATION_POOLS, MODULE_EVENT], vector[]);
     }
 
     #[test_only]