fix: prevent GraphTallyCollector from accepting non QueryFee payments

Signed-off-by: Tomás Migone <tomas@edgeandnode.com>

Author: tmigone

packages/horizon/contracts/interfaces/IGraphTallyCollector.sol

@@ -99,6 +99,12 @@ interface IGraphTallyCollector is IPaymentsCollector {
      */
     error GraphTallyCollectorInvalidTokensToCollectAmount(uint256 tokensToCollect, uint256 maxTokensToCollect);
 
+    /**
+     * @notice Thrown when the payment type is invalid
+     * @param paymentType The payment type
+     */
+    error GraphTallyCollectorInvalidPaymentType(IGraphPayments.PaymentTypes paymentType);
+
     /**
      * @notice See {IPaymentsCollector.collect}
      * This variant adds the ability to partially collect a RAV by specifying the amount of tokens to collect.

packages/horizon/contracts/payments/collectors/GraphTallyCollector.sol

@@ -102,6 +102,8 @@ contract GraphTallyCollector is EIP712, GraphDirectory, Authorizable, IGraphTall
         bytes calldata _data,
         uint256 _tokensToCollect
     ) private returns (uint256) {
+        require(_paymentType == IGraphPayments.PaymentTypes.QueryFee, GraphTallyCollectorInvalidPaymentType(_paymentType));
+
         (SignedRAV memory signedRAV, uint256 dataServiceCut, address receiverDestination) = abi.decode(
             _data,
             (SignedRAV, uint256, address)

packages/horizon/test/unit/payments/graph-tally-collector/collect/collect.t.sol

@@ -483,4 +483,31 @@ contract GraphTallyCollectTest is GraphTallyTest {
         bytes memory allocation0Data = _getQueryFeeEncodedData(signerPrivateKey, collectTestParams[0]);
         _collect(IGraphPayments.PaymentTypes.QueryFee, allocation0Data);
     }
+
+    function testGraphTally_Collect_RevertWhen_IncorrectPaymentType(
+        uint256 tokens
+    ) public useIndexer useProvisionDataService(users.verifier, 100, 0, 0) useGateway useSigner {
+        tokens = bound(tokens, 1, type(uint128).max);
+
+        _depositTokens(address(graphTallyCollector), users.indexer, tokens);
+
+        CollectTestParams memory params = CollectTestParams({
+            tokens: tokens,
+            allocationId: _allocationId,
+            payer: users.gateway,
+            indexer: users.indexer,
+            collector: users.verifier
+        });
+
+        bytes memory data = _getQueryFeeEncodedData(signerPrivateKey, params);
+
+        resetPrank(users.verifier);
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IGraphTallyCollector.GraphTallyCollectorInvalidPaymentType.selector,
+                IGraphPayments.PaymentTypes.IndexingRewards
+            )
+        );
+        graphTallyCollector.collect(IGraphPayments.PaymentTypes.IndexingRewards, data);
+    }
 }

packages/subgraph-service/test/unit/subgraphService/collect/indexing/indexing.t.sol

@@ -171,4 +171,14 @@ contract SubgraphServiceCollectIndexingTest is SubgraphServiceTest {
         );
         subgraphService.collect(newIndexer, paymentType, data);
     }
+
+    function test_SubgraphService_Collect_Indexing_RevertWhen_IncorrectPaymentType(uint256 tokens) public useIndexer useAllocation(tokens) {
+        bytes memory data = abi.encode(allocationID, bytes32("POI"), _getHardcodedPOIMetadata());
+
+        // skip time to ensure allocation gets rewards
+        vm.roll(block.number + EPOCH_LENGTH);
+
+        vm.expectRevert();
+        subgraphService.collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data);
+    }
 }

packages/subgraph-service/test/unit/subgraphService/collect/query/query.t.sol

@@ -11,6 +11,7 @@ import { MessageHashUtils } from "@openzeppelin/contracts/utils/cryptography/Mes
 
 import { ISubgraphService } from "../../../../../contracts/interfaces/ISubgraphService.sol";
 import { SubgraphServiceTest } from "../../SubgraphService.t.sol";
+import { Allocation } from "../../../../../contracts/libraries/Allocation.sol";
 
 contract SubgraphServiceRegisterTest is SubgraphServiceTest {
     using PPMMath for uint128;
@@ -279,4 +280,50 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
         );
         assertEq(afterTokensCollected, intermediateTokensCollected + tokensToCollect + (oddTokensPayment ? 1 : 0));
     }
+
+    function testCollect_QueryFees_ClosedAllocation(
+        uint256 tokensAllocated,
+        uint256 tokensPayment
+    ) public useIndexer useAllocation(tokensAllocated) {
+        vm.assume(tokensAllocated > minimumProvisionTokens * stakeToFeesRatio);
+        uint256 maxTokensPayment = tokensAllocated / stakeToFeesRatio > type(uint128).max
+            ? type(uint128).max
+            : tokensAllocated / stakeToFeesRatio;
+        tokensPayment = bound(tokensPayment, minimumProvisionTokens, maxTokensPayment);
+
+        resetPrank(users.gateway);
+        _deposit(tokensPayment);
+        _authorizeSigner();
+
+        // Close the allocation
+        resetPrank(users.indexer);
+        bytes memory closeAlloData = abi.encode(allocationID);
+        _stopService(users.indexer, closeAlloData);
+
+        // Collect the RAV
+        resetPrank(users.indexer);
+        bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokensPayment), 0);
+        _collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data);
+    }
+
+    function testCollect_QueryFees_RevertWhen_IncorrectPaymentType(
+        uint256 tokensAllocated,
+        uint256 tokensPayment
+    ) public useIndexer useAllocation(tokensAllocated) {
+        vm.assume(tokensAllocated > minimumProvisionTokens * stakeToFeesRatio);
+        uint256 maxTokensPayment = tokensAllocated / stakeToFeesRatio > type(uint128).max
+            ? type(uint128).max
+            : tokensAllocated / stakeToFeesRatio;
+        tokensPayment = bound(tokensPayment, minimumProvisionTokens, maxTokensPayment);
+
+        resetPrank(users.gateway);
+        _deposit(tokensPayment);
+        _authorizeSigner();
+
+        resetPrank(users.indexer);
+        bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokensPayment), 0);
+
+        vm.expectRevert();
+        subgraphService.collect(users.indexer, IGraphPayments.PaymentTypes.IndexingRewards, data);
+    }
 }