Pass through extraArgs from user to VRFV2PlusWrapper (#10816)

* Pass through extraArgs from user to VRFV2PlusWrapper via VRFV2PlusWrapperConsumerBase * Generate go wrappers and prettify * Remove abi.decode, removeDomainSeparator and add a simpler checkPaymentMode * Fix issue in VRFV2PlusWrapperLoadTestConsumer::makeRequests * Fix comments --------- Co-authored-by: Sri Kidambi <[email protected]> Co-authored-by: Makram <[email protected]>
smartcontractkit · Oct 2, 2023 · 07d4038 · 07d4038
1 parent 6ae09dc
commit 07d4038
Show file tree

Hide file tree

Showing 10 changed files with 95 additions and 30 deletions.
diff --git a/contracts/src/v0.8/dev/interfaces/VRFV2PlusWrapperInterface.sol b/contracts/src/v0.8/dev/interfaces/VRFV2PlusWrapperInterface.sol
@@ -65,6 +65,7 @@ interface VRFV2PlusWrapperInterface {
   function requestRandomWordsInNative(
     uint32 _callbackGasLimit,
     uint16 _requestConfirmations,
-    uint32 _numWords
+    uint32 _numWords,
+    bytes memory extraArgs
   ) external payable returns (uint256 requestId);
 }
diff --git a/contracts/src/v0.8/dev/vrf/VRFV2PlusWrapper.sol b/contracts/src/v0.8/dev/vrf/VRFV2PlusWrapper.sol
@@ -20,6 +20,9 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
 
   error LinkAlreadySet();
   error FailedToTransferLink();
+  error IncorrectExtraArgsLength(uint16 expectedMinimumLength, uint16 actualLength);
+  error NativePaymentInOnTokenTransfer();
+  error LINKPaymentInRequestRandomWordsInNative();
 
   /* Storage Slot 1: BEGIN */
   // s_keyHash is the key hash to use when requesting randomness. Fees are paid based on current gas
@@ -104,6 +107,8 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
 
   // s_maxNumWords is the max number of words that can be requested in a single wrapped VRF request.
   uint8 s_maxNumWords;
+
+  uint16 private constant EXPECTED_MIN_LENGTH = 36;
   /* Storage Slot 8: END */
 
   struct Callback {
@@ -358,10 +363,11 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
   function onTokenTransfer(address _sender, uint256 _amount, bytes calldata _data) external onlyConfiguredNotDisabled {
     require(msg.sender == address(s_link), "only callable from LINK");
 
-    (uint32 callbackGasLimit, uint16 requestConfirmations, uint32 numWords) = abi.decode(
+    (uint32 callbackGasLimit, uint16 requestConfirmations, uint32 numWords, bytes memory extraArgs) = abi.decode(
       _data,
-      (uint32, uint16, uint32)
+      (uint32, uint16, uint32, bytes)
     );
+    checkPaymentMode(extraArgs, true);
     uint32 eip150Overhead = getEIP150Overhead(callbackGasLimit);
     int256 weiPerUnitLink = getFeedData();
     uint256 price = calculateRequestPriceInternal(callbackGasLimit, tx.gasprice, weiPerUnitLink);
@@ -373,7 +379,7 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
       requestConfirmations: requestConfirmations,
       callbackGasLimit: callbackGasLimit + eip150Overhead + s_wrapperGasOverhead,
       numWords: numWords,
-      extraArgs: "" // empty extraArgs defaults to link payment
+      extraArgs: extraArgs // empty extraArgs defaults to link payment
     });
     uint256 requestId = s_vrfCoordinator.requestRandomWords(req);
     s_callbacks[requestId] = Callback({
@@ -384,11 +390,38 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
     lastRequestId = requestId;
   }
 
+  function checkPaymentMode(bytes memory extraArgs, bool isLinkMode) public pure {
+    // If extraArgs is empty, payment mode is LINK by default
+    if (extraArgs.length == 0) {
+      if (!isLinkMode) {
+        revert LINKPaymentInRequestRandomWordsInNative();
+      }
+      return;
+    }
+    if (extraArgs.length < EXPECTED_MIN_LENGTH) {
+      revert IncorrectExtraArgsLength(EXPECTED_MIN_LENGTH, uint16(extraArgs.length));
+    }
+    // ExtraArgsV1 only has struct {bool nativePayment} as of now
+    // The following condition checks if nativePayment in abi.encode of
+    // ExtraArgsV1 matches the appropriate function call (onTokenTransfer
+    // for LINK and requestRandomWordsInNative for Native payment)
+    bool nativePayment = extraArgs[35] == hex"01";
+    if (nativePayment && isLinkMode) {
+      revert NativePaymentInOnTokenTransfer();
+    }
+    if (!nativePayment && !isLinkMode) {
+      revert LINKPaymentInRequestRandomWordsInNative();
+    }
+  }
+
   function requestRandomWordsInNative(
     uint32 _callbackGasLimit,
     uint16 _requestConfirmations,
-    uint32 _numWords
+    uint32 _numWords,
+    bytes calldata extraArgs
   ) external payable override returns (uint256 requestId) {
+    checkPaymentMode(extraArgs, false);
+
     uint32 eip150Overhead = getEIP150Overhead(_callbackGasLimit);
     uint256 price = calculateRequestPriceNativeInternal(_callbackGasLimit, tx.gasprice);
     require(msg.value >= price, "fee too low");
@@ -399,7 +432,7 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
       requestConfirmations: _requestConfirmations,
       callbackGasLimit: _callbackGasLimit + eip150Overhead + s_wrapperGasOverhead,
       numWords: _numWords,
-      extraArgs: VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: true}))
+      extraArgs: extraArgs
     });
     requestId = s_vrfCoordinator.requestRandomWords(req);
     s_callbacks[requestId] = Callback({

diff --git a/contracts/src/v0.8/dev/vrf/VRFV2PlusWrapperConsumerBase.sol b/contracts/src/v0.8/dev/vrf/VRFV2PlusWrapperConsumerBase.sol
@@ -73,27 +73,30 @@ abstract contract VRFV2PlusWrapperConsumerBase {
   function requestRandomness(
     uint32 _callbackGasLimit,
     uint16 _requestConfirmations,
-    uint32 _numWords
+    uint32 _numWords,
+    bytes memory extraArgs
   ) internal returns (uint256 requestId) {
     LINK.transferAndCall(
       address(VRF_V2_PLUS_WRAPPER),
       VRF_V2_PLUS_WRAPPER.calculateRequestPrice(_callbackGasLimit),
-      abi.encode(_callbackGasLimit, _requestConfirmations, _numWords)
+      abi.encode(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs)
     );
     return VRF_V2_PLUS_WRAPPER.lastRequestId();
   }
 
   function requestRandomnessPayInNative(
     uint32 _callbackGasLimit,
     uint16 _requestConfirmations,
-    uint32 _numWords
+    uint32 _numWords,
+    bytes memory extraArgs
   ) internal returns (uint256 requestId) {
     uint256 requestPrice = VRF_V2_PLUS_WRAPPER.calculateRequestPriceNative(_callbackGasLimit);
     return
       VRF_V2_PLUS_WRAPPER.requestRandomWordsInNative{value: requestPrice}(
         _callbackGasLimit,
         _requestConfirmations,
-        _numWords
+        _numWords,
+        extraArgs
       );
   }
 

diff --git a/contracts/src/v0.8/dev/vrf/testhelpers/VRFV2PlusWrapperConsumerExample.sol b/contracts/src/v0.8/dev/vrf/testhelpers/VRFV2PlusWrapperConsumerExample.sol
@@ -3,6 +3,7 @@ pragma solidity ^0.8.6;
 
 import "../VRFV2PlusWrapperConsumerBase.sol";
 import "../../../shared/access/ConfirmedOwner.sol";
+import {VRFV2PlusClient} from "../libraries/VRFV2PlusClient.sol";
 
 contract VRFV2PlusWrapperConsumerExample is VRFV2PlusWrapperConsumerBase, ConfirmedOwner {
   event WrappedRequestFulfilled(uint256 requestId, uint256[] randomWords, uint256 payment);
@@ -27,7 +28,8 @@ contract VRFV2PlusWrapperConsumerExample is VRFV2PlusWrapperConsumerBase, Confir
     uint16 _requestConfirmations,
     uint32 _numWords
   ) external onlyOwner returns (uint256 requestId) {
-    requestId = requestRandomness(_callbackGasLimit, _requestConfirmations, _numWords);
+    bytes memory extraArgs = VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: false}));
+    requestId = requestRandomness(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs);
     uint256 paid = VRF_V2_PLUS_WRAPPER.calculateRequestPrice(_callbackGasLimit);
     s_requests[requestId] = RequestStatus({paid: paid, randomWords: new uint256[](0), fulfilled: false, native: false});
     emit WrapperRequestMade(requestId, paid);
@@ -39,7 +41,8 @@ contract VRFV2PlusWrapperConsumerExample is VRFV2PlusWrapperConsumerBase, Confir
     uint16 _requestConfirmations,
     uint32 _numWords
   ) external onlyOwner returns (uint256 requestId) {
-    requestId = requestRandomnessPayInNative(_callbackGasLimit, _requestConfirmations, _numWords);
+    bytes memory extraArgs = VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: true}));
+    requestId = requestRandomnessPayInNative(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs);
     uint256 paid = VRF_V2_PLUS_WRAPPER.calculateRequestPriceNative(_callbackGasLimit);
     s_requests[requestId] = RequestStatus({paid: paid, randomWords: new uint256[](0), fulfilled: false, native: true});
     emit WrapperRequestMade(requestId, paid);

diff --git a/contracts/src/v0.8/dev/vrf/testhelpers/VRFV2PlusWrapperLoadTestConsumer.sol b/contracts/src/v0.8/dev/vrf/testhelpers/VRFV2PlusWrapperLoadTestConsumer.sol
@@ -4,6 +4,7 @@ pragma solidity ^0.8.6;
 import "../VRFV2PlusWrapperConsumerBase.sol";
 import "../../../shared/access/ConfirmedOwner.sol";
 import "../../../ChainSpecificUtil.sol";
+import {VRFV2PlusClient} from "../libraries/VRFV2PlusClient.sol";
 
 contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, ConfirmedOwner {
   uint256 public s_responseCount;
@@ -42,7 +43,8 @@ contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, Confi
     uint16 _requestCount
   ) external onlyOwner {
     for (uint16 i = 0; i < _requestCount; i++) {
-      uint256 requestId = requestRandomness(_callbackGasLimit, _requestConfirmations, _numWords);
+      bytes memory extraArgs = VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: false}));
+      uint256 requestId = requestRandomness(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs);
       s_lastRequestId = requestId;
 
       uint256 requestBlockNumber = ChainSpecificUtil.getBlockNumber();
@@ -70,7 +72,8 @@ contract VRFV2PlusWrapperLoadTestConsumer is VRFV2PlusWrapperConsumerBase, Confi
     uint16 _requestCount
   ) external onlyOwner {
     for (uint16 i = 0; i < _requestCount; i++) {
-      uint256 requestId = requestRandomnessPayInNative(_callbackGasLimit, _requestConfirmations, _numWords);
+      bytes memory extraArgs = VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: true}));
+      uint256 requestId = requestRandomnessPayInNative(_callbackGasLimit, _requestConfirmations, _numWords, extraArgs);
       s_lastRequestId = requestId;
 
       uint256 requestBlockNumber = ChainSpecificUtil.getBlockNumber();

diff --git a/core/gethwrappers/generated/vrfv2plus_client/vrfv2plus_client.go b/core/gethwrappers/generated/vrfv2plus_client/vrfv2plus_client.go
diff --git a/core/gethwrappers/generated/vrfv2plus_wrapper/vrfv2plus_wrapper.go b/core/gethwrappers/generated/vrfv2plus_wrapper/vrfv2plus_wrapper.go