Skip to content

Commit 9a74331

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 caa7fa9 commit 9a74331

File tree

3 files changed

+70
-14
lines changed

3 files changed

+70
-14
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/subgraphService/SubgraphService.t.sol

+3-3
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,7 @@ 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(_data, (IGraphTallyCollector.SignedRAV, uint256));
282282
address allocationId = address(uint160(uint256(signedRav.rav.collectionId)));
283283
Allocation.State memory allocation = subgraphService.getAllocation(allocationId);
284284
bytes32 subgraphDeploymentId = allocation.subgraphDeploymentId;
@@ -293,7 +293,7 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
293293
_indexer,
294294
payer
295295
);
296-
paymentCollected = signedRav.rav.valueAggregate - tokensCollected;
296+
paymentCollected = tokensToCollect == 0 ? signedRav.rav.valueAggregate - tokensCollected : tokensToCollect;
297297

298298
QueryFeeData memory queryFeeData = _queryFeeData(allocation.subgraphDeploymentId);
299299
uint256 tokensProtocol = paymentCollected.mulPPMRoundUp(queryFeeData.protocolPaymentCut);
@@ -370,7 +370,7 @@ contract SubgraphServiceTest is SubgraphServiceSharedTest {
370370
CollectPaymentData memory collectPaymentDataBefore,
371371
CollectPaymentData memory collectPaymentDataAfter
372372
) private view {
373-
IGraphTallyCollector.SignedRAV memory signedRav = abi.decode(_data, (IGraphTallyCollector.SignedRAV));
373+
(IGraphTallyCollector.SignedRAV memory signedRav, uint256 tokensToCollect) = abi.decode(_data, (IGraphTallyCollector.SignedRAV, uint256));
374374
Allocation.State memory allocation = subgraphService.getAllocation(
375375
address(uint160(uint256(signedRav.rav.collectionId)))
376376
);

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

+54-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,47 @@ 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(address(subgraphService), bytes32(uint256(uint160(allocationID))), users.indexer, users.gateway);
217+
218+
// Collect the RAV in two steps
219+
resetPrank(users.indexer);
220+
uint256 tokensToCollect = tokensPayment / 2;
221+
bool oddTokensPayment = tokensPayment % 2 == 1;
222+
bytes memory data = _getQueryFeeEncodedData(users.indexer, uint128(tokensPayment), tokensToCollect);
223+
_collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data);
224+
225+
uint256 intermediateGatewayBalance = escrow.getBalance(users.gateway, address(graphTallyCollector), users.indexer);
226+
assertEq(intermediateGatewayBalance, beforeGatewayBalance - tokensToCollect);
227+
uint256 intermediateTokensCollected = graphTallyCollector.tokensCollected(address(subgraphService), bytes32(uint256(uint160(allocationID))), users.indexer, users.gateway);
228+
assertEq(intermediateTokensCollected, beforeTokensCollected + tokensToCollect);
229+
230+
bytes memory data2 = _getQueryFeeEncodedData(
231+
users.indexer,
232+
uint128(tokensPayment),
233+
oddTokensPayment ? tokensToCollect + 1 : tokensToCollect
234+
);
235+
_collect(users.indexer, IGraphPayments.PaymentTypes.QueryFee, data2);
236+
237+
// Check the indexer received the correct amount of tokens
238+
uint256 afterGatewayBalance = escrow.getBalance(users.gateway, address(graphTallyCollector), users.indexer);
239+
assertEq(afterGatewayBalance, beforeGatewayBalance - tokensPayment);
240+
uint256 afterTokensCollected = graphTallyCollector.tokensCollected(address(subgraphService), bytes32(uint256(uint160(allocationID))), users.indexer, users.gateway);
241+
assertEq(afterTokensCollected, intermediateTokensCollected + tokensToCollect + (oddTokensPayment ? 1 : 0));
242+
}
196243
}

0 commit comments

Comments
 (0)