Allows operators to set beneficiaries

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

@@ -37,7 +37,7 @@ module aptos_framework::staking_contract {
     use aptos_framework::aptos_account;
     use aptos_framework::aptos_coin::AptosCoin;
     use aptos_framework::coin::{Self, Coin};
-    use aptos_framework::event::{EventHandle, emit_event};
+    use aptos_framework::event::{EventHandle, emit, emit_event};
     use aptos_framework::stake::{Self, OwnerCapability};
     use aptos_framework::staking_config;
 
@@ -57,8 +57,8 @@ module aptos_framework::staking_contract {
     const ESTAKING_CONTRACT_ALREADY_EXISTS: u64 = 6;
     /// Not enough active stake to withdraw. Some stake might still pending and will be active in the next epoch.
     const EINSUFFICIENT_ACTIVE_STAKE_TO_WITHDRAW: u64 = 7;
-    /// Caller must be either the staker or operator.
-    const ENOT_STAKER_OR_OPERATOR: u64 = 8;
+    /// Caller must be either the staker, operator, or beneficiary.
+    const ENOT_STAKER_OR_OPERATOR_OR_BENEFICIARY: u64 = 8;
 
     /// Maximum number of distributions a stake pool can support.
     const MAXIMUM_PENDING_DISTRIBUTIONS: u64 = 20;
@@ -95,6 +95,10 @@ module aptos_framework::staking_contract {
         distribute_events: EventHandle<DistributeEvent>,
     }
 
+    struct Beneficiary has key {
+        beneficiary: address,
+    }
+
     struct UpdateCommissionEvent has drop, store {
         staker: address,
         operator: address,
@@ -166,6 +170,13 @@ module aptos_framework::staking_contract {
         amount: u64,
     }
 
+    #[event]
+    struct SetBeneficiaryForOperatorEvent has drop, store {
+        operator: address,
+        old_beneficiary: address,
+        new_beneficiary: address,
+    }
+
     #[view]
     /// Return the address of the underlying stake pool for the staking contract between the provided staker and
     /// operator.
@@ -234,6 +245,16 @@ module aptos_framework::staking_contract {
         simple_map::contains_key(&store.staking_contracts, &operator)
     }
 
+    #[view]
+    /// Return the beneficiary address of the operator.
+    public fun beneficiary_for_operator(operator: address): address acquires Beneficiary {
+        if (exists<Beneficiary>(operator)) {
+            return borrow_global<Beneficiary>(operator).beneficiary
+        } else {
+            operator
+        }
+    }
+
     #[view]
     /// Return the address of the stake pool to be created with the provided staker, operator and seed.
     public fun get_expected_stake_pool_address(
@@ -375,7 +396,7 @@ module aptos_framework::staking_contract {
 
     /// Convenience function to allow a staker to update the commission percentage paid to the operator.
     /// TODO: fix the typo in function name. commision -> commission
-    public entry fun update_commision(staker: &signer, operator: address, new_commission_percentage: u64) acquires Store, StakingGroupUpdateCommissionEvent {
+    public entry fun update_commision(staker: &signer, operator: address, new_commission_percentage: u64) acquires Store, StakingGroupUpdateCommissionEvent, Beneficiary {
         assert!(
             new_commission_percentage >= 0 && new_commission_percentage <= 100,
             error::invalid_argument(EINVALID_COMMISSION_PERCENTAGE),
@@ -407,10 +428,13 @@ module aptos_framework::staking_contract {
     /// Unlock commission amount from the stake pool. Operator needs to wait for the amount to become withdrawable
     /// at the end of the stake pool's lockup period before they can actually can withdraw_commission.
     ///
-    /// Only staker or operator can call this.
-    public entry fun request_commission(account: &signer, staker: address, operator: address) acquires Store {
+    /// Only staker, operator or beneficiary can call this.
+    public entry fun request_commission(account: &signer, staker: address, operator: address) acquires Store, Beneficiary {
         let account_addr = signer::address_of(account);
-        assert!(account_addr == staker || account_addr == operator, error::unauthenticated(ENOT_STAKER_OR_OPERATOR));
+        assert!(
+            account_addr == staker || account_addr == operator || account_addr == beneficiary_for_operator(operator),
+            error::unauthenticated(ENOT_STAKER_OR_OPERATOR_OR_BENEFICIARY)
+        );
         assert_staking_contract_exists(staker, operator);
 
         let store = borrow_global_mut<Store>(staker);
@@ -465,7 +489,7 @@ module aptos_framework::staking_contract {
 
     /// Staker can call this to request withdrawal of part or all of their staking_contract.
     /// This also triggers paying commission to the operator for accounting simplicity.
-    public entry fun unlock_stake(staker: &signer, operator: address, amount: u64) acquires Store {
+    public entry fun unlock_stake(staker: &signer, operator: address, amount: u64) acquires Store, Beneficiary {
         // Short-circuit if amount is 0.
         if (amount == 0) return;
 
@@ -511,7 +535,7 @@ module aptos_framework::staking_contract {
     }
 
     /// Unlock all accumulated rewards since the last recorded principals.
-    public entry fun unlock_rewards(staker: &signer, operator: address) acquires Store {
+    public entry fun unlock_rewards(staker: &signer, operator: address) acquires Store, Beneficiary {
         let staker_address = signer::address_of(staker);
         assert_staking_contract_exists(staker_address, operator);
 
@@ -526,7 +550,7 @@ module aptos_framework::staking_contract {
         staker: &signer,
         old_operator: address,
         new_operator: address,
-    ) acquires Store {
+    ) acquires Store, Beneficiary {
         let staker_address = signer::address_of(staker);
         assert_staking_contract_exists(staker_address, old_operator);
 
@@ -540,7 +564,7 @@ module aptos_framework::staking_contract {
         old_operator: address,
         new_operator: address,
         new_commission_percentage: u64,
-    ) acquires Store {
+    ) acquires Store, Beneficiary {
         let staker_address = signer::address_of(staker);
         assert_staking_contract_exists(staker_address, old_operator);
 
@@ -577,9 +601,26 @@ module aptos_framework::staking_contract {
         );
     }
 
+    /// Allows an operator to change its beneficiary.
+    public entry fun set_beneficiary_for_operator(operator: &signer, new_beneficiary: address) acquires Beneficiary {
+        let operator_addr = signer::address_of(operator);
+        let old_beneficiary = beneficiary_for_operator(operator_addr);
+        if (exists<Beneficiary>(operator_addr)) {
+            borrow_global_mut<Beneficiary>(operator_addr).beneficiary = new_beneficiary;
+        } else {
+            move_to(operator, Beneficiary { beneficiary: new_beneficiary });
+        };
+
+        emit(SetBeneficiaryForOperatorEvent {
+            operator: operator_addr,
+            old_beneficiary,
+            new_beneficiary,
+        });
+    }
+
     /// Allow anyone to distribute already unlocked funds. This does not affect reward compounding and therefore does
     /// not need to be restricted to just the staker or operator.
-    public entry fun distribute(staker: address, operator: address) acquires Store {
+    public entry fun distribute(staker: address, operator: address) acquires Store, Beneficiary {
         assert_staking_contract_exists(staker, operator);
         let store = borrow_global_mut<Store>(staker);
         let staking_contract = simple_map::borrow_mut(&mut store.staking_contracts, &operator);
@@ -592,7 +633,7 @@ module aptos_framework::staking_contract {
         operator: address,
         staking_contract: &mut StakingContract,
         distribute_events: &mut EventHandle<DistributeEvent>,
-    ) {
+    ) acquires Beneficiary {
         let pool_address = staking_contract.pool_address;
         let (_, inactive, _, pending_inactive) = stake::get_stake(pool_address);
         let total_potential_withdrawable = inactive + pending_inactive;
@@ -613,6 +654,10 @@ module aptos_framework::staking_contract {
             let recipient = *vector::borrow(&mut recipients, 0);
             let current_shares = pool_u64::shares(distribution_pool, recipient);
             let amount_to_distribute = pool_u64::redeem_shares(distribution_pool, recipient, current_shares);
+            // If the recipient is the operator, send the commission to the beneficiary instead.
+            if (recipient == operator) {
+                recipient = beneficiary_for_operator(operator);
+            };
             aptos_account::deposit_coins(recipient, coin::extract(&mut coins, amount_to_distribute));
 
             emit_event(
@@ -776,6 +821,9 @@ module aptos_framework::staking_contract {
     #[test_only]
     const MAXIMUM_STAKE: u64 = 100000000000000000; // 1B APT coins with 8 decimals.
 
+    #[test_only]
+    const MODULE_EVENT: u64 = 26;
+
     #[test_only]
     public fun setup(aptos_framework: &signer, staker: &signer, operator: &signer, initial_balance: u64) {
         // Reward rate of 0.1% per epoch.
@@ -806,10 +854,11 @@ module aptos_framework::staking_contract {
 
         // Voter is initially set to operator but then updated to be staker.
         create_staking_contract(staker, operator_address, operator_address, amount, commission, vector::empty<u8>());
+        std::features::change_feature_flags(aptos_framework, vector[MODULE_EVENT], vector[]);
     }
 
     #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
-    public entry fun test_end_to_end(aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store {
+    public entry fun test_end_to_end(aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, Beneficiary {
         setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
         let staker_address = signer::address_of(staker);
         let operator_address = signer::address_of(operator);
@@ -912,7 +961,7 @@ module aptos_framework::staking_contract {
 
     #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
     public entry fun test_operator_cannot_request_same_commission_multiple_times(
-        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store {
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, Beneficiary {
         setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
         let staker_address = signer::address_of(staker);
         let operator_address = signer::address_of(operator);
@@ -940,7 +989,7 @@ module aptos_framework::staking_contract {
 
     #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
     public entry fun test_unlock_rewards(
-        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store {
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, Beneficiary {
         setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
         let staker_address = signer::address_of(staker);
         let operator_address = signer::address_of(operator);
@@ -1037,7 +1086,7 @@ module aptos_framework::staking_contract {
         staker: &signer,
         operator_1: &signer,
         operator_2: &signer,
-    ) acquires Store {
+    ) acquires Store, Beneficiary {
         setup_staking_contract(aptos_framework, staker, operator_1, INITIAL_BALANCE, 10);
         account::create_account_for_test(signer::address_of(operator_2));
         stake::mint(operator_2, INITIAL_BALANCE);
@@ -1114,7 +1163,7 @@ module aptos_framework::staking_contract {
         staker: &signer,
         operator_1: &signer,
         operator_2: &signer,
-    ) acquires Store {
+    ) acquires Store, Beneficiary {
         setup_staking_contract(aptos_framework, staker, operator_1, INITIAL_BALANCE, 10);
         let staker_address = signer::address_of(staker);
         let operator_1_address = signer::address_of(operator_1);
@@ -1128,9 +1177,70 @@ module aptos_framework::staking_contract {
         assert!(commission_percentage(staker_address, operator_2_address) == 10, 2);
     }
 
+    #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234, beneficiary = @0x345)]
+    public entry fun test_operator_can_set_beneficiary(
+        aptos_framework: &signer,
+        staker: &signer,
+        operator: &signer,
+        beneficiary: &signer,
+    ) acquires Store, Beneficiary {
+        setup_staking_contract(aptos_framework, staker, operator, INITIAL_BALANCE, 10);
+        let staker_address = signer::address_of(staker);
+        let operator_address = signer::address_of(operator);
+        let beneficiary_address = signer::address_of(beneficiary);
+        // account::create_account_for_test(beneficiary_address);
+        aptos_framework::aptos_account::create_account(beneficiary_address);
+        assert_staking_contract_exists(staker_address, operator_address);
+        assert_staking_contract(staker_address, operator_address, INITIAL_BALANCE, 10);
+
+        // Verify that the stake pool has been set up properly.
+        let pool_address = stake_pool_address(staker_address, operator_address);
+        stake::assert_stake_pool(pool_address, INITIAL_BALANCE, 0, 0, 0);
+        assert!(last_recorded_principal(staker_address, operator_address) == INITIAL_BALANCE, 0);
+        assert!(stake::get_operator(pool_address) == operator_address, 0);
+        assert!(beneficiary_for_operator(operator_address) == operator_address, 0);
+
+        // Operator joins the validator set.
+        let (_sk, pk, pop) = stake::generate_identity();
+        stake::join_validator_set_for_test(&pk, &pop, operator, pool_address, true);
+        assert!(stake::get_validator_state(pool_address) == VALIDATOR_STATUS_ACTIVE, 1);
+
+        // Set beneficiary.
+        set_beneficiary_for_operator(operator, beneficiary_address);
+        assert!(beneficiary_for_operator(operator_address) == beneficiary_address, 0);
+
+        // Fast forward to generate rewards.
+        stake::end_epoch();
+        let new_balance = with_rewards(INITIAL_BALANCE);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+
+        // Operator claims 10% of rewards so far as commissions.
+        let expected_commission_1 = (new_balance - last_recorded_principal(staker_address, operator_address)) / 10;
+        new_balance = new_balance - expected_commission_1;
+        request_commission(operator, staker_address, operator_address);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, expected_commission_1);
+        assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
+        assert_distribution(staker_address, operator_address, operator_address, expected_commission_1);
+        stake::fast_forward_to_unlock(pool_address);
+
+        // Both original stake and operator commissions have received rewards.
+        expected_commission_1 = with_rewards(expected_commission_1);
+        new_balance = with_rewards(new_balance);
+        stake::assert_stake_pool(pool_address, new_balance, expected_commission_1, 0, 0);
+        distribute(staker_address, operator_address);
+        let operator_balance = coin::balance<AptosCoin>(operator_address);
+        let beneficiary_balance = coin::balance<AptosCoin>(beneficiary_address);
+        let expected_operator_balance = INITIAL_BALANCE;
+        let expected_beneficiary_balance = expected_commission_1;
+        assert!(operator_balance == expected_operator_balance, operator_balance);
+        assert!(beneficiary_balance == expected_beneficiary_balance, beneficiary_balance);
+        stake::assert_stake_pool(pool_address, new_balance, 0, 0, 0);
+        assert_no_pending_distributions(staker_address, operator_address);
+    }
+
     #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
     public entry fun test_staker_can_withdraw_partial_stake(
-        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store {
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, Beneficiary {
         let initial_balance = INITIAL_BALANCE * 2;
         setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
         let staker_address = signer::address_of(staker);
@@ -1180,7 +1290,7 @@ module aptos_framework::staking_contract {
 
     #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
     public entry fun test_staker_can_withdraw_partial_stake_if_operator_never_joined_validator_set(
-        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store {
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, Beneficiary {
         let initial_balance = INITIAL_BALANCE * 2;
         setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
         let staker_address = signer::address_of(staker);
@@ -1213,7 +1323,7 @@ module aptos_framework::staking_contract {
 
     #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
     public entry fun test_multiple_distributions_added_before_distribute(
-        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store {
+        aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, Beneficiary {
         let initial_balance = INITIAL_BALANCE * 2;
         setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
         let staker_address = signer::address_of(staker);
@@ -1258,7 +1368,7 @@ module aptos_framework::staking_contract {
         assert!(last_recorded_principal(staker_address, operator_address) == new_balance, 0);
     }
     #[test(aptos_framework = @0x1, staker = @0x123, operator = @0x234)]
-    public entry fun test_update_commission(aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, StakingGroupUpdateCommissionEvent {
+    public entry fun test_update_commission(aptos_framework: &signer, staker: &signer, operator: &signer) acquires Store, StakingGroupUpdateCommissionEvent, Beneficiary {
         let initial_balance = INITIAL_BALANCE * 2;
         setup_staking_contract(aptos_framework, staker, operator, initial_balance, 10);
         let staker_address = signer::address_of(staker);

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

@@ -192,6 +192,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`.