Allow anyone to call the fulfillRandomWords function with a valid ora…

…cle signature
PaintSwap · Feb 19, 2024 · 131a04e · 131a04e
1 parent 3c19629
commit 131a04e
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 42 deletions.
diff --git a/contracts/ISamWitchRNGConsumer.sol b/contracts/ISamWitchRNGConsumer.sol
@@ -5,7 +5,7 @@ interface ISamWitchRNGConsumer {
   /**
    * @notice fulfillRandomness handles the RNG response. Your contract must
    * @notice implement it.
-   * 
+   *
    * @param requestId The Id initially returned by requestRandomness
    * @param data the RNG output expanded to the requested number of words
    */

diff --git a/contracts/SamWitchRNG.sol b/contracts/SamWitchRNG.sol
@@ -3,49 +3,43 @@ pragma solidity ^0.8.24;
 
 import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
 import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {SignatureChecker} from "@openzeppelin/contracts/utils/cryptography/SignatureChecker.sol";
+import {MessageHashUtils} from "@openzeppelin/contracts/utils/cryptography/MessageHashUtils.sol";
 
 import {ISamWitchRNGConsumer} from "./ISamWitchRNGConsumer.sol";
 
 /// @title SamWitchRNG - Random Number Generator
 /// @author Sam Witch (SamWitchRNG & Estfor Kingdom)
-/// @notice This contract listens for requests for RNG, and allows the caller to fulfill random numbers
+/// @notice This contract listens for requests for RNG, and allows the oracle to fulfill random numbers
 contract SamWitchRNG is UUPSUpgradeable, OwnableUpgradeable {
   event ConsumerRegistered(address consumer);
   event RandomWordsRequested(bytes32 requestId, address fulfillAddress, uint numWords);
   event RandomWordsFulfilled(bytes32 requestId, bytes data);
 
   error FulfillmentFailed(bytes32 requestId);
   error InvalidConsumer(address consumer);
-  error OnlyCaller();
+  error OnlyOracle();
   error RequestAlreadyFulfilled();
   error RequestIdDoesNotExist(bytes32 requestId);
 
+  mapping(address consumer => uint64 nonce) public consumers;
+  address private oracle;
+
   // 5k is plenty for an EXTCODESIZE call (2600) + warm CALL (100)
   // and some arithmetic operations.
   uint private constant GAS_FOR_CALL_EXACT_CHECK = 5_000;
 
-  mapping(address consumer => uint64 nonce) public consumers;
-
-  address private caller;
-
-  modifier onlyCaller() {
-    if (msg.sender != caller) {
-      revert OnlyCaller();
-    }
-    _;
-  }
-
   /// @custom:oz-upgrades-unsafe-allow constructor
   constructor() {
     _disableInitializers();
   }
 
   /// @notice Initialize the contract as part of the proxy contract deployment
-  function initialize(address _caller) external payable initializer {
+  function initialize(address _oracle) external payable initializer {
     __UUPSUpgradeable_init();
     __Ownable_init(_msgSender());
 
-    caller = _caller;
+    oracle = _oracle;
   }
 
   /// @notice Called by the requester to make a full request, which provides
@@ -60,27 +54,40 @@ contract SamWitchRNG is UUPSUpgradeable, OwnableUpgradeable {
 
     uint64 nonce = ++currentNonce;
     consumers[msg.sender] = currentNonce;
-    requestId = computeRequestId(msg.sender, nonce);
+    requestId = _computeRequestId(msg.sender, nonce);
 
     emit RandomWordsRequested(
       requestId,
-      msg.sender, // fulfillAddress
+      msg.sender,
       numWords
     );
   }
 
-  /// @notice Called by the allowed caller to fulfill the request
+  /// @notice Called by the allowed oracle to fulfill the request
   /// @param requestId Request ID
-  /// @param randomWordsData The random words to assign abi encoded
+  /// @param randomWordsData The random words to assign (abi encoded)
   /// @param fulfillAddress Address that will be called to fulfill
   /// @return callSuccess If the fulfillment call succeeded
   function fulfillRandomWords(
     bytes32 requestId,
     bytes calldata randomWordsData,
+    bytes calldata signature,
     address fulfillAddress,
     uint gasAmount
-  ) external onlyCaller returns (bool callSuccess) {
-    bytes memory data = abi.encodeWithSelector(ISamWitchRNGConsumer.fulfillRandomWords.selector, requestId, randomWordsData);
+  ) external returns (bool callSuccess) {
+    // Verify it was created by the oracle
+    bytes32 signedDataHash = keccak256(abi.encodePacked(requestId, randomWordsData));
+    bytes32 message = MessageHashUtils.toEthSignedMessageHash(signedDataHash);
+    if (!SignatureChecker.isValidSignatureNow(oracle, message, signature)) {
+      revert OnlyOracle();
+    }
+
+    // Call the consumer contract callback
+    bytes memory data = abi.encodeWithSelector(
+      ISamWitchRNGConsumer.fulfillRandomWords.selector,
+      requestId,
+      randomWordsData
+    );
     callSuccess = _callWithExactGas(gasAmount, fulfillAddress, data);
     if (callSuccess) {
       emit RandomWordsFulfilled(requestId, randomWordsData);
@@ -94,7 +101,7 @@ contract SamWitchRNG is UUPSUpgradeable, OwnableUpgradeable {
     emit ConsumerRegistered(_consumer);
   }
 
-  function computeRequestId(address sender, uint64 nonce) private pure returns (bytes32) {
+  function _computeRequestId(address sender, uint64 nonce) private pure returns (bytes32) {
     return keccak256(abi.encodePacked(sender, nonce));
   }
 

diff --git a/contracts/test/TestRNGConsumer.sol b/contracts/test/TestRNGConsumer.sol
@@ -28,11 +28,8 @@ contract TestRNGConsumer is ISamWitchRNGConsumer, Ownable {
     requestId = samWitchRNG.requestRandomWords(numWords);
   }
 
-// Called by the RNG contract to fulfill a random number request
-  function fulfillRandomWords(
-    bytes32 requestId,
-    bytes calldata data
-  ) external onlySamWitchRNG {
+  // Called by the RNG contract to fulfill a random number request
+  function fulfillRandomWords(bytes32 requestId, bytes calldata data) external onlySamWitchRNG {
     uint[] memory randomWords = abi.decode(data, (uint[]));
     allRandomWords[requestId] = randomWords;
     if (shouldRevert) {

diff --git a/test/SamWitchRNG.ts b/test/SamWitchRNG.ts
@@ -7,9 +7,9 @@ describe("SamWitchRNG", function () {
   async function deployContractsFixture() {
     const [owner, alice, bob, charlie, dev] = await ethers.getSigners();
 
-    const caller = bob.address;
+    const oracle = bob.address;
     const SamWitchRNG = await ethers.getContractFactory("SamWitchRNG");
-    const samWitchRNG = (await upgrades.deployProxy(SamWitchRNG, [caller], {
+    const samWitchRNG = (await upgrades.deployProxy(SamWitchRNG, [oracle], {
       kind: "uups",
     })) as unknown as SamWitchRNG;
 
@@ -43,7 +43,10 @@ describe("SamWitchRNG", function () {
     const randomNum = ethers.hexlify(ethers.randomBytes(32));
     const gasLimit = 1_000_000;
     const data = ethers.AbiCoder.defaultAbiCoder().encode(["uint256[]"], [[randomNum]]);
-    tx = await samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, rngConsumer, gasLimit);
+    const signature = await bob.signMessage(
+      ethers.toBeArray(ethers.solidityPackedKeccak256(["bytes32", "bytes"], [requestId, data])),
+    );
+    tx = await samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, signature, rngConsumer, gasLimit);
     receipt = await tx.wait();
     log = samWitchRNG.interface.parseLog(receipt?.logs[0]);
 
@@ -73,26 +76,42 @@ describe("SamWitchRNG", function () {
 
     let log = samWitchRNG.interface.parseLog(receipt?.logs[0]);
     const requestId = log.args[0];
-
     const randomNum = ethers.hexlify(ethers.randomBytes(32));
-    const gasLimit = 1_000_000;
     const data = ethers.AbiCoder.defaultAbiCoder().encode(["uint256[]"], [[randomNum]]);
-    tx = await samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, rngConsumer, gasLimit);
+    const signature = await bob.signMessage(
+      ethers.toBeArray(ethers.solidityPackedKeccak256(["bytes32", "bytes"], [requestId, data])),
+    );
+
+    const gasLimit = 1_000_000;
+    tx = await samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, signature, rngConsumer, gasLimit);
     receipt = await tx.wait();
     log = samWitchRNG.interface.parseLog(receipt?.logs[0]);
 
     // Can call fulfill again, it is a requirement that it is reverted on the consumer side if the requestId is already fulfilled
-    await expect(samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, rngConsumer, gasLimit)).to.not.be
-      .reverted;
+    await expect(samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, signature, rngConsumer, gasLimit)).to.not
+      .be.reverted;
   });
 
-  it("Only the caller can fulfill the random words", async function () {
-    const {samWitchRNG, rngConsumer, alice} = await loadFixture(deployContractsFixture);
+  it("Anyone can call fulfill as long as the signature is signed by the oracle", async function () {
+    const {samWitchRNG, rngConsumer, alice, bob} = await loadFixture(deployContractsFixture);
     const gasLimit = 1_000_000;
     const data = ethers.AbiCoder.defaultAbiCoder().encode(["uint256[]"], [[1]]);
+
+    // Create a hash of your data
+    const requestId = ethers.encodeBytes32String("1");
+    const hash = ethers.solidityPackedKeccak256(["bytes32", "bytes"], [requestId, data]);
+
+    // Sign the hash
+    const signatureWrongSigner = await alice.signMessage(ethers.toBeArray(hash));
+
+    await expect(
+      samWitchRNG.connect(alice).fulfillRandomWords(requestId, data, signatureWrongSigner, rngConsumer, gasLimit),
+    ).to.be.revertedWithCustomError(samWitchRNG, "OnlyOracle");
+
+    const signatureCorrectSigner = await bob.signMessage(ethers.toBeArray(hash));
     await expect(
-      samWitchRNG.connect(alice).fulfillRandomWords(ethers.encodeBytes32String("1"), data, rngConsumer, gasLimit),
-    ).to.be.revertedWithCustomError(samWitchRNG, "OnlyCaller");
+      samWitchRNG.connect(alice).fulfillRandomWords(requestId, data, signatureCorrectSigner, rngConsumer, gasLimit),
+    ).to.not.be.reverted;
   });
 
   it("Consumer ran out of gas", async function () {
@@ -112,8 +131,11 @@ describe("SamWitchRNG", function () {
     const randomNum = ethers.hexlify(ethers.randomBytes(32));
     const gasLimit = 1;
     const data = ethers.AbiCoder.defaultAbiCoder().encode(["uint256[]"], [[randomNum]]);
+    const signature = await bob.signMessage(
+      ethers.toBeArray(ethers.solidityPackedKeccak256(["bytes32", "bytes"], [requestId, data])),
+    );
     await expect(
-      samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, rngConsumer, gasLimit),
+      samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, signature, rngConsumer, gasLimit),
     ).to.be.revertedWithCustomError(samWitchRNG, "FulfillmentFailed");
   });
 
@@ -135,8 +157,11 @@ describe("SamWitchRNG", function () {
     const gasLimit = 1_000_000;
     await rngConsumer.setShouldRevert(true);
     const data = ethers.AbiCoder.defaultAbiCoder().encode(["uint256[]"], [[randomNum]]);
+    const signature = await bob.signMessage(
+      ethers.toBeArray(ethers.solidityPackedKeccak256(["bytes32", "bytes"], [requestId, data])),
+    );
     await expect(
-      samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, rngConsumer, gasLimit),
+      samWitchRNG.connect(bob).fulfillRandomWords(requestId, data, signature, rngConsumer, gasLimit),
     ).to.be.revertedWithCustomError(samWitchRNG, "FulfillmentFailed");
   });