move all after checks to CWToken

cloudwalk · Jul 3, 2024 · ac8109d · ac8109d
1 parent 80bdba3
commit ac8109d
Show file tree

Hide file tree

Showing 8 changed files with 20 additions and 93 deletions.
diff --git a/contracts/BRLCTokenBridgeable.sol b/contracts/BRLCTokenBridgeable.sol
@@ -95,11 +95,4 @@ contract BRLCTokenBridgeable is ERC20Base, ERC20Bridgeable, ERC20Freezable {
     ) internal virtual override(ERC20Base, ERC20Freezable) {
         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
@@ -123,6 +123,26 @@ contract CWToken is
         uint256 amount
     ) internal virtual override(ERC20Base, ERC20Mintable, ERC20Freezable, 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 (balancePreminted != 0 && balanceTotal < balancePreminted) {
+                revert TransferExceededPremintedAmount();
+            } else if (balanceFrozen != 0 &&
+                msg.sig != this.transferFrozen.selector &&
+                balanceFreezable < balanceFrozen
+            ) {
+                revert TransferExceededFrozenAmount();
+            } else if (balanceRestricted != 0) {
+                revert TransferExceededRestrictedAmount();
+            }
+        }
     }
 
     /**
@@ -135,31 +155,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,24 +131,11 @@ 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();
-            }
-        }
     }
 
     /**

diff --git a/contracts/base/ERC20Mintable.sol b/contracts/base/ERC20Mintable.sol
@@ -420,12 +420,6 @@ abstract contract ERC20Mintable is ERC20Base, IERC20Mintable {
      */
     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) {
@@ -563,11 +557,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);
-    }
 }