Merge pull request #49 from PacificYield/updates-comp

feat: Updates Comp.sol

contracts/governance/Comp.sol

@@ -3,12 +3,14 @@ pragma solidity ^0.8.24;
 
 import "fhevm/lib/TFHE.sol";
 import { Ownable2Step, Ownable } from "@openzeppelin/contracts/access/Ownable2Step.sol";
+import { EIP712 } from "@openzeppelin/contracts/utils/cryptography/EIP712.sol";
+import { SignatureChecker } from "@openzeppelin/contracts/utils/cryptography/SignatureChecker.sol";
 import { EncryptedERC20 } from "../token/ERC20/EncryptedERC20.sol";
 import { IComp } from "./IComp.sol";
 
 /**
  * @title       Comp
- * @notice      This contract inherits EncryptedERC20 and Ownable2Step.
+ * @notice      This contract inherits EncryptedERC20, EIP712, and Ownable2Step.
  *              This is based on the Comp.sol contract written by Compound Labs.
  *              see: compound-finance/compound-protocol/blob/master/contracts/Governance/Comp.sol
  *              It is a governance token used to delegate votes, which can be used by contracts such as
@@ -17,14 +19,24 @@ import { IComp } from "./IComp.sol";
  *              with an account's balance.
  * @dev         The delegation of votes leaks information about the account's encrypted balance to the `delegatee`.
  */
-abstract contract Comp is IComp, EncryptedERC20, Ownable2Step {
+abstract contract Comp is IComp, EncryptedERC20, EIP712, Ownable2Step {
     /// @notice Returned if the `blockNumber` is higher or equal to the (current) `block.number`.
     /// @dev    It is returned in requests to access votes.
     error BlockNumberEqualOrHigherThanCurrentBlock();
 
     /// @notice Returned if the `msg.sender` is not the `governor` contract.
     error GovernorInvalid();
 
+    /// @notice Returned if the signature has expired.
+    error SignatureExpired();
+
+    /// @notice Returned if the signature's nonce is invalid.
+    error SignatureNonceInvalid();
+
+    /// @notice Returned if the signature's verification has failed.
+    /// @dev    See {SignatureChecker} for potential reasons.
+    error SignatureVerificationFail();
+
     /// @notice Emitted when an `account` (i.e. `delegator`) changes its delegate.
     event DelegateChanged(address indexed delegator, address indexed fromDelegate, address indexed toDelegate);
 
@@ -35,6 +47,9 @@ abstract contract Comp is IComp, EncryptedERC20, Ownable2Step {
     /// @dev    WARNING: it can be set to a malicious contract, which could reencrypt all user votes.
     event NewGovernor(address indexed governor);
 
+    /// @notice Emitted when the account cancels a signature.
+    event NonceIncremented(address account, uint256 newNonce);
+
     /// @notice          A checkpoint for marking number of votes from a given block.
     /// @param fromBlock Block from where the checkpoint applies.
     /// @param votes     Total number of votes for the account power.
@@ -46,10 +61,6 @@ abstract contract Comp is IComp, EncryptedERC20, Ownable2Step {
         euint64 votes;
     }
 
-    /// @notice The EIP-712 typehash for the contract's domain.
-    bytes32 public constant DOMAIN_TYPEHASH =
-        keccak256("EIP712Domain(string name,uint256 chainId,address verifyingContract)");
-
     /// @notice The EIP-712 typehash for the `Delegation` struct.
     bytes32 public constant DELEGATION_TYPEHASH =
         keccak256("Delegation(address delegatee,uint256 nonce,uint256 expiry)");
@@ -77,11 +88,21 @@ abstract contract Comp is IComp, EncryptedERC20, Ownable2Step {
     euint64 private _EUINT64_ZERO;
 
     /**
-     * @param owner Owner address
+     * @param owner_        Owner address.
+     * @param name_         Token name.
+     * @param symbol_       Token symbol.
+     * @param version_      Version (e.g. "0.1", "1.0").
+     * @param totalSupply_  Total supply to mint.
      */
-    constructor(address owner) EncryptedERC20("Compound", "COMP") Ownable(owner) {
-        _unsafeMint(owner, TFHE.asEuint64(10000000e6)); /// 10 million Comp
-        _totalSupply = 10000000e6;
+    constructor(
+        address owner_,
+        string memory name_,
+        string memory symbol_,
+        string memory version_,
+        uint64 totalSupply_
+    ) EncryptedERC20(name_, symbol_) EIP712(name_, version_) Ownable(owner_) {
+        _unsafeMint(owner_, TFHE.asEuint64(totalSupply_));
+        _totalSupply = totalSupply_;
 
         /// @dev Define the constant in the storage.
         _EUINT64_ZERO = TFHE.asEuint64(0);
@@ -98,32 +119,48 @@ abstract contract Comp is IComp, EncryptedERC20, Ownable2Step {
 
     /**
      * @notice          Delegate votes from signatory to `delegatee`.
+     * @param delegator The account that delegates its votes. It must be the signer.
      * @param delegatee The address to delegate votes to.
      * @param nonce     The contract state required to match the signature.
      * @param expiry    The time at which to expire the signature.
-     * @param v         The recovery byte of the signature.
-     * @param r         Half of the ECDSA signature pair.
-     * @param s         Half of the ECDSA signature pair.
+     * @param signature The signature.
+     * @dev             Signature can be either 64-byte or 65-byte long if it is from an EOA.
+     *                  Else, it must adhere to ERC1271. See {https://eips.ethereum.org/EIPS/eip-1271}
      */
     function delegateBySig(
+        address delegator,
         address delegatee,
         uint256 nonce,
         uint256 expiry,
-        uint8 v,
-        bytes32 r,
-        bytes32 s
+        bytes memory signature
     ) public virtual {
-        bytes32 domainSeparator = keccak256(
-            abi.encode(DOMAIN_TYPEHASH, keccak256(bytes(name())), block.chainid, address(this))
-        );
         bytes32 structHash = keccak256(abi.encode(DELEGATION_TYPEHASH, delegatee, nonce, expiry));
-        bytes32 digest = keccak256(abi.encodePacked("\x19\x01", domainSeparator, structHash));
-        address signatory = ecrecover(digest, v, r, s);
+        bytes32 digest = keccak256(abi.encodePacked("\x19\x01", _domainSeparatorV4(), structHash));
+
+        if (!SignatureChecker.isValidSignatureNow(delegator, digest, signature)) {
+            revert SignatureVerificationFail();
+        }
+
+        if (nonce != nonces[delegator]++) {
+            revert SignatureNonceInvalid();
+        }
+
+        if (block.timestamp > expiry) {
+            revert SignatureExpired();
+        }
+
+        return _delegate(delegator, delegatee);
+    }
+
+    /**
+     * @notice          Increment the nonce.
+     * @dev             This function enables the sender to cancel a signature.
+     */
+    function incrementNonce() public virtual {
+        uint256 currentNonce = nonces[msg.sender];
+        nonces[msg.sender] = ++currentNonce;
 
-        require(signatory != address(0), "Comp::delegateBySig: invalid signature");
-        require(nonce == nonces[signatory]++, "Comp::delegateBySig: invalid nonce");
-        require(block.timestamp <= expiry, "Comp::delegateBySig: signature expired");
-        return _delegate(signatory, delegatee);
+        emit NonceIncremented(msg.sender, currentNonce);
     }
 
     /**

contracts/test/governance/TestComp.sol

@@ -5,7 +5,13 @@ import { Comp } from "../../governance/Comp.sol";
 import { DefaultFHEVMConfig } from "../DefaultFHEVMConfig.sol";
 
 contract TestComp is DefaultFHEVMConfig, Comp {
-    constructor(address owner_) Comp(owner_) {
+    constructor(
+        address owner_,
+        string memory name_,
+        string memory symbol_,
+        string memory version_,
+        uint64 totalSupply_
+    ) Comp(owner_, name_, symbol_, version_, totalSupply_) {
         //
     }
 }

test/governance/Comp.fixture.ts

@@ -1,3 +1,4 @@
+import { parseUnits } from "ethers";
 import { ethers } from "hardhat";
 
 import type { TestComp } from "../../types";
@@ -7,7 +8,9 @@ import { FhevmInstances } from "../types";
 
 export async function deployCompFixture(signers: Signers): Promise<TestComp> {
   const contractFactory = await ethers.getContractFactory("TestComp");
-  const contract = await contractFactory.connect(signers.alice).deploy(signers.alice.address);
+  const contract = await contractFactory
+    .connect(signers.alice)
+    .deploy(signers.alice.address, "CompoundZama", "COMP", "1.0", parseUnits("10000000", 6));
   await contract.waitForDeployment();
   return contract;
 }

test/governance/Comp.test.ts

@@ -67,14 +67,17 @@ describe("Comp", function () {
   });
 
   it("can delegate votes via delegateBySig if signature is valid", async function () {
+    const delegator = this.signers.alice;
     const delegatee = this.signers.bob;
     const nonce = 0;
     let latestBlockNumber = await ethers.provider.getBlockNumber();
     const block = await ethers.provider.getBlock(latestBlockNumber);
     const expiry = block!.timestamp + 100;
-    const [v, r, s] = await delegateBySig(this.signers.alice, delegatee.address, this.comp, nonce, expiry);
+    const signature = await delegateBySig(delegator, delegatee.address, this.comp, nonce, expiry, false);
 
-    const tx = await this.comp.connect(this.signers.alice).delegateBySig(delegatee, nonce, expiry, v, r, s);
+    const tx = await this.comp
+      .connect(this.signers.alice)
+      .delegateBySig(delegator, delegatee, nonce, expiry, signature);
     await tx.wait();
 
     latestBlockNumber = await ethers.provider.getBlockNumber();
@@ -125,49 +128,75 @@ describe("Comp", function () {
   });
 
   it("cannot delegate votes if nonce is invalid", async function () {
+    const delegator = this.signers.alice;
     const delegatee = this.signers.bob;
     const nonce = 0;
     let latestBlockNumber = await ethers.provider.getBlockNumber();
     const block = await ethers.provider.getBlock(latestBlockNumber);
     const expiry = block!.timestamp + 100;
-    const [v, r, s] = await delegateBySig(this.signers.alice, delegatee.address, this.comp, nonce, expiry);
+    const signature = await delegateBySig(delegator, delegatee.address, this.comp, nonce, expiry);
 
-    const tx = await this.comp.connect(this.signers.alice).delegateBySig(delegatee, nonce, expiry, v, r, s);
+    const tx = await this.comp
+      .connect(this.signers.alice)
+      .delegateBySig(delegator, delegatee, nonce, expiry, signature);
     await tx.wait();
 
     // Cannot reuse same nonce when delegating by sig
-    await expect(this.comp.delegateBySig(delegatee, nonce, expiry, v, r, s)).to.be.revertedWith(
-      "Comp::delegateBySig: invalid nonce",
+    await expect(this.comp.delegateBySig(delegator, delegatee, nonce, expiry, signature)).to.be.revertedWithCustomError(
+      this.comp,
+      "SignatureNonceInvalid",
+    );
+  });
+
+  it("cannot delegate votes if nonce is invalid due to the delegator incrementing her nonce", async function () {
+    const delegator = this.signers.alice;
+    const delegatee = this.signers.bob;
+    const nonce = 0;
+    let latestBlockNumber = await ethers.provider.getBlockNumber();
+    const block = await ethers.provider.getBlock(latestBlockNumber);
+    const expiry = block!.timestamp + 100;
+    const signature = await delegateBySig(delegator, delegatee.address, this.comp, nonce, expiry);
+
+    const tx = await this.comp.connect(delegator).incrementNonce();
+    await tx.wait();
+
+    // Cannot reuse same nonce when delegating by sig
+    await expect(this.comp.delegateBySig(delegator, delegatee, nonce, expiry, signature)).to.be.revertedWithCustomError(
+      this.comp,
+      "SignatureNonceInvalid",
     );
   });
 
   it("cannot delegate votes if signer is invalid", async function () {
+    const delegator = this.signers.alice;
     const delegatee = this.signers.bob;
     const nonce = 0;
     let latestBlockNumber = await ethers.provider.getBlockNumber();
     const block = await ethers.provider.getBlock(latestBlockNumber);
     const expiry = block!.timestamp + 100;
-    const [v, r, s] = await delegateBySig(this.signers.alice, delegatee.address, this.comp, nonce, expiry);
 
-    // Cannot use invalid signature when delegating by sig
-    await expect(this.comp.delegateBySig(delegatee, nonce, expiry, 30, r, s)).to.be.revertedWith(
-      "Comp::delegateBySig: invalid signature",
+    // Signer is not the delegator
+    const signature = await delegateBySig(this.signers.carol, delegatee.address, this.comp, nonce, expiry);
+    await expect(this.comp.delegateBySig(delegator, delegatee, nonce, expiry, signature)).to.be.revertedWithCustomError(
+      this.comp,
+      "SignatureVerificationFail",
     );
   });
 
   it("cannot delegate votes if signature has expired", async function () {
+    const delegator = this.signers.alice;
     const delegatee = this.signers.bob;
     const nonce = 0;
     let latestBlockNumber = await ethers.provider.getBlockNumber();
     const block = await ethers.provider.getBlock(latestBlockNumber);
     const expiry = block!.timestamp + 100;
-    const [v, r, s] = await delegateBySig(this.signers.alice, delegatee.address, this.comp, nonce, expiry);
+    const signature = await delegateBySig(delegator, delegatee.address, this.comp, nonce, expiry);
 
     ethers.provider.send("evm_increaseTime", ["0xffff"]);
 
-    await expect(this.comp.connect(delegatee).delegateBySig(delegatee, nonce, expiry, v, r, s)).to.be.revertedWith(
-      "Comp::delegateBySig: signature expired",
-    );
+    await expect(
+      this.comp.connect(delegatee).delegateBySig(delegator, delegatee, nonce, expiry, signature),
+    ).to.be.revertedWithCustomError(this.comp, "SignatureExpired");
   });
 
   it("cannot request votes if blocktime is equal to current blocktime", async function () {
@@ -343,7 +372,7 @@ describe("Comp", function () {
           await reencryptPriorVotes(
             this.signers,
             this.instances,
-            "alice",
+            "carol",
             blockNumbers[i],
             this.comp,
             this.compAddress,
@@ -354,7 +383,7 @@ describe("Comp", function () {
           await reencryptPriorVotes(
             this.signers,
             this.instances,
-            "carol",
+            "alice",
             blockNumbers[i],
             this.comp,
             this.compAddress,

test/governance/DelegateBySig.ts

@@ -4,26 +4,38 @@ import { Address } from "hardhat-deploy/types";
 
 import type { Comp } from "../../types";
 
+/**
+ *
+ * @param _signer     Signer from ethers.
+ * @param _delegatee  Delegatee address.
+ * @param _comp       Comp token.
+ * @param _nonce      Nonce to sign.
+ * @param _expiry     Expiry timestamp.
+ * @returns           The signature.
+ */
 export const delegateBySig = async (
   _signer: HardhatEthersSigner,
   _delegatee: Address,
   _comp: Comp,
   _nonce: number,
   _expiry: number,
-): Promise<[BigInt, string, string]> => {
+): Promise<string> => {
   const compAddress_ = await _comp.getAddress();
   const delegatee_ = _delegatee;
   const nonce_ = _nonce;
   const expiry_ = _expiry;
-
   const network = await ethers.provider.getNetwork();
   const chainId = network.chainId;
+
   const domain = {
     name: await _comp.name(),
+    version: "1.0",
     chainId: chainId,
     verifyingContract: compAddress_,
   };
+
   // Delegation(address delegatee,uint256 nonce,uint256 expiry)
+
   const types = {
     Delegation: [
       {
@@ -48,9 +60,5 @@ export const delegateBySig = async (
   };
 
   const signature = await _signer.signTypedData(domain, types, message);
-  const sigRSV = ethers.Signature.from(signature);
-  const v = 27 + sigRSV.yParity;
-  const r = sigRSV.r;
-  const s = sigRSV.s;
-  return [BigInt(v), r, s];
+  return signature;
 };