Skip to content

Commit

Permalink
soft delete nonce in s_consumers so that request IDs do not repeat. W…
Browse files Browse the repository at this point in the history
…IP (#12405)

* soft delete nonce in s_consumers so that request IDs do not repeat. WIP

* forge unit tests passing except for exact billing amount

* add forge tests and optimize gas

* regenerate wrappers

* remove comment

* address comments

* fix test

* add changeset
  • Loading branch information
jinhoonbang committed Mar 14, 2024
1 parent b3c0fa4 commit 2bd210b
Show file tree
Hide file tree
Showing 11 changed files with 288 additions and 84 deletions.
5 changes: 5 additions & 0 deletions .changeset/gorgeous-crabs-repeat.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"chainlink": patch
---

Soft delete consumer nonce in VRF coordinator v2.5
25 changes: 16 additions & 9 deletions contracts/src/v0.8/vrf/dev/SubscriptionAPI.sol
Original file line number Diff line number Diff line change
Expand Up @@ -60,8 +60,13 @@ abstract contract SubscriptionAPI is ConfirmedOwner, IERC677Receiver, IVRFSubscr
// consumer is valid without reading all the consumers from storage.
address[] consumers;
}
struct ConsumerConfig {
bool active;
uint64 nonce;
}
// Note a nonce of 0 indicates an the consumer is not assigned to that subscription.
mapping(address => mapping(uint256 => uint64)) /* consumer */ /* subId */ /* nonce */ internal s_consumers;
mapping(address => mapping(uint256 => ConsumerConfig)) /* consumerAddress */ /* subId */ /* consumerConfig */
internal s_consumers;
mapping(uint256 => SubscriptionConfig) /* subId */ /* subscriptionConfig */ internal s_subscriptionConfigs;
mapping(uint256 => Subscription) /* subId */ /* subscription */ internal s_subscriptions;
// subscription nonce used to construct subId. Rises monotonically
Expand Down Expand Up @@ -400,19 +405,21 @@ abstract contract SubscriptionAPI is ConfirmedOwner, IERC677Receiver, IVRFSubscr
* @inheritdoc IVRFSubscriptionV2Plus
*/
function addConsumer(uint256 subId, address consumer) external override onlySubOwner(subId) nonReentrant {
ConsumerConfig storage consumerConfig = s_consumers[consumer][subId];
if (consumerConfig.active) {
// Idempotence - do nothing if already added.
// Ensures uniqueness in s_subscriptions[subId].consumers.
return;
}
// Already maxed, cannot add any more consumers.
address[] storage consumers = s_subscriptionConfigs[subId].consumers;
if (consumers.length == MAX_CONSUMERS) {
revert TooManyConsumers();
}
mapping(uint256 => uint64) storage nonces = s_consumers[consumer];
if (nonces[subId] != 0) {
// Idempotence - do nothing if already added.
// Ensures uniqueness in s_subscriptions[subId].consumers.
return;
}
// Initialize the nonce to 1, indicating the consumer is allocated.
nonces[subId] = 1;
// consumerConfig.nonce is 0 if the consumer had never sent a request to this subscription
// otherwise, consumerConfig.nonce is non-zero
// in both cases, use consumerConfig.nonce as is and set active status to true
consumerConfig.active = true;
consumers.push(consumer);

emit SubscriptionConsumerAdded(subId, consumer);
Expand Down
24 changes: 14 additions & 10 deletions contracts/src/v0.8/vrf/dev/VRFCoordinatorV2_5.sol
Original file line number Diff line number Diff line change
Expand Up @@ -263,10 +263,9 @@ contract VRFCoordinatorV2_5 is VRF, SubscriptionAPI, IVRFCoordinatorV2Plus {
}
// Its important to ensure that the consumer is in fact who they say they
// are, otherwise they could use someone else's subscription balance.
// A nonce of 0 indicates consumer is not allocated to the sub.
mapping(uint256 => uint64) storage nonces = s_consumers[msg.sender];
uint64 nonce = nonces[subId];
if (nonce == 0) {
mapping(uint256 => ConsumerConfig) storage consumerConfigs = s_consumers[msg.sender];
ConsumerConfig memory consumerConfig = consumerConfigs[subId];
if (!consumerConfig.active) {
revert InvalidConsumer(subId, msg.sender);
}
// Input validation using the config storage word.
Expand All @@ -293,9 +292,9 @@ contract VRFCoordinatorV2_5 is VRF, SubscriptionAPI, IVRFCoordinatorV2Plus {
// Note we do not check whether the keyHash is valid to save gas.
// The consequence for users is that they can send requests
// for invalid keyHashes which will simply not be fulfilled.
++nonce;
++consumerConfig.nonce;
uint256 preSeed;
(requestId, preSeed) = _computeRequestId(req.keyHash, msg.sender, subId, nonce);
(requestId, preSeed) = _computeRequestId(req.keyHash, msg.sender, subId, consumerConfig.nonce);

bytes memory extraArgsBytes = VRFV2PlusClient._argsToBytes(_fromBytes(req.extraArgs));
s_requestCommitments[requestId] = keccak256(
Expand All @@ -320,7 +319,7 @@ contract VRFCoordinatorV2_5 is VRF, SubscriptionAPI, IVRFCoordinatorV2Plus {
extraArgsBytes,
msg.sender
);
nonces[subId] = nonce;
consumerConfigs[subId] = consumerConfig;

return requestId;
}
Expand Down Expand Up @@ -630,7 +629,12 @@ contract VRFCoordinatorV2_5 is VRF, SubscriptionAPI, IVRFCoordinatorV2Plus {
for (uint256 i = 0; i < consumersLength; ++i) {
address consumer = consumers[i];
for (uint256 j = 0; j < provingKeyHashesLength; ++j) {
(uint256 reqId, ) = _computeRequestId(s_provingKeyHashes[j], consumer, subId, s_consumers[consumer][subId]);
(uint256 reqId, ) = _computeRequestId(
s_provingKeyHashes[j],
consumer,
subId,
s_consumers[consumer][subId].nonce
);
if (s_requestCommitments[reqId] != 0) {
return true;
}
Expand All @@ -646,7 +650,7 @@ contract VRFCoordinatorV2_5 is VRF, SubscriptionAPI, IVRFCoordinatorV2Plus {
if (pendingRequestExists(subId)) {
revert PendingRequestExists();
}
if (s_consumers[consumer][subId] == 0) {
if (!s_consumers[consumer][subId].active) {
revert InvalidConsumer(subId, consumer);
}
// Note bounded by MAX_CONSUMERS
Expand All @@ -662,7 +666,7 @@ contract VRFCoordinatorV2_5 is VRF, SubscriptionAPI, IVRFCoordinatorV2Plus {
break;
}
}
delete s_consumers[consumer][subId];
s_consumers[consumer][subId].active = false;
emit SubscriptionConsumerRemoved(subId, consumer);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -240,8 +240,9 @@ contract VRFCoordinatorV2PlusUpgradedVersion is
// Its important to ensure that the consumer is in fact who they say they
// are, otherwise they could use someone else's subscription balance.
// A nonce of 0 indicates consumer is not allocated to the sub.
uint64 currentNonce = s_consumers[msg.sender][req.subId];
if (currentNonce == 0) {
mapping(uint256 => ConsumerConfig) storage consumerConfigs = s_consumers[msg.sender];
ConsumerConfig memory consumerConfig = consumerConfigs[req.subId];
if (!consumerConfig.active) {
revert InvalidConsumer(req.subId, msg.sender);
}
// Input validation using the config storage word.
Expand All @@ -267,8 +268,8 @@ contract VRFCoordinatorV2PlusUpgradedVersion is
// Note we do not check whether the keyHash is valid to save gas.
// The consequence for users is that they can send requests
// for invalid keyHashes which will simply not be fulfilled.
uint64 nonce = currentNonce + 1;
(uint256 requestId, uint256 preSeed) = _computeRequestId(req.keyHash, msg.sender, req.subId, nonce);
++consumerConfig.nonce;
(uint256 requestId, uint256 preSeed) = _computeRequestId(req.keyHash, msg.sender, req.subId, consumerConfig.nonce);

VRFV2PlusClient.ExtraArgsV1 memory extraArgs = _fromBytes(req.extraArgs);
bytes memory extraArgsBytes = VRFV2PlusClient._argsToBytes(extraArgs);
Expand All @@ -294,7 +295,7 @@ contract VRFCoordinatorV2PlusUpgradedVersion is
extraArgsBytes,
msg.sender
);
s_consumers[msg.sender][req.subId] = nonce;
s_consumers[msg.sender][req.subId] = consumerConfig;

return requestId;
}
Expand Down Expand Up @@ -548,7 +549,7 @@ contract VRFCoordinatorV2PlusUpgradedVersion is
s_provingKeyHashes[j],
subConfig.consumers[i],
subId,
s_consumers[subConfig.consumers[i]][subId]
s_consumers[subConfig.consumers[i]][subId].nonce
);
if (s_requestCommitments[reqId] != 0) {
return true;
Expand All @@ -565,7 +566,7 @@ contract VRFCoordinatorV2PlusUpgradedVersion is
if (pendingRequestExists(subId)) {
revert PendingRequestExists();
}
if (s_consumers[consumer][subId] == 0) {
if (!s_consumers[consumer][subId].active) {
revert InvalidConsumer(subId, consumer);
}
// Note bounded by MAX_CONSUMERS
Expand Down Expand Up @@ -712,7 +713,7 @@ contract VRFCoordinatorV2PlusUpgradedVersion is
}

for (uint256 i = 0; i < migrationData.consumers.length; i++) {
s_consumers[migrationData.consumers[i]][migrationData.subId] = 1;
s_consumers[migrationData.consumers[i]][migrationData.subId] = ConsumerConfig({active: true, nonce: 0});
}

s_subscriptions[migrationData.subId] = Subscription({
Expand Down
Loading

0 comments on commit 2bd210b

Please sign in to comment.