Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add reentrancy guard to VRFCoordinatorV2Mock to align its behaviour w… #10585

Merged
merged 3 commits into from
Sep 12, 2023
Merged

Add reentrancy guard to VRFCoordinatorV2Mock to align its behaviour w… #10585

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7f42838
Add reentrancy guard to VRFCoordinatorV2Mock to align its behaviour w…
kidambisrinivas Sep 11, 2023
914bf65
Prettier fix
kidambisrinivas Sep 11, 2023
899f5a6
Added nonReentrant for cancelSubscription, and removeConsumer
kidambisrinivas Sep 11, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions contracts/scripts/native_solc_compile_all_vrf
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ compileContract vrf/BatchBlockhashStore.sol
compileContract vrf/BatchVRFCoordinatorV2.sol
compileContract vrf/testhelpers/VRFCoordinatorV2TestHelper.sol
compileContractAltOpts vrf/VRFCoordinatorV2.sol 10000
compileContract mocks/VRFCoordinatorV2Mock.sol
compileContract vrf/VRFOwner.sol

# VRF V2Plus
Expand Down
38 changes: 32 additions & 6 deletions contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,9 @@ pragma solidity ^0.8.4;
import "../shared/interfaces/LinkTokenInterface.sol";
import "../interfaces/VRFCoordinatorV2Interface.sol";
import "../vrf/VRFConsumerBaseV2.sol";
import "../shared/access/ConfirmedOwner.sol";

contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface, ConfirmedOwner {
uint96 public immutable BASE_FEE;
uint96 public immutable GAS_PRICE_LINK;
uint16 public immutable MAX_CONSUMERS = 100;
Expand All @@ -17,6 +18,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
error TooManyConsumers();
error InvalidConsumer();
error InvalidRandomWords();
error Reentrant();

event RandomWordsRequested(
bytes32 indexed keyHash,
Expand All @@ -34,7 +36,13 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
event SubscriptionCanceled(uint64 indexed subId, address to, uint256 amount);
event ConsumerAdded(uint64 indexed subId, address consumer);
event ConsumerRemoved(uint64 indexed subId, address consumer);
event ConfigSet();

struct Config {
// Reentrancy protection.
bool reentrancyLock;
}
Config private s_config;
uint64 s_currentSubId;
uint256 s_nextRequestId = 1;
uint256 s_nextPreSeed = 100;
Expand All @@ -52,9 +60,18 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
}
mapping(uint256 => Request) s_requests; /* requestId */ /* request */

constructor(uint96 _baseFee, uint96 _gasPriceLink) {
constructor(uint96 _baseFee, uint96 _gasPriceLink) ConfirmedOwner(msg.sender) {
BASE_FEE = _baseFee;
GAS_PRICE_LINK = _gasPriceLink;
setConfig();
}

/**
* @notice Sets the configuration of the vrfv2 mock coordinator
*/
function setConfig() public onlyOwner {
kidambisrinivas marked this conversation as resolved.
Show resolved Hide resolved
s_config = Config({reentrancyLock: false});
emit ConfigSet();
}

function consumerIsAdded(uint64 _subId, address _consumer) public view returns (bool) {
Expand Down Expand Up @@ -85,7 +102,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
* @param _requestId the request to fulfill
* @param _consumer the VRF randomness consumer to send the result to
*/
function fulfillRandomWords(uint256 _requestId, address _consumer) external {
function fulfillRandomWords(uint256 _requestId, address _consumer) external nonReentrant {
fulfillRandomWordsWithOverride(_requestId, _consumer, new uint256[](0));
}

Expand Down Expand Up @@ -114,7 +131,9 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {

VRFConsumerBaseV2 v;
bytes memory callReq = abi.encodeWithSelector(v.rawFulfillRandomWords.selector, _requestId, _words);
s_config.reentrancyLock = true;
(bool success, ) = _consumer.call{gas: req.callbackGasLimit}(callReq);
s_config.reentrancyLock = false;

uint96 payment = uint96(BASE_FEE + ((startGas - gasleft()) * GAS_PRICE_LINK));
if (s_subscriptions[req.subId].balance < payment) {
Expand Down Expand Up @@ -146,7 +165,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
uint16 _minimumRequestConfirmations,
uint32 _callbackGasLimit,
uint32 _numWords
) external override onlyValidConsumer(_subId, msg.sender) returns (uint256) {
) external override nonReentrant onlyValidConsumer(_subId, msg.sender) returns (uint256) {
kidambisrinivas marked this conversation as resolved.
Show resolved Hide resolved
if (s_subscriptions[_subId].owner == address(0)) {
revert InvalidSubscription();
}
Expand Down Expand Up @@ -185,7 +204,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
return (s_subscriptions[_subId].balance, 0, s_subscriptions[_subId].owner, s_consumers[_subId]);
}

function cancelSubscription(uint64 _subId, address _to) external override onlySubOwner(_subId) {
function cancelSubscription(uint64 _subId, address _to) external override onlySubOwner(_subId) nonReentrant {
emit SubscriptionCanceled(_subId, _to, s_subscriptions[_subId].balance);
delete (s_subscriptions[_subId]);
}
Expand Down Expand Up @@ -221,7 +240,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
function removeConsumer(
uint64 _subId,
address _consumer
) external override onlySubOwner(_subId) onlyValidConsumer(_subId, _consumer) {
) external override onlySubOwner(_subId) onlyValidConsumer(_subId, _consumer) nonReentrant {
address[] storage consumers = s_consumers[_subId];
for (uint256 i = 0; i < consumers.length; i++) {
if (consumers[i] == _consumer) {
Expand Down Expand Up @@ -276,6 +295,13 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface {
);
}

modifier nonReentrant() {
if (s_config.reentrancyLock) {
revert Reentrant();
}
_;
}

function getFallbackWeiPerUnitLink() external view returns (int256) {
return 4000000000000000; // 0.004 Ether
}
Expand Down
4 changes: 2 additions & 2 deletions contracts/test/v0.8/dev/VRFCoordinatorV2Mock.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -262,7 +262,7 @@ describe('VRFCoordinatorV2Mock', () => {
expect(receipt.events[0].args['success']).to.equal(true)
assert(
receipt.events[0].args['payment']
.sub(BigNumber.from('100119017000000000'))
.sub(BigNumber.from('100119403000000000'))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If you have the time, some foundry tests on the mockfile could be nice. However, I don't think it's worth the effort ever adding typescript/hardhat tests, since we don't really plan on working with that software going forward.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure, I'll try to add Foundry tests for VRFCoordinatorV2Mock. Shall I add it as a follow up PR to this?

.lt(BigNumber.from('10000000000')),
)

Expand Down Expand Up @@ -315,7 +315,7 @@ describe('VRFCoordinatorV2Mock', () => {
expect(receipt.events[0].args['success']).to.equal(true)
assert(
receipt.events[0].args['payment']
.sub(BigNumber.from('100119017000000000'))
.sub(BigNumber.from('100120516000000000'))
.lt(BigNumber.from('10000000000')),
)

Expand Down