refactor: Contracts

contracts/DAO/Comp.sol

@@ -11,24 +11,34 @@ import { IComp } from "./IComp.sol";
  * @notice      This contract inherits EncryptedERC20 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
+ *              GovernorAlphaZama.sol.
  *              It uses encrypted votes to delegate the voting power associated
  *              with an account's balance.
  * @dev         The delegation of votes leaks information about the account's encrypted balance to the delegate.
  */
 contract Comp is IComp, EncryptedERC20, Ownable2Step {
+    /// @notice Returned if the `blockNumber` is higher or equal to the (current) `block.number`.
+    /// @dev    It is returned for requests to access votes.
+    error BlockNumberEqualOrHigherThanCurrentBlock();
+
+    /// @notice Returned if the `msg.sender` is not the `governor` contract.
+    error GovernorInvalid();
+
     /// @notice Emitted when an account changes its delegate.
     event DelegateChanged(address indexed delegator, address indexed fromDelegate, address indexed toDelegate);
 
     /// @notice Emitted when a delegate account's vote balance changes.
     event DelegateVotesChanged(address indexed delegate);
 
-    /// @notice Emitted when the contract that can reencrypt changes.
-    event NewAllowedContract(address indexed allowedContract);
+    /// @notice Emitted when the governor contract that can reencrypt votes changes.
+    /// @dev    It can be set to a malicious contract, which could reencrypt all user votes.
+    event NewGovernor(address indexed governor);
 
     /// @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.
-    /// @dev             In Compound's implementation, `fromBlock` is defined as uint32 to allow tight-packing
+    /// @dev             In Compound's implementation, `fromBlock` is defined as uint32 to allow tight-packing.
     ///                  However, in this implementations `votes` is uint256-based.
     ///                  `fromBlock`'s type is set to uint256, which simplifies the codebase.
     struct Checkpoint {
@@ -44,8 +54,9 @@ contract Comp is IComp, EncryptedERC20, Ownable2Step {
     bytes32 public constant DELEGATION_TYPEHASH =
         keccak256("Delegation(address delegatee,uint256 nonce,uint256 expiry)");
 
-    /// @notice The smart contract that can access votes.
-    address public allowedContract;
+    /// @notice The smart contract that can access encrypted votes.
+    /// @dev    The contract is expected to be a governor contract.
+    address public governor;
 
     /// @notice A record of each account's delegate.
     mapping(address account => address delegate) public delegates;
@@ -111,15 +122,21 @@ contract Comp is IComp, EncryptedERC20, Ownable2Step {
     /**
      * @notice              Determine the prior number of votes for an account as of a block number.
      * @dev                 Block number must be a finalized block or else this function will revert.
-     *                      This function can change the state since the allowedContract needs access in the ACL
+     *                      This function can change the state since the governor needs access in the ACL
      *                      contract.
      * @param account       Account address.
      * @param blockNumber   The block number to get the vote balance at.
      * @return votes        Number of votes the account as of the given block number.
      */
-    function getPriorVotesForAllowedContract(address account, uint256 blockNumber) external returns (euint64 votes) {
-        require(msg.sender == allowedContract, "Caller not allowed to call this function");
-        require(blockNumber < block.number, "Comp::getPriorVotes: not yet determined");
+    function getPriorVotesForGovernor(address account, uint256 blockNumber) external returns (euint64 votes) {
+        if (msg.sender != governor) {
+            revert GovernorInvalid();
+        }
+
+        if (blockNumber >= block.number) {
+            revert BlockNumberEqualOrHigherThanCurrentBlock();
+        }
+
         votes = _getPriorVote(account, blockNumber);
         TFHE.allow(votes, msg.sender);
     }
@@ -142,17 +159,20 @@ contract Comp is IComp, EncryptedERC20, Ownable2Step {
      * @return votes        Number of votes the account as of the given block.
      */
     function getPriorVotes(address account, uint256 blockNumber) external view returns (euint64 votes) {
-        require(blockNumber < block.number, "Comp::getPriorVotes: not yet determined");
+        if (blockNumber >= block.number) {
+            revert BlockNumberEqualOrHigherThanCurrentBlock();
+        }
+
         return _getPriorVote(account, blockNumber);
     }
 
     /**
      * @notice                  Set an allowed contract that can access votes.
-     * @param contractAddress   The address of the smart contract that may access votes.
+     * @param newGovernor       New governor contract that can reencrypt/access votes.
      */
-    function setAllowedContract(address contractAddress) public onlyOwner {
-        allowedContract = contractAddress;
-        emit NewAllowedContract(contractAddress);
+    function setGovernor(address newGovernor) public onlyOwner {
+        governor = newGovernor;
+        emit NewGovernor(newGovernor);
     }
 
     function _delegate(address delegator, address delegatee) internal {

contracts/DAO/CompoundTimelock.sol

@@ -5,36 +5,13 @@ import { ICompoundTimelock } from "./ICompoundTimelock.sol";
 
 /**
  * @title       CompoundTimelock
+ * @notice      This contract allows the admin to set a delay period before executing transactions.
+ *              Transactions must be queued before execution. No transaction can be executed during this period,
+ *              which offers time to verify the validity of pending transactions.
+ *              It also has a grace period to allow for transactions
+ *              not to be executed after a specific period following the queuing.
  */
 contract CompoundTimelock is ICompoundTimelock {
-    event NewAdmin(address indexed newAdmin);
-    event NewPendingAdmin(address indexed newPendingAdmin);
-    event NewDelay(uint256 indexed newDelay);
-    event CancelTransaction(
-        bytes32 indexed txHash,
-        address indexed target,
-        uint256 value,
-        string signature,
-        bytes data,
-        uint256 eta
-    );
-    event ExecuteTransaction(
-        bytes32 indexed txHash,
-        address indexed target,
-        uint256 value,
-        string signature,
-        bytes data,
-        uint256 eta
-    );
-    event QueueTransaction(
-        bytes32 indexed txHash,
-        address indexed target,
-        uint256 value,
-        string signature,
-        bytes data,
-        uint256 eta
-    );
-
     uint256 public constant GRACE_PERIOD = 14 days;
     uint256 public constant MINIMUM_DELAY = 2 days;
     uint256 public constant MAXIMUM_DELAY = 30 days;

contracts/DAO/GovernorAlphaZama.sol

@@ -12,6 +12,12 @@ import { ICompoundTimelock } from "./ICompoundTimelock.sol";
  * @title       GovernorAlphaZama
  * @notice      This is based on the GovernorAlpha.sol contract written by Compound Labs.
  *              see: compound-finance/compound-protocol/blob/master/contracts/Governance/GovernorAlpha.sol
+ *              This decentralized governance system allows users to propose and vote on changes to the protocol.
+ *              The contract is responsible for:
+ *              - Proposal: A new proposal is made to introduce a change.
+ *              - Voting: Users can vote on the proposal, either in favor or against it.
+ *              - Quorum: A minimum number of votes (quorum) must be reached for the proposal to pass.
+ *              - Execution: Once a proposal passes, it is executed and takes effect on the Compound protocol.
  */
 contract GovernorAlphaZama is Ownable2Step, GatewayCaller {
     /// @notice Returned if proposal contains too many changes.
@@ -427,7 +433,7 @@ contract GovernorAlphaZama is Ownable2Step, GatewayCaller {
 
         ebool canPropose = TFHE.lt(
             _EUINT64_PROPOSAL_THRESHOLD,
-            COMP.getPriorVotesForAllowedContract(msg.sender, block.number - 1)
+            COMP.getPriorVotesForGovernor(msg.sender, block.number - 1)
         );
 
         uint256[] memory cts = new uint256[](1);
@@ -640,7 +646,7 @@ contract GovernorAlphaZama is Ownable2Step, GatewayCaller {
             revert VoterHasAlreadyVoted();
         }
 
-        euint64 votes = COMP.getPriorVotesForAllowedContract(voter, proposal.startBlock);
+        euint64 votes = COMP.getPriorVotesForGovernor(voter, proposal.startBlock);
         proposal.forVotes = TFHE.select(support, TFHE.add(proposal.forVotes, votes), proposal.forVotes);
         proposal.againstVotes = TFHE.select(support, proposal.againstVotes, TFHE.add(proposal.againstVotes, votes));
 

contracts/DAO/IComp.sol

@@ -8,5 +8,5 @@ import "fhevm/lib/TFHE.sol";
  * @dev         The GovernorAlphaZama relies on this interface.
  */
 interface IComp {
-    function getPriorVotesForAllowedContract(address account, uint256 blockNumber) external returns (euint64 votes);
+    function getPriorVotesForGovernor(address account, uint256 blockNumber) external returns (euint64 votes);
 }

contracts/DAO/ICompoundTimelock.sol

@@ -5,6 +5,34 @@ pragma solidity ^0.8.24;
  * @title       ICompoundTimelock
  */
 interface ICompoundTimelock {
+    event NewAdmin(address indexed newAdmin);
+    event NewPendingAdmin(address indexed newPendingAdmin);
+    event NewDelay(uint256 indexed newDelay);
+    event CancelTransaction(
+        bytes32 indexed txHash,
+        address indexed target,
+        uint256 value,
+        string signature,
+        bytes data,
+        uint256 eta
+    );
+    event ExecuteTransaction(
+        bytes32 indexed txHash,
+        address indexed target,
+        uint256 value,
+        string signature,
+        bytes data,
+        uint256 eta
+    );
+    event QueueTransaction(
+        bytes32 indexed txHash,
+        address indexed target,
+        uint256 value,
+        string signature,
+        bytes data,
+        uint256 eta
+    );
+
     function delay() external view returns (uint256);
 
     function GRACE_PERIOD() external view returns (uint256);