fix: refactor payments cut distribution (TRST-L12)

Signed-off-by: Tomás Migone <[email protected]>

packages/horizon/contracts/interfaces/IGraphPayments.sol

@@ -23,48 +23,57 @@ interface IGraphPayments {
      * @param payer The address of the payer
      * @param receiver The address of the receiver
      * @param dataService The address of the data service
-     * @param tokensReceiver Amount of tokens for the receiver
-     * @param tokensDelegationPool Amount of tokens for delegators
-     * @param tokensDataService Amount of tokens for the data service
+     * @param tokens The total amount of tokens being collected
      * @param tokensProtocol Amount of tokens charged as protocol tax
+     * @param tokensDataService Amount of tokens for the data service
+     * @param tokensDelegationPool Amount of tokens for delegators
+     * @param tokensReceiver Amount of tokens for the receiver
      */
-    event PaymentCollected(
+    event GraphPaymentCollected(
         address indexed payer,
         address indexed receiver,
         address indexed dataService,
-        uint256 tokensReceiver,
-        uint256 tokensDelegationPool,
+        uint256 tokens,
+        uint256 tokensProtocol,
         uint256 tokensDataService,
-        uint256 tokensProtocol
+        uint256 tokensDelegationPool,
+        uint256 tokensReceiver
     );
 
     /**
-     * @notice Thrown when there are insufficient tokens to pay the required amount
-     * @param tokens The amount of tokens available
-     * @param minTokens The amount of tokens being collected
+     * @notice Thrown when the calculated amount of tokens to be paid out to all parties is
+     * not the same as the amount of tokens being collected
+     * @param tokens The amount of tokens being collected
+     * @param tokensCalculated The sum of all the tokens to be paid out
      */
-    error GraphPaymentsInsufficientTokens(uint256 tokens, uint256 minTokens);
+    error GraphPaymentsBadAccounting(uint256 tokens, uint256 tokensCalculated);
 
     /**
      * @notice Thrown when the protocol payment cut is invalid
      * @param protocolPaymentCut The protocol payment cut
      */
     error GraphPaymentsInvalidProtocolPaymentCut(uint256 protocolPaymentCut);
 
+    /**
+     * @notice Thrown when trying to use a cut that is not expressed in PPM
+     * @param cut The cut
+     */
+    error GraphPaymentsInvalidCut(uint256 cut);
+
     /**
      * @notice Collects funds from a payer.
      * It will pay cuts to all relevant parties and forward the rest to the receiver.
      * @param paymentType The type of payment as defined in {IGraphPayments}
      * @param receiver The address of the receiver
      * @param tokens The amount of tokens being collected
      * @param dataService The address of the data service
-     * @param tokensDataService The amount of tokens that should be sent to the data service
+     * @param dataServiceCut The data service cut in PPM
      */
     function collect(
         PaymentTypes paymentType,
         address receiver,
         uint256 tokens,
         address dataService,
-        uint256 tokensDataService
+        uint256 dataServiceCut
     ) external;
 }

packages/horizon/contracts/interfaces/IPaymentsCollector.sol

@@ -19,17 +19,15 @@ interface IPaymentsCollector {
      * @param paymentType The payment type collected as defined by {IGraphPayments}
      * @param payer The address of the payer
      * @param receiver The address of the receiver
-     * @param tokensReceiver The amount of tokens received by the receiver
      * @param dataService The address of the data service
-     * @param tokensDataService The amount of tokens received by the data service
+     * @param tokens The amount of tokens being collected
      */
     event PaymentCollected(
         IGraphPayments.PaymentTypes indexed paymentType,
         address indexed payer,
         address receiver,
-        uint256 tokensReceiver,
         address indexed dataService,
-        uint256 tokensDataService
+        uint256 tokens
     );
 
     /**

packages/horizon/contracts/interfaces/IPaymentsEscrow.sol

@@ -186,15 +186,15 @@ interface IPaymentsEscrow {
      * @param receiver The address of the receiver
      * @param tokens The amount of tokens to collect
      * @param dataService The address of the data service
-     * @param tokensDataService The amount of tokens that {GraphPayments} should send to the data service
+     * @param dataServiceCut The data service cut in PPM that {GraphPayments} should send
      */
     function collect(
         IGraphPayments.PaymentTypes paymentType,
         address payer,
         address receiver,
         uint256 tokens,
         address dataService,
-        uint256 tokensDataService
+        uint256 dataServiceCut
     ) external;
 
     /**

packages/horizon/contracts/payments/GraphPayments.sol

@@ -32,7 +32,7 @@ contract GraphPayments is Initializable, MulticallUpgradeable, GraphDirectory, I
      * @param protocolPaymentCut The protocol tax in PPM
      */
     constructor(address controller, uint256 protocolPaymentCut) GraphDirectory(controller) {
-        require(PPMMath.isValidPPM(protocolPaymentCut), GraphPaymentsInvalidProtocolPaymentCut(protocolPaymentCut));
+        require(PPMMath.isValidPPM(protocolPaymentCut), GraphPaymentsInvalidCut(protocolPaymentCut));
         PROTOCOL_PAYMENT_CUT = protocolPaymentCut;
         _disableInitializers();
     }
@@ -52,41 +52,54 @@ contract GraphPayments is Initializable, MulticallUpgradeable, GraphDirectory, I
         address receiver,
         uint256 tokens,
         address dataService,
-        uint256 tokensDataService
+        uint256 dataServiceCut
     ) external {
+        require(PPMMath.isValidPPM(dataServiceCut), GraphPaymentsInvalidCut(dataServiceCut));
+
+        // Pull tokens from the sender
         _graphToken().pullTokens(msg.sender, tokens);
 
-        // Calculate cuts
-        uint256 tokensProtocol = tokens.mulPPM(PROTOCOL_PAYMENT_CUT);
-        uint256 delegationFeeCut = _graphStaking().getDelegationFeeCut(receiver, dataService, paymentType);
-        uint256 tokensDelegationPool = tokens.mulPPM(delegationFeeCut);
-        uint256 totalCut = tokensProtocol + tokensDataService + tokensDelegationPool;
-        require(tokens >= totalCut, GraphPaymentsInsufficientTokens(tokens, totalCut));
+        // Calculate token amounts for each party
+        // Order matters: protocol -> data service -> delegators -> receiver
+        // Note the substractions should not underflow as we are only deducting a percentage of the remainder
+        uint256 tokensRemaining = tokens;
+
+        uint256 tokensProtocol = tokensRemaining.mulPPMRoundUp(PROTOCOL_PAYMENT_CUT);
+        tokensRemaining = tokensRemaining - tokensProtocol;
+
+        uint256 tokensDataService = tokensRemaining.mulPPMRoundUp(dataServiceCut);
+        tokensRemaining = tokensRemaining - tokensDataService;
+
+        uint256 tokensDelegationPool = tokensRemaining.mulPPMRoundUp(
+            _graphStaking().getDelegationFeeCut(receiver, dataService, paymentType)
+        );
+        tokensRemaining = tokensRemaining - tokensDelegationPool;
+
+        // Ensure accounting is correct
+        uint256 tokensTotal = tokensProtocol + tokensDataService + tokensDelegationPool + tokensRemaining;
+        require(tokens == tokensTotal, GraphPaymentsBadAccounting(tokens, tokensTotal));
 
-        // Pay protocol cut
+        // Pay all parties
         _graphToken().burnTokens(tokensProtocol);
 
-        // Pay data service cut
         _graphToken().pushTokens(dataService, tokensDataService);
 
-        // Pay delegators
         if (tokensDelegationPool > 0) {
             _graphToken().approve(address(_graphStaking()), tokensDelegationPool);
             _graphStaking().addToDelegationPool(receiver, dataService, tokensDelegationPool);
         }
 
-        // Pay receiver
-        uint256 tokensReceiverRemaining = tokens - totalCut;
-        _graphToken().pushTokens(receiver, tokensReceiverRemaining);
+        _graphToken().pushTokens(receiver, tokensRemaining);
 
-        emit PaymentCollected(
+        emit GraphPaymentCollected(
             msg.sender,
             receiver,
             dataService,
-            tokensReceiverRemaining,
-            tokensDelegationPool,
+            tokens,
+            tokensProtocol,
             tokensDataService,
-            tokensProtocol
+            tokensDelegationPool,
+            tokensRemaining
         );
     }
 }

packages/horizon/contracts/payments/PaymentsEscrow.sol

@@ -129,7 +129,7 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,
         address receiver,
         uint256 tokens,
         address dataService,
-        uint256 tokensDataService
+        uint256 dataServiceCut
     ) external override notPaused {
         // Check if there are enough funds in the escrow account
         EscrowAccount storage account = escrowAccounts[payer][msg.sender][receiver];
@@ -141,7 +141,7 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,
         uint256 balanceBefore = _graphToken().balanceOf(address(this));
 
         _graphToken().approve(address(_graphPayments()), tokens);
-        _graphPayments().collect(paymentType, receiver, tokens, dataService, tokensDataService);
+        _graphPayments().collect(paymentType, receiver, tokens, dataService, dataServiceCut);
 
         uint256 balanceAfter = _graphToken().balanceOf(address(this));
         require(

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

@@ -214,21 +214,12 @@ contract TAPCollector is EIP712, GraphDirectory, ITAPCollector {
             }
         }
 
-        uint256 tokensDataService = tokensToCollect.mulPPM(dataServiceCut);
-
         if (tokensToCollect > 0) {
             tokensCollected[dataService][receiver][payer] += tokensToCollect;
-            _graphPaymentsEscrow().collect(
-                _paymentType,
-                payer,
-                receiver,
-                tokensToCollect,
-                dataService,
-                tokensDataService
-            );
+            _graphPaymentsEscrow().collect(_paymentType, payer, receiver, tokensToCollect, dataService, dataServiceCut);
         }
 
-        emit PaymentCollected(_paymentType, payer, receiver, tokensToCollect, dataService, tokensDataService);
+        emit PaymentCollected(_paymentType, payer, receiver, dataService, tokensToCollect);
         emit RAVCollected(
             payer,
             dataService,

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

@@ -7,8 +7,11 @@ import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
 
 import { HorizonStakingSharedTest } from "../shared/horizon-staking/HorizonStakingShared.t.sol";
 import { PaymentsEscrowSharedTest } from "../shared/payments-escrow/PaymentsEscrowShared.t.sol";
+import { PPMMath } from "../../contracts/libraries/PPMMath.sol";
 
 contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
+    using PPMMath for uint256;
+
     /*
      * MODIFIERS
      */
@@ -59,6 +62,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
         uint256 receiverBalance;
         uint256 delegationPoolBalance;
         uint256 dataServiceBalance;
+        uint256 payerEscrowBalance;
     }
 
     function _collectEscrow(
@@ -67,53 +71,62 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
         address _receiver,
         uint256 _tokens,
         address _dataService,
-        uint256 _tokensDataService
+        uint256 _dataServiceCut
     ) internal {
         (, address _collector, ) = vm.readCallers();
 
         // Previous balances
-        (uint256 previousPayerEscrowBalance, , ) = escrow.escrowAccounts(_payer, _collector, _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)
+            dataServiceBalance: token.balanceOf(_dataService),
+            payerEscrowBalance: 0
         });
 
+        {
+            (uint256 payerEscrowBalance, , ) = escrow.escrowAccounts(_payer, _collector, _receiver);
+            previousBalances.payerEscrowBalance = payerEscrowBalance;
+        }
+
         vm.expectEmit(address(escrow));
         emit IPaymentsEscrow.EscrowCollected(_payer, _collector, _receiver, _tokens);
-        escrow.collect(_paymentType, _payer, _receiver, _tokens, _dataService, _tokensDataService);
+        escrow.collect(_paymentType, _payer, _receiver, _tokens, _dataService, _dataServiceCut);
 
         // Calculate cuts
-        uint256 protocolPaymentCut = payments.PROTOCOL_PAYMENT_CUT();
-        uint256 delegatorCut = staking.getDelegationFeeCut(_receiver, _dataService, _paymentType);
+        // this is nasty but stack is indeed too deep
+        uint256 tokensDataService = (_tokens - _tokens.mulPPMRoundUp(payments.PROTOCOL_PAYMENT_CUT())).mulPPMRoundUp(
+            _dataServiceCut
+        );
+        uint256 tokensDelegation = (_tokens -
+            _tokens.mulPPMRoundUp(payments.PROTOCOL_PAYMENT_CUT()) -
+            tokensDataService).mulPPMRoundUp(staking.getDelegationFeeCut(_receiver, _dataService, _paymentType));
+        uint256 receiverExpectedPayment = _tokens -
+            _tokens.mulPPMRoundUp(payments.PROTOCOL_PAYMENT_CUT()) -
+            tokensDataService -
+            tokensDelegation;
 
         // After balances
-        (uint256 afterPayerEscrowBalance, , ) = escrow.escrowAccounts(_payer, _collector, _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)
+            dataServiceBalance: token.balanceOf(_dataService),
+            payerEscrowBalance: 0
         });
+        {
+            (uint256 afterPayerEscrowBalance, , ) = escrow.escrowAccounts(_payer, _collector, _receiver);
+            afterBalances.payerEscrowBalance = afterPayerEscrowBalance;
+        }
 
         // Check receiver balance after payment
-        uint256 receiverExpectedPayment = _tokens -
-            _tokensDataService -
-            (_tokens * protocolPaymentCut) /
-            MAX_PPM -
-            (_tokens * delegatorCut) /
-            MAX_PPM;
         assertEq(afterBalances.receiverBalance - previousBalances.receiverBalance, receiverExpectedPayment);
         assertEq(token.balanceOf(address(payments)), 0);
 
         // Check delegation pool balance after payment
-        assertEq(
-            afterBalances.delegationPoolBalance - previousBalances.delegationPoolBalance,
-            (_tokens * delegatorCut) / MAX_PPM
-        );
+        assertEq(afterBalances.delegationPoolBalance - previousBalances.delegationPoolBalance, tokensDelegation);
 
         // Check that the escrow account has been updated
         assertEq(previousBalances.escrowBalance, afterBalances.escrowBalance + _tokens);
@@ -122,9 +135,9 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
         assertEq(previousBalances.paymentsBalance, afterBalances.paymentsBalance);
 
         // Check data service balance after payment
-        assertEq(afterBalances.dataServiceBalance - previousBalances.dataServiceBalance, _tokensDataService);
+        assertEq(afterBalances.dataServiceBalance - previousBalances.dataServiceBalance, tokensDataService);
 
         // Check payers escrow balance after payment
-        assertEq(previousPayerEscrowBalance - _tokens, afterPayerEscrowBalance);
+        assertEq(previousBalances.payerEscrowBalance - _tokens, afterBalances.payerEscrowBalance);
     }
 }

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

@@ -16,18 +16,16 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
     function testCollect_Tokens(
         uint256 tokens,
         uint256 delegationTokens,
-        uint256 tokensDataService
+        uint256 dataServiceCut
     )
         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);
-
+        dataServiceCut = bound(dataServiceCut, 0, MAX_PPM);
         delegationTokens = bound(delegationTokens, 1, MAX_STAKING_TOKENS);
+
         resetPrank(users.delegator);
         _delegate(users.indexer, subgraphDataServiceAddress, delegationTokens, 0);
 
@@ -41,7 +39,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
             users.indexer,
             tokens,
             subgraphDataServiceAddress,
-            tokensDataService
+            dataServiceCut
         );
     }