From f0564100e0175014504a6d3164446f9bdd748d61 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Migone?= Date: Fri, 13 Dec 2024 17:19:04 -0300 Subject: [PATCH] fix: couple minor improvments to PaymentsEscrow and some new tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Tomás Migone --- .../contracts/interfaces/IGraphPayments.sol | 4 +- .../contracts/interfaces/IPaymentsEscrow.sol | 18 ++++++++- .../contracts/payments/PaymentsEscrow.sol | 30 +++++++------- .../horizon/test/escrow/GraphEscrow.t.sol | 30 +++++++++++++- packages/horizon/test/escrow/collect.t.sol | 40 +++++++++++++++++++ packages/horizon/test/escrow/deposit.t.sol | 25 ++++++++++-- packages/horizon/test/escrow/thaw.t.sol | 22 ++++++++-- packages/horizon/test/escrow/withdraw.t.sol | 37 +++++++++++++++-- .../PaymentsEscrowShared.t.sol | 12 ++++++ 9 files changed, 188 insertions(+), 30 deletions(-) diff --git a/packages/horizon/contracts/interfaces/IGraphPayments.sol b/packages/horizon/contracts/interfaces/IGraphPayments.sol index ebb27a71a..f62828868 100644 --- a/packages/horizon/contracts/interfaces/IGraphPayments.sol +++ b/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, diff --git a/packages/horizon/contracts/interfaces/IPaymentsEscrow.sol b/packages/horizon/contracts/interfaces/IPaymentsEscrow.sol index 1c44c4e42..da05d8c9f 100644 --- a/packages/horizon/contracts/interfaces/IPaymentsEscrow.sol +++ b/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 -- diff --git a/packages/horizon/contracts/payments/PaymentsEscrow.sol b/packages/horizon/contracts/payments/PaymentsEscrow.sol index 5e2408365..c1de3bce1 100644 --- a/packages/horizon/contracts/payments/PaymentsEscrow.sol +++ b/packages/horizon/contracts/payments/PaymentsEscrow.sol @@ -23,10 +23,6 @@ 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; @@ -34,6 +30,14 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory, /// @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,18 +147,19 @@ 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); } /** @@ -162,10 +167,7 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory, */ 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; } /** diff --git a/packages/horizon/test/escrow/GraphEscrow.t.sol b/packages/horizon/test/escrow/GraphEscrow.t.sol index 6bb56c97c..ee41a7909 100644 --- a/packages/horizon/test/escrow/GraphEscrow.t.sol +++ b/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,7 +64,7 @@ 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); @@ -70,6 +72,30 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest { 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 diff --git a/packages/horizon/test/escrow/collect.t.sol b/packages/horizon/test/escrow/collect.t.sol index 3c18c93b6..f4357d213 100644 --- a/packages/horizon/test/escrow/collect.t.sol +++ b/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 + // ); + } } diff --git a/packages/horizon/test/escrow/deposit.t.sol b/packages/horizon/test/escrow/deposit.t.sol index 002d2b1a4..fba3b7c0d 100644 --- a/packages/horizon/test/escrow/deposit.t.sol +++ b/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); } -} \ No newline at end of file + + 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); + } +} diff --git a/packages/horizon/test/escrow/thaw.t.sol b/packages/horizon/test/escrow/thaw.t.sol index e8e915b22..f7c23371b 100644 --- a/packages/horizon/test/escrow/thaw.t.sol +++ b/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); } diff --git a/packages/horizon/test/escrow/withdraw.t.sol b/packages/horizon/test/escrow/withdraw.t.sol index a141d039b..db7001fd8 100644 --- a/packages/horizon/test/escrow/withdraw.t.sol +++ b/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); + } } \ No newline at end of file diff --git a/packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol b/packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol index c3714dfd6..28b0327ae 100644 --- a/packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol +++ b/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); + } }