fix: allow custodian to freeze full balance, add edge tests, and clean unused code

contracts/token/ERC20/extensions/ERC20Custodian.sol

@@ -14,7 +14,7 @@ import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
  * of a user.
  *
  * The frozen balance is not available for transfers or approvals
- * to other entities to operate on its behalf if. The frozen balance
+ * to other entities to operate on its behalf. The frozen balance
  * can be reduced by calling {freeze} again with a lower amount.
  */
 abstract contract ERC20Custodian is ERC20 {
@@ -30,23 +30,11 @@ abstract contract ERC20Custodian is ERC20 {
      */
     event TokensFrozen(address indexed user, uint256 amount);
 
-    /**
-     * @dev Emitted when tokens are unfrozen for a user.
-     * @param user The address of the user whose tokens were unfrozen.
-     * @param amount The amount of tokens that were unfrozen.
-     */
-    event TokensUnfrozen(address indexed user, uint256 amount);
-
     /**
      * @dev The operation failed because the user has insufficient unfrozen balance.
      */
     error ERC20InsufficientUnfrozenBalance(address user);
 
-    /**
-     * @dev The operation failed because the user has insufficient frozen balance.
-     */
-    error ERC20InsufficientFrozenBalance(address user);
-
     /**
      * @dev Error thrown when a non-custodian account attempts to perform a custodian-only operation.
      */
@@ -70,14 +58,14 @@ abstract contract ERC20Custodian is ERC20 {
     /**
      * @dev Adjusts the amount of tokens frozen for a user.
      * @param user The address of the user whose tokens to freeze.
-     * @param amount The amount of tokens frozen.
+     * @param amount The total amount of tokens frozen for the user.
      *
      * Requirements:
      *
-     * - The user must have sufficient unfrozen balance.
+     * - The amount must not exceed the user's total balance.
      */
     function freeze(address user, uint256 amount) external virtual onlyCustodian {
-        if (availableBalance(user) < amount) revert ERC20InsufficientUnfrozenBalance(user);
+        if (balanceOf(user) < amount) revert ERC20InsufficientUnfrozenBalance(user);
         _frozen[user] = amount;
         emit TokensFrozen(user, amount);
     }
@@ -98,6 +86,10 @@ abstract contract ERC20Custodian is ERC20 {
      */
     function _isCustodian(address user) internal view virtual returns (bool);
 
+    /**
+     * @dev Override of {ERC20-_update} that enforces frozen balance restrictions.
+     * Prevents transfers when the sender doesn't have sufficient unfrozen balance.
+     */
     function _update(address from, address to, uint256 value) internal virtual override {
         if (from != address(0) && availableBalance(from) < value) revert ERC20InsufficientUnfrozenBalance(from);
         super._update(from, to, value);

docs/index.md

@@ -541,16 +541,15 @@ _See {ERC20-_update}._
 
 _Extension of {ERC20} that allows to implement a custodian
 mechanism that can be managed by an authorized account with the
-{freeze} and {unfreeze} functions.
+{freeze} function.
 
 This mechanism allows a custodian (e.g. a DAO or a
 well-configured multisig) to freeze and unfreeze the balance
 of a user.
 
 The frozen balance is not available for transfers or approvals
-to other entities to operate on its behalf if {freeze} was not
-called with such account as an argument. Similarly, the account
-will be unfrozen again if {unfreeze} is called._
+to other entities to operate on its behalf. The frozen balance
+can be reduced by calling {freeze} again with a lower amount.
 
 ### _frozen
 
@@ -573,22 +572,7 @@ _Emitted when tokens are frozen for a user._
 | Name | Type | Description |
 | ---- | ---- | ----------- |
 | user | address | The address of the user whose tokens were frozen. |
-| amount | uint256 | The amount of tokens that were frozen. |
-
-### TokensUnfrozen
-
-```solidity
-event TokensUnfrozen(address user, uint256 amount)
-```
-
-_Emitted when tokens are unfrozen for a user._
-
-#### Parameters
-
-| Name | Type | Description |
-| ---- | ---- | ----------- |
-| user | address | The address of the user whose tokens were unfrozen. |
-| amount | uint256 | The amount of tokens that were unfrozen. |
+| amount | uint256 | The total amount of tokens frozen for the user. |
 
 ### ERC20InsufficientUnfrozenBalance
 
@@ -598,14 +582,6 @@ error ERC20InsufficientUnfrozenBalance(address user)
 
 _The operation failed because the user has insufficient unfrozen balance._
 
-### ERC20InsufficientFrozenBalance
-
-```solidity
-error ERC20InsufficientFrozenBalance(address user)
-```
-
-_The operation failed because the user has insufficient frozen balance._
-
 ### ERC20NotCustodian
 
 ```solidity
@@ -643,7 +619,7 @@ _Adjusts the amount of tokens frozen for a user._
 | Name | Type | Description |
 | ---- | ---- | ----------- |
 | user | address | The address of the user whose tokens to freeze. |
-| amount | uint256 | The amount of tokens frozen. Requirements: - The user must have sufficient unfrozen balance. |
+| amount | uint256 | The total amount of tokens frozen for the user. Requirements: - The amount must not exceed the user's total balance. |
 
 ### availableBalance
 
@@ -1034,4 +1010,3 @@ any of the following bytes , )\x00
 
 If the `contentsType` is invalid, this returns an empty string. Otherwise, the return string has non-zero
 length._
-

test/token/ERC20/extensions/ERC20Custodian.test.js

@@ -132,6 +132,153 @@ describe('ERC20CustodianMock', function () {
           'ERC20InsufficientUnfrozenBalance',
         );
       });
+
+      it('should allow reducing frozen amount', async function () {
+        // First, freeze 80 tokens out of 100
+        await this.token.freeze(this.holder, 80n);
+        expect(await this.token.frozen(this.holder)).to.equal(80n);
+
+        // Now reduce frozen amount to 50 (this should work with the fix)
+        await this.token.freeze(this.holder, 50n);
+        expect(await this.token.frozen(this.holder)).to.equal(50n);
+
+        // Verify available balance is now 100 - 50 = 50
+        expect(await this.token.availableBalance(this.holder)).to.equal(50n);
+      });
+    });
+
+    describe('edge cases', function () {
+      it('should revert when non-custodian tries to freeze', async function () {
+        await expect(this.token.connect(this.recipient).freeze(this.holder, 50n)).to.be.revertedWithCustomError(
+          this.token,
+          'ERC20NotCustodian',
+        );
+      });
+
+      it('should revert when trying to freeze zero address', async function () {
+        await expect(this.token.freeze(ethers.ZeroAddress, 50n)).to.be.revertedWithCustomError(
+          this.token,
+          'ERC20InsufficientUnfrozenBalance',
+        );
+      });
+
+      it('should allow freezing zero amount (complete unfreeze)', async function () {
+        // First freeze some tokens
+        await this.token.freeze(this.holder, 80n);
+        expect(await this.token.frozen(this.holder)).to.equal(80n);
+
+        // Then unfreeze completely by setting to 0
+        await this.token.freeze(this.holder, 0n);
+        expect(await this.token.frozen(this.holder)).to.equal(0n);
+        expect(await this.token.availableBalance(this.holder)).to.equal(initialSupply);
+      });
+
+      it('should allow freezing exact balance amount', async function () {
+        await this.token.freeze(this.holder, initialSupply);
+
+        expect(await this.token.frozen(this.holder)).to.equal(initialSupply);
+        expect(await this.token.availableBalance(this.holder)).to.equal(0n);
+
+        // Should prevent any transfers
+        await expect(this.token.connect(this.holder).transfer(this.recipient, 1n)).to.be.revertedWithCustomError(
+          this.token,
+          'ERC20InsufficientUnfrozenBalance',
+        );
+      });
+
+      it('should allow partial transfers up to available balance', async function () {
+        await this.token.freeze(this.holder, 70n); // Freeze 70, available = 30
+
+        // Should allow transfer of exactly available amount
+        await this.token.connect(this.holder).transfer(this.recipient, 30n);
+        expect(await this.token.balanceOf(this.holder)).to.equal(70n);
+
+        // Should prevent transfer of even 1 more token
+        await expect(this.token.connect(this.holder).transfer(this.recipient, 1n)).to.be.revertedWithCustomError(
+          this.token,
+          'ERC20InsufficientUnfrozenBalance',
+        );
+      });
+
+      it('should allow minting to accounts with frozen balance', async function () {
+        await this.token.freeze(this.holder, 50n);
+
+        // Mint 20 more tokens
+        await this.token.$_mint(this.holder, 20n);
+
+        expect(await this.token.balanceOf(this.holder)).to.equal(120n);
+        expect(await this.token.frozen(this.holder)).to.equal(50n);
+        expect(await this.token.availableBalance(this.holder)).to.equal(70n);
+      });
+
+      it('should respect frozen balance in transferFrom', async function () {
+        const allowance = 50n;
+        await this.token.connect(this.holder).approve(this.approved, allowance);
+        await this.token.freeze(this.holder, 80n); // Available = 20
+
+        // Should fail even with approval when trying to transfer more than available
+        await expect(
+          this.token.connect(this.approved).transferFrom(this.holder, this.recipient, 30n),
+        ).to.be.revertedWithCustomError(this.token, 'ERC20InsufficientUnfrozenBalance');
+
+        // Should work for available amount
+        await this.token.connect(this.approved).transferFrom(this.holder, this.recipient, 20n);
+        expect(await this.token.balanceOf(this.holder)).to.equal(80n);
+      });
+
+      it('should allow increasing frozen amount', async function () {
+        await this.token.freeze(this.holder, 50n);
+        expect(await this.token.availableBalance(this.holder)).to.equal(50n);
+
+        // Increase frozen amount
+        await this.token.freeze(this.holder, 80n);
+        expect(await this.token.frozen(this.holder)).to.equal(80n);
+        expect(await this.token.availableBalance(this.holder)).to.equal(20n);
+      });
+
+      it('should emit TokensFrozen event with correct parameters', async function () {
+        await expect(this.token.freeze(this.holder, 50n))
+          .to.emit(this.token, 'TokensFrozen')
+          .withArgs(this.holder, 50n);
+      });
+
+      it('should handle accounts with zero balance', async function () {
+        // Test with recipient who has 0 balance
+        expect(await this.token.balanceOf(this.recipient)).to.equal(0n);
+
+        // Should allow freezing 0 amount
+        await this.token.freeze(this.recipient, 0n);
+        expect(await this.token.frozen(this.recipient)).to.equal(0n);
+
+        // Should revert when trying to freeze non-zero amount
+        await expect(this.token.freeze(this.recipient, 1n)).to.be.revertedWithCustomError(
+          this.token,
+          'ERC20InsufficientUnfrozenBalance',
+        );
+      });
+
+      it('should allow burning frozen tokens', async function () {
+        await this.token.freeze(this.holder, 50n);
+
+        // Burning should work even for frozen tokens
+        await expect(this.token.$_burn(this.holder, 30n)).to.changeTokenBalance(this.token, this.holder, -30n);
+
+        // Balance should now be 70, frozen still 50, available = 20
+        expect(await this.token.balanceOf(this.holder)).to.equal(70n);
+        expect(await this.token.frozen(this.holder)).to.equal(50n);
+        expect(await this.token.availableBalance(this.holder)).to.equal(20n);
+      });
+
+      it('should handle multiple freeze operations correctly', async function () {
+        // Start with multiple freeze operations
+        await this.token.freeze(this.holder, 30n);
+        await this.token.freeze(this.holder, 60n); // Increase
+        await this.token.freeze(this.holder, 40n); // Decrease
+        await this.token.freeze(this.holder, 0n); // Complete unfreeze
+
+        expect(await this.token.frozen(this.holder)).to.equal(0n);
+        expect(await this.token.availableBalance(this.holder)).to.equal(initialSupply);
+      });
     });
   });
 });