fix: check limit after upgrade

VenusProtocol · Sep 21, 2023 · b98dc82 · b98dc82
1 parent 5a04aa0
commit b98dc82
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 7 deletions.
diff --git a/contracts/Tokens/Prime/Prime.sol b/contracts/Tokens/Prime/Prime.sol
@@ -271,12 +271,7 @@ contract Prime is IIncomeDestination, AccessControlledV8, PausableUpgradeable, M
             for (uint256 i = 0; i < users.length; ) {
                 Token storage userToken = tokens[users[i]];
                 if (userToken.exists && !userToken.isIrrevocable) {
-                    //upgrade to irrevocable token
-                    userToken.isIrrevocable = true;      
-                    totalIrrevocable++;
-                    totalRevocable--; 
-
-                    emit TokenUpgraded(users[i]);
+                    _upgrade(users[i]);
                 } else {
                     _mint(true, users[i]);
                     _initializeMarkets(users[i]);
@@ -575,6 +570,22 @@ contract Prime is IIncomeDestination, AccessControlledV8, PausableUpgradeable, M
         emit Burn(user);
     }
 
+    /**
+     * @notice Used to upgrade an token
+     * @param user owner whose prime token to upgrade
+     */
+    function _upgrade(address user) internal {
+        Token storage userToken = tokens[user];
+
+        userToken.isIrrevocable = true;      
+        totalIrrevocable++;
+        totalRevocable--; 
+
+        if (totalIrrevocable > irrevocableLimit) revert InvalidLimit();
+
+        emit TokenUpgraded(user);
+    }
+
     /**
      * @notice Used to get if the XVS balance is eligible for prime token
      * @param amount amount of XVS

diff --git a/contracts/Tokens/Prime/PrimeLiquidityProvider.sol b/contracts/Tokens/Prime/PrimeLiquidityProvider.sol
@@ -342,7 +342,7 @@ contract PrimeLiquidityProvider is AccessControlledV8, PausableUpgradeable {
     function _ensureTokenInitialized(address token_) internal view {
         uint256 lastBlockAccrued = lastAccruedBlock[token_];
 
-        if (lastBlockAccrued == 0)) {
+        if (lastBlockAccrued == 0) {
             revert TokenNotInitialized(token_);
         }
     }