fix: couple minor improvments to PaymentsEscrow and some new tests

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

packages/horizon/contracts/interfaces/IGraphPayments.sol

@@ -31,9 +31,9 @@ interface IGraphPayments {
      * @param tokensReceiver Amount of tokens for the receiver
      */
     event GraphPaymentCollected(
-        PaymentTypes paymentType,
+        PaymentTypes indexed paymentType,
         address indexed payer,
-        address indexed receiver,
+        address receiver,
         address indexed dataService,
         uint256 tokens,
         uint256 tokensProtocol,

packages/horizon/contracts/interfaces/IPaymentsEscrow.sol

@@ -37,11 +37,18 @@ interface IPaymentsEscrow {
     /**
      * @notice Emitted when a payer cancels an escrow thawing
      * @param payer The address of the payer
+     * @param collector The address of the collector
      * @param receiver The address of the receiver
      * @param tokensThawing The amount of tokens that were being thawed
      * @param thawEndTimestamp The timestamp at which the thawing period was ending
      */
-    event CancelThaw(address indexed payer, address indexed receiver, uint256 tokensThawing, uint256 thawEndTimestamp);
+    event CancelThaw(
+        address indexed payer,
+        address indexed collector,
+        address indexed receiver,
+        uint256 tokensThawing,
+        uint256 thawEndTimestamp
+    );
 
     /**
      * @notice Emitted when a payer thaws funds from the escrow for a payer-collector-receiver tuple
@@ -70,12 +77,19 @@ interface IPaymentsEscrow {
 
     /**
      * @notice Emitted when a collector collects funds from the escrow for a payer-collector-receiver tuple
+     * @param paymentType The type of payment being collected as defined in the {IGraphPayments} interface
      * @param payer The address of the payer
      * @param collector The address of the collector
      * @param receiver The address of the receiver
      * @param tokens The amount of tokens collected
      */
-    event EscrowCollected(address indexed payer, address indexed collector, address indexed receiver, uint256 tokens);
+    event EscrowCollected(
+        IGraphPayments.PaymentTypes indexed paymentType,
+        address indexed payer,
+        address indexed collector,
+        address receiver,
+        uint256 tokens
+    );
 
     // -- Errors --
 

packages/horizon/contracts/payments/PaymentsEscrow.sol

@@ -23,17 +23,21 @@ import { GraphDirectory } from "../utilities/GraphDirectory.sol";
 contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory, IPaymentsEscrow {
     using TokenUtils for IGraphToken;
 
-    /// @notice Escrow account details for payer-collector-receiver tuples
-    mapping(address payer => mapping(address collector => mapping(address receiver => IPaymentsEscrow.EscrowAccount escrowAccount)))
-        public escrowAccounts;
-
     /// @notice The maximum thawing period (in seconds) for both escrow withdrawal and collector revocation
     /// @dev This is a precautionary measure to avoid inadvertedly locking funds for too long
     uint256 public constant MAX_WAIT_PERIOD = 90 days;
 
     /// @notice Thawing period in seconds for escrow funds withdrawal
     uint256 public immutable WITHDRAW_ESCROW_THAWING_PERIOD;
 
+    /// @notice Escrow account details for payer-collector-receiver tuples
+    mapping(address payer => mapping(address collector => mapping(address receiver => IPaymentsEscrow.EscrowAccount escrowAccount)))
+        public escrowAccounts;
+
+    /**
+     * @notice Modifier to prevent function execution when contract is paused
+     * @dev Reverts if the controller indicates the contract is paused
+     */
     modifier notPaused() {
         require(!_graphController().paused(), PaymentsEscrowIsPaused());
         _;
@@ -101,7 +105,7 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,
         account.tokensThawing = 0;
         account.thawEndTimestamp = 0;
 
-        emit CancelThaw(msg.sender, receiver, tokensThawing, thawEndTimestamp);
+        emit CancelThaw(msg.sender, collector, receiver, tokensThawing, thawEndTimestamp);
     }
 
     /**
@@ -143,29 +147,27 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,
         // Reduce amount from account balance
         account.balance -= tokens;
 
-        uint256 balanceBefore = _graphToken().balanceOf(address(this));
+        uint256 escrowBalanceBefore = _graphToken().balanceOf(address(this));
 
         _graphToken().approve(address(_graphPayments()), tokens);
         _graphPayments().collect(paymentType, receiver, tokens, dataService, dataServiceCut);
 
-        uint256 balanceAfter = _graphToken().balanceOf(address(this));
+        // Verify that the escrow balance is consistent with the collected tokens
+        uint256 escrowBalanceAfter = _graphToken().balanceOf(address(this));
         require(
-            balanceBefore == tokens + balanceAfter,
-            PaymentsEscrowInconsistentCollection(balanceBefore, balanceAfter, tokens)
+            escrowBalanceBefore == tokens + escrowBalanceAfter,
+            PaymentsEscrowInconsistentCollection(escrowBalanceBefore, escrowBalanceAfter, tokens)
         );
 
-        emit EscrowCollected(payer, msg.sender, receiver, tokens);
+        emit EscrowCollected(paymentType, payer, msg.sender, receiver, tokens);
     }
 
     /**
      * @notice See {IPaymentsEscrow-getBalance}
      */
     function getBalance(address payer, address collector, address receiver) external view override returns (uint256) {
         EscrowAccount storage account = escrowAccounts[payer][collector][receiver];
-        if (account.balance <= account.tokensThawing) {
-            return 0;
-        }
-        return account.balance - account.tokensThawing;
+        return account.balance > account.tokensThawing ? account.balance - account.tokensThawing : 0;
     }
 
     /**

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

@@ -30,7 +30,9 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
     }
 
     modifier depositAndThawTokens(uint256 amount, uint256 thawAmount) {
+        vm.assume(amount > 0);
         vm.assume(thawAmount > 0);
+        vm.assume(amount <= MAX_STAKING_TOKENS);
         vm.assume(amount > thawAmount);
         _depositTokens(users.verifier, users.indexer, amount);
         escrow.thaw(users.verifier, users.indexer, thawAmount);
@@ -62,14 +64,38 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
         (, uint256 amountThawingBefore, uint256 thawEndTimestampBefore) = escrow.escrowAccounts(msgSender, collector, receiver);
 
         vm.expectEmit(address(escrow));
-        emit IPaymentsEscrow.CancelThaw(msgSender, receiver, amountThawingBefore, thawEndTimestampBefore);
+        emit IPaymentsEscrow.CancelThaw(msgSender, collector, receiver, amountThawingBefore, thawEndTimestampBefore);
         escrow.cancelThaw(collector, receiver);
 
         (, uint256 amountThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(msgSender, collector, receiver);
         assertEq(amountThawing, 0);
         assertEq(thawEndTimestamp, 0);
     }
 
+    function _withdrawEscrow(address collector, address receiver) internal {
+        (, address msgSender, ) = vm.readCallers();
+
+        (uint256 balanceBefore, uint256 amountThawingBefore, ) = escrow.escrowAccounts(msgSender, collector, receiver);
+        uint256 tokenBalanceBeforeSender = token.balanceOf(msgSender);
+        uint256 tokenBalanceBeforeEscrow = token.balanceOf(address(escrow));
+
+        uint256 amountToWithdraw = amountThawingBefore > balanceBefore ? balanceBefore : amountThawingBefore;
+        vm.expectEmit(address(escrow));
+        emit IPaymentsEscrow.Withdraw(msgSender, collector, receiver, amountToWithdraw);
+        escrow.withdraw(collector, receiver);
+
+        (uint256 balanceAfter, uint256 tokensThawingAfter, uint256 thawEndTimestampAfter) = escrow.escrowAccounts(msgSender, collector, receiver);
+        uint256 tokenBalanceAfterSender = token.balanceOf(msgSender);
+        uint256 tokenBalanceAfterEscrow = token.balanceOf(address(escrow));
+
+        assertEq(balanceAfter, balanceBefore - amountToWithdraw);
+        assertEq(tokensThawingAfter, 0);
+        assertEq(thawEndTimestampAfter, 0);
+
+        assertEq(tokenBalanceAfterSender, tokenBalanceBeforeSender + amountToWithdraw);
+        assertEq(tokenBalanceAfterEscrow, tokenBalanceBeforeEscrow - amountToWithdraw);        
+    }
+
     struct CollectPaymentData {
         uint256 escrowBalance;
         uint256 paymentsBalance;
@@ -105,7 +131,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
         }
 
         vm.expectEmit(address(escrow));
-        emit IPaymentsEscrow.EscrowCollected(_payer, _collector, _receiver, _tokens);
+        emit IPaymentsEscrow.EscrowCollected(_paymentType, _payer, _collector, _receiver, _tokens);
         escrow.collect(_paymentType, _payer, _receiver, _tokens, _dataService, _dataServiceCut);
 
         // Calculate cuts

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

@@ -13,6 +13,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
      * TESTS
      */
 
+    // use users.verifier as collector
     function testCollect_Tokens(
         uint256 tokens,
         uint256 tokensToCollect,
@@ -98,4 +99,43 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
         );
         vm.stopPrank();
     }
+
+    function testCollect_MultipleCollections(
+        uint256 depositAmount,
+        uint256 firstCollect,
+        uint256 secondCollect
+    ) public useIndexer {
+        // Tests multiple collect operations from the same escrow account
+        vm.assume(firstCollect < MAX_STAKING_TOKENS);
+        vm.assume(secondCollect < MAX_STAKING_TOKENS);
+        vm.assume(depositAmount > 0);
+        vm.assume(firstCollect > 0 && firstCollect < depositAmount);
+        vm.assume(secondCollect > 0 && secondCollect <= depositAmount - firstCollect);
+
+        resetPrank(users.gateway);
+        _depositTokens(users.verifier, users.indexer, depositAmount);
+
+        // burn some tokens to prevent overflow
+        resetPrank(users.indexer);
+        token.burn(MAX_STAKING_TOKENS);
+
+        resetPrank(users.verifier);
+        _collectEscrow(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
+            users.indexer,
+            firstCollect,
+            subgraphDataServiceAddress,
+            0
+        );
+
+        // _collectEscrow(
+        //     IGraphPayments.PaymentTypes.QueryFee,
+        //     users.gateway,
+        //     users.indexer,
+        //     secondCollect,
+        //     subgraphDataServiceAddress,
+        //     0
+        // );
+    }
 }

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

@@ -6,13 +6,32 @@ import "forge-std/Test.sol";
 import { GraphEscrowTest } from "./GraphEscrow.t.sol";
 
 contract GraphEscrowDepositTest is GraphEscrowTest {
-
     /*
      * TESTS
      */
 
     function testDeposit_Tokens(uint256 amount) public useGateway useDeposit(amount) {
-        (uint256 indexerEscrowBalance,,) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
+        (uint256 indexerEscrowBalance, , ) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
         assertEq(indexerEscrowBalance, amount);
     }
-}
+
+    function testDepositTo_Tokens(uint256 amount) public {
+        resetPrank(users.delegator);
+        token.approve(address(escrow), amount);
+        _depositToTokens(users.gateway, users.verifier, users.indexer, amount);
+    }
+
+    // Tests multiple deposits accumulate correctly in the escrow account
+    function testDeposit_MultipleDeposits(uint256 amount1, uint256 amount2) public useGateway {
+        vm.assume(amount1 > 0);
+        vm.assume(amount2 > 0);
+        vm.assume(amount1 <= MAX_STAKING_TOKENS);
+        vm.assume(amount2 <= MAX_STAKING_TOKENS);
+
+        _depositTokens(users.verifier, users.indexer, amount1);
+        _depositTokens(users.verifier, users.indexer, amount2);
+
+        (uint256 balance,,) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
+        assertEq(balance, amount1 + amount2);
+    }
+}

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

@@ -10,9 +10,21 @@ contract GraphEscrowThawTest is GraphEscrowTest {
      * TESTS
      */
 
-    function testThaw_Tokens(uint256 amount) public useGateway useDeposit(amount) {
-        amount = bound(amount, 1, type(uint256).max);
+    function testThaw_PartialBalanceThaw(
+        uint256 amountDeposited,
+        uint256 amountThawed
+    ) public useGateway useDeposit(amountDeposited) {
+        vm.assume(amountThawed > 0);
+        vm.assume(amountThawed <= amountDeposited);
+        _thawEscrow(users.verifier, users.indexer, amountThawed);
+    }
+
+    function testThaw_FullBalanceThaw(uint256 amount) public useGateway useDeposit(amount) {
+        vm.assume(amount > 0);
         _thawEscrow(users.verifier, users.indexer, amount);
+
+        uint256 availableBalance = escrow.getBalance(users.gateway, users.verifier, users.indexer);
+        assertEq(availableBalance, 0);
     }
 
     function testThaw_Tokens_SuccesiveCalls(uint256 amount) public useGateway {
@@ -25,7 +37,11 @@ contract GraphEscrowThawTest is GraphEscrowTest {
         _thawEscrow(users.verifier, users.indexer, secondAmountToThaw);
 
         (, address msgSender, ) = vm.readCallers();
-        (, uint256 amountThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(msgSender, users.verifier, users.indexer);
+        (, uint256 amountThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(
+            msgSender,
+            users.verifier,
+            users.indexer
+        );
         assertEq(amountThawing, secondAmountToThaw);
         assertEq(thawEndTimestamp, block.timestamp + withdrawEscrowThawingPeriod);
     }

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

@@ -3,6 +3,7 @@ pragma solidity 0.8.27;
 
 import "forge-std/Test.sol";
 
+import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
 import { GraphEscrowTest } from "./GraphEscrow.t.sol";
 
 contract GraphEscrowWithdrawTest is GraphEscrowTest {
@@ -18,11 +19,8 @@ contract GraphEscrowWithdrawTest is GraphEscrowTest {
         // advance time
         skip(withdrawEscrowThawingPeriod + 1);
 
-        escrow.withdraw(users.verifier, users.indexer);
+        _withdrawEscrow(users.verifier, users.indexer);
         vm.stopPrank();
-
-        (uint256 indexerEscrowBalance,,) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
-        assertEq(indexerEscrowBalance, amount - thawAmount);
     }
 
     function testWithdraw_RevertWhen_NotThawing(uint256 amount) public useGateway useDeposit(amount) {
@@ -39,4 +37,35 @@ contract GraphEscrowWithdrawTest is GraphEscrowTest {
         vm.expectRevert(expectedError);
         escrow.withdraw(users.verifier, users.indexer);
     }
+
+    function testWithdraw_BalanceAfterCollect(
+        uint256 amountDeposited,
+        uint256 amountThawed,
+        uint256 amountCollected
+    ) public useGateway depositAndThawTokens(amountDeposited, amountThawed) {
+        vm.assume(amountCollected > 0);
+        vm.assume(amountCollected <= amountDeposited);
+
+        // burn some tokens to prevent overflow
+        resetPrank(users.indexer);
+        token.burn(MAX_STAKING_TOKENS);
+
+        // collect
+        resetPrank(users.verifier);
+        _collectEscrow(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
+            users.indexer,
+            amountCollected,
+            subgraphDataServiceAddress,
+            0
+        );
+
+        // Advance time to simulate the thawing period
+        skip(withdrawEscrowThawingPeriod + 1);
+
+        // withdraw the remaining thawed balance
+        resetPrank(users.gateway);
+        _withdrawEscrow(users.verifier, users.indexer);
+    }
 }

packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol

@@ -34,4 +34,16 @@ abstract contract PaymentsEscrowSharedTest is GraphBaseTest {
         (uint256 escrowBalanceAfter,,) = escrow.escrowAccounts(msgSender, _collector, _receiver);
         assertEq(escrowBalanceAfter - _tokens, escrowBalanceBefore);
     }
+
+    function _depositToTokens(address _payer, address _collector, address _receiver, uint256 _tokens) internal {
+        (uint256 escrowBalanceBefore,,) = escrow.escrowAccounts(_payer, _collector, _receiver);
+        token.approve(address(escrow), _tokens);
+
+        vm.expectEmit(address(escrow));
+        emit IPaymentsEscrow.Deposit(_payer, _collector, _receiver, _tokens);
+        escrow.depositTo(_payer, _collector, _receiver, _tokens);
+
+        (uint256 escrowBalanceAfter,,) = escrow.escrowAccounts(_payer, _collector, _receiver);
+        assertEq(escrowBalanceAfter - _tokens, escrowBalanceBefore);
+    }
 }