chore(Horizon): add a beneficiary address to undelegate (#1052)

packages/horizon/contracts/interfaces/internal/IHorizonStakingMain.sol

@@ -409,6 +409,11 @@ interface IHorizonStakingMain {
      */
     error HorizonStakingInvalidDelegationPool(address serviceProvider, address verifier);
 
+    /**
+     * @notice Thrown when attempting to undelegate with a beneficiary that is the zero address.
+     */
+    error HorizonStakingInvalidBeneficiaryZeroAddress();
+
     // -- Errors: thaw requests --
 
     error HorizonStakingNothingThawing();
@@ -706,6 +711,33 @@ interface IHorizonStakingMain {
      */
     function undelegate(address serviceProvider, address verifier, uint256 shares) external returns (bytes32);
 
+    /**
+     * @notice Undelegate tokens from a provision and start thawing them.
+     * The tokens will be withdrawable by the `beneficiary` after the thawing period.
+     *
+     * Note that undelegating tokens from a provision is a two step process:
+     * - First the tokens are thawed using this function.
+     * - Then after the thawing period, the tokens are removed from the provision using {withdrawDelegated}.
+     *
+     * Requirements:
+     * - `shares` cannot be zero.
+     * - `beneficiary` cannot be the zero address.
+     *
+     * Emits a {TokensUndelegated} and {ThawRequestCreated} event.
+     *
+     * @param serviceProvider The service provider address
+     * @param verifier The verifier address
+     * @param shares The amount of shares to undelegate
+     * @param beneficiary The address where the tokens will be withdrawn after thawing
+     * @return The ID of the thaw request
+     */
+    function undelegate(
+        address serviceProvider,
+        address verifier,
+        uint256 shares,
+        address beneficiary
+    ) external returns (bytes32);
+
     /**
      * @notice Withdraw undelegated tokens from a provision after thawing.
      * Tokens can be automatically re-delegated to another provision by setting `newServiceProvider`.

packages/horizon/contracts/staking/HorizonStaking.sol

@@ -302,7 +302,20 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
         address verifier,
         uint256 shares
     ) external override notPaused returns (bytes32) {
-        return _undelegate(serviceProvider, verifier, shares);
+        return _undelegate(serviceProvider, verifier, shares, msg.sender);
+    }
+
+    /**
+     * @notice See {IHorizonStakingMain-undelegate}.
+     */
+    function undelegate(
+        address serviceProvider,
+        address verifier,
+        uint256 shares,
+        address beneficiary
+    ) external override notPaused returns (bytes32) {
+        require(beneficiary != address(0), HorizonStakingInvalidBeneficiaryZeroAddress());
+        return _undelegate(serviceProvider, verifier, shares, beneficiary);
     }
 
     /**
@@ -345,7 +358,7 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
      * @notice See {IHorizonStakingMain-undelegate}.
      */
     function undelegate(address serviceProvider, uint256 shares) external override notPaused {
-        _undelegate(serviceProvider, SUBGRAPH_DATA_SERVICE_ADDRESS, shares);
+        _undelegate(serviceProvider, SUBGRAPH_DATA_SERVICE_ADDRESS, shares, msg.sender);
     }
 
     /**
@@ -762,7 +775,12 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
      * @dev To allow delegation to be slashable even while thawing without breaking accounting
      * the delegation pool shares are burned and replaced with thawing pool shares.
      */
-    function _undelegate(address _serviceProvider, address _verifier, uint256 _shares) private returns (bytes32) {
+    function _undelegate(
+        address _serviceProvider,
+        address _verifier,
+        uint256 _shares,
+        address beneficiary
+    ) private returns (bytes32) {
         require(_shares > 0, HorizonStakingInvalidZeroShares());
         DelegationPoolInternal storage pool = _getDelegationPool(_serviceProvider, _verifier);
         DelegationInternal storage delegation = pool.delegators[msg.sender];
@@ -789,7 +807,7 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
         bytes32 thawRequestId = _createThawRequest(
             _serviceProvider,
             _verifier,
-            msg.sender,
+            beneficiary,
             thawingShares,
             thawingUntil
         );

packages/horizon/test/shared/horizon-staking/HorizonStakingShared.t.sol

@@ -853,11 +853,17 @@ abstract contract HorizonStakingSharedTest is GraphBaseTest {
     }
 
     function _undelegate(address serviceProvider, address verifier, uint256 shares) internal {
-        __undelegate(serviceProvider, verifier, shares, false);
+        (, address caller, ) = vm.readCallers();
+        __undelegate(serviceProvider, verifier, shares, false, caller);
+    }
+
+    function _undelegate(address serviceProvider, address verifier, uint256 shares, address beneficiary) internal {
+        __undelegate(serviceProvider, verifier, shares, false, beneficiary);
     }
 
     function _undelegate(address serviceProvider, uint256 shares) internal {
-        __undelegate(serviceProvider, subgraphDataServiceLegacyAddress, shares, true);
+        (, address caller, ) = vm.readCallers();
+        __undelegate(serviceProvider, subgraphDataServiceLegacyAddress, shares, true, caller);
     }
 
     struct BeforeValues_Undelegate {
@@ -873,7 +879,7 @@ abstract contract HorizonStakingSharedTest is GraphBaseTest {
         bytes32 thawRequestId;
     }
 
-    function __undelegate(address serviceProvider, address verifier, uint256 shares, bool legacy) private {
+    function __undelegate(address serviceProvider, address verifier, uint256 shares, bool legacy, address beneficiary) private {
         (, address delegator, ) = vm.readCallers();
 
         // before
@@ -893,15 +899,15 @@ abstract contract HorizonStakingSharedTest is GraphBaseTest {
             staking.getProvision(serviceProvider, verifier).thawingPeriod +
             uint64(block.timestamp);
         calcValues.thawRequestId = keccak256(
-            abi.encodePacked(serviceProvider, verifier, delegator, beforeValues.thawRequestList.nonce)
+            abi.encodePacked(serviceProvider, verifier, beneficiary, beforeValues.thawRequestList.nonce)
         );
 
         // undelegate
         vm.expectEmit();
         emit IHorizonStakingMain.ThawRequestCreated(
             serviceProvider,
             verifier,
-            delegator,
+            beneficiary,
             calcValues.thawingShares,
             calcValues.thawingUntil,
             calcValues.thawRequestId
@@ -911,7 +917,7 @@ abstract contract HorizonStakingSharedTest is GraphBaseTest {
         if (legacy) {
             staking.undelegate(serviceProvider, shares);
         } else {
-            staking.undelegate(serviceProvider, verifier, shares);
+            staking.undelegate(serviceProvider, verifier, shares, beneficiary);
         }
 
         // after
@@ -923,10 +929,10 @@ abstract contract HorizonStakingSharedTest is GraphBaseTest {
         DelegationInternal memory afterDelegation = _getStorage_Delegation(
             serviceProvider,
             verifier,
-            delegator,
+            beneficiary,
             legacy
         );
-        LinkedList.List memory afterThawRequestList = staking.getThawRequestList(serviceProvider, verifier, delegator);
+        LinkedList.List memory afterThawRequestList = staking.getThawRequestList(serviceProvider, verifier, beneficiary);
         ThawRequest memory afterThawRequest = staking.getThawRequest(calcValues.thawRequestId);
         uint256 afterDelegatedTokens = staking.getDelegatedTokensAvailable(serviceProvider, verifier);
 

packages/horizon/test/staking/delegation/undelegate.t.sol

@@ -40,6 +40,17 @@ contract HorizonStakingUndelegateTest is HorizonStakingTest {
         }
     }
 
+    function testUndelegate_WithBeneficiary(
+        uint256 amount,
+        uint256 delegationAmount,
+        address beneficiary
+    ) public useIndexer useProvision(amount, 0, 0) useDelegation(delegationAmount) {
+        vm.assume(beneficiary != address(0));
+        resetPrank(users.delegator);
+        DelegationInternal memory delegation = _getStorage_Delegation(users.indexer, subgraphDataServiceAddress, users.delegator, false);
+        _undelegate(users.indexer, subgraphDataServiceAddress, delegation.shares, beneficiary);
+    }
+
     function testUndelegate_RevertWhen_TooManyUndelegations()
         public
         useIndexer
@@ -133,4 +144,15 @@ contract HorizonStakingUndelegateTest is HorizonStakingTest {
         ));
         staking.undelegate(users.indexer, subgraphDataServiceAddress, delegation.shares);
     }
+
+    function testUndelegate_RevertIf_BeneficiaryIsZero(
+        uint256 amount,
+        uint256 delegationAmount
+    ) public useIndexer useProvision(amount, 0, 0) useDelegation(delegationAmount) {
+        resetPrank(users.delegator);
+        DelegationInternal memory delegation = _getStorage_Delegation(users.indexer, subgraphDataServiceAddress, users.delegator, false);
+        bytes memory expectedError = abi.encodeWithSelector(IHorizonStakingMain.HorizonStakingInvalidBeneficiaryZeroAddress.selector);
+        vm.expectRevert(expectedError);
+        staking.undelegate(users.indexer, subgraphDataServiceAddress, delegation.shares, address(0));
+    }
 }

packages/horizon/test/staking/delegation/withdraw.t.sol

@@ -176,4 +176,57 @@ contract HorizonStakingWithdrawDelegationTest is HorizonStakingTest {
         ));
         staking.withdrawDelegated(users.indexer, subgraphDataServiceAddress, address(0), 0, 0);
     }
+
+    function testWithdrawDelegation_WithBeneficiary(
+        uint256 delegationAmount,
+        address beneficiary
+    )
+        public
+        useIndexer
+        useProvision(10_000_000 ether, 0, MAX_THAWING_PERIOD)
+        useDelegation(delegationAmount)
+     {
+        vm.assume(beneficiary != address(0));
+
+        // Delegator undelegates to beneficiary
+        resetPrank(users.delegator);
+        DelegationInternal memory delegation = _getStorage_Delegation(users.indexer, subgraphDataServiceAddress, users.delegator, false);
+        _undelegate(users.indexer, subgraphDataServiceAddress, delegation.shares, beneficiary);
+
+        // Thawing period ends
+        LinkedList.List memory thawingRequests = staking.getThawRequestList(users.indexer, subgraphDataServiceAddress, beneficiary);
+        ThawRequest memory thawRequest = staking.getThawRequest(thawingRequests.tail);
+        skip(thawRequest.thawingUntil + 1);
+
+        // Beneficiary withdraws delegated tokens
+        resetPrank(beneficiary);
+        _withdrawDelegated(users.indexer, subgraphDataServiceAddress, address(0), 0, 1);
+    }
+
+    function testWithdrawDelegation_RevertWhen_PreviousOwnerAttemptsToWithdraw(
+        uint256 delegationAmount,
+        address beneficiary
+    )
+        public
+        useIndexer
+        useProvision(10_000_000 ether, 0, MAX_THAWING_PERIOD)
+        useDelegation(delegationAmount)
+    {
+        vm.assume(beneficiary != address(0));
+
+        // Delegator undelegates to beneficiary
+        resetPrank(users.delegator);
+        DelegationInternal memory delegation = _getStorage_Delegation(users.indexer, subgraphDataServiceAddress, users.delegator, false);
+        _undelegate(users.indexer, subgraphDataServiceAddress, delegation.shares, beneficiary);
+
+        // Thawing period ends
+        LinkedList.List memory thawingRequests = staking.getThawRequestList(users.indexer, subgraphDataServiceAddress, users.delegator);
+        ThawRequest memory thawRequest = staking.getThawRequest(thawingRequests.tail);
+        skip(thawRequest.thawingUntil + 1);
+
+        // Delegator attempts to withdraw delegated tokens, should revert since beneficiary is the thaw request owner
+        bytes memory expectedError = abi.encodeWithSelector(IHorizonStakingMain.HorizonStakingNothingThawing.selector);
+        vm.expectRevert(expectedError);
+        staking.withdrawDelegated(users.indexer, subgraphDataServiceAddress, address(0), 0, 1);
+    }
 }