fix(Horizon): only allow addToDelegationPool if provision exists and pool has shares

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

@@ -402,6 +402,13 @@ interface IHorizonStakingMain {
      */
     error HorizonStakingInvalidDelegationPoolState(address serviceProvider, address verifier);
 
+    /**
+     * @notice Thrown when attempting to operate with a delegation pool that does not exist.
+     * @param serviceProvider The service provider address
+     * @param verifier The verifier address
+     */
+    error HorizonStakingInvalidDelegationPool(address serviceProvider, address verifier);
+
     // -- Errors: thaw requests --
 
     error HorizonStakingNothingThawing();

packages/horizon/contracts/staking/HorizonStaking.sol

@@ -279,7 +279,15 @@ contract HorizonStaking is HorizonStakingBase, IHorizonStakingMain {
         uint256 tokens
     ) external override notPaused {
         require(tokens != 0, HorizonStakingInvalidZeroTokens());
+
+        // Provision must exist before adding to delegation pool
+        Provision memory prov = _provisions[serviceProvider][verifier];
+        require(prov.createdAt != 0, HorizonStakingInvalidProvision(serviceProvider, verifier));
+
+        // Delegation pool must exist before adding tokens
         DelegationPoolInternal storage pool = _getDelegationPool(serviceProvider, verifier);
+        require(pool.shares > 0, HorizonStakingInvalidDelegationPool(serviceProvider, verifier));
+
         pool.tokens = pool.tokens + tokens;
         _graphToken().pullTokens(msg.sender, tokens);
         emit TokensToDelegationPoolAdded(serviceProvider, verifier, tokens);

packages/horizon/test/escrow/collect.t.sol

@@ -3,37 +3,115 @@ pragma solidity 0.8.27;
 
 import "forge-std/Test.sol";
 
+import { IHorizonStakingMain } from "../../contracts/interfaces/internal/IHorizonStakingMain.sol";
 import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
 import { IPaymentsEscrow } from "../../contracts/interfaces/IPaymentsEscrow.sol";
 
 import { GraphEscrowTest } from "./GraphEscrow.t.sol";
 
 contract GraphEscrowCollectTest is GraphEscrowTest {
 
+    struct CollectPaymentData {
+        uint256 escrowBalance;
+        uint256 paymentsBalance;
+        uint256 receiverBalance;
+        uint256 delegationPoolBalance;
+        uint256 dataServiceBalance;
+    }
+
+    function _collect(
+        IGraphPayments.PaymentTypes _paymentType,
+        address _payer,
+        address _receiver,
+        uint256 _tokens,
+        address _dataService,
+        uint256 _tokensDataService
+    ) private {
+        // Previous balances
+        (uint256 previousPayerEscrowBalance,,) = escrow.escrowAccounts(_payer, _receiver);
+        CollectPaymentData memory previousBalances = CollectPaymentData({
+            escrowBalance: token.balanceOf(address(escrow)),
+            paymentsBalance: token.balanceOf(address(payments)),
+            receiverBalance: token.balanceOf(_receiver),
+            delegationPoolBalance: staking.getDelegatedTokensAvailable(
+                _receiver,
+                _dataService
+            ),
+            dataServiceBalance: token.balanceOf(_dataService)
+        });
+
+        vm.expectEmit(address(escrow));
+        emit IPaymentsEscrow.EscrowCollected(_payer, _receiver, _tokens);
+        escrow.collect(_paymentType, _payer, _receiver, _tokens, _dataService, _tokensDataService);
+
+        // Calculate cuts
+        uint256 protocolPaymentCut = payments.PROTOCOL_PAYMENT_CUT();
+        uint256 delegatorCut = staking.getDelegationFeeCut(
+            _receiver,
+            _dataService,
+            _paymentType
+        );
+        uint256 tokensProtocol = _tokens * protocolPaymentCut / MAX_PPM;
+        uint256 tokensDelegation = _tokens * delegatorCut / MAX_PPM;
+
+        // After balances
+        (uint256 afterPayerEscrowBalance,,) = escrow.escrowAccounts(_payer, _receiver);
+        CollectPaymentData memory afterBalances = CollectPaymentData({
+            escrowBalance: token.balanceOf(address(escrow)),
+            paymentsBalance: token.balanceOf(address(payments)),
+            receiverBalance: token.balanceOf(_receiver),
+            delegationPoolBalance: staking.getDelegatedTokensAvailable(
+                _receiver,
+                _dataService
+            ),
+            dataServiceBalance: token.balanceOf(_dataService)
+        });
+
+        // Check receiver balance after payment
+        uint256 receiverExpectedPayment = _tokens - _tokensDataService - tokensProtocol - tokensDelegation;
+        assertEq(afterBalances.receiverBalance - previousBalances.receiverBalance, receiverExpectedPayment);
+        assertEq(token.balanceOf(address(payments)), 0);
+
+        // Check delegation pool balance after payment
+        assertEq(afterBalances.delegationPoolBalance - previousBalances.delegationPoolBalance, tokensDelegation);
+
+        // Check that the escrow account has been updated
+        assertEq(previousBalances.escrowBalance, afterBalances.escrowBalance + _tokens);
+
+        // Check that payments balance didn't change
+        assertEq(previousBalances.paymentsBalance, afterBalances.paymentsBalance);
+
+        // Check data service balance after payment
+        assertEq(afterBalances.dataServiceBalance - previousBalances.dataServiceBalance, _tokensDataService);
+
+        // Check payers escrow balance after payment
+        assertEq(previousPayerEscrowBalance - _tokens, afterPayerEscrowBalance);
+    }
+
     /*
      * TESTS
      */
 
     function testCollect_Tokens(
-        uint256 amount,
+        uint256 tokens,
+        uint256 delegationTokens,
         uint256 tokensDataService
-    ) public useIndexer useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
-        uint256 tokensProtocol = amount * protocolPaymentCut / MAX_PPM;
-        uint256 tokensDelegatoion = amount * delegationFeeCut / MAX_PPM;
-        vm.assume(tokensDataService < amount - tokensProtocol - tokensDelegatoion);
-
-        vm.startPrank(users.gateway);
-        escrow.approveCollector(users.verifier, amount);
-        _depositTokens(amount);
+    ) public useIndexer useProvision(tokens, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+        uint256 tokensProtocol = tokens * protocolPaymentCut / MAX_PPM;
+        uint256 tokensDelegatoion = tokens * delegationFeeCut / MAX_PPM;
+        vm.assume(tokensDataService < tokens - tokensProtocol - tokensDelegatoion);
 
-        uint256 indexerPreviousBalance = token.balanceOf(users.indexer);
-        vm.startPrank(users.verifier);
-        escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, tokensDataService);
+        vm.assume(delegationTokens > MIN_DELEGATION);
+        vm.assume(delegationTokens <= MAX_STAKING_TOKENS);
+        resetPrank(users.delegator);
+        _delegate(users.indexer, subgraphDataServiceAddress, delegationTokens, 0);
 
-        uint256 indexerBalance = token.balanceOf(users.indexer);
-        uint256 indexerExpectedPayment = amount - tokensDataService - tokensProtocol - tokensDelegatoion;
-        assertEq(indexerBalance - indexerPreviousBalance, indexerExpectedPayment);
-        assertEq(token.balanceOf(address(payments)), 0);
+        resetPrank(users.gateway);
+        escrow.approveCollector(users.verifier, tokens);
+        _depositTokens(tokens);
+
+        resetPrank(users.verifier);
+        _collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, tokens, subgraphDataServiceAddress, tokensDataService);
     }
 
     function testCollect_RevertWhen_CollectorNotAuthorized(uint256 amount) public {
@@ -78,4 +156,41 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
         escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 0);
         vm.stopPrank();
     }
+
+    function testCollect_RevertWhen_InvalidPool(
+        uint256 amount
+    ) public useIndexer useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+        vm.assume(amount > 1 ether);
+
+        resetPrank(users.gateway);
+        escrow.approveCollector(users.verifier, amount);
+        _depositTokens(amount);
+
+        resetPrank(users.verifier);
+        vm.expectRevert(abi.encodeWithSelector(
+            IHorizonStakingMain.HorizonStakingInvalidDelegationPool.selector,
+            users.indexer,
+            subgraphDataServiceAddress
+        ));
+        escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 1);
+    }
+
+    function testCollect_RevertWhen_InvalidProvision(
+        uint256 amount
+    ) public useIndexer useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+        vm.assume(amount > 1 ether);
+        vm.assume(amount <= MAX_STAKING_TOKENS);
+
+        resetPrank(users.gateway);
+        escrow.approveCollector(users.verifier, amount);
+        _depositTokens(amount);
+
+        resetPrank(users.verifier);
+        vm.expectRevert(abi.encodeWithSelector(
+            IHorizonStakingMain.HorizonStakingInvalidProvision.selector,
+            users.indexer,
+            subgraphDataServiceAddress
+        ));
+        escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 1);
+    }
 }

packages/horizon/test/payments/GraphPayments.t.sol

@@ -3,13 +3,95 @@ pragma solidity 0.8.27;
 
 import "forge-std/Test.sol";
 
+import { IHorizonStakingMain } from "../../contracts/interfaces/internal/IHorizonStakingMain.sol";
 import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
 import { GraphPayments } from "../../contracts/payments/GraphPayments.sol";
 
 import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
 
 contract GraphPaymentsTest is HorizonStakingSharedTest {
 
+    struct CollectPaymentData {
+        uint256 escrowBalance;
+        uint256 paymentsBalance;
+        uint256 receiverBalance;
+        uint256 delegationPoolBalance;
+        uint256 dataServiceBalance;
+    }
+
+    function _collect(
+        IGraphPayments.PaymentTypes _paymentType,
+        address _receiver,
+        uint256 _tokens,
+        address _dataService,
+        uint256 _tokensDataService
+    ) private {
+        // Previous balances
+        CollectPaymentData memory previousBalances = CollectPaymentData({
+            escrowBalance: token.balanceOf(address(escrow)),
+            paymentsBalance: token.balanceOf(address(payments)),
+            receiverBalance: token.balanceOf(_receiver),
+            delegationPoolBalance: staking.getDelegatedTokensAvailable(
+                _receiver,
+                _dataService
+            ),
+            dataServiceBalance: token.balanceOf(_dataService)
+        });
+
+        // Calculate cuts
+        uint256 protocolPaymentCut = payments.PROTOCOL_PAYMENT_CUT();
+        uint256 delegatorCut = staking.getDelegationFeeCut(
+            _receiver,
+            _dataService,
+            _paymentType
+        );
+        uint256 tokensProtocol = _tokens * protocolPaymentCut / MAX_PPM;
+        uint256 tokensDelegation = _tokens * delegatorCut / MAX_PPM;
+
+        uint256 receiverExpectedPayment = _tokens - _tokensDataService - tokensProtocol - tokensDelegation;
+
+        (,address msgSender, ) = vm.readCallers();
+        vm.expectEmit(address(payments));
+        emit IGraphPayments.PaymentCollected(
+            msgSender,
+            _receiver,
+            _dataService,
+            receiverExpectedPayment,
+            tokensDelegation,
+            _tokensDataService,
+            tokensProtocol
+        );
+        payments.collect(_paymentType, _receiver, _tokens, _dataService, _tokensDataService);
+
+        // After balances
+        CollectPaymentData memory afterBalances = CollectPaymentData({
+            escrowBalance: token.balanceOf(address(escrow)),
+            paymentsBalance: token.balanceOf(address(payments)),
+            receiverBalance: token.balanceOf(_receiver),
+            delegationPoolBalance: staking.getDelegatedTokensAvailable(
+                _receiver,
+                _dataService
+            ),
+            dataServiceBalance: token.balanceOf(_dataService)
+        });
+
+        // Check receiver balance after payment
+        assertEq(afterBalances.receiverBalance - previousBalances.receiverBalance, receiverExpectedPayment);
+        assertEq(token.balanceOf(address(payments)), 0);
+
+        // Check delegation pool balance after payment
+        assertEq(afterBalances.delegationPoolBalance - previousBalances.delegationPoolBalance, tokensDelegation);
+
+        // Check that the escrow account has been updated
+        assertEq(previousBalances.escrowBalance, afterBalances.escrowBalance + _tokens);
+
+        // Check that payments balance didn't change
+        assertEq(previousBalances.paymentsBalance, afterBalances.paymentsBalance);
+
+        // Check data service balance after payment
+        assertEq(afterBalances.dataServiceBalance - previousBalances.dataServiceBalance, _tokensDataService);
+    }
+
     /*
      * TESTS
      */
@@ -28,32 +110,28 @@ contract GraphPaymentsTest is HorizonStakingSharedTest {
 
     function testCollect(
         uint256 amount,
-        uint256 tokensDataService
+        uint256 tokensDataService,
+        uint256 tokensDelegate
     ) public useIndexer useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
         uint256 tokensProtocol = amount * protocolPaymentCut / MAX_PPM;
-        uint256 tokensDelegatoion = amount * delegationFeeCut / MAX_PPM;
-        vm.assume(tokensDataService < amount - tokensProtocol - tokensDelegatoion);
+        uint256 tokensDelegation = amount * delegationFeeCut / MAX_PPM;
+        vm.assume(tokensDataService < amount - tokensProtocol - tokensDelegation);
         address escrowAddress = address(escrow);
 
+        // Delegate tokens
+        vm.assume(tokensDelegate > MIN_DELEGATION);
+        vm.assume(tokensDelegate <= MAX_STAKING_TOKENS);
+        vm.startPrank(users.delegator);
+        _delegate(users.indexer, subgraphDataServiceAddress, tokensDelegate, 0);
+
         // Add tokens in escrow
         mint(escrowAddress, amount);
         vm.startPrank(escrowAddress);
         approve(address(payments), amount);
 
         // Collect payments through GraphPayments
-        uint256 indexerPreviousBalance = token.balanceOf(users.indexer);
-        payments.collect(IGraphPayments.PaymentTypes.QueryFee, users.indexer, amount, subgraphDataServiceAddress, tokensDataService);
+        _collect(IGraphPayments.PaymentTypes.QueryFee, users.indexer, amount, subgraphDataServiceAddress, tokensDataService);
         vm.stopPrank();
-
-        uint256 indexerBalance = token.balanceOf(users.indexer);
-        uint256 expectedPayment = amount - tokensDataService - tokensProtocol - tokensDelegatoion;
-        assertEq(indexerBalance - indexerPreviousBalance, expectedPayment);
-
-        uint256 dataServiceBalance = token.balanceOf(subgraphDataServiceAddress);
-        assertEq(dataServiceBalance, tokensDataService);
-
-        uint256 delegatorBalance = staking.getDelegatedTokensAvailable(users.indexer, subgraphDataServiceAddress);
-        assertEq(delegatorBalance, tokensDelegatoion);
     }
 
     function testCollect_RevertWhen_InsufficientAmount(
@@ -77,4 +155,45 @@ contract GraphPaymentsTest is HorizonStakingSharedTest {
         vm.expectRevert(expectedError);
         payments.collect(IGraphPayments.PaymentTypes.QueryFee, users.indexer, amount, subgraphDataServiceAddress, tokensDataService);
     }
+
+    function testCollect_RevertWhen_InvalidPool(
+        uint256 amount
+    ) public useIndexer useProvision(amount, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+        vm.assume(amount > 1 ether);
+        address escrowAddress = address(escrow);
+
+        // Add tokens in escrow
+        mint(escrowAddress, amount);
+        vm.startPrank(escrowAddress);
+        approve(address(payments), amount);
+
+        // Collect payments through GraphPayments
+        vm.expectRevert(abi.encodeWithSelector(
+            IHorizonStakingMain.HorizonStakingInvalidDelegationPool.selector,
+            users.indexer,
+            subgraphDataServiceAddress
+        ));
+        payments.collect(IGraphPayments.PaymentTypes.QueryFee, users.indexer, amount, subgraphDataServiceAddress, 1);
+    }
+
+    function testCollect_RevertWhen_InvalidProvision(
+        uint256 amount
+    ) public useIndexer useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
+        vm.assume(amount > 1 ether);
+        vm.assume(amount <= MAX_STAKING_TOKENS);
+        address escrowAddress = address(escrow);
+
+        // Add tokens in escrow
+        mint(escrowAddress, amount);
+        vm.startPrank(escrowAddress);
+        approve(address(payments), amount);
+
+        // Collect payments through GraphPayments
+        vm.expectRevert(abi.encodeWithSelector(
+            IHorizonStakingMain.HorizonStakingInvalidProvision.selector,
+            users.indexer,
+            subgraphDataServiceAddress
+        ));
+        payments.collect(IGraphPayments.PaymentTypes.QueryFee, users.indexer, amount, subgraphDataServiceAddress, 1);
+    }
 }