From dc1f494710c79258d89f52e0e1cd4df9dee72144 Mon Sep 17 00:00:00 2001
From: YikJiun Lee <yikjiun.lee@smartcontract.com>
Date: Tue, 11 Jun 2024 15:55:58 +0800
Subject: [PATCH 1/2] vrfv2plus: Add batch coordinator tests

---
 .../src/v0.8/vrf/dev/VRFV2PlusWrapper.sol     |   1 +
 .../vrf/test/BatchVRFCoordinatorV2Plus.t.sol  | 178 ++++++++++++++++++
 .../vrf/test/FixtureVRFCoordinatorV2_5.t.sol  | 120 ++++++++++++
 3 files changed, 299 insertions(+)
 create mode 100644 contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol
 create mode 100644 contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol

diff --git a/contracts/src/v0.8/vrf/dev/VRFV2PlusWrapper.sol b/contracts/src/v0.8/vrf/dev/VRFV2PlusWrapper.sol
index 1d741280851..86578fb1400 100644
--- a/contracts/src/v0.8/vrf/dev/VRFV2PlusWrapper.sol
+++ b/contracts/src/v0.8/vrf/dev/VRFV2PlusWrapper.sol
@@ -232,6 +232,7 @@ contract VRFV2PlusWrapper is ConfirmedOwner, TypeAndVersionInterface, VRFConsume
    * @param _fulfillmentFlatFeeLinkDiscountPPM is the flat fee discount in millionths of native that VRFCoordinatorV2Plus
    *        charges for link payment.
    */
+  /// @dev This function while having only 12 parameters is causing a Stack too deep error when running forge coverage.
   function setConfig(
     uint32 _wrapperGasOverhead,
     uint32 _coordinatorGasOverheadNative,
diff --git a/contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol b/contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol
new file mode 100644
index 00000000000..8f323f6450c
--- /dev/null
+++ b/contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol
@@ -0,0 +1,178 @@
+pragma solidity 0.8.19;
+
+import {console} from "forge-std/console.sol";
+import {VRF} from "../VRF.sol";
+import {VRFTypes} from "../VRFTypes.sol";
+import {BatchVRFCoordinatorV2Plus} from "../dev/BatchVRFCoordinatorV2Plus.sol";
+import {VRFV2PlusClient} from "../dev/libraries/VRFV2PlusClient.sol";
+import {VRFCoordinatorV2_5} from "../dev/VRFCoordinatorV2_5.sol";
+import "./BaseTest.t.sol";
+import {FixtureVRFCoordinatorV2_5} from "./FixtureVRFCoordinatorV2_5.t.sol";
+
+contract BatchVRFCoordinatorV2PlusTest is FixtureVRFCoordinatorV2_5 {
+  BatchVRFCoordinatorV2Plus private s_batchCoordinator;
+
+  event RandomWordsFulfilled(
+    uint256 indexed requestId,
+    uint256 outputSeed,
+    uint256 indexed subId,
+    uint96 payment,
+    bool nativePayment,
+    bool success,
+    bool onlyPremium
+  );
+
+  function setUp() public override {
+    FixtureVRFCoordinatorV2_5.setUp();
+
+    s_batchCoordinator = new BatchVRFCoordinatorV2Plus(address(s_coordinator));
+  }
+
+  function test_fulfillRandomWords() public {
+    _setUpConfig();
+    _setUpProvingKey();
+    _setUpSubscription();
+
+    uint32 requestBlock = 10;
+    vm.roll(requestBlock);
+
+    vm.startPrank(SUBSCRIPTION_OWNER);
+    vm.deal(SUBSCRIPTION_OWNER, 10 ether);
+    s_coordinator.fundSubscriptionWithNative{value: 10 ether}(s_subId);
+
+    // Request random words.
+    s_consumer.requestRandomWords(CALLBACK_GAS_LIMIT, MIN_CONFIRMATIONS, NUM_WORDS, VRF_KEY_HASH, true);
+    vm.stopPrank();
+
+    // Move on to the next block.
+    // Store the previous block's blockhash.
+    vm.roll(requestBlock + 1);
+    s_bhs.store(requestBlock);
+
+    VRFTypes.Proof[] memory proofs = new VRFTypes.Proof[](2);
+    VRFTypes.RequestCommitmentV2Plus[] memory rcs = new VRFTypes.RequestCommitmentV2Plus[](2);
+
+    // Proof generated via the generate-proof-v2-plus script command. Example usage:
+    // _printGenerateProofV2PlusCommand(address(s_consumer), 1, requestBlock, true);
+    /*
+       go run . generate-proof-v2-plus \
+         -key-hash 0x9f2353bde94264dbc3d554a94cceba2d7d2b4fdce4304d3e09a1fea9fbeb1528 \
+         -pre-seed 33855227690351884611579800220581891477580182035146587491531555927634180294480 \
+         -block-hash 0x0a \
+         -block-num 10 \
+         -sender 0xdc90e8ce61c1af8a638b95264037c8e67ee5765c \
+         -native-payment true
+
+    */
+    proofs[0] = VRFTypes.Proof({
+      pk: [
+        72488970228380509287422715226575535698893157273063074627791787432852706183111,
+        62070622898698443831883535403436258712770888294397026493185421712108624767191
+      ],
+      gamma: [
+        80420391742429647505172101941811820476888293644816377569181566466584288434705,
+        24046736031266889997051641830469514057863365715722268340801477580836256044582
+      ],
+      c: 74775128390693502914275156263410881155583102046081919417827483535122161050585,
+      s: 69563235412360165148368009853509434870917653835330501139204071967997764190111,
+      seed: 33855227690351884611579800220581891477580182035146587491531555927634180294480,
+      uWitness: 0xfB0663eaf48785540dE0FD0F837FD9c09BF4B80A,
+      cGammaWitness: [
+        53711159452748734758194447734939737695995909567499536035707522847057731697403,
+        113650002631484103366420937668971311744887820666944514581352028601506700116835
+      ],
+      sHashWitness: [
+        89656531714223714144489731263049239277719465105516547297952288438117443488525,
+        90859682705760125677895017864538514058733199985667976488434404721197234427011
+      ],
+      zInv: 97275608653505690744303242942631893944856831559408852202478373762878300587548
+    });
+    rcs[0] = VRFTypes.RequestCommitmentV2Plus({
+      blockNum: requestBlock,
+      subId: s_subId,
+      callbackGasLimit: CALLBACK_GAS_LIMIT,
+      numWords: 1,
+      sender: address(s_consumer),
+      extraArgs: VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: true}))
+    });
+
+    VRFCoordinatorV2_5.Output memory output = s_coordinator.getRandomnessFromProofExternal(
+      abi.decode(abi.encode(proofs[0]), (VRF.Proof)),
+      rcs[0]
+    );
+
+    requestBlock = 20;
+    vm.roll(requestBlock);
+
+    vm.startPrank(SUBSCRIPTION_OWNER);
+    s_linkToken.setBalance(address(SUBSCRIPTION_OWNER), 10 ether);
+    s_linkToken.transferAndCall(address(s_coordinator), 10 ether, abi.encode(s_subId));
+
+    // Request random words.
+    s_consumer1.requestRandomWords(CALLBACK_GAS_LIMIT, MIN_CONFIRMATIONS, NUM_WORDS, VRF_KEY_HASH, false);
+    vm.stopPrank();
+
+    // Move on to the next block.
+    // Store the previous block's blockhash.
+    vm.roll(requestBlock + 1);
+    s_bhs.store(requestBlock);
+
+    // Proof generated via the generate-proof-v2-plus script command. Example usage:
+    // _printGenerateProofV2PlusCommand(address(s_consumer1), 1, requestBlock, false);
+    /*
+       go run . generate-proof-v2-plus \
+        -key-hash 0x9f2353bde94264dbc3d554a94cceba2d7d2b4fdce4304d3e09a1fea9fbeb1528 \
+        -pre-seed 76568185840201037774581758921393822690942290841865097674309745036496166431060 \
+        -block-hash 0x14 \
+        -block-num 20 \
+        -sender 0x2f1c0761d6e4b1e5f01968d6c746f695e5f3e25d \
+        -native-payment false
+    */
+    proofs[1] = VRFTypes.Proof({
+      pk: [
+        72488970228380509287422715226575535698893157273063074627791787432852706183111,
+        62070622898698443831883535403436258712770888294397026493185421712108624767191
+      ],
+      gamma: [
+        21323932463597506192387578758854201988004673105893105492473194972397109828006,
+        96834737826889397196571646974355352644437196500310392203712129010026003355112
+      ],
+      c: 8775807990949224376582975115621037245862755412370175152581490650310350359728,
+      s: 6805708577951013810918872616271445638109899206333819877111740872779453350091,
+      seed: 76568185840201037774581758921393822690942290841865097674309745036496166431060,
+      uWitness: 0xE82fF24Fecfbe73d682f38308bE3E039Dfabdf5c,
+      cGammaWitness: [
+        92810770919624535241476539842820168209710445519252592382122118536598338376923,
+        17271305664006119131434661141858450289379246199095231636439133258170648418554
+      ],
+      sHashWitness: [
+        29540023305939374439696120003978246982707698669656874393367212257432197207536,
+        93902323936532381028323379401739289810874348405259732508442252936582467730050
+      ],
+      zInv: 88845170436601946907659333156418518556235340365885668267853966404617557948692
+    });
+    rcs[1] = VRFTypes.RequestCommitmentV2Plus({
+      blockNum: requestBlock,
+      subId: s_subId,
+      callbackGasLimit: CALLBACK_GAS_LIMIT,
+      numWords: 1,
+      sender: address(s_consumer1),
+      extraArgs: VRFV2PlusClient._argsToBytes(VRFV2PlusClient.ExtraArgsV1({nativePayment: false}))
+    });
+
+    VRFCoordinatorV2_5.Output memory output1 = s_coordinator.getRandomnessFromProofExternal(
+      abi.decode(abi.encode(proofs[1]), (VRF.Proof)),
+      rcs[1]
+    );
+
+    // The payments are NOT pre-calculated and simply copied from the actual event.
+    // We can assert and ignore the payment field but the code will be considerably longer.
+    vm.expectEmit(true, true, false, true, address(s_coordinator));
+    emit RandomWordsFulfilled(output.requestId, output.randomness, s_subId, 500000000000142215, true, true, false);
+    vm.expectEmit(true, true, false, true, address(s_coordinator));
+    emit RandomWordsFulfilled(output1.requestId, output1.randomness, s_subId, 800000000000304103, false, true, false);
+
+    // Fulfill the requests.
+    s_batchCoordinator.fulfillRandomWords(proofs, rcs);
+  }
+}
diff --git a/contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol b/contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol
new file mode 100644
index 00000000000..8144e22548d
--- /dev/null
+++ b/contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol
@@ -0,0 +1,120 @@
+pragma solidity ^0.8.19;
+
+import {console} from "forge-std/console.sol";
+import "@openzeppelin/contracts/utils/Strings.sol";
+import {VRF} from "../VRF.sol";
+import {VRFTypes} from "../VRFTypes.sol";
+import {BlockhashStore} from "../dev/BlockhashStore.sol";
+import {VRFV2PlusClient} from "../dev/libraries/VRFV2PlusClient.sol";
+import {ExposedVRFCoordinatorV2_5} from "../dev/testhelpers/ExposedVRFCoordinatorV2_5.sol";
+import {VRFV2PlusConsumerExample} from "../dev/testhelpers/VRFV2PlusConsumerExample.sol";
+import {MockLinkToken} from "../../mocks/MockLinkToken.sol";
+import {MockV3Aggregator} from "../../tests/MockV3Aggregator.sol";
+import "./BaseTest.t.sol";
+
+contract FixtureVRFCoordinatorV2_5 is BaseTest, VRF {
+  address internal SUBSCRIPTION_OWNER = makeAddr("SUBSCRIPTION_OWNER");
+
+  uint64 internal constant GAS_LANE_MAX_GAS = 5000 gwei;
+  uint16 internal constant MIN_CONFIRMATIONS = 0;
+  uint32 internal constant CALLBACK_GAS_LIMIT = 1_000_000;
+  uint32 internal constant NUM_WORDS = 1;
+
+  // VRF KeyV2 generated from a node; not sensitive information.
+  // The secret key used to generate this key is: 10.
+  bytes internal constant VRF_UNCOMPRESSED_PUBLIC_KEY =
+    hex"a0434d9e47f3c86235477c7b1ae6ae5d3442d49b1943c2b752a68e2a47e247c7893aba425419bc27a3b6c7e693a24c696f794c2ed877a1593cbee53b037368d7";
+  bytes internal constant VRF_COMPRESSED_PUBLIC_KEY =
+    hex"a0434d9e47f3c86235477c7b1ae6ae5d3442d49b1943c2b752a68e2a47e247c701";
+  bytes32 internal constant VRF_KEY_HASH = hex"9f2353bde94264dbc3d554a94cceba2d7d2b4fdce4304d3e09a1fea9fbeb1528";
+
+  BlockhashStore internal s_bhs;
+  ExposedVRFCoordinatorV2_5 internal s_coordinator;
+
+  // Use multiple consumers because VRFV2PlusConsumerExample cannot have multiple pending requests.
+  uint256 internal s_subId;
+  VRFV2PlusConsumerExample internal s_consumer;
+  VRFV2PlusConsumerExample internal s_consumer1;
+
+  MockLinkToken internal s_linkToken;
+  MockV3Aggregator internal s_linkNativeFeed;
+
+  function setUp() public virtual override {
+    BaseTest.setUp();
+    vm.stopPrank();
+
+    vm.startPrank(OWNER);
+    s_bhs = new BlockhashStore();
+
+    // Deploy coordinator.
+    s_coordinator = new ExposedVRFCoordinatorV2_5(address(s_bhs));
+    s_linkToken = new MockLinkToken();
+    s_linkNativeFeed = new MockV3Aggregator(18, 500000000000000000); // .5 ETH (good for testing)
+
+    // Configure the coordinator.
+    s_coordinator.setLINKAndLINKNativeFeed(address(s_linkToken), address(s_linkNativeFeed));
+    vm.stopPrank();
+
+    // Deploy consumers.
+    vm.startPrank(SUBSCRIPTION_OWNER);
+    s_consumer = new VRFV2PlusConsumerExample(address(s_coordinator), address(s_linkToken));
+    s_consumer1 = new VRFV2PlusConsumerExample(address(s_coordinator), address(s_linkToken));
+    vm.stopPrank();
+  }
+
+  function _setUpConfig() internal {
+    vm.prank(OWNER);
+    s_coordinator.setConfig(
+      0, // minRequestConfirmations
+      2_500_000, // maxGasLimit
+      1, // stalenessSeconds
+      50_000, // gasAfterPaymentCalculation
+      50000000000000000, // fallbackWeiPerUnitLink
+      500_000, // fulfillmentFlatFeeNativePPM
+      100_000, // fulfillmentFlatFeeLinkDiscountPPM
+      15, // nativePremiumPercentage
+      10 // linkPremiumPercentage
+    );
+  }
+
+  function _setUpProvingKey() internal {
+    uint256[2] memory uncompressedKeyParts = this._getProvingKeyParts(VRF_UNCOMPRESSED_PUBLIC_KEY);
+    vm.prank(OWNER);
+    s_coordinator.registerProvingKey(uncompressedKeyParts, GAS_LANE_MAX_GAS);
+  }
+
+  function _setUpSubscription() internal {
+    vm.startPrank(SUBSCRIPTION_OWNER);
+    s_subId = s_coordinator.createSubscription();
+    s_coordinator.addConsumer(s_subId, address(s_consumer));
+    s_consumer.setSubId(s_subId);
+    s_coordinator.addConsumer(s_subId, address(s_consumer1));
+    s_consumer1.setSubId(s_subId);
+    vm.stopPrank();
+  }
+
+  // note: Call this function via this.getProvingKeyParts to be able to pass memory as calldata and
+  // index over the byte array.
+  function _getProvingKeyParts(bytes calldata uncompressedKey) public pure returns (uint256[2] memory) {
+    uint256 keyPart1 = uint256(bytes32(uncompressedKey[0:32]));
+    uint256 keyPart2 = uint256(bytes32(uncompressedKey[32:64]));
+    return [keyPart1, keyPart2];
+  }
+
+  function _printGenerateProofV2PlusCommand(
+    address sender,
+    uint64 nonce,
+    uint256 requestBlock,
+    bool nativePayment
+  ) internal {
+    (, uint256 preSeed) = s_coordinator.computeRequestIdExternal(VRF_KEY_HASH, sender, s_subId, nonce);
+
+    console.log("go run . generate-proof-v2-plus \\");
+    console.log(string.concat("  -key-hash ", Strings.toHexString(uint256(VRF_KEY_HASH)), " \\"));
+    console.log(string.concat("  -pre-seed ", Strings.toString(preSeed), " \\"));
+    console.log(string.concat("  -block-hash ", Strings.toHexString(uint256(blockhash(requestBlock))), " \\"));
+    console.log(string.concat("  -block-num ", Strings.toString(requestBlock), " \\"));
+    console.log(string.concat("  -sender ", Strings.toHexString(sender), " \\"));
+    console.log(string.concat("  -native-payment ", nativePayment ? "true" : "false"));
+  }
+}

From e5e21e8c345c0590c3a246b58f336989ea58e048 Mon Sep 17 00:00:00 2001
From: YikJiun Lee <yikjiun.lee@smartcontract.com>
Date: Thu, 13 Jun 2024 11:33:19 +0800
Subject: [PATCH 2/2] Add function comment for print generate proof v2 plus

---
 .../v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol  | 12 ++++++------
 .../v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol  | 14 ++++++++++++++
 2 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol b/contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol
index 8f323f6450c..298af8c76e3 100644
--- a/contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol
+++ b/contracts/src/v0.8/vrf/test/BatchVRFCoordinatorV2Plus.t.sol
@@ -121,12 +121,12 @@ contract BatchVRFCoordinatorV2PlusTest is FixtureVRFCoordinatorV2_5 {
     // _printGenerateProofV2PlusCommand(address(s_consumer1), 1, requestBlock, false);
     /*
        go run . generate-proof-v2-plus \
-        -key-hash 0x9f2353bde94264dbc3d554a94cceba2d7d2b4fdce4304d3e09a1fea9fbeb1528 \
-        -pre-seed 76568185840201037774581758921393822690942290841865097674309745036496166431060 \
-        -block-hash 0x14 \
-        -block-num 20 \
-        -sender 0x2f1c0761d6e4b1e5f01968d6c746f695e5f3e25d \
-        -native-payment false
+         -key-hash 0x9f2353bde94264dbc3d554a94cceba2d7d2b4fdce4304d3e09a1fea9fbeb1528 \
+         -pre-seed 76568185840201037774581758921393822690942290841865097674309745036496166431060 \
+         -block-hash 0x14 \
+         -block-num 20 \
+         -sender 0x2f1c0761d6e4b1e5f01968d6c746f695e5f3e25d \
+         -native-payment false
     */
     proofs[1] = VRFTypes.Proof({
       pk: [
diff --git a/contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol b/contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol
index 8144e22548d..c1c2c7eb27c 100644
--- a/contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol
+++ b/contracts/src/v0.8/vrf/test/FixtureVRFCoordinatorV2_5.t.sol
@@ -101,6 +101,20 @@ contract FixtureVRFCoordinatorV2_5 is BaseTest, VRF {
     return [keyPart1, keyPart2];
   }
 
+  /**
+   * Prints the command to generate a proof for a VRF request.
+   *
+   * This function provides a convenient way to generate the proof off-chain to be copied into the tests.
+   *
+   * An example of the command looks like this:
+   * go run . generate-proof-v2-plus \
+   *   -key-hash 0x9f2353bde94264dbc3d554a94cceba2d7d2b4fdce4304d3e09a1fea9fbeb1528 \
+   *   -pre-seed 76568185840201037774581758921393822690942290841865097674309745036496166431060 \
+   *   -block-hash 0x14 \
+   *   -block-num 20 \
+   *   -sender 0x2f1c0761d6e4b1e5f01968d6c746f695e5f3e25d \
+   *   -native-payment false
+   */
   function _printGenerateProofV2PlusCommand(
     address sender,
     uint64 nonce,