From 7f428380a86892c76d5ee6d32b2ce5dcae07bb7a Mon Sep 17 00:00:00 2001 From: Sri Kidambi Date: Mon, 11 Sep 2023 18:12:40 +0100 Subject: [PATCH 1/3] Add reentrancy guard to VRFCoordinatorV2Mock to align its behaviour with VRFCoordinatorV2 --- contracts/scripts/native_solc_compile_all_vrf | 1 + .../src/v0.8/mocks/VRFCoordinatorV2Mock.sol | 36 ++++++++++++++++--- .../v0.8/dev/VRFCoordinatorV2Mock.test.ts | 4 +-- 3 files changed, 35 insertions(+), 6 deletions(-) diff --git a/contracts/scripts/native_solc_compile_all_vrf b/contracts/scripts/native_solc_compile_all_vrf index 72be6eebf8d..feebbdcf145 100755 --- a/contracts/scripts/native_solc_compile_all_vrf +++ b/contracts/scripts/native_solc_compile_all_vrf @@ -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 diff --git a/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol b/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol index 8dfb0dca7c5..05714a549ad 100644 --- a/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol +++ b/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol @@ -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; @@ -17,6 +18,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface { error TooManyConsumers(); error InvalidConsumer(); error InvalidRandomWords(); + error Reentrant(); event RandomWordsRequested( bytes32 indexed keyHash, @@ -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; @@ -52,9 +60,20 @@ 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 { + s_config = Config({ + reentrancyLock: false + }); + emit ConfigSet(); } function consumerIsAdded(uint64 _subId, address _consumer) public view returns (bool) { @@ -85,7 +104,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)); } @@ -114,7 +133,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) { @@ -146,7 +167,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) { if (s_subscriptions[_subId].owner == address(0)) { revert InvalidSubscription(); } @@ -276,6 +297,13 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface { ); } + modifier nonReentrant() { + if (s_config.reentrancyLock) { + revert Reentrant(); + } + _; + } + function getFallbackWeiPerUnitLink() external view returns (int256) { return 4000000000000000; // 0.004 Ether } diff --git a/contracts/test/v0.8/dev/VRFCoordinatorV2Mock.test.ts b/contracts/test/v0.8/dev/VRFCoordinatorV2Mock.test.ts index b3446ddb154..b0a2a10b201 100644 --- a/contracts/test/v0.8/dev/VRFCoordinatorV2Mock.test.ts +++ b/contracts/test/v0.8/dev/VRFCoordinatorV2Mock.test.ts @@ -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')) .lt(BigNumber.from('10000000000')), ) @@ -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')), ) From 914bf65b5725a667a3c959b11287b66823f4d831 Mon Sep 17 00:00:00 2001 From: Sri Kidambi Date: Mon, 11 Sep 2023 18:22:44 +0100 Subject: [PATCH 2/3] Prettier fix --- contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol b/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol index 05714a549ad..5b23eabbcec 100644 --- a/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol +++ b/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol @@ -70,9 +70,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface, ConfirmedOwner { * @notice Sets the configuration of the vrfv2 mock coordinator */ function setConfig() public onlyOwner { - s_config = Config({ - reentrancyLock: false - }); + s_config = Config({reentrancyLock: false}); emit ConfigSet(); } From 899f5a64c30e7db9c55a7dbb3df719d8a6987c63 Mon Sep 17 00:00:00 2001 From: Sri Kidambi Date: Tue, 12 Sep 2023 00:11:44 +0100 Subject: [PATCH 3/3] Added nonReentrant for cancelSubscription, and removeConsumer --- contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol b/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol index 5b23eabbcec..34a3334189d 100644 --- a/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol +++ b/contracts/src/v0.8/mocks/VRFCoordinatorV2Mock.sol @@ -204,7 +204,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface, ConfirmedOwner { 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]); } @@ -240,7 +240,7 @@ contract VRFCoordinatorV2Mock is VRFCoordinatorV2Interface, ConfirmedOwner { 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) {