fix: make getBalance return 0 if balance is less than thawing amount (TRST-L10)

tmigone · tmigone · commit 55ab1b635ac3 · 2024-11-27T14:43:48.000-03:00
Signed-off-by: Tomás Migone &lt;tomas@edgeandnode.com&gt;
diff --git a/packages/horizon/contracts/interfaces/IPaymentsEscrow.sol b/packages/horizon/contracts/interfaces/IPaymentsEscrow.sol
@@ -199,6 +199,7 @@ interface IPaymentsEscrow {
 
     /**
      * @notice Get the balance of a payer-collector-receiver tuple
+     * This function will return 0 if the current balance is less than the amount of funds being thawed.
      * @param payer The address of the payer
      * @param collector The address of the collector
      * @param receiver The address of the receiver
diff --git a/packages/horizon/contracts/payments/PaymentsEscrow.sol b/packages/horizon/contracts/payments/PaymentsEscrow.sol
@@ -157,6 +157,9 @@ 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;
     }
 
diff --git a/packages/horizon/test/escrow/GraphEscrow.t.sol b/packages/horizon/test/escrow/GraphEscrow.t.sol
@@ -2,12 +2,13 @@
 pragma solidity 0.8.27;
 
 import "forge-std/Test.sol";
+import { IPaymentsEscrow } from "../../contracts/interfaces/IPaymentsEscrow.sol";
+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";
 
 contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
-
     /*
      * MODIFIERS
      */
@@ -39,4 +40,91 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
     function _approveEscrow(uint256 tokens) internal {
         token.approve(address(escrow), tokens);
     }
-}
+
+    function _thawEscrow(address collector, address receiver, uint256 amount) internal {
+        (, address msgSender, ) = vm.readCallers();
+        uint256 expectedThawEndTimestamp = block.timestamp + withdrawEscrowThawingPeriod;
+        vm.expectEmit(address(escrow));
+        emit IPaymentsEscrow.Thaw(msgSender, collector, receiver, amount, expectedThawEndTimestamp);
+        escrow.thaw(collector, receiver, amount);
+
+        (, uint256 amountThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(msgSender, collector, receiver);
+        assertEq(amountThawing, amount);
+        assertEq(thawEndTimestamp, expectedThawEndTimestamp);
+    }
+
+    struct CollectPaymentData {
+        uint256 escrowBalance;
+        uint256 paymentsBalance;
+        uint256 receiverBalance;
+        uint256 delegationPoolBalance;
+        uint256 dataServiceBalance;
+    }
+
+    function _collectEscrow(
+        IGraphPayments.PaymentTypes _paymentType,
+        address _payer,
+        address _receiver,
+        uint256 _tokens,
+        address _dataService,
+        uint256 _tokensDataService
+    ) 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)
+        });
+
+        vm.expectEmit(address(escrow));
+        emit IPaymentsEscrow.EscrowCollected(_payer, _collector, _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);
+
+        // 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)
+        });
+
+        // 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
+        );
+
+        // 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);
+    }
+}
diff --git a/packages/horizon/test/escrow/collect.t.sol b/packages/horizon/test/escrow/collect.t.sol
@@ -5,89 +5,10 @@ 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 {
-        (, 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)
-        });
-
-        vm.expectEmit(address(escrow));
-        emit IPaymentsEscrow.EscrowCollected(_payer, _collector, _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
-        );
-
-        // 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)
-        });
-
-        // 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);
-
-        // 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
      */
@@ -96,9 +17,14 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
         uint256 tokens,
         uint256 delegationTokens,
         uint256 tokensDataService
-    ) public useIndexer useProvision(tokens, 0, 0) useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut) {
-        uint256 tokensProtocol = tokens * protocolPaymentCut / MAX_PPM;
-        uint256 tokensDelegatoion = tokens * delegationFeeCut / MAX_PPM;
+    )
+        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);
 
         delegationTokens = bound(delegationTokens, 1, MAX_STAKING_TOKENS);
@@ -109,55 +35,96 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
         _depositTokens(users.verifier, users.indexer, tokens);
 
         resetPrank(users.verifier);
-        _collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, tokens, subgraphDataServiceAddress, tokensDataService);
+        _collectEscrow(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
+            users.indexer,
+            tokens,
+            subgraphDataServiceAddress,
+            tokensDataService
+        );
     }
 
     function testCollect_RevertWhen_SenderHasInsufficientAmountInEscrow(
-        uint256 amount, 
+        uint256 amount,
         uint256 insufficientAmount
-    ) public useGateway useDeposit(insufficientAmount)  {
+    ) public useGateway useDeposit(insufficientAmount) {
         vm.assume(amount > 0);
         vm.assume(insufficientAmount < amount);
 
         vm.startPrank(users.verifier);
-        bytes memory expectedError = abi.encodeWithSignature("PaymentsEscrowInsufficientBalance(uint256,uint256)", insufficientAmount, amount);
+        bytes memory expectedError = abi.encodeWithSignature(
+            "PaymentsEscrowInsufficientBalance(uint256,uint256)",
+            insufficientAmount,
+            amount
+        );
         vm.expectRevert(expectedError);
-        escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 0);
+        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) {
+    )
+        public
+        useIndexer
+        useProvision(amount, 0, 0)
+        useDelegationFeeCut(IGraphPayments.PaymentTypes.QueryFee, delegationFeeCut)
+    {
         vm.assume(amount > 1 ether);
 
         resetPrank(users.gateway);
         _depositTokens(users.verifier, users.indexer, amount);
 
         resetPrank(users.verifier);
-        vm.expectRevert(abi.encodeWithSelector(
-            IHorizonStakingMain.HorizonStakingInvalidDelegationPool.selector,
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IHorizonStakingMain.HorizonStakingInvalidDelegationPool.selector,
+                users.indexer,
+                subgraphDataServiceAddress
+            )
+        );
+        escrow.collect(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
             users.indexer,
-            subgraphDataServiceAddress
-        ));
-        escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 1);
+            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);
         _depositTokens(users.verifier, users.indexer, amount);
 
         resetPrank(users.verifier);
-        vm.expectRevert(abi.encodeWithSelector(
-            IHorizonStakingMain.HorizonStakingInvalidProvision.selector,
+        vm.expectRevert(
+            abi.encodeWithSelector(
+                IHorizonStakingMain.HorizonStakingInvalidProvision.selector,
+                users.indexer,
+                subgraphDataServiceAddress
+            )
+        );
+        escrow.collect(
+            IGraphPayments.PaymentTypes.QueryFee,
+            users.gateway,
             users.indexer,
-            subgraphDataServiceAddress
-        ));
-        escrow.collect(IGraphPayments.PaymentTypes.QueryFee, users.gateway, users.indexer, amount, subgraphDataServiceAddress, 1);
+            amount,
+            subgraphDataServiceAddress,
+            1
+        );
     }
-}
+}
diff --git a/packages/horizon/test/escrow/getters.t.sol b/packages/horizon/test/escrow/getters.t.sol
diff --git a/packages/horizon/test/escrow/thaw.t.sol b/packages/horizon/test/escrow/thaw.t.sol

Original file line number	Diff line number	Diff line change
`@@ -157,6 +157,9 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,`
`157`	`157`	`*/`
`158`	`158`	`function getBalance(address payer, address collector, address receiver) external view override returns (uint256) {`
`159`	`159`	`EscrowAccount storage account = escrowAccounts[payer][collector][receiver];`
	`160`	`+ if (account.balance <= account.tokensThawing) {`
	`161`	`+ return 0;`
	`162`	`+ }`
`160`	`163`	`return account.balance - account.tokensThawing;`
`161`	`164`	`}`
`162`	`165`