Skip to content

Commit d3318eb

Browse files
tmigonetmigone
committed
fix: couple minor improvments to PaymentsEscrow and some new tests
Signed-off-by: Tomás Migone <[email protected]>
1 parent 8694e06 commit d3318eb

File tree

9 files changed

+188
-30
lines changed

9 files changed

+188
-30
lines changed

Diff for: ‎packages/horizon/contracts/interfaces/IGraphPayments.sol

+2-2
Original file line numberDiff line numberDiff line change
@@ -31,9 +31,9 @@ interface IGraphPayments {
3131
* @param tokensReceiver Amount of tokens for the receiver
3232
*/
3333
event GraphPaymentCollected(
34-
PaymentTypes paymentType,
34+
PaymentTypes indexed paymentType,
3535
address indexed payer,
36-
address indexed receiver,
36+
address receiver,
3737
address indexed dataService,
3838
uint256 tokens,
3939
uint256 tokensProtocol,

Diff for: ‎packages/horizon/contracts/interfaces/IPaymentsEscrow.sol

+16-2
Original file line numberDiff line numberDiff line change
@@ -37,11 +37,18 @@ interface IPaymentsEscrow {
3737
/**
3838
* @notice Emitted when a payer cancels an escrow thawing
3939
* @param payer The address of the payer
40+
* @param collector The address of the collector
4041
* @param receiver The address of the receiver
4142
* @param tokensThawing The amount of tokens that were being thawed
4243
* @param thawEndTimestamp The timestamp at which the thawing period was ending
4344
*/
44-
event CancelThaw(address indexed payer, address indexed receiver, uint256 tokensThawing, uint256 thawEndTimestamp);
45+
event CancelThaw(
46+
address indexed payer,
47+
address indexed collector,
48+
address indexed receiver,
49+
uint256 tokensThawing,
50+
uint256 thawEndTimestamp
51+
);
4552

4653
/**
4754
* @notice Emitted when a payer thaws funds from the escrow for a payer-collector-receiver tuple
@@ -70,12 +77,19 @@ interface IPaymentsEscrow {
7077

7178
/**
7279
* @notice Emitted when a collector collects funds from the escrow for a payer-collector-receiver tuple
80+
* @param paymentType The type of payment being collected as defined in the {IGraphPayments} interface
7381
* @param payer The address of the payer
7482
* @param collector The address of the collector
7583
* @param receiver The address of the receiver
7684
* @param tokens The amount of tokens collected
7785
*/
78-
event EscrowCollected(address indexed payer, address indexed collector, address indexed receiver, uint256 tokens);
86+
event EscrowCollected(
87+
IGraphPayments.PaymentTypes indexed paymentType,
88+
address indexed payer,
89+
address indexed collector,
90+
address receiver,
91+
uint256 tokens
92+
);
7993

8094
// -- Errors --
8195

Diff for: ‎packages/horizon/contracts/payments/PaymentsEscrow.sol

+16-14
Original file line numberDiff line numberDiff line change
@@ -23,17 +23,21 @@ import { GraphDirectory } from "../utilities/GraphDirectory.sol";
2323
contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory, IPaymentsEscrow {
2424
using TokenUtils for IGraphToken;
2525

26-
/// @notice Escrow account details for payer-collector-receiver tuples
27-
mapping(address payer => mapping(address collector => mapping(address receiver => IPaymentsEscrow.EscrowAccount escrowAccount)))
28-
public escrowAccounts;
29-
3026
/// @notice The maximum thawing period (in seconds) for both escrow withdrawal and collector revocation
3127
/// @dev This is a precautionary measure to avoid inadvertedly locking funds for too long
3228
uint256 public constant MAX_WAIT_PERIOD = 90 days;
3329

3430
/// @notice Thawing period in seconds for escrow funds withdrawal
3531
uint256 public immutable WITHDRAW_ESCROW_THAWING_PERIOD;
3632

33+
/// @notice Escrow account details for payer-collector-receiver tuples
34+
mapping(address payer => mapping(address collector => mapping(address receiver => IPaymentsEscrow.EscrowAccount escrowAccount)))
35+
public escrowAccounts;
36+
37+
/**
38+
* @notice Modifier to prevent function execution when contract is paused
39+
* @dev Reverts if the controller indicates the contract is paused
40+
*/
3741
modifier notPaused() {
3842
require(!_graphController().paused(), PaymentsEscrowIsPaused());
3943
_;
@@ -101,7 +105,7 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,
101105
account.tokensThawing = 0;
102106
account.thawEndTimestamp = 0;
103107

104-
emit CancelThaw(msg.sender, receiver, tokensThawing, thawEndTimestamp);
108+
emit CancelThaw(msg.sender, collector, receiver, tokensThawing, thawEndTimestamp);
105109
}
106110

107111
/**
@@ -143,29 +147,27 @@ contract PaymentsEscrow is Initializable, MulticallUpgradeable, GraphDirectory,
143147
// Reduce amount from account balance
144148
account.balance -= tokens;
145149

146-
uint256 balanceBefore = _graphToken().balanceOf(address(this));
150+
uint256 escrowBalanceBefore = _graphToken().balanceOf(address(this));
147151

148152
_graphToken().approve(address(_graphPayments()), tokens);
149153
_graphPayments().collect(paymentType, receiver, tokens, dataService, dataServiceCut);
150154

151-
uint256 balanceAfter = _graphToken().balanceOf(address(this));
155+
// Verify that the escrow balance is consistent with the collected tokens
156+
uint256 escrowBalanceAfter = _graphToken().balanceOf(address(this));
152157
require(
153-
balanceBefore == tokens + balanceAfter,
154-
PaymentsEscrowInconsistentCollection(balanceBefore, balanceAfter, tokens)
158+
escrowBalanceBefore == tokens + escrowBalanceAfter,
159+
PaymentsEscrowInconsistentCollection(escrowBalanceBefore, escrowBalanceAfter, tokens)
155160
);
156161

157-
emit EscrowCollected(payer, msg.sender, receiver, tokens);
162+
emit EscrowCollected(paymentType, payer, msg.sender, receiver, tokens);
158163
}
159164

160165
/**
161166
* @notice See {IPaymentsEscrow-getBalance}
162167
*/
163168
function getBalance(address payer, address collector, address receiver) external view override returns (uint256) {
164169
EscrowAccount storage account = escrowAccounts[payer][collector][receiver];
165-
if (account.balance <= account.tokensThawing) {
166-
return 0;
167-
}
168-
return account.balance - account.tokensThawing;
170+
return account.balance > account.tokensThawing ? account.balance - account.tokensThawing : 0;
169171
}
170172

171173
/**

Diff for: ‎packages/horizon/test/escrow/GraphEscrow.t.sol

+28-2
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,9 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
3030
}
3131

3232
modifier depositAndThawTokens(uint256 amount, uint256 thawAmount) {
33+
vm.assume(amount > 0);
3334
vm.assume(thawAmount > 0);
35+
vm.assume(amount <= MAX_STAKING_TOKENS);
3436
vm.assume(amount > thawAmount);
3537
_depositTokens(users.verifier, users.indexer, amount);
3638
escrow.thaw(users.verifier, users.indexer, thawAmount);
@@ -62,14 +64,38 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
6264
(, uint256 amountThawingBefore, uint256 thawEndTimestampBefore) = escrow.escrowAccounts(msgSender, collector, receiver);
6365

6466
vm.expectEmit(address(escrow));
65-
emit IPaymentsEscrow.CancelThaw(msgSender, receiver, amountThawingBefore, thawEndTimestampBefore);
67+
emit IPaymentsEscrow.CancelThaw(msgSender, collector, receiver, amountThawingBefore, thawEndTimestampBefore);
6668
escrow.cancelThaw(collector, receiver);
6769

6870
(, uint256 amountThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(msgSender, collector, receiver);
6971
assertEq(amountThawing, 0);
7072
assertEq(thawEndTimestamp, 0);
7173
}
7274

75+
function _withdrawEscrow(address collector, address receiver) internal {
76+
(, address msgSender, ) = vm.readCallers();
77+
78+
(uint256 balanceBefore, uint256 amountThawingBefore, ) = escrow.escrowAccounts(msgSender, collector, receiver);
79+
uint256 tokenBalanceBeforeSender = token.balanceOf(msgSender);
80+
uint256 tokenBalanceBeforeEscrow = token.balanceOf(address(escrow));
81+
82+
uint256 amountToWithdraw = amountThawingBefore > balanceBefore ? balanceBefore : amountThawingBefore;
83+
vm.expectEmit(address(escrow));
84+
emit IPaymentsEscrow.Withdraw(msgSender, collector, receiver, amountToWithdraw);
85+
escrow.withdraw(collector, receiver);
86+
87+
(uint256 balanceAfter, uint256 tokensThawingAfter, uint256 thawEndTimestampAfter) = escrow.escrowAccounts(msgSender, collector, receiver);
88+
uint256 tokenBalanceAfterSender = token.balanceOf(msgSender);
89+
uint256 tokenBalanceAfterEscrow = token.balanceOf(address(escrow));
90+
91+
assertEq(balanceAfter, balanceBefore - amountToWithdraw);
92+
assertEq(tokensThawingAfter, 0);
93+
assertEq(thawEndTimestampAfter, 0);
94+
95+
assertEq(tokenBalanceAfterSender, tokenBalanceBeforeSender + amountToWithdraw);
96+
assertEq(tokenBalanceAfterEscrow, tokenBalanceBeforeEscrow - amountToWithdraw);
97+
}
98+
7399
struct CollectPaymentData {
74100
uint256 escrowBalance;
75101
uint256 paymentsBalance;
@@ -105,7 +131,7 @@ contract GraphEscrowTest is HorizonStakingSharedTest, PaymentsEscrowSharedTest {
105131
}
106132

107133
vm.expectEmit(address(escrow));
108-
emit IPaymentsEscrow.EscrowCollected(_payer, _collector, _receiver, _tokens);
134+
emit IPaymentsEscrow.EscrowCollected(_paymentType, _payer, _collector, _receiver, _tokens);
109135
escrow.collect(_paymentType, _payer, _receiver, _tokens, _dataService, _dataServiceCut);
110136

111137
// Calculate cuts

Diff for: ‎packages/horizon/test/escrow/collect.t.sol

+40
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
1313
* TESTS
1414
*/
1515

16+
// use users.verifier as collector
1617
function testCollect_Tokens(
1718
uint256 tokens,
1819
uint256 tokensToCollect,
@@ -98,4 +99,43 @@ contract GraphEscrowCollectTest is GraphEscrowTest {
9899
);
99100
vm.stopPrank();
100101
}
102+
103+
function testCollect_MultipleCollections(
104+
uint256 depositAmount,
105+
uint256 firstCollect,
106+
uint256 secondCollect
107+
) public useIndexer {
108+
// Tests multiple collect operations from the same escrow account
109+
vm.assume(firstCollect < MAX_STAKING_TOKENS);
110+
vm.assume(secondCollect < MAX_STAKING_TOKENS);
111+
vm.assume(depositAmount > 0);
112+
vm.assume(firstCollect > 0 && firstCollect < depositAmount);
113+
vm.assume(secondCollect > 0 && secondCollect <= depositAmount - firstCollect);
114+
115+
resetPrank(users.gateway);
116+
_depositTokens(users.verifier, users.indexer, depositAmount);
117+
118+
// burn some tokens to prevent overflow
119+
resetPrank(users.indexer);
120+
token.burn(MAX_STAKING_TOKENS);
121+
122+
resetPrank(users.verifier);
123+
_collectEscrow(
124+
IGraphPayments.PaymentTypes.QueryFee,
125+
users.gateway,
126+
users.indexer,
127+
firstCollect,
128+
subgraphDataServiceAddress,
129+
0
130+
);
131+
132+
// _collectEscrow(
133+
// IGraphPayments.PaymentTypes.QueryFee,
134+
// users.gateway,
135+
// users.indexer,
136+
// secondCollect,
137+
// subgraphDataServiceAddress,
138+
// 0
139+
// );
140+
}
101141
}

Diff for: ‎packages/horizon/test/escrow/deposit.t.sol

+22-3
Original file line numberDiff line numberDiff line change
@@ -6,13 +6,32 @@ import "forge-std/Test.sol";
66
import { GraphEscrowTest } from "./GraphEscrow.t.sol";
77

88
contract GraphEscrowDepositTest is GraphEscrowTest {
9-
109
/*
1110
* TESTS
1211
*/
1312

1413
function testDeposit_Tokens(uint256 amount) public useGateway useDeposit(amount) {
15-
(uint256 indexerEscrowBalance,,) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
14+
(uint256 indexerEscrowBalance, , ) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
1615
assertEq(indexerEscrowBalance, amount);
1716
}
18-
}
17+
18+
function testDepositTo_Tokens(uint256 amount) public {
19+
resetPrank(users.delegator);
20+
token.approve(address(escrow), amount);
21+
_depositToTokens(users.gateway, users.verifier, users.indexer, amount);
22+
}
23+
24+
// Tests multiple deposits accumulate correctly in the escrow account
25+
function testDeposit_MultipleDeposits(uint256 amount1, uint256 amount2) public useGateway {
26+
vm.assume(amount1 > 0);
27+
vm.assume(amount2 > 0);
28+
vm.assume(amount1 <= MAX_STAKING_TOKENS);
29+
vm.assume(amount2 <= MAX_STAKING_TOKENS);
30+
31+
_depositTokens(users.verifier, users.indexer, amount1);
32+
_depositTokens(users.verifier, users.indexer, amount2);
33+
34+
(uint256 balance,,) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
35+
assertEq(balance, amount1 + amount2);
36+
}
37+
}

Diff for: ‎packages/horizon/test/escrow/thaw.t.sol

+19-3
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,21 @@ contract GraphEscrowThawTest is GraphEscrowTest {
1010
* TESTS
1111
*/
1212

13-
function testThaw_Tokens(uint256 amount) public useGateway useDeposit(amount) {
14-
amount = bound(amount, 1, type(uint256).max);
13+
function testThaw_PartialBalanceThaw(
14+
uint256 amountDeposited,
15+
uint256 amountThawed
16+
) public useGateway useDeposit(amountDeposited) {
17+
vm.assume(amountThawed > 0);
18+
vm.assume(amountThawed <= amountDeposited);
19+
_thawEscrow(users.verifier, users.indexer, amountThawed);
20+
}
21+
22+
function testThaw_FullBalanceThaw(uint256 amount) public useGateway useDeposit(amount) {
23+
vm.assume(amount > 0);
1524
_thawEscrow(users.verifier, users.indexer, amount);
25+
26+
uint256 availableBalance = escrow.getBalance(users.gateway, users.verifier, users.indexer);
27+
assertEq(availableBalance, 0);
1628
}
1729

1830
function testThaw_Tokens_SuccesiveCalls(uint256 amount) public useGateway {
@@ -25,7 +37,11 @@ contract GraphEscrowThawTest is GraphEscrowTest {
2537
_thawEscrow(users.verifier, users.indexer, secondAmountToThaw);
2638

2739
(, address msgSender, ) = vm.readCallers();
28-
(, uint256 amountThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(msgSender, users.verifier, users.indexer);
40+
(, uint256 amountThawing, uint256 thawEndTimestamp) = escrow.escrowAccounts(
41+
msgSender,
42+
users.verifier,
43+
users.indexer
44+
);
2945
assertEq(amountThawing, secondAmountToThaw);
3046
assertEq(thawEndTimestamp, block.timestamp + withdrawEscrowThawingPeriod);
3147
}

Diff for: ‎packages/horizon/test/escrow/withdraw.t.sol

+33-4
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ pragma solidity 0.8.27;
33

44
import "forge-std/Test.sol";
55

6+
import { IGraphPayments } from "../../contracts/interfaces/IGraphPayments.sol";
67
import { GraphEscrowTest } from "./GraphEscrow.t.sol";
78

89
contract GraphEscrowWithdrawTest is GraphEscrowTest {
@@ -18,11 +19,8 @@ contract GraphEscrowWithdrawTest is GraphEscrowTest {
1819
// advance time
1920
skip(withdrawEscrowThawingPeriod + 1);
2021

21-
escrow.withdraw(users.verifier, users.indexer);
22+
_withdrawEscrow(users.verifier, users.indexer);
2223
vm.stopPrank();
23-
24-
(uint256 indexerEscrowBalance,,) = escrow.escrowAccounts(users.gateway, users.verifier, users.indexer);
25-
assertEq(indexerEscrowBalance, amount - thawAmount);
2624
}
2725

2826
function testWithdraw_RevertWhen_NotThawing(uint256 amount) public useGateway useDeposit(amount) {
@@ -39,4 +37,35 @@ contract GraphEscrowWithdrawTest is GraphEscrowTest {
3937
vm.expectRevert(expectedError);
4038
escrow.withdraw(users.verifier, users.indexer);
4139
}
40+
41+
function testWithdraw_BalanceAfterCollect(
42+
uint256 amountDeposited,
43+
uint256 amountThawed,
44+
uint256 amountCollected
45+
) public useGateway depositAndThawTokens(amountDeposited, amountThawed) {
46+
vm.assume(amountCollected > 0);
47+
vm.assume(amountCollected <= amountDeposited);
48+
49+
// burn some tokens to prevent overflow
50+
resetPrank(users.indexer);
51+
token.burn(MAX_STAKING_TOKENS);
52+
53+
// collect
54+
resetPrank(users.verifier);
55+
_collectEscrow(
56+
IGraphPayments.PaymentTypes.QueryFee,
57+
users.gateway,
58+
users.indexer,
59+
amountCollected,
60+
subgraphDataServiceAddress,
61+
0
62+
);
63+
64+
// Advance time to simulate the thawing period
65+
skip(withdrawEscrowThawingPeriod + 1);
66+
67+
// withdraw the remaining thawed balance
68+
resetPrank(users.gateway);
69+
_withdrawEscrow(users.verifier, users.indexer);
70+
}
4271
}

Diff for: ‎packages/horizon/test/shared/payments-escrow/PaymentsEscrowShared.t.sol

+12
Original file line numberDiff line numberDiff line change
@@ -34,4 +34,16 @@ abstract contract PaymentsEscrowSharedTest is GraphBaseTest {
3434
(uint256 escrowBalanceAfter,,) = escrow.escrowAccounts(msgSender, _collector, _receiver);
3535
assertEq(escrowBalanceAfter - _tokens, escrowBalanceBefore);
3636
}
37+
38+
function _depositToTokens(address _payer, address _collector, address _receiver, uint256 _tokens) internal {
39+
(uint256 escrowBalanceBefore,,) = escrow.escrowAccounts(_payer, _collector, _receiver);
40+
token.approve(address(escrow), _tokens);
41+
42+
vm.expectEmit(address(escrow));
43+
emit IPaymentsEscrow.Deposit(_payer, _collector, _receiver, _tokens);
44+
escrow.depositTo(_payer, _collector, _receiver, _tokens);
45+
46+
(uint256 escrowBalanceAfter,,) = escrow.escrowAccounts(_payer, _collector, _receiver);
47+
assertEq(escrowBalanceAfter - _tokens, escrowBalanceBefore);
48+
}
3749
}

0 commit comments

Comments
 (0)