Skip to content

Commit f9bee71

Browse files
tmigonetmigone
committed
fix: allow partial RAV collection in subgraph service (OZ L-13)
Signed-off-by: Tomás Migone <[email protected]>
1 parent e053527 commit f9bee71

File tree

4 files changed

+105
-19
lines changed

4 files changed

+105
-19
lines changed

‎packages/subgraph-service/contracts/SubgraphService.sol

+13-4
Original file line numberDiff line numberDiff line change
@@ -268,12 +268,15 @@ contract SubgraphService is
268268
uint256 paymentCollected = 0;
269269

270270
if (paymentType == IGraphPayments.PaymentTypes.QueryFee) {
271-
IGraphTallyCollector.SignedRAV memory signedRav = abi.decode(data, (IGraphTallyCollector.SignedRAV));
271+
(IGraphTallyCollector.SignedRAV memory signedRav, uint256 tokensToCollect) = abi.decode(
272+
data,
273+
(IGraphTallyCollector.SignedRAV, uint256)
274+
);
272275
require(
273276
signedRav.rav.serviceProvider == indexer,
274277
SubgraphServiceIndexerMismatch(signedRav.rav.serviceProvider, indexer)
275278
);
276-
paymentCollected = _collectQueryFees(signedRav);
279+
paymentCollected = _collectQueryFees(signedRav, tokensToCollect);
277280
} else if (paymentType == IGraphPayments.PaymentTypes.IndexingRewards) {
278281
(address allocationId, bytes32 poi) = abi.decode(data, (address, bytes32));
279282
require(
@@ -475,9 +478,14 @@ contract SubgraphService is
475478
* Emits a {QueryFeesCollected} event.
476479
*
477480
* @param _signedRav Signed RAV
481+
* @param tokensToCollect The amount of tokens to collect. Allows partially collecting a RAV. If 0, the entire RAV will
482+
* be collected.
478483
* @return The amount of fees collected
479484
*/
480-
function _collectQueryFees(IGraphTallyCollector.SignedRAV memory _signedRav) private returns (uint256) {
485+
function _collectQueryFees(
486+
IGraphTallyCollector.SignedRAV memory _signedRav,
487+
uint256 tokensToCollect
488+
) private returns (uint256) {
481489
address indexer = _signedRav.rav.serviceProvider;
482490

483491
// Check that collectionId (256 bits) is a valid address (160 bits)
@@ -502,7 +510,8 @@ contract SubgraphService is
502510
uint256 curationCut = _curation().isCurated(subgraphDeploymentId) ? curationFeesCut : 0;
503511
uint256 tokensCollected = _graphTallyCollector().collect(
504512
IGraphPayments.PaymentTypes.QueryFee,
505-
abi.encode(_signedRav, curationCut)
513+
abi.encode(_signedRav, curationCut),
514+
tokensToCollect
506515
);
507516

508517
uint256 balanceAfter = _graphToken().balanceOf(address(this));

‎packages/subgraph-service/test/mocks/MockRewardsManager.sol

+10-5
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,13 @@ interface IRewardsIssuer {
1414
)
1515
external
1616
view
17-
returns (bool isActive, address indexer, bytes32 subgraphDeploymentId, uint256 tokens, uint256 accRewardsPerAllocatedToken);
17+
returns (
18+
bool isActive,
19+
address indexer,
20+
bytes32 subgraphDeploymentId,
21+
uint256 tokens,
22+
uint256 accRewardsPerAllocatedToken
23+
);
1824
}
1925

2026
contract MockRewardsManager is IRewardsManager {
@@ -71,13 +77,12 @@ contract MockRewardsManager is IRewardsManager {
7177

7278
function takeRewards(address _allocationID) external returns (uint256) {
7379
address rewardsIssuer = msg.sender;
74-
(bool isActive, , , uint256 tokens, uint256 accRewardsPerAllocatedToken) = IRewardsIssuer(rewardsIssuer).getAllocationData(
75-
_allocationID
76-
);
80+
(bool isActive, , , uint256 tokens, uint256 accRewardsPerAllocatedToken) = IRewardsIssuer(rewardsIssuer)
81+
.getAllocationData(_allocationID);
7782

7883
if (!isActive) {
7984
return 0;
80-
}
85+
}
8186

8287
uint256 accRewardsPerTokens = tokens.mulPPM(rewardsPerSignal);
8388
uint256 rewards = accRewardsPerTokens - accRewardsPerAllocatedToken;

‎packages/subgraph-service/test/subgraphService/SubgraphService.t.sol

+9-3
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,10 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
278278
address _indexer,
279279
bytes memory _data
280280
) private returns (uint256 paymentCollected) {
281-
IGraphTallyCollector.SignedRAV memory signedRav = abi.decode(_data, (IGraphTallyCollector.SignedRAV));
281+
(IGraphTallyCollector.SignedRAV memory signedRav, uint256 tokensToCollect) = abi.decode(
282+
_data,
283+
(IGraphTallyCollector.SignedRAV, uint256)
284+
);
282285
address allocationId = address(uint160(uint256(signedRav.rav.collectionId)));
283286
Allocation.State memory allocation = subgraphService.getAllocation(allocationId);
284287
bytes32 subgraphDeploymentId = allocation.subgraphDeploymentId;
@@ -293,7 +296,7 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
293296
_indexer,
294297
payer
295298
);
296-
paymentCollected = signedRav.rav.valueAggregate - tokensCollected;
299+
paymentCollected = tokensToCollect == 0 ? signedRav.rav.valueAggregate - tokensCollected : tokensToCollect;
297300

298301
QueryFeeData memory queryFeeData = _queryFeeData(allocation.subgraphDeploymentId);
299302
uint256 tokensProtocol = paymentCollected.mulPPMRoundUp(queryFeeData.protocolPaymentCut);
@@ -370,7 +373,10 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
370373
CollectPaymentData memory collectPaymentDataBefore,
371374
CollectPaymentData memory collectPaymentDataAfter
372375
) private view {
373-
IGraphTallyCollector.SignedRAV memory signedRav = abi.decode(_data, (IGraphTallyCollector.SignedRAV));
376+
(IGraphTallyCollector.SignedRAV memory signedRav, uint256 tokensToCollect) = abi.decode(
377+
_data,
378+
(IGraphTallyCollector.SignedRAV, uint256)
379+
);
374380
Allocation.State memory allocation = subgraphService.getAllocation(
375381
address(uint160(uint256(signedRav.rav.collectionId)))
376382
);

‎packages/subgraph-service/test/subgraphService/collect/query/query.t.sol

+73-7
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,11 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
3939
return abi.encodePacked(r, s, v);
4040
}
4141

42-
function _getQueryFeeEncodedData(address indexer, uint128 tokens) private view returns (bytes memory) {
42+
function _getQueryFeeEncodedData(
43+
address indexer,
44+
uint128 tokens,
45+
uint256 tokensToCollect
46+
) private view returns (bytes memory) {
4347
IGraphTallyCollector.ReceiptAggregateVoucher memory rav = _getRAV(
4448
indexer,
4549
bytes32(uint256(uint160(allocationID))),
@@ -49,7 +53,7 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
4953
(uint8 v, bytes32 r, bytes32 s) = vm.sign(signerPrivateKey, messageHash);
5054
bytes memory signature = abi.encodePacked(r, s, v);
5155
IGraphTallyCollector.SignedRAV memory signedRAV = IGraphTallyCollector.SignedRAV(rav, signature);
52-
return abi.encode(signedRAV);
56+
return abi.encode(signedRAV, tokensToCollect);
5357
}
5458

5559
function _getRAV(
@@ -109,7 +113,7 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
109113
_authorizeSigner();
110114

111115
resetPrank(users.indexer);
112-
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokensPayment));
116+
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokensPayment), 0);
113117
_collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data);
114118
}
115119

@@ -129,14 +133,14 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
129133
uint256 accTokensPayment = 0;
130134
for (uint i = 0; i < numPayments; i++) {
131135
accTokensPayment = accTokensPayment + tokensPayment;
132-
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(accTokensPayment));
136+
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(accTokensPayment), 0);
133137
_collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data);
134138
}
135139
}
136140

137141
function testCollect_RevertWhen_NotAuthorized(uint256 tokens) public useIndexer useAllocation(tokens) {
138142
IGraphPayments.PaymentTypes paymentType = IGraphPayments.PaymentTypes.QueryFee;
139-
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokens));
143+
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokens), 0);
140144
resetPrank(users.operator);
141145
vm.expectRevert(
142146
abi.encodeWithSelector(
@@ -157,7 +161,7 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
157161
_createAndStartAllocation(newIndexer, tokens);
158162

159163
// This data is for user.indexer allocationId
160-
bytes memory data = _getQueryFeeEncodedData(newIndexer, uint128(tokens));
164+
bytes memory data = _getQueryFeeEncodedData(newIndexer, uint128(tokens), 0);
161165

162166
resetPrank(newIndexer);
163167
vm.expectRevert(
@@ -173,7 +177,7 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
173177
// Setup new indexer
174178
address newIndexer = makeAddr("newIndexer");
175179
_createAndStartAllocation(newIndexer, tokens);
176-
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokens));
180+
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokens), 0);
177181
vm.expectRevert(
178182
abi.encodeWithSelector(ISubgraphService.SubgraphServiceIndexerMismatch.selector, users.indexer, newIndexer)
179183
);
@@ -193,4 +197,66 @@ contract SubgraphServiceRegisterTest is SubgraphServiceTest {
193197
);
194198
subgraphService.collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data);
195199
}
200+
201+
function testCollect_QueryFees_PartialCollect(
202+
uint256 tokensAllocated,
203+
uint256 tokensPayment
204+
) public useIndexer useAllocation(tokensAllocated) {
205+
vm.assume(tokensAllocated > minimumProvisionTokens * stakeToFeesRatio);
206+
uint256 maxTokensPayment = tokensAllocated / stakeToFeesRatio > type(uint128).max
207+
? type(uint128).max
208+
: tokensAllocated / stakeToFeesRatio;
209+
tokensPayment = bound(tokensPayment, minimumProvisionTokens, maxTokensPayment);
210+
211+
resetPrank(users.gateway);
212+
_deposit(tokensPayment);
213+
_authorizeSigner();
214+
215+
uint256 beforeGatewayBalance = escrow.getBalance(users.gateway, address(graphTallyCollector), users.indexer);
216+
uint256 beforeTokensCollected = graphTallyCollector.tokensCollected(
217+
address(subgraphService),
218+
bytes32(uint256(uint160(allocationID))),
219+
users.indexer,
220+
users.gateway
221+
);
222+
223+
// Collect the RAV in two steps
224+
resetPrank(users.indexer);
225+
uint256 tokensToCollect = tokensPayment / 2;
226+
bool oddTokensPayment = tokensPayment % 2 == 1;
227+
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokensPayment), tokensToCollect);
228+
_collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data);
229+
230+
uint256 intermediateGatewayBalance = escrow.getBalance(
231+
users.gateway,
232+
address(graphTallyCollector),
233+
users.indexer
234+
);
235+
assertEq(intermediateGatewayBalance, beforeGatewayBalance - tokensToCollect);
236+
uint256 intermediateTokensCollected = graphTallyCollector.tokensCollected(
237+
address(subgraphService),
238+
bytes32(uint256(uint160(allocationID))),
239+
users.indexer,
240+
users.gateway
241+
);
242+
assertEq(intermediateTokensCollected, beforeTokensCollected + tokensToCollect);
243+
244+
bytes memory data2 = _getQueryFeeEncodedData(
245+
users.indexer,
246+
uint128(tokensPayment),
247+
oddTokensPayment ? tokensToCollect + 1 : tokensToCollect
248+
);
249+
_collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data2);
250+
251+
// Check the indexer received the correct amount of tokens
252+
uint256 afterGatewayBalance = escrow.getBalance(users.gateway, address(graphTallyCollector), users.indexer);
253+
assertEq(afterGatewayBalance, beforeGatewayBalance - tokensPayment);
254+
uint256 afterTokensCollected = graphTallyCollector.tokensCollected(
255+
address(subgraphService),
256+
bytes32(uint256(uint160(allocationID))),
257+
users.indexer,
258+
users.gateway
259+
);
260+
assertEq(afterTokensCollected, intermediateTokensCollected + tokensToCollect + (oddTokensPayment ? 1 : 0));
261+
}
196262
}

0 commit comments

Comments
 (0)