Skip to content

Commit

Permalink
Fix example consumer contracts to be migratable (#10982)
Browse files Browse the repository at this point in the history 
* Fix example consumer contracts to be migratable

* Generate go wrappers after merge

* Generate go wrappers- fix

---------

Co-authored-by: Sri Kidambi <[email protected]>
  • Loading branch information
@kidambisrinivas
kidambisrinivas and kidambisrinivas authored Oct 17, 2023
1 parent c89866d commit edaefc6
Show file tree
Hide file tree
Showing 10 changed files with 30 additions and 44 deletions.
2 changes: 2 additions & 0 deletions contracts/src/v0.8/vrf/dev/VRFConsumerBaseV2Plus.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,8 @@ abstract contract VRFConsumerBaseV2Plus is IVRFMigratableConsumerV2Plus, Confirm
error OnlyOwnerOrCoordinator(address have, address owner, address coordinator);
error ZeroAddress();

// s_vrfCoordinator should be used by consumers to make requests to vrfCoordinator
// so that coordinator reference is updated after migration
IVRFCoordinatorV2Plus public s_vrfCoordinator;

/**
Expand Down
16 changes: 6 additions & 10 deletions contracts/src/v0.8/vrf/dev/testhelpers/VRFMaliciousConsumerV2Plus.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,23 +2,19 @@
pragma solidity ^0.8.0;

import {LinkTokenInterface} from "../../../shared/interfaces/LinkTokenInterface.sol";
import {IVRFCoordinatorV2Plus} from "../interfaces/IVRFCoordinatorV2Plus.sol";
import {VRFConsumerBaseV2Plus} from "../VRFConsumerBaseV2Plus.sol";
import {VRFV2PlusClient} from "../libraries/VRFV2PlusClient.sol";

contract VRFMaliciousConsumerV2Plus is VRFConsumerBaseV2Plus {
uint256[] public s_randomWords;
uint256 public s_requestId;
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
IVRFCoordinatorV2Plus internal COORDINATOR;
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
LinkTokenInterface internal LINKTOKEN;
uint256 public s_gasAvailable;
uint256 internal s_subId;
bytes32 internal s_keyHash;

constructor(address vrfCoordinator, address link) VRFConsumerBaseV2Plus(vrfCoordinator) {
COORDINATOR = IVRFCoordinatorV2Plus(vrfCoordinator);
LINKTOKEN = LinkTokenInterface(link);
}

Expand All @@ -36,23 +32,23 @@ contract VRFMaliciousConsumerV2Plus is VRFConsumerBaseV2Plus {
extraArgs: "" // empty extraArgs defaults to link payment
});
// Should revert
COORDINATOR.requestRandomWords(req);
s_vrfCoordinator.requestRandomWords(req);
}

function createSubscriptionAndFund(uint96 amount) external {
if (s_subId == 0) {
s_subId = COORDINATOR.createSubscription();
COORDINATOR.addConsumer(s_subId, address(this));
s_subId = s_vrfCoordinator.createSubscription();
s_vrfCoordinator.addConsumer(s_subId, address(this));
}
// Approve the link transfer.
LINKTOKEN.transferAndCall(address(COORDINATOR), amount, abi.encode(s_subId));
LINKTOKEN.transferAndCall(address(s_vrfCoordinator), amount, abi.encode(s_subId));
}

function updateSubscription(address[] memory consumers) external {
// solhint-disable-next-line custom-errors
require(s_subId != 0, "subID not set");
for (uint256 i = 0; i < consumers.length; i++) {
COORDINATOR.addConsumer(s_subId, consumers[i]);
s_vrfCoordinator.addConsumer(s_subId, consumers[i]);
}
}

Expand All @@ -66,6 +62,6 @@ contract VRFMaliciousConsumerV2Plus is VRFConsumerBaseV2Plus {
numWords: 1,
extraArgs: "" // empty extraArgs defaults to link payment
});
return COORDINATOR.requestRandomWords(req);
return s_vrfCoordinator.requestRandomWords(req);
}
}
6 changes: 1 addition & 5 deletions contracts/src/v0.8/vrf/dev/testhelpers/VRFV2PlusExternalSubOwnerExample.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,14 +2,11 @@
pragma solidity ^0.8.0;

import {LinkTokenInterface} from "../../../shared/interfaces/LinkTokenInterface.sol";
import {IVRFCoordinatorV2Plus} from "../interfaces/IVRFCoordinatorV2Plus.sol";
import {VRFConsumerBaseV2Plus} from "../VRFConsumerBaseV2Plus.sol";
import {VRFV2PlusClient} from "../libraries/VRFV2PlusClient.sol";

/// @notice This contract is used for testing only and should not be used for production.
contract VRFV2PlusExternalSubOwnerExample is VRFConsumerBaseV2Plus {
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
IVRFCoordinatorV2Plus internal COORDINATOR;
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
LinkTokenInterface internal LINKTOKEN;

Expand All @@ -18,7 +15,6 @@ contract VRFV2PlusExternalSubOwnerExample is VRFConsumerBaseV2Plus {
address internal s_owner;

constructor(address vrfCoordinator, address link) VRFConsumerBaseV2Plus(vrfCoordinator) {
COORDINATOR = IVRFCoordinatorV2Plus(vrfCoordinator);
LINKTOKEN = LinkTokenInterface(link);
s_owner = msg.sender;
}
Expand Down Expand Up @@ -47,6 +43,6 @@ contract VRFV2PlusExternalSubOwnerExample is VRFConsumerBaseV2Plus {
extraArgs: VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: nativePayment}))
});
// Will revert if subscription is not funded.
s_requestId = COORDINATOR.requestRandomWords(req);
s_requestId = s_vrfCoordinator.requestRandomWords(req);
}
}
16 changes: 6 additions & 10 deletions contracts/src/v0.8/vrf/dev/testhelpers/VRFV2PlusRevertingExample.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
pragma solidity ^0.8.0;

import {LinkTokenInterface} from "../../../shared/interfaces/LinkTokenInterface.sol";
import {IVRFCoordinatorV2Plus} from "../interfaces/IVRFCoordinatorV2Plus.sol";
import {VRFConsumerBaseV2Plus} from "../VRFConsumerBaseV2Plus.sol";
import {VRFV2PlusClient} from "../libraries/VRFV2PlusClient.sol";

Expand All @@ -11,14 +10,11 @@ contract VRFV2PlusRevertingExample is VRFConsumerBaseV2Plus {
uint256[] public s_randomWords;
uint256 public s_requestId;
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
IVRFCoordinatorV2Plus internal COORDINATOR;
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
LinkTokenInterface internal LINKTOKEN;
uint256 public s_subId;
uint256 public s_gasAvailable;

constructor(address vrfCoordinator, address link) VRFConsumerBaseV2Plus(vrfCoordinator) {
COORDINATOR = IVRFCoordinatorV2Plus(vrfCoordinator);
LINKTOKEN = LinkTokenInterface(link);
}

Expand All @@ -30,25 +26,25 @@ contract VRFV2PlusRevertingExample is VRFConsumerBaseV2Plus {

function createSubscriptionAndFund(uint96 amount) external {
if (s_subId == 0) {
s_subId = COORDINATOR.createSubscription();
COORDINATOR.addConsumer(s_subId, address(this));
s_subId = s_vrfCoordinator.createSubscription();
s_vrfCoordinator.addConsumer(s_subId, address(this));
}
// Approve the link transfer.
LINKTOKEN.transferAndCall(address(COORDINATOR), amount, abi.encode(s_subId));
LINKTOKEN.transferAndCall(address(s_vrfCoordinator), amount, abi.encode(s_subId));
}

function topUpSubscription(uint96 amount) external {
// solhint-disable-next-line custom-errors
require(s_subId != 0, "sub not set");
// Approve the link transfer.
LINKTOKEN.transferAndCall(address(COORDINATOR), amount, abi.encode(s_subId));
LINKTOKEN.transferAndCall(address(s_vrfCoordinator), amount, abi.encode(s_subId));
}

function updateSubscription(address[] memory consumers) external {
// solhint-disable-next-line custom-errors
require(s_subId != 0, "subID not set");
for (uint256 i = 0; i < consumers.length; i++) {
COORDINATOR.addConsumer(s_subId, consumers[i]);
s_vrfCoordinator.addConsumer(s_subId, consumers[i]);
}
}

Expand All @@ -67,7 +63,7 @@ contract VRFV2PlusRevertingExample is VRFConsumerBaseV2Plus {
numWords: numWords,
extraArgs: "" // empty extraArgs defaults to link payment
});
s_requestId = COORDINATOR.requestRandomWords(req);
s_requestId = s_vrfCoordinator.requestRandomWords(req);
return s_requestId;
}
}
18 changes: 7 additions & 11 deletions contracts/src/v0.8/vrf/dev/testhelpers/VRFV2PlusSingleConsumerExample.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,11 @@
pragma solidity ^0.8.0;

import {LinkTokenInterface} from "../../../shared/interfaces/LinkTokenInterface.sol";
import {IVRFCoordinatorV2Plus} from "../interfaces/IVRFCoordinatorV2Plus.sol";
import {VRFConsumerBaseV2Plus} from "../VRFConsumerBaseV2Plus.sol";
import {VRFV2PlusClient} from "../libraries/VRFV2PlusClient.sol";

/// @notice This contract is used for testing only and should not be used for production.
contract VRFV2PlusSingleConsumerExample is VRFConsumerBaseV2Plus {
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
IVRFCoordinatorV2Plus internal COORDINATOR;
// solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
LinkTokenInterface internal LINKTOKEN;

Expand All @@ -36,7 +33,6 @@ contract VRFV2PlusSingleConsumerExample is VRFConsumerBaseV2Plus {
bytes32 keyHash,
bool nativePayment
) VRFConsumerBaseV2Plus(vrfCoordinator) {
COORDINATOR = IVRFCoordinatorV2Plus(vrfCoordinator);
LINKTOKEN = LinkTokenInterface(link);
s_owner = msg.sender;
s_requestConfig = RequestConfig({
Expand Down Expand Up @@ -69,7 +65,7 @@ contract VRFV2PlusSingleConsumerExample is VRFConsumerBaseV2Plus {
extraArgs: VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: rc.nativePayment}))
});
// Will revert if subscription is not set and funded.
s_requestId = COORDINATOR.requestRandomWords(req);
s_requestId = s_vrfCoordinator.requestRandomWords(req);
}

// Assumes this contract owns link
Expand All @@ -78,7 +74,7 @@ contract VRFV2PlusSingleConsumerExample is VRFConsumerBaseV2Plus {
// with different link costs).
function fundAndRequestRandomWords(uint256 amount) external onlyOwner {
RequestConfig memory rc = s_requestConfig;
LINKTOKEN.transferAndCall(address(COORDINATOR), amount, abi.encode(s_requestConfig.subId));
LINKTOKEN.transferAndCall(address(s_vrfCoordinator), amount, abi.encode(s_requestConfig.subId));
VRFV2PlusClient.RandomWordsRequest memory req = VRFV2PlusClient.RandomWordsRequest({
keyHash: rc.keyHash,
subId: rc.subId,
Expand All @@ -88,12 +84,12 @@ contract VRFV2PlusSingleConsumerExample is VRFConsumerBaseV2Plus {
extraArgs: VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: rc.nativePayment}))
});
// Will revert if subscription is not set and funded.
s_requestId = COORDINATOR.requestRandomWords(req);
s_requestId = s_vrfCoordinator.requestRandomWords(req);
}

// Assumes this contract owns link
function topUpSubscription(uint256 amount) external onlyOwner {
LINKTOKEN.transferAndCall(address(COORDINATOR), amount, abi.encode(s_requestConfig.subId));
LINKTOKEN.transferAndCall(address(s_vrfCoordinator), amount, abi.encode(s_requestConfig.subId));
}

function withdraw(uint256 amount, address to) external onlyOwner {
Expand All @@ -102,7 +98,7 @@ contract VRFV2PlusSingleConsumerExample is VRFConsumerBaseV2Plus {

function unsubscribe(address to) external onlyOwner {
// Returns funds to this address
COORDINATOR.cancelSubscription(s_requestConfig.subId, to);
s_vrfCoordinator.cancelSubscription(s_requestConfig.subId, to);
s_requestConfig.subId = 0;
}

Expand All @@ -112,7 +108,7 @@ contract VRFV2PlusSingleConsumerExample is VRFConsumerBaseV2Plus {
// Create a subscription, current subId
address[] memory consumers = new address[](1);
consumers[0] = address(this);
s_requestConfig.subId = COORDINATOR.createSubscription();
COORDINATOR.addConsumer(s_requestConfig.subId, consumers[0]);
s_requestConfig.subId = s_vrfCoordinator.createSubscription();
s_vrfCoordinator.addConsumer(s_requestConfig.subId, consumers[0]);
}
}
Loading

0 comments on commit edaefc6

Please sign in to comment.