feat: improve the logic of the complex balance (#80)

* move all after checks to CWToken * remove unnecessary tests * comparison change * remove redundant functions * add condition * update tests * add tests * add tests * add tests with greater restriction * test: make additional changes in tests * codacy notification * build: upgrade BRLC and USJIM token contracts in Testnet * build: prepare BRLC and USJIM token contracts upgrade in Mainnet --------- Co-authored-by: Evgenii Zaitsev <evgenii@cloudwalk.io> Co-authored-by: Igor Senych <igor.senych@cloudwalk.io>
cloudwalk · Jul 29, 2024 · 63a0831 · 63a0831
1 parent 80bdba3
commit 63a0831
Show file tree

Hide file tree

Showing 14 changed files with 2,525 additions and 158 deletions.
diff --git a/.openzeppelin/unknown-2008.json b/.openzeppelin/unknown-2008.json
diff --git a/.openzeppelin/unknown-2009.json b/.openzeppelin/unknown-2009.json
diff --git a/contracts/BRLCTokenBridgeable.sol b/contracts/BRLCTokenBridgeable.sol
@@ -92,14 +92,7 @@ contract BRLCTokenBridgeable is ERC20Base, ERC20Bridgeable, ERC20Freezable {
         address from,
         address to,
         uint256 amount
-    ) internal virtual override(ERC20Base, ERC20Freezable) {
+    ) internal virtual override(ERC20Base) {
         super._afterTokenTransfer(from, to, amount);
     }
-
-    /**
-     * @inheritdoc ERC20Freezable
-     */
-    function _balanceOf_ERC20Freezable(address account) internal view override returns (uint256) {
-        return balanceOf(account);
-    }
 }
diff --git a/contracts/base/CWToken.sol b/contracts/base/CWToken.sol
@@ -112,17 +112,32 @@ contract CWToken is
 
     /**
      * @dev See {ERC20Base-_afterTokenTransfer}
-     * @dev See {ERC20Mintable-_afterTokenTransfer}
-     * @dev See {ERC20Freezable-_afterTokenTransfer}
      * @dev See {ERC20Restrictable-_afterTokenTransfer}
      * @dev See {ERC20Hookable-_afterTokenTransfer}
      */
     function _afterTokenTransfer(
         address from,
         address to,
         uint256 amount
-    ) internal virtual override(ERC20Base, ERC20Mintable, ERC20Freezable, ERC20Restrictable, ERC20Hookable) {
+    ) internal virtual override(ERC20Base, ERC20Restrictable, ERC20Hookable) {
         super._afterTokenTransfer(from, to, amount);
+
+        uint256 balanceTotal = balanceOf(from);
+        uint256 balanceFrozen = balanceOfFrozen(from);
+        uint256 balancePreminted = balanceOfPremint(from);
+        uint256 balanceRestricted = balanceOfRestricted(from, bytes32(0));
+
+        if (balanceTotal < balanceFrozen + balancePreminted + balanceRestricted) {
+            uint256 balanceFreezable = (balanceTotal >= balancePreminted) ? balanceTotal - balancePreminted : 0;
+
+            if (balanceTotal < balancePreminted) {
+                revert TransferExceededPremintedAmount();
+            } else if (balanceFreezable < balanceFrozen && msg.sig != this.transferFrozen.selector) {
+                revert TransferExceededFrozenAmount();
+            } else if (balanceRestricted != 0 && msg.sig != this.transferFrozen.selector) {
+                revert TransferExceededRestrictedAmount();
+            }
+        }
     }
 
     /**
@@ -135,31 +150,4 @@ contract CWToken is
     ) public view override(ERC20Base, ERC20Trustable) returns (uint256) {
         return super.allowance(owner, spender);
     }
-
-    /**
-     * @inheritdoc ERC20Mintable
-     */
-    function _balanceOf_ERC20Mintable(address account) internal view virtual override returns (uint256) {
-        return balanceOf(account);
-    }
-
-    /**
-     * @inheritdoc ERC20Freezable
-     */
-    function _balanceOf_ERC20Freezable(address account) internal view virtual override returns (uint256) {
-        return balanceOf(account) - balanceOfPremint(account);
-    }
-
-    /**
-     * @inheritdoc ERC20Restrictable
-     */
-    function _balanceOf_ERC20Restrictable(address account) internal view virtual override returns (uint256) {
-        uint256 frozenBalance = balanceOfFrozen(account);
-        uint256 restBalance = balanceOf(account) - balanceOfPremint(account);
-        if (frozenBalance < restBalance) {
-            return restBalance - frozenBalance;
-        } else {
-            return 0;
-        }
-    }
 }
diff --git a/contracts/base/ERC20Freezable.sol b/contracts/base/ERC20Freezable.sol
@@ -131,26 +131,6 @@ abstract contract ERC20Freezable is ERC20Base, IERC20Freezable {
         return balanceOfFrozen(account);
     }
 
-    /**
-     * @notice Returns the transferable amount of tokens owned by account
-     *
-     * @param account The account to get the balance of
-     */
-    function _balanceOf_ERC20Freezable(address account) internal view virtual returns (uint256);
-
-    /**
-     * @inheritdoc ERC20Base
-     */
-    function _afterTokenTransfer(address from, address to, uint256 amount) internal virtual override {
-        super._afterTokenTransfer(from, to, amount);
-        uint256 frozen = _frozenBalances[from];
-        if (frozen != 0 && msg.sig != this.transferFrozen.selector) {
-            if (_balanceOf_ERC20Freezable(from) < frozen) {
-                revert TransferExceededFrozenAmount();
-            }
-        }
-    }
-
     /**
      * @dev This empty reserved space is put in place to allow future versions
      * to add new variables without shifting down storage in the inheritance chain

diff --git a/contracts/base/ERC20Mintable.sol b/contracts/base/ERC20Mintable.sol
@@ -415,19 +415,6 @@ abstract contract ERC20Mintable is ERC20Base, IERC20Mintable {
         return true;
     }
 
-    /**
-     * @inheritdoc ERC20Base
-     */
-    function _afterTokenTransfer(address from, address to, uint256 amount) internal virtual override {
-        super._afterTokenTransfer(from, to, amount);
-        uint256 preminted = balanceOfPremint(from);
-        if (preminted != 0) {
-            if (_balanceOf_ERC20Mintable(from) < preminted) {
-                revert TransferExceededPremintedAmount();
-            }
-        }
-    }
-
     function _getExtendedStorageSlot() internal pure returns (ExtendedStorageSlot storage r) {
         /// @solidity memory-safe-assembly
         assembly {
@@ -563,11 +550,4 @@ abstract contract ERC20Mintable is ERC20Base, IERC20Mintable {
         }
         premintRecords.pop();
     }
-
-    /**
-     * @notice Returns the transferable amount of tokens owned by account
-     *
-     * @param account The account to get the balance of
-     */
-    function _balanceOf_ERC20Mintable(address account) internal view virtual returns (uint256);
 }
diff --git a/contracts/base/ERC20Restrictable.sol b/contracts/base/ERC20Restrictable.sol
@@ -163,21 +163,10 @@ abstract contract ERC20Restrictable is ERC20Base, IERC20Restrictable {
                 }
             }
 
-            if (_balanceOf_ERC20Restrictable(from) < restrictedBalance) {
-                revert TransferExceededRestrictedAmount();
-            }
-
             _totalRestrictedBalances[from] = restrictedBalance;
         }
     }
 
-    /**
-     * @notice Returns the transferable amount of tokens owned by account
-     *
-     * @param account The account to get the balance of
-     */
-    function _balanceOf_ERC20Restrictable(address account) internal view virtual returns (uint256);
-
     /**
      * @dev This empty reserved space is put in place to allow future versions
      * to add new variables without shifting down storage in the inheritance chain

diff --git a/contracts/mocks/base/ERC20FreezableMock.sol b/contracts/mocks/base/ERC20FreezableMock.sol
@@ -63,11 +63,4 @@ contract ERC20FreezableMock is ERC20Freezable {
         _mint(account, amount);
         return true;
     }
-
-    /**
-     * @inheritdoc ERC20Freezable
-     */
-    function _balanceOf_ERC20Freezable(address account) internal view virtual override returns (uint256) {
-        return balanceOf(account);
-    }
 }
diff --git a/contracts/mocks/base/ERC20MintableMock.sol b/contracts/mocks/base/ERC20MintableMock.sol
@@ -32,7 +32,6 @@ contract ERC20MintableMock is ERC20Mintable {
      */
     function initialize(string memory name_, string memory symbol_) public initializer {
         __ERC20Mintable_init(name_, symbol_);
-        _balanceOf_ERC20Mintable(address(0)); // To ensure 100% coverage
     }
 
     /**
@@ -53,11 +52,4 @@ contract ERC20MintableMock is ERC20Mintable {
     function call_parent_initialize_unchained() public {
         __ERC20Mintable_init_unchained();
     }
-
-    /**
-     * @inheritdoc ERC20Mintable
-     */
-    function _balanceOf_ERC20Mintable(address account) internal view virtual override returns (uint256) {
-        return balanceOf(account);
-    }
 }
diff --git a/contracts/mocks/base/ERC20RestrictableMock.sol b/contracts/mocks/base/ERC20RestrictableMock.sol
@@ -63,11 +63,4 @@ contract ERC20RestrictableMock is ERC20Restrictable {
         _mint(account, amount);
         return true;
     }
-
-    /**
-     * @inheritdoc ERC20Restrictable
-     */
-    function _balanceOf_ERC20Restrictable(address account) internal view virtual override returns (uint256) {
-        return balanceOf(account);
-    }
 }