implement migrate method in coordinator v2 plus (#9807)

* implement migrate method in coordinator v2 plus

* add extra args field to requestRandomWords()

* add missing client file

* prettier

* prettier

* add forge tests for migrations

* address comments

* add missing file

* remove s_requestPayment()

* fix failing tests

* run prettier

contracts/src/v0.8/dev/interfaces/IVRFCoordinatorV2Plus.sol

@@ -3,6 +3,9 @@ pragma solidity ^0.8.0;
 
 import "../vrf/libraries/VRFV2PlusClient.sol";
 
+// Interface for initial version of VRFCoordinatorV2Plus
+// Functions in this interface may not be supported when VRFCoordinatorV2Plus is upgraded to a new version
+// TODO: Revisit these functions and decide which functions need to be backwards compatible
 interface IVRFCoordinatorV2Plus {
   /**
    * @notice Get configuration relevant for making requests
@@ -107,11 +110,4 @@ interface IVRFCoordinatorV2Plus {
    * otherwise.
    */
   function pendingRequestExists(uint64 subId) external view returns (bool);
-
-  /*
-   * @notice Check to see the payment made for the provided request id.
-   * @param requestId - ID of the request
-   * @return amountPaid - amount paid for the request
-   */
-  function s_requestPayments(uint256 requestId) external view returns (uint96);
 }

contracts/src/v0.8/dev/interfaces/IVRFCoordinatorV2PlusMigration.sol

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+// Future versions of VRFCoordinatorV2Plus must implement IVRFCoordinatorV2PlusMigration
+// to support migrations from previous versions
+interface IVRFCoordinatorV2PlusMigration {
+  /**
+   * @notice called by older versions of coordinator for migration.
+   * @notice only callable by older versions of coordinator
+   * @notice supports transfer of native currency
+   * @param encodedData - user data from older version of coordinator
+   * @return subId
+   */
+  function onMigration(bytes calldata encodedData) external payable returns (uint64 subId);
+}

contracts/src/v0.8/dev/interfaces/IVRFMigratableConsumerV2Plus.sol

@@ -5,7 +5,7 @@ pragma solidity ^0.8.0;
 /// @notice method required to be implemented by all V2Plus consumers.
 /// @dev This interface is designed to be used in VRFConsumerBaseV2Plus.
 interface IVRFMigratableConsumerV2Plus {
-  /// @notice Set the VRF Coordinator address for the consumer.
-  /// @notice This method is should only be callable by the subscription admin.
-  function setVRFCoordinator(address vrfCoordinator) external;
+  /// @notice Sets the VRF Coordinator address and sub ID
+  /// @notice This method is should only be callable by the coordinator or contract owner
+  function setConfig(address vrfCoordinator, uint64 subId) external;
 }

contracts/src/v0.8/dev/interfaces/IVRFMigratableCoordinatorV2Plus.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+import "../vrf/libraries/VRFV2PlusClient.sol";
+
+// Interface that enables consumers of VRFCoordinatorV2Plus to be future-proof for upgrades
+// This interface is supported by subsequent versions of VRFCoordinatorV2Plus
+interface IVRFMigratableCoordinatorV2Plus {
+  /**
+   * @notice Request a set of random words.
+   * @param req - a struct containing following fiels for randomness request:
+   * keyHash - Corresponds to a particular oracle job which uses
+   * that key for generating the VRF proof. Different keyHash's have different gas price
+   * ceilings, so you can select a specific one to bound your maximum per request cost.
+   * subId  - The ID of the VRF subscription. Must be funded
+   * with the minimum subscription balance required for the selected keyHash.
+   * minimumRequestConfirmations - How many blocks you'd like the
+   * oracle to wait before responding to the request. See SECURITY CONSIDERATIONS
+   * for why you may want to request more. The acceptable range is
+   * [minimumRequestBlockConfirmations, 200].
+   * callbackGasLimit - How much gas you'd like to receive in your
+   * fulfillRandomWords callback. Note that gasleft() inside fulfillRandomWords
+   * may be slightly less than this amount because of gas used calling the function
+   * (argument decoding etc.), so you may need to request slightly more than you expect
+   * to have inside fulfillRandomWords. The acceptable range is
+   * [0, maxGasLimit]
+   * numWords - The number of uint256 random values you'd like to receive
+   * in your fulfillRandomWords callback. Note these numbers are expanded in a
+   * secure way by the VRFCoordinator from a single random value supplied by the oracle.
+   * nativePayment - Whether payment should be made in ETH or LINK.
+   * @return requestId - A unique identifier of the request. Can be used to match
+   * a request to a response in fulfillRandomWords.
+   */
+  function requestRandomWords(VRFV2PlusClient.RandomWordsRequest calldata req) external returns (uint256 requestId);
+}

contracts/src/v0.8/dev/vrf/SubscriptionAPI.sol

@@ -218,7 +218,7 @@ abstract contract SubscriptionAPI is ConfirmedOwner, ReentrancyGuard, ERC677Rece
    */
   function getSubscription(
     uint64 subId
-  ) external view override returns (uint96 balance, uint96 ethBalance, address owner, address[] memory consumers) {
+  ) public view override returns (uint96 balance, uint96 ethBalance, address owner, address[] memory consumers) {
     if (s_subscriptionConfigs[subId].owner == address(0)) {
       revert InvalidSubscription();
     }
@@ -298,11 +298,11 @@ abstract contract SubscriptionAPI is ConfirmedOwner, ReentrancyGuard, ERC677Rece
     emit SubscriptionConsumerAdded(subId, consumer);
   }
 
-  function cancelSubscriptionHelper(uint64 subId, address to) internal {
+  function deleteSubscription(uint64 subId) internal returns (uint96 balance, uint96 ethBalance) {
     SubscriptionConfig memory subConfig = s_subscriptionConfigs[subId];
     Subscription memory sub = s_subscriptions[subId];
-    uint96 balance = sub.balance;
-    uint96 ethBalance = sub.ethBalance;
+    balance = sub.balance;
+    ethBalance = sub.ethBalance;
     // Note bounded by MAX_CONSUMERS;
     // If no consumers, does nothing.
     for (uint256 i = 0; i < subConfig.consumers.length; i++) {
@@ -312,6 +312,11 @@ abstract contract SubscriptionAPI is ConfirmedOwner, ReentrancyGuard, ERC677Rece
     delete s_subscriptions[subId];
     s_totalBalance -= balance;
     s_totalEthBalance -= ethBalance;
+    return (balance, ethBalance);
+  }
+
+  function cancelSubscriptionHelper(uint64 subId, address to) internal {
+    (uint96 balance, uint96 ethBalance) = deleteSubscription(subId);
     if (!LINK.transfer(to, uint256(balance))) {
       revert InsufficientBalance();
     }

contracts/src/v0.8/dev/vrf/VRFConsumerBaseV2Plus.sol

@@ -1,8 +1,9 @@
 // SPDX-License-Identifier: MIT
 pragma solidity ^0.8.4;
 
-import "../interfaces/IVRFCoordinatorV2Plus.sol";
+import "../interfaces/IVRFMigratableCoordinatorV2Plus.sol";
 import "../interfaces/IVRFMigratableConsumerV2Plus.sol";
+import "../../ConfirmedOwner.sol";
 
 /** ****************************************************************************
  * @notice Interface for contracts using VRF randomness
@@ -97,20 +98,19 @@ import "../interfaces/IVRFMigratableConsumerV2Plus.sol";
  * @dev responding to the request (however this is not enforced in the contract
  * @dev and so remains effective only in the case of unmodified oracle software).
  */
-abstract contract VRFConsumerBaseV2Plus is IVRFMigratableConsumerV2Plus {
+abstract contract VRFConsumerBaseV2Plus is IVRFMigratableConsumerV2Plus, ConfirmedOwner {
   error OnlyCoordinatorCanFulfill(address have, address want);
-  error OnlySubOwnerCanSetVRFCoordinator(address have, address want);
+  error OnlyOwnerOrCoordinator(address have, address owner, address coordinator);
   error ZeroAddress();
 
-  IVRFCoordinatorV2Plus private vrfCoordinator;
-  address private subOwner;
+  IVRFMigratableCoordinatorV2Plus internal s_vrfCoordinator;
+  uint64 internal s_subId;
 
   /**
    * @param _vrfCoordinator address of VRFCoordinator contract
    */
-  constructor(address _vrfCoordinator) {
-    require(_vrfCoordinator != address(0), "zero address");
-    vrfCoordinator = IVRFCoordinatorV2Plus(_vrfCoordinator);
+  constructor(address _vrfCoordinator) ConfirmedOwner(msg.sender) {
+    s_vrfCoordinator = IVRFMigratableCoordinatorV2Plus(_vrfCoordinator);
   }
 
   /**
@@ -133,26 +133,24 @@ abstract contract VRFConsumerBaseV2Plus is IVRFMigratableConsumerV2Plus {
   // proof. rawFulfillRandomness then calls fulfillRandomness, after validating
   // the origin of the call
   function rawFulfillRandomWords(uint256 requestId, uint256[] memory randomWords) external {
-    if (msg.sender != address(vrfCoordinator)) {
-      revert OnlyCoordinatorCanFulfill(msg.sender, address(vrfCoordinator));
+    if (msg.sender != address(s_vrfCoordinator)) {
+      revert OnlyCoordinatorCanFulfill(msg.sender, address(s_vrfCoordinator));
     }
     fulfillRandomWords(requestId, randomWords);
   }
 
   /**
    * @inheritdoc IVRFMigratableConsumerV2Plus
    */
-  function setVRFCoordinator(address _vrfCoordinator) external override {
-    if (msg.sender != subOwner) {
-      revert OnlySubOwnerCanSetVRFCoordinator(msg.sender, subOwner);
-    }
-    vrfCoordinator = IVRFCoordinatorV2Plus(_vrfCoordinator);
+  function setConfig(address _vrfCoordinator, uint64 _subId) public override onlyOwnerOrCoordinator {
+    s_vrfCoordinator = IVRFMigratableCoordinatorV2Plus(_vrfCoordinator);
+    s_subId = _subId;
   }
 
-  function _setSubOwner(address _subOwner) internal {
-    if (_subOwner == address(0)) {
-      revert ZeroAddress();
+  modifier onlyOwnerOrCoordinator() {
+    if (msg.sender != owner() && msg.sender != address(s_vrfCoordinator)) {
+      revert OnlyOwnerOrCoordinator(msg.sender, owner(), address(s_vrfCoordinator));
     }
-    subOwner = _subOwner;
+    _;
   }
 }

contracts/src/v0.8/dev/vrf/VRFCoordinatorV2Plus.sol

@@ -10,8 +10,9 @@ import "./VRFConsumerBaseV2Plus.sol";
 import "../../ChainSpecificUtil.sol";
 import "./SubscriptionAPI.sol";
 import "./libraries/VRFV2PlusClient.sol";
+import "../interfaces/IVRFCoordinatorV2PlusMigration.sol";
 
-contract VRFCoordinatorV2Plus is VRF, TypeAndVersionInterface, SubscriptionAPI {
+contract VRFCoordinatorV2Plus is VRF, SubscriptionAPI {
   /// @dev may not be provided upon construction on some chains due to lack of availability
   AggregatorV3Interface public LINK_ETH_FEED;
   /// @dev should always be available
@@ -607,11 +608,97 @@ contract VRFCoordinatorV2Plus is VRF, TypeAndVersionInterface, SubscriptionAPI {
     cancelSubscriptionHelper(subId, to);
   }
 
-  /**
-   * @notice The type and version of this contract
-   * @return Type and version string
-   */
-  function typeAndVersion() external pure virtual override returns (string memory) {
-    return "VRFCoordinatorV2Plus 1.0.0";
+  /***************************************************************************
+   * Section: Migration
+   ***************************************************************************/
+
+  address[] internal s_migrationTargets;
+
+  /// @dev Emitted when new coordinator is registered as migratable target
+  event CoordinatorRegistered(address coordinatorAddress);
+
+  /// @dev Emitted when new coordinator is deregistered
+  event CoordinatorDeregistered(address coordinatorAddress);
+
+  /// @notice emitted when migration to new coordinator completes successfully
+  /// @param newCoordinator coordinator address after migration
+  /// @param prevSubId old subscription ID
+  /// @param newSubId new subscription ID
+  event MigrationCompleted(address newCoordinator, uint64 prevSubId, uint64 newSubId);
+
+  /// @notice emitted when migrate() is called and given coordinator is not registered as migratable target
+  error CoordinatorNotRegistered(address coordinatorAddress);
+
+  /// @notice emitted when migrate() is called and given coordinator is registered as migratable target
+  error CoordinatorAlreadyRegistered(address coordinatorAddress);
+
+  /// @dev encapsulates data to be migrated from current coordinator
+  struct V1MigrationData {
+    uint8 fromVersion;
+    address subOwner;
+    address[] consumers;
+    uint96 linkBalance;
+    uint96 ethBalance;
+  }
+
+  function isTargetRegistered(address target) internal view returns (bool) {
+    for (uint256 i = 0; i < s_migrationTargets.length; i++) {
+      if (s_migrationTargets[i] == target) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  function registerMigratableCoordinator(address target) external onlyOwner {
+    if (isTargetRegistered(target)) {
+      revert CoordinatorAlreadyRegistered(target);
+    }
+    s_migrationTargets.push(target);
+    emit CoordinatorRegistered(target);
+  }
+
+  function deregisterMigratableCoordinator(address target) external onlyOwner {
+    uint256 nTargets = s_migrationTargets.length;
+    for (uint256 i = 0; i < nTargets; i++) {
+      if (s_migrationTargets[i] == target) {
+        s_migrationTargets[i] = s_migrationTargets[nTargets - 1];
+        s_migrationTargets[nTargets - 1] = target;
+        s_migrationTargets.pop();
+        emit CoordinatorDeregistered(target);
+        return;
+      }
+    }
+    revert CoordinatorNotRegistered(target);
+  }
+
+  function migrate(uint64 subId, address newCoordinator) external returns (uint64) {
+    if (!isTargetRegistered(newCoordinator)) {
+      revert CoordinatorNotRegistered(newCoordinator);
+    }
+    (uint96 balance, uint96 ethBalance, address owner, address[] memory consumers) = getSubscription(subId);
+    require(owner == msg.sender, "Not subscription owner");
+    require(!pendingRequestExists(subId), "Pending request exists");
+
+    V1MigrationData memory migrationData = V1MigrationData({
+      fromVersion: migrationVersion(),
+      subOwner: owner,
+      consumers: consumers,
+      linkBalance: balance,
+      ethBalance: ethBalance
+    });
+    bytes memory encodedData = abi.encode(migrationData);
+    deleteSubscription(subId);
+    uint64 newSubId = IVRFCoordinatorV2PlusMigration(newCoordinator).onMigration{value: ethBalance}(encodedData);
+    require(LINK.transfer(address(newCoordinator), balance), "insufficient funds");
+    for (uint256 i = 0; i < consumers.length; i++) {
+      IVRFMigratableConsumerV2Plus(consumers[i]).setConfig(newCoordinator, newSubId);
+    }
+    emit MigrationCompleted(newCoordinator, subId, newSubId);
+    return newSubId;
+  }
+
+  function migrationVersion() public pure returns (uint8 version) {
+    return 1;
   }
 }

contracts/src/v0.8/dev/vrf/VRFV2PlusSubscriptionManager.sol

@@ -116,7 +116,7 @@ contract VRFV2PlusSubscriptionManager is ConfirmedOwner {
     // note that this is bounded by MAX_CONSUMERS in the coordinator
     for (uint256 i = 0; i < consumers.length; i++) {
       newCoord.addConsumer(newSubId, consumers[i]);
-      IVRFMigratableConsumerV2Plus(consumers[i]).setVRFCoordinator(newCoordinator);
+      IVRFMigratableConsumerV2Plus(consumers[i]).setConfig(newCoordinator, newSubId);
     }
     // set the subscription id and the vrf coordinator in this owner contract
     s_subId = newSubId;

contracts/src/v0.8/dev/vrf/VRFV2PlusWrapper.sol

@@ -96,11 +96,7 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
   }
   mapping(uint256 => Callback) /* requestID */ /* callback */ public s_callbacks;
 
-  constructor(
-    address _link,
-    address _linkEthFeed,
-    address _coordinator
-  ) ConfirmedOwner(msg.sender) VRFConsumerBaseV2Plus(_coordinator) {
+  constructor(address _link, address _linkEthFeed, address _coordinator) VRFConsumerBaseV2Plus(_coordinator) {
     if (_link != address(0)) {
       s_link = LinkTokenInterface(_link);
     }

contracts/src/v0.8/dev/vrf/testhelpers/ExposedVRFCoordinatorV2Plus.sol

@@ -14,4 +14,8 @@ contract ExposedVRFCoordinatorV2Plus is VRFCoordinatorV2Plus {
   ) external pure returns (uint256, uint256) {
     return computeRequestId(keyHash, sender, subId, nonce);
   }
+
+  function isTargetRegisteredExternal(address target) external view returns (bool) {
+    return isTargetRegistered(target);
+  }
 }