fix: addToDelegationPool access control and token transfer (OZ C-01)

graphprotocol · Aug 7, 2024 · e414021 · e414021
1 parent f47cf91
commit e414021
Show file tree

Hide file tree

Showing 7 changed files with 93 additions and 2 deletions.
diff --git a/packages/horizon/contracts/interfaces/internal/IHorizonStakingMain.sol b/packages/horizon/contracts/interfaces/internal/IHorizonStakingMain.sol
@@ -338,6 +338,12 @@ interface IHorizonStakingMain {
      */
     error HorizonStakingNotAuthorized(address caller, address serviceProvider, address verifier);
 
+    /**
+     * @notice Thrown when an unauthorized sender attempts to deposit tokens into the delegation pool.
+     * @param sender The message sender address
+     */
+    error HorizonStakingInvalidDelegationPoolSender(address sender);
+
     /**
      * @notice Thrown when attempting to create a provision with an invalid maximum verifier cut.
      * @param maxVerifierCut The maximum verifier cut

diff --git a/packages/horizon/contracts/mocks/MockGRTToken.sol b/packages/horizon/contracts/mocks/MockGRTToken.sol
@@ -7,7 +7,9 @@ import { IGraphToken } from "@graphprotocol/contracts/contracts/token/IGraphToke
 contract MockGRTToken is ERC20, IGraphToken {
     constructor() ERC20("Graph Token", "GRT") {}
 
-    function burn(uint256 tokens) external {}
+    function burn(uint256 tokens) external {
+        _burn(msg.sender, tokens);
+    }
 
     function burnFrom(address from, uint256 tokens) external {
         _burn(from, tokens);

diff --git a/packages/horizon/contracts/payments/GraphPayments.sol b/packages/horizon/contracts/payments/GraphPayments.sol
@@ -68,6 +68,7 @@ contract GraphPayments is Initializable, MulticallUpgradeable, GraphDirectory, I
 
         // Pay delegators
         if (tokensDelegationPool > 0) {
+            _graphToken().approve(address(_graphStaking()), tokensDelegationPool);
             _graphStaking().addToDelegationPool(receiver, dataService, tokensDelegationPool);
         }
 

diff --git a/packages/horizon/contracts/staking/HorizonStaking.sol b/packages/horizon/contracts/staking/HorizonStaking.sol
@@ -278,6 +278,8 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
         uint256 tokens
     ) external override notPaused {
         require(tokens != 0, HorizonStakingInvalidZeroTokens());
+        require(msg.sender == verifier || msg.sender == address(_graphPayments()), HorizonStakingInvalidDelegationPoolSender(msg.sender));
+        _graphToken().pullTokens(msg.sender, tokens);
         DelegationPoolInternal storage pool = _getDelegationPool(serviceProvider, verifier);
         pool.tokens = pool.tokens + tokens;
         emit TokensToDelegationPoolAdded(serviceProvider, verifier, tokens);

diff --git a/packages/horizon/test/escrow/collect.t.sol b/packages/horizon/test/escrow/collect.t.sol
@@ -31,7 +31,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
         uint256 indexerBalance = token.balanceOf(users.indexer);
         uint256 indexerExpectedPayment = amount - tokensDataService - tokensProtocol - tokensDelegatoion;
         assertEq(indexerBalance - indexerPreviousBalance, indexerExpectedPayment);
-        assertTrue(true);
+        assertEq(token.balanceOf(address(payments)), 0);
     }
 
     function testCollect_RevertWhen_CollectorNotAuthorized(uint256 amount) public {

diff --git a/packages/horizon/test/staking/allocation/collect.t.sol b/packages/horizon/test/staking/allocation/collect.t.sol
@@ -141,6 +141,7 @@ contract HorizonStakingCollectAllocationTest is HorizonStakingExtensionTest {
         assertEq(staking.getStake(address(users.indexer)), provisionTokens + payment);
         assertEq(curation.curation(_subgraphDeploymentID), curationTokens + curationCutTokens);
         assertEq(staking.getDelegationPool(users.indexer, subgraphDataServiceLegacyAddress).tokens, delegationTokens + delegationFeeCut);
+        assertEq(token.balanceOf(address(payments)), 0);
     }
 
     function testCollect_WithBeneficiaryAddress(

diff --git a/packages/horizon/test/staking/delegation/addToPool.t.sol b/packages/horizon/test/staking/delegation/addToPool.t.sol
@@ -0,0 +1,79 @@
+// SPDX-License-Identifier: MIT
+pragma solidity 0.8.26;
+
+import "forge-std/Test.sol";
+
+import { IHorizonStakingMain } from "../../../contracts/interfaces/internal/IHorizonStakingMain.sol";
+import { HorizonStakingTest } from "../HorizonStaking.t.sol";
+
+contract HorizonStakingDelegationAddToPoolTest is HorizonStakingTest {
+
+    modifier useValidDelegationAmount(uint256 tokens) {
+        vm.assume(tokens > 0);
+        vm.assume(tokens <= MAX_STAKING_TOKENS);
+        _;
+    }
+
+    /*
+     * TESTS
+     */
+
+    function test_Delegation_AddToPool_Verifier(
+        uint256 amount,
+        uint256 delegationAmount
+    ) public useIndexer useProvision(amount, 0, 0) useValidDelegationAmount(delegationAmount) {
+        uint256 stakingPreviousBalance = token.balanceOf(address(staking));
+
+        resetPrank(subgraphDataServiceAddress);
+        mint(subgraphDataServiceAddress, delegationAmount);
+        token.approve(address(staking), delegationAmount);
+        vm.expectEmit(address(staking));
+        emit IHorizonStakingMain.TokensToDelegationPoolAdded(users.indexer, subgraphDataServiceAddress, delegationAmount);
+        staking.addToDelegationPool(users.indexer, subgraphDataServiceAddress, delegationAmount);
+
+        uint256 delegatedTokens = staking.getDelegatedTokensAvailable(users.indexer, subgraphDataServiceAddress);
+        assertEq(delegatedTokens, delegationAmount);
+        assertEq(token.balanceOf(subgraphDataServiceAddress), 0);
+        assertEq(token.balanceOf(address(staking)), stakingPreviousBalance + delegationAmount);
+    }
+
+    function test_Delegation_AddToPool_Payments(
+        uint256 amount,
+        uint256 delegationAmount
+    ) public useIndexer useProvision(amount, 0, 0) useValidDelegationAmount(delegationAmount) {
+        uint256 stakingPreviousBalance = token.balanceOf(address(staking));
+
+        resetPrank(address(payments));
+        mint(address(payments), delegationAmount);
+        token.approve(address(staking), delegationAmount);
+        staking.addToDelegationPool(users.indexer, subgraphDataServiceAddress, delegationAmount);
+
+        uint256 delegatedTokens = staking.getDelegatedTokensAvailable(users.indexer, subgraphDataServiceAddress);
+        assertEq(delegatedTokens, delegationAmount);
+        assertEq(token.balanceOf(subgraphDataServiceAddress), 0);
+        assertEq(token.balanceOf(address(staking)), stakingPreviousBalance + delegationAmount);
+    }
+
+    function test_Delegation_AddToPool_RevertWhen_InvalidSender(
+        uint256 amount,
+        uint256 delegationAmount
+    ) public useIndexer useProvision(amount, 0, 0) useValidDelegationAmount(delegationAmount) {
+        vm.assume(delegationAmount > 0);
+        vm.startPrank(users.delegator);
+        bytes memory expectedError = abi.encodeWithSelector(
+            IHorizonStakingMain.HorizonStakingInvalidDelegationPoolSender.selector,
+            users.delegator
+        );
+        vm.expectRevert(expectedError);
+        staking.addToDelegationPool(users.indexer, subgraphDataServiceAddress, delegationAmount);
+    }
+
+    function test_Delegation_AddToPool_RevertWhen_ZeroTokens(
+        uint256 amount
+    ) public useIndexer useProvision(amount, 0, 0) {
+        vm.startPrank(subgraphDataServiceAddress);
+        bytes memory expectedError = abi.encodeWithSelector(IHorizonStakingMain.HorizonStakingInvalidZeroTokens.selector);
+        vm.expectRevert(expectedError);
+        staking.addToDelegationPool(users.indexer, subgraphDataServiceAddress, 0);
+    }
+}