From d7d9aa0a476e5a4664cb90e3fb747cec470a3920 Mon Sep 17 00:00:00 2001
From: 0xRaccoon <raccoon@defi.sucks>
Date: Mon, 18 Dec 2023 16:03:47 -0300
Subject: [PATCH] fix: add proposalType validity check

Signed-off-by: 0xRaccoon <raccoon@defi.sucks>
---
 .../governance/utils/WonderVotes.sol          | 21 +++++++++++++++----
 .../governance/utils/IWonderVotes.sol         |  5 +++++
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/solidity/contracts/governance/utils/WonderVotes.sol b/solidity/contracts/governance/utils/WonderVotes.sol
index c48dcd1..6d85965 100644
--- a/solidity/contracts/governance/utils/WonderVotes.sol
+++ b/solidity/contracts/governance/utils/WonderVotes.sol
@@ -132,7 +132,7 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
   /**
    * @dev Delegates votes from the sender to `delegatee`.
    */
-  function delegate(Delegate[] calldata delegatees, uint8 proposalType) public virtual {
+  function delegate(Delegate[] calldata delegatees, uint8 proposalType) public virtual validProposalType(proposalType) {
     address account = _msgSender();
     _delegate(account, proposalType, delegatees);
   }
@@ -140,7 +140,7 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
   /**
    * @dev See {IWonderVotes-delegate}.
    */
-  function delegate(address delegatee, uint8 proposalType) public virtual {
+  function delegate(address delegatee, uint8 proposalType) public virtual validProposalType(proposalType) {
     address account = _msgSender();
     Delegate[] memory _singleDelegate = new Delegate[](1);
     _singleDelegate[0] = Delegate({account: delegatee, weight: _totalWeight()});
@@ -187,7 +187,7 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
     uint8 v,
     bytes32 r,
     bytes32 s
-  ) public virtual {
+  ) public virtual validProposalType(proposalType) {
     if (block.timestamp > expiry) {
       revert VotesExpiredSignature(expiry);
     }
@@ -209,7 +209,7 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
     uint8 v,
     bytes32 r,
     bytes32 s
-  ) public virtual {
+  ) public virtual validProposalType(proposalType) {
     Delegate[] memory _singleDelegate = new Delegate[](1);
     _singleDelegate[0] = Delegate({account: delegatee, weight: _totalWeight()});
     delegateBySig(_singleDelegate, proposalType, nonce, expiry, v, r, s);
@@ -364,4 +364,17 @@ abstract contract WonderVotes is Context, EIP712, Nonces, IERC6372, IWonderVotes
    * @dev Returns the maximum number of delegates that `proposalType` can delegate to.
    */
   function _maxDelegates() internal view virtual returns (uint8);
+
+  /**
+   * @dev Returns true if the `proposalType` is valid, false otherwise.
+   */
+  function _validProposalType(uint8 proposalType) internal view virtual returns (bool);
+
+  /**
+   * @dev checks the `proposalType` validity
+   */
+  modifier validProposalType(uint8 proposalType) {
+    if (!_validProposalType(proposalType)) revert InvalidProposalType(proposalType);
+    _;
+  }
 }
diff --git a/solidity/interfaces/governance/utils/IWonderVotes.sol b/solidity/interfaces/governance/utils/IWonderVotes.sol
index 3483db9..3989e0d 100644
--- a/solidity/interfaces/governance/utils/IWonderVotes.sol
+++ b/solidity/interfaces/governance/utils/IWonderVotes.sol
@@ -26,6 +26,11 @@ interface IWonderVotes {
    */
   error ZeroWeight();
 
+  /**
+   * @dev The proposal type is invalid.
+   */
+  error InvalidProposalType(uint8 proposalType);
+
   /**
    * @dev The delegates number for a `proposalType` exceeds the maximum number of delegates.
    */