feat: add increase/decrease restriction functions

cloudwalk · May 2, 2024 · b31e270 · b31e270
1 parent 40a93a4
commit b31e270
Show file tree

Hide file tree

Showing 4 changed files with 188 additions and 3 deletions.
diff --git a/contracts/base/ERC20Base.sol b/contracts/base/ERC20Base.sol
@@ -21,6 +21,12 @@ abstract contract ERC20Base is
     BlocklistableUpgradeable,
     ERC20Upgradeable
 {
+    /// @dev Throws if the zero address is passed to the function
+    error ZeroAddress();
+
+    /// @dev Throws if the zero amount is passed to the function
+    error ZeroAmount();
+
     /**
      * @notice The internal initializer of the upgradable contract
      *

diff --git a/contracts/base/ERC20Restrictable.sol b/contracts/base/ERC20Restrictable.sol
@@ -84,6 +84,52 @@ abstract contract ERC20Restrictable is ERC20Base, IERC20Restrictable {
         emit UpdateRestriction(account, purpose, balance, oldBalance);
     }
 
+    /**
+     * @inheritdoc IERC20Restrictable
+     */
+    function restrictionIncrease(address account, bytes32 purpose, uint256 amount) external onlyBlocklister {
+        if (account == address(0)) {
+            revert ZeroAddress();
+        }
+        if (purpose == bytes32(0)) {
+            revert ZeroPurpose();
+        }
+        if (amount == 0) {
+            revert ZeroAmount();
+        }
+
+        uint256 oldBalance = _restrictedPurposeBalances[account][purpose];
+        uint256 newBalance = oldBalance + amount;
+
+        _restrictedPurposeBalances[account][purpose] = newBalance;
+        _totalRestrictedBalances[account] += amount;
+
+        emit UpdateRestriction(account, purpose, newBalance, oldBalance);
+    }
+
+    /**
+     * @inheritdoc IERC20Restrictable
+     */
+    function restrictionDecrease(address account, bytes32 purpose, uint256 amount) external onlyBlocklister {
+        if (account == address(0)) {
+            revert ZeroAddress();
+        }
+        if (purpose == bytes32(0)) {
+            revert ZeroPurpose();
+        }
+        if (amount == 0) {
+            revert ZeroAmount();
+        }
+
+        uint256 oldBalance = _restrictedPurposeBalances[account][purpose];
+        uint256 newBalance = oldBalance - amount;
+
+        _restrictedPurposeBalances[account][purpose] = newBalance;
+        _totalRestrictedBalances[account] -= amount;
+
+        emit UpdateRestriction(account, purpose, newBalance, oldBalance);
+    }
+
     /**
      * @inheritdoc IERC20Restrictable
      */

diff --git a/contracts/base/interfaces/IERC20Restrictable.sol b/contracts/base/interfaces/IERC20Restrictable.sol
@@ -51,6 +51,24 @@ interface IERC20Restrictable {
      */
     function updateRestriction(address account, bytes32 purpose, uint256 balance) external;
 
+    /**
+     * @notice Increases the restriction balance for an account
+     *
+     * @param account The account to increase restriction for
+     * @param purpose The restriction purpose
+     * @param amount The amount to increase the restriction balance by
+     */
+    function restrictionIncrease(address account, bytes32 purpose, uint256 amount) external;
+
+    /**
+     * @notice Decreases the restriction balance for an account
+     *
+     * @param account The account to decrease restriction for
+     * @param purpose The restriction purpose
+     * @param amount The amount to decrease the restriction balance by
+     */
+    function restrictionDecrease(address account, bytes32 purpose, uint256 amount) external;
+
     /**
      * @notice Returns the restricted balance for the account and the restriction purpose
      *

diff --git a/test/base/ERC20Restrictable.test.ts b/test/base/ERC20Restrictable.test.ts
@@ -20,6 +20,8 @@ describe("Contract 'ERC20Restrictable'", async () => {
   const REVERT_MESSAGE_INITIALIZABLE_CONTRACT_IS_ALREADY_INITIALIZED = "Initializable: contract is already initialized";
   const REVERT_MESSAGE_INITIALIZABLE_CONTRACT_IS_NOT_INITIALIZING = "Initializable: contract is not initializing";
 
+  const REVERT_ERROR_ZERO_ADDRESS = "ZeroAddress";
+  const REVERT_ERROR_ZERO_AMOUNT = "ZeroAmount";
   const REVERT_ERROR_ZERO_PURPOSE = "ZeroPurpose";
   const REVERT_ERROR_UNAUTHORIZED_BLOCKLISTER = "UnauthorizedBlocklister";
   const REVERT_ERROR_TRANSFER_EXCEEDED_RESTRICTED_AMOUNT = "TransferExceededRestrictedAmount";
@@ -178,9 +180,6 @@ describe("Contract 'ERC20Restrictable'", async () => {
 
     it("Is reverted if the caller is not a blocklister", async () => {
       const { token } = await setUpFixture(deployAndConfigureToken);
-      await expect(
-        token.updateRestriction(purposeAccount1.address, PURPOSE_1, 100)
-      ).to.be.revertedWithCustomError(token, REVERT_ERROR_UNAUTHORIZED_BLOCKLISTER);
       await expect(
         connect(token, user1).updateRestriction(purposeAccount1.address, PURPOSE_1, 100)
       ).to.be.revertedWithCustomError(token, REVERT_ERROR_UNAUTHORIZED_BLOCKLISTER);
@@ -194,6 +193,122 @@ describe("Contract 'ERC20Restrictable'", async () => {
     });
   });
 
+  describe("Function 'restrictionIncrease()'", async () => {
+    it("Increase restriction and emits the correct event", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_1)).to.eq(0);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_2)).to.eq(0);
+
+      await expect(connect(token, blocklister).restrictionIncrease(user1.address, PURPOSE_1, 100))
+        .to.emit(token, "UpdateRestriction")
+        .withArgs(user1.address, PURPOSE_1, 100, 0);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_1)).to.eq(100);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_2)).to.eq(0);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_ZERO)).to.eq(100);
+
+      await expect(connect(token, blocklister).restrictionIncrease(user1.address, PURPOSE_2, 200))
+        .to.emit(token, "UpdateRestriction")
+        .withArgs(user1.address, PURPOSE_2, 200, 0);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_1)).to.eq(100);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_2)).to.eq(200);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_ZERO)).to.eq(300);
+
+      await expect(connect(token, blocklister).restrictionIncrease(user1.address, PURPOSE_1, 100))
+        .to.emit(token, "UpdateRestriction")
+        .withArgs(user1.address, PURPOSE_1, 200, 100);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_1)).to.eq(200);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_2)).to.eq(200);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_ZERO)).to.eq(400);
+    });
+
+    it("Is reverted if the caller is not a blocklister", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, user1).restrictionIncrease(purposeAccount1.address, PURPOSE_1, 100)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_UNAUTHORIZED_BLOCKLISTER);
+    });
+
+    it("Is reverted if the provided account is zero", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, blocklister).restrictionIncrease(ethers.ZeroAddress, PURPOSE_1, 100)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_ZERO_ADDRESS);
+    });
+
+    it("Is reverted if the provided amount is zero", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, blocklister).restrictionIncrease(purposeAccount1.address, PURPOSE_1, 0)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_ZERO_AMOUNT);
+    });
+
+    it("Is reverted if the provided purpose is zero", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, blocklister).restrictionIncrease(purposeAccount1.address, PURPOSE_ZERO, 100)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_ZERO_PURPOSE);
+    });
+  });
+
+  describe("Function 'restrictionDecrease()'", async () => {
+    it("Increase restriction and emits the correct event", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+
+      await proveTx(connect(token, blocklister).restrictionIncrease(user1.address, PURPOSE_1, 200));
+      await proveTx(connect(token, blocklister).restrictionIncrease(user1.address, PURPOSE_2, 200));
+
+      await expect(connect(token, blocklister).restrictionDecrease(user1.address, PURPOSE_1, 100))
+        .to.emit(token, "UpdateRestriction")
+        .withArgs(user1.address, PURPOSE_1, 100, 200);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_1)).to.eq(100);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_2)).to.eq(200);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_ZERO)).to.eq(300);
+
+      await expect(connect(token, blocklister).restrictionDecrease(user1.address, PURPOSE_2, 200))
+        .to.emit(token, "UpdateRestriction")
+        .withArgs(user1.address, PURPOSE_2, 0, 200);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_1)).to.eq(100);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_2)).to.eq(0);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_ZERO)).to.eq(100);
+
+      await expect(connect(token, blocklister).restrictionDecrease(user1.address, PURPOSE_1, 100))
+        .to.emit(token, "UpdateRestriction")
+        .withArgs(user1.address, PURPOSE_1, 0, 100);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_1)).to.eq(0);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_2)).to.eq(0);
+      expect(await token.balanceOfRestricted(user1.address, PURPOSE_ZERO)).to.eq(0);
+    });
+
+    it("Is reverted if the caller is not a blocklister", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, user1).restrictionDecrease(purposeAccount1.address, PURPOSE_1, 100)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_UNAUTHORIZED_BLOCKLISTER);
+    });
+
+    it("Is reverted if the provided account is zero", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, blocklister).restrictionDecrease(ethers.ZeroAddress, PURPOSE_1, 100)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_ZERO_ADDRESS);
+    });
+
+    it("Is reverted if the provided amount is zero", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, blocklister).restrictionDecrease(purposeAccount1.address, PURPOSE_1, 0)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_ZERO_AMOUNT);
+    });
+
+    it("Is reverted if the provided purpose is zero", async () => {
+      const { token } = await setUpFixture(deployAndConfigureToken);
+      await expect(
+        connect(token, blocklister).restrictionDecrease(purposeAccount1.address, PURPOSE_ZERO, 100)
+      ).to.be.revertedWithCustomError(token, REVERT_ERROR_ZERO_PURPOSE);
+    });
+  });
+
   describe("Function 'balanceOfRestricted()'", async () => {
     it("Returns the correct value", async () => {
       const { token } = await setUpFixture(deployAndConfigureToken);