Skip to content

Commit

Permalink
fix: make getBalance return 0 if balance is less than thawing amount …
Browse files Browse the repository at this point in the history
…(TRST-L10)

Signed-off-by: Tomás Migone <tomas@edgeandnode.com>
  • Loading branch information
tmigone committed Nov 28, 2024
1 parent ec69bf9 commit d55cdf4
Show file tree
Hide file tree
Showing 6 changed files with 229 additions and 112 deletions.
1 change: 1 addition & 0 deletions packages/horizon/contracts/interfaces/IPaymentsEscrow.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 3 additions & 0 deletions packages/horizon/contracts/payments/PaymentsEscrow.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}

Expand Down
92 changes: 90 additions & 2 deletions packages/horizon/test/escrow/GraphEscrow.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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
*/
Expand Down Expand Up @@ -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);
}
}
167 changes: 67 additions & 100 deletions packages/horizon/test/escrow/collect.t.sol
Original file line number Diff line number Diff line change
Expand Up @@ -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
*/
Expand All @@ -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);
Expand All @@ -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
);
}
}
}
Loading

0 comments on commit d55cdf4

Please sign in to comment.