Enhance VRFV2PlusWrapperConsumerBase (#10966)

* VRFV2PlusWrapperConsumerBas to provide VRFWrapper and native balance * Made VRF_WRAPPER public immutable * Prettier fix * Refactor LINK token * Cascading fixes * Fix lint issues --------- Co-authored-by: Sri Kidambi <[email protected]>
smartcontractkit · Oct 17, 2023 · 93c75b1 · 93c75b1
1 parent f74c15a
commit 93c75b1
Show file tree

Hide file tree

Showing 7 changed files with 143 additions and 45 deletions.
diff --git a/contracts/src/v0.8/vrf/dev/VRFV2PlusWrapperConsumerBase.sol b/contracts/src/v0.8/vrf/dev/VRFV2PlusWrapperConsumerBase.sol
@@ -30,34 +30,33 @@ import {IVRFV2PlusWrapper} from "./interfaces/IVRFV2PlusWrapper.sol";
  */
 abstract contract VRFV2PlusWrapperConsumerBase {
   error LINKAlreadySet();
+  error OnlyVRFWrapperCanFulfill(address have, address want);
 
-  // solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
-  LinkTokenInterface internal LINK;
-  // solhint-disable-next-line chainlink-solidity/prefix-storage-variables-with-s-underscore
-  IVRFV2PlusWrapper internal VRF_V2_PLUS_WRAPPER;
+  LinkTokenInterface internal s_linkToken;
+  IVRFV2PlusWrapper public immutable i_vrfV2PlusWrapper;
 
   /**
    * @param _link is the address of LinkToken
    * @param _vrfV2PlusWrapper is the address of the VRFV2Wrapper contract
    */
   constructor(address _link, address _vrfV2PlusWrapper) {
     if (_link != address(0)) {
-      LINK = LinkTokenInterface(_link);
+      s_linkToken = LinkTokenInterface(_link);
     }
 
-    VRF_V2_PLUS_WRAPPER = IVRFV2PlusWrapper(_vrfV2PlusWrapper);
+    i_vrfV2PlusWrapper = IVRFV2PlusWrapper(_vrfV2PlusWrapper);
   }
 
   /**
    * @notice setLinkToken changes the LINK token address.
    * @param _link is the address of the new LINK token contract
    */
   function setLinkToken(address _link) external {
-    if (address(LINK) != address(0)) {
+    if (address(s_linkToken) != address(0)) {
       revert LINKAlreadySet();
     }
 
-    LINK = LinkTokenInterface(_link);
+    s_linkToken = LinkTokenInterface(_link);
   }
 
   /**
@@ -79,12 +78,12 @@ abstract contract VRFV2PlusWrapperConsumerBase {
     uint32 _numWords,
     bytes memory extraArgs
   ) internal returns (uint256 requestId) {
-    LINK.transferAndCall(
-      address(VRF_V2_PLUS_WRAPPER),
-      VRF_V2_PLUS_WRAPPER.calculateRequestPrice(_callbackGasLimit),
+    s_linkToken.transferAndCall(
+      address(i_vrfV2PlusWrapper),
+      i_vrfV2PlusWrapper.calculateRequestPrice(_callbackGasLimit),
       abi.encode(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs)
     );
-    return VRF_V2_PLUS_WRAPPER.lastRequestId();
+    return i_vrfV2PlusWrapper.lastRequestId();
   }
 
   // solhint-disable-next-line chainlink-solidity/prefix-internal-functions-with-underscore
@@ -94,9 +93,9 @@ abstract contract VRFV2PlusWrapperConsumerBase {
     uint32 _numWords,
     bytes memory extraArgs
   ) internal returns (uint256 requestId) {
-    uint256 requestPrice = VRF_V2_PLUS_WRAPPER.calculateRequestPriceNative(_callbackGasLimit);
+    uint256 requestPrice = i_vrfV2PlusWrapper.calculateRequestPriceNative(_callbackGasLimit);
     return
-      VRF_V2_PLUS_WRAPPER.requestRandomWordsInNative{value: requestPrice}(
+      i_vrfV2PlusWrapper.requestRandomWordsInNative{value: requestPrice}(
         _callbackGasLimit,
         _requestConfirmations,
         _numWords,
@@ -115,8 +114,20 @@ abstract contract VRFV2PlusWrapperConsumerBase {
   function fulfillRandomWords(uint256 _requestId, uint256[] memory _randomWords) internal virtual;
 
   function rawFulfillRandomWords(uint256 _requestId, uint256[] memory _randomWords) external {
-    // solhint-disable-next-line custom-errors
-    require(msg.sender == address(VRF_V2_PLUS_WRAPPER), "only VRF V2 Plus wrapper can fulfill");
+    address vrfWrapperAddr = address(i_vrfV2PlusWrapper);
+    if (msg.sender != vrfWrapperAddr) {
+      revert OnlyVRFWrapperCanFulfill(msg.sender, vrfWrapperAddr);
+    }
     fulfillRandomWords(_requestId, _randomWords);
   }
+
+  /// @notice getBalance returns the native balance of the consumer contract
+  function getBalance() public view returns (uint256) {
+    return address(this).balance;
+  }
+
+  /// @notice getLinkToken returns the link token contract
+  function getLinkToken() public view returns (LinkTokenInterface) {
+    return s_linkToken;
+  }
 }
diff --git a/contracts/src/v0.8/vrf/dev/testhelpers/VRFV2PlusWrapperConsumerExample.sol b/contracts/src/v0.8/vrf/dev/testhelpers/VRFV2PlusWrapperConsumerExample.sol
@@ -30,7 +30,7 @@ contract VRFV2PlusWrapperConsumerExample is VRFV2PlusWrapperConsumerBase, Confir
   ) external onlyOwner returns (uint256 requestId) {
     bytes memory extraArgs = VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: false}));
     requestId = requestRandomness(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs);
-    uint256 paid = VRF_V2_PLUS_WRAPPER.calculateRequestPrice(_callbackGasLimit);
+    uint256 paid = i_vrfV2PlusWrapper.calculateRequestPrice(_callbackGasLimit);
     s_requests[requestId] = RequestStatus({paid: paid, randomWords: new uint256[](0), fulfilled: false, native: false});
     emit WrapperRequestMade(requestId, paid);
     return requestId;
@@ -43,7 +43,7 @@ contract VRFV2PlusWrapperConsumerExample is VRFV2PlusWrapperConsumerBase, Confir
   ) external onlyOwner returns (uint256 requestId) {
     bytes memory extraArgs = VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: true}));
     requestId = requestRandomnessPayInNative(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs);
-    uint256 paid = VRF_V2_PLUS_WRAPPER.calculateRequestPriceNative(_callbackGasLimit);
+    uint256 paid = i_vrfV2PlusWrapper.calculateRequestPriceNative(_callbackGasLimit);
     s_requests[requestId] = RequestStatus({paid: paid, randomWords: new uint256[](0), fulfilled: false, native: true});
     emit WrapperRequestMade(requestId, paid);
     return requestId;
@@ -70,7 +70,7 @@ contract VRFV2PlusWrapperConsumerExample is VRFV2PlusWrapperConsumerBase, Confir
   /// @notice withdrawLink withdraws the amount specified in amount to the owner
   /// @param amount the amount to withdraw, in juels
   function withdrawLink(uint256 amount) external onlyOwner {
-    LINK.transfer(owner(), amount);
+    s_linkToken.transfer(owner(), amount);
   }
 
   /// @notice withdrawNative withdraws the amount specified in amount to the owner

diff --git a/contracts/src/v0.8/vrf/dev/testhelpers/VRFV2PlusWrapperLoadTestConsumer.sol b/contracts/src/v0.8/vrf/dev/testhelpers/VRFV2PlusWrapperLoadTestConsumer.sol
@@ -5,7 +5,6 @@ import {VRFV2PlusWrapperConsumerBase} from "../VRFV2PlusWrapperConsumerBase.sol"
 import {ConfirmedOwner} from "../../../shared/access/ConfirmedOwner.sol";
 import {ChainSpecificUtil} from "../../../ChainSpecificUtil.sol";
 import {VRFV2PlusClient} from "../libraries/VRFV2PlusClient.sol";
-import {IVRFV2PlusWrapper} from "../interfaces/IVRFV2PlusWrapper.sol";
 
 contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, ConfirmedOwner {
   uint256 public s_responseCount;
@@ -50,7 +49,7 @@ contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, Confi
       s_lastRequestId = requestId;
 
       uint256 requestBlockNumber = ChainSpecificUtil._getBlockNumber();
-      uint256 paid = VRF_V2_PLUS_WRAPPER.calculateRequestPrice(_callbackGasLimit);
+      uint256 paid = i_vrfV2PlusWrapper.calculateRequestPrice(_callbackGasLimit);
       s_requests[requestId] = RequestStatus({
         paid: paid,
         fulfilled: false,
@@ -79,7 +78,7 @@ contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, Confi
       s_lastRequestId = requestId;
 
       uint256 requestBlockNumber = ChainSpecificUtil._getBlockNumber();
-      uint256 paid = VRF_V2_PLUS_WRAPPER.calculateRequestPriceNative(_callbackGasLimit);
+      uint256 paid = i_vrfV2PlusWrapper.calculateRequestPriceNative(_callbackGasLimit);
       s_requests[requestId] = RequestStatus({
         paid: paid,
         fulfilled: false,
@@ -161,7 +160,7 @@ contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, Confi
   /// @notice withdrawLink withdraws the amount specified in amount to the owner
   /// @param amount the amount to withdraw, in juels
   function withdrawLink(uint256 amount) external onlyOwner {
-    LINK.transfer(owner(), amount);
+    s_linkToken.transfer(owner(), amount);
   }
 
   /// @notice withdrawNative withdraws the amount specified in amount to the owner
@@ -172,13 +171,5 @@ contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, Confi
     require(success, "withdrawNative failed");
   }
 
-  function getWrapper() external view returns (IVRFV2PlusWrapper) {
-    return VRF_V2_PLUS_WRAPPER;
-  }
-
   receive() external payable {}
-
-  function getBalance() public view returns (uint256) {
-    return address(this).balance;
-  }
 }