Merge pull request #43 from CMTA/whitelist-improvement

Whitelist improvement
CMTA · Jun 17, 2024 · 431fdf4 · 431fdf4
2 parents 061f26c + 1d0b481
commit 431fdf4
Show file tree

Hide file tree

Showing 9 changed files with 79 additions and 99 deletions.
diff --git a/doc/schema/rule/IncomeVault-claimDividend.drawio.png b/doc/schema/rule/IncomeVault-claimDividend.drawio.png
diff --git a/src/rules/validation/RuleWhitelist.sol b/src/rules/validation/RuleWhitelist.sol
@@ -2,18 +2,13 @@
 
 pragma solidity ^0.8.20;
 
-import "./abstract/RuleAddressList/invariantStorage/RuleWhitelistInvariantStorage.sol";
 import "./abstract/RuleAddressList/RuleAddressList.sol";
-import "./abstract/RuleValidateTransfer.sol";
-
+import "./abstract/RuleWhitelistCommon.sol";
 /**
  * @title a whitelist manager
  */
 
-contract RuleWhitelist is
-    RuleValidateTransfer,
-    RuleAddressList,
-    RuleWhitelistInvariantStorage
+contract RuleWhitelist is RuleAddressList, RuleWhitelistCommon
 {
     /**
      * @param admin Address of the contract (Access Control)
@@ -43,34 +38,4 @@ contract RuleWhitelist is
             return uint8(REJECTED_CODE_BASE.TRANSFER_OK);
         }
     }
-
-    /**
-     * @notice To know if the restriction code is valid for this rule or not
-     * @param _restrictionCode The target restriction code
-     * @return true if the restriction code is known, false otherwise
-     **/
-    function canReturnTransferRestrictionCode(
-        uint8 _restrictionCode
-    ) external pure override returns (bool) {
-        return
-            _restrictionCode == CODE_ADDRESS_FROM_NOT_WHITELISTED ||
-            _restrictionCode == CODE_ADDRESS_TO_NOT_WHITELISTED;
-    }
-
-    /**
-     * @notice Return the corresponding message
-     * @param _restrictionCode The target restriction code
-     * @return true if the transfer is valid, false otherwise
-     **/
-    function messageForTransferRestriction(
-        uint8 _restrictionCode
-    ) external pure override returns (string memory) {
-        if (_restrictionCode == CODE_ADDRESS_FROM_NOT_WHITELISTED) {
-            return TEXT_ADDRESS_FROM_NOT_WHITELISTED;
-        } else if (_restrictionCode == CODE_ADDRESS_TO_NOT_WHITELISTED) {
-            return TEXT_ADDRESS_TO_NOT_WHITELISTED;
-        } else {
-            return TEXT_CODE_NOT_FOUND;
-        }
-    }
 }
diff --git a/src/rules/validation/RuleWhitelistWrapper.sol b/src/rules/validation/RuleWhitelistWrapper.sol
@@ -5,18 +5,16 @@ pragma solidity ^0.8.20;
 import "OZ/access/AccessControl.sol";
 import "../../modules/RuleEngineValidationCommon.sol";
 import "../../modules/MetaTxModuleStandalone.sol";
-import "./abstract/RuleValidateTransfer.sol";
-import "./abstract/RuleAddressList/invariantStorage/RuleWhitelistInvariantStorage.sol";
 import "./abstract/RuleAddressList/RuleAddressList.sol";
+import "./abstract/RuleWhitelistCommon.sol";
 
 /**
  * @title Wrapper to call several different whitelist rules
  */
 contract RuleWhitelistWrapper is
     RuleEngineValidationCommon,
     MetaTxModuleStandalone,
-    RuleValidateTransfer,
-    RuleWhitelistInvariantStorage
+    RuleWhitelistCommon
 {
     /**
      * @param admin Address of the contract (Access Control)
@@ -73,36 +71,6 @@ contract RuleWhitelistWrapper is
         }
     }
 
-    /**
-     * @notice To know if the restriction code is valid for this rule or not
-     * @param _restrictionCode The target restriction code
-     * @return true if the restriction code is known, false otherwise
-     **/
-    function canReturnTransferRestrictionCode(
-        uint8 _restrictionCode
-    ) external pure override returns (bool) {
-        return
-            _restrictionCode == CODE_ADDRESS_FROM_NOT_WHITELISTED ||
-            _restrictionCode == CODE_ADDRESS_TO_NOT_WHITELISTED;
-    }
-
-    /**
-     * @notice Return the corresponding message
-     * @param _restrictionCode The target restriction code
-     * @return true if the transfer is valid, false otherwise
-     **/
-    function messageForTransferRestriction(
-        uint8 _restrictionCode
-    ) external pure override returns (string memory) {
-        if (_restrictionCode == CODE_ADDRESS_FROM_NOT_WHITELISTED) {
-            return TEXT_ADDRESS_FROM_NOT_WHITELISTED;
-        } else if (_restrictionCode == CODE_ADDRESS_TO_NOT_WHITELISTED) {
-            return TEXT_ADDRESS_TO_NOT_WHITELISTED;
-        } else {
-            return TEXT_CODE_NOT_FOUND;
-        }
-    }
-
     /**
      * @dev This surcharge is not necessary if you do not use the MetaTxModule
      */

diff --git a/src/rules/validation/abstract/RuleAddressList/RuleAddressList.sol b/src/rules/validation/abstract/RuleAddressList/RuleAddressList.sol
@@ -32,7 +32,8 @@ abstract contract RuleAddressList is
             revert RuleAddressList_AdminWithAddressZeroNotAllowed();
         }
         _grantRole(DEFAULT_ADMIN_ROLE, admin);
-        _grantRole(ADDRESS_LIST_ROLE, admin);
+        _grantRole(ADDRESS_LIST_ADD_ROLE, admin);
+        _grantRole(ADDRESS_LIST_REMOVE_ROLE, admin);
     }
 
     /**
@@ -42,7 +43,7 @@ abstract contract RuleAddressList is
      */
     function addAddressesToTheList(
         address[] calldata listWhitelistedAddress
-    ) public onlyRole(ADDRESS_LIST_ROLE) {
+    ) public onlyRole(ADDRESS_LIST_ADD_ROLE) {
         _addAddressesToThelist(listWhitelistedAddress);
     }
 
@@ -54,7 +55,7 @@ abstract contract RuleAddressList is
      */
     function removeAddressesFromTheList(
         address[] calldata listWhitelistedAddress
-    ) public onlyRole(ADDRESS_LIST_ROLE) {
+    ) public onlyRole(ADDRESS_LIST_REMOVE_ROLE) {
         _removeAddressesFromThelist(listWhitelistedAddress);
     }
 
@@ -65,7 +66,7 @@ abstract contract RuleAddressList is
      */
     function addAddressToTheList(
         address _newWhitelistAddress
-    ) public onlyRole(ADDRESS_LIST_ROLE) {
+    ) public onlyRole(ADDRESS_LIST_ADD_ROLE) {
         _addAddressToThelist(_newWhitelistAddress);
     }
 
@@ -77,7 +78,7 @@ abstract contract RuleAddressList is
      */
     function removeAddressFromTheList(
         address _removeWhitelistAddress
-    ) public onlyRole(ADDRESS_LIST_ROLE) {
+    ) public onlyRole(ADDRESS_LIST_REMOVE_ROLE) {
         _removeAddressFromThelist(_removeWhitelistAddress);
     }
 

diff --git a/.../validation/abstract/RuleAddressList/invariantStorage/RuleAddressListInvariantStorage.sol b/.../validation/abstract/RuleAddressList/invariantStorage/RuleAddressListInvariantStorage.sol
@@ -6,6 +6,7 @@ abstract contract RuleAddressListInvariantStorage {
     // custom errors
     error RuleAddressList_AdminWithAddressZeroNotAllowed();
 
-    // Role
-    bytes32 public constant ADDRESS_LIST_ROLE = keccak256("ADDRESS_LIST_ROLE");
+    // Add/Remove
+    bytes32 public constant ADDRESS_LIST_REMOVE_ROLE = keccak256("ADDRESS_LIST_REMOVE_ROLE");
+    bytes32 public constant ADDRESS_LIST_ADD_ROLE = keccak256("ADDRESS_LIST_ADD_ROLE");
 }
diff --git a/src/rules/validation/abstract/RuleWhitelistCommon.sol b/src/rules/validation/abstract/RuleWhitelistCommon.sol
@@ -0,0 +1,40 @@
+// SPDX-License-Identifier: MPL-2.0
+
+pragma solidity ^0.8.20;
+
+import "./RuleAddressList/invariantStorage/RuleWhitelistInvariantStorage.sol";
+import "./RuleValidateTransfer.sol";
+
+abstract contract RuleWhitelistCommon is
+    RuleValidateTransfer,
+    RuleWhitelistInvariantStorage {
+     /**
+     * @notice To know if the restriction code is valid for this rule or not
+     * @param _restrictionCode The target restriction code
+     * @return true if the restriction code is known, false otherwise
+     **/
+    function canReturnTransferRestrictionCode(
+        uint8 _restrictionCode
+    ) external pure override returns (bool) {
+        return
+            _restrictionCode == CODE_ADDRESS_FROM_NOT_WHITELISTED ||
+            _restrictionCode == CODE_ADDRESS_TO_NOT_WHITELISTED;
+    }
+
+    /**
+     * @notice Return the corresponding message
+     * @param _restrictionCode The target restriction code
+     * @return true if the transfer is valid, false otherwise
+     **/
+    function messageForTransferRestriction(
+        uint8 _restrictionCode
+    ) external pure override returns (string memory) {
+        if (_restrictionCode == CODE_ADDRESS_FROM_NOT_WHITELISTED) {
+            return TEXT_ADDRESS_FROM_NOT_WHITELISTED;
+        } else if (_restrictionCode == CODE_ADDRESS_TO_NOT_WHITELISTED) {
+            return TEXT_ADDRESS_TO_NOT_WHITELISTED;
+        } else {
+            return TEXT_CODE_NOT_FOUND;
+        }
+    }
+}
diff --git a/test/RuleWhitelist/AccessControl/RuleWhitelistAccessControl.t.sol b/test/RuleWhitelist/AccessControl/RuleWhitelistAccessControl.t.sol
@@ -35,7 +35,7 @@ contract RuleWhitelistAccessControl is Test, HelperContract {
             abi.encodeWithSelector(
                 AccessControlUnauthorizedAccount.selector,
                 ATTACKER,
-                ADDRESS_LIST_ROLE
+                ADDRESS_LIST_ADD_ROLE
             )
         );
         vm.prank(ATTACKER);
@@ -61,7 +61,7 @@ contract RuleWhitelistAccessControl is Test, HelperContract {
             abi.encodeWithSelector(
                 AccessControlUnauthorizedAccount.selector,
                 ATTACKER,
-                ADDRESS_LIST_ROLE
+                ADDRESS_LIST_ADD_ROLE
             )
         );
         vm.prank(ATTACKER);
@@ -98,7 +98,7 @@ contract RuleWhitelistAccessControl is Test, HelperContract {
             abi.encodeWithSelector(
                 AccessControlUnauthorizedAccount.selector,
                 ATTACKER,
-                ADDRESS_LIST_ROLE
+                ADDRESS_LIST_REMOVE_ROLE
             )
         );
         vm.prank(ATTACKER);
@@ -135,7 +135,7 @@ contract RuleWhitelistAccessControl is Test, HelperContract {
             abi.encodeWithSelector(
                 AccessControlUnauthorizedAccount.selector,
                 ATTACKER,
-                ADDRESS_LIST_ROLE
+                ADDRESS_LIST_REMOVE_ROLE
             )
         );
         vm.prank(ATTACKER);

diff --git a/test/RuleWhitelist/AccessControl/RuleWhitelistAccessControlOZ.t.sol b/test/RuleWhitelist/AccessControl/RuleWhitelistAccessControlOZ.t.sol
@@ -34,41 +34,41 @@ contract RuleWhitelistAccessControlOZ is Test, HelperContract, AccessControl {
         vm.prank(WHITELIST_OPERATOR_ADDRESS);
         vm.expectEmit(true, true, false, true);
         emit RoleGranted(
-            ADDRESS_LIST_ROLE,
+            ADDRESS_LIST_ADD_ROLE,
             ADDRESS1,
             WHITELIST_OPERATOR_ADDRESS
         );
-        ruleWhitelist.grantRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        ruleWhitelist.grantRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         // Assert
-        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertEq(res1, true);
     }
 
     function testRevokeRoleAsAdmin() public {
         // Arrange
         vm.prank(WHITELIST_OPERATOR_ADDRESS);
-        ruleWhitelist.grantRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        ruleWhitelist.grantRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         // Arrange - Assert
-        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertEq(res1, true);
 
         // Act
         vm.prank(WHITELIST_OPERATOR_ADDRESS);
         vm.expectEmit(true, true, false, true);
         emit RoleRevoked(
-            ADDRESS_LIST_ROLE,
+            ADDRESS_LIST_ADD_ROLE,
             ADDRESS1,
             WHITELIST_OPERATOR_ADDRESS
         );
-        ruleWhitelist.revokeRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        ruleWhitelist.revokeRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         // Assert
-        bool res2 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res2 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertFalse(res2);
     }
 
     function testCannotGrantFromNonAdmin() public {
         // Arrange - Assert
-        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertFalse(res1);
 
         // Act
@@ -80,22 +80,22 @@ contract RuleWhitelistAccessControlOZ is Test, HelperContract, AccessControl {
             )
         );
         vm.prank(ADDRESS2);
-        ruleWhitelist.grantRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        ruleWhitelist.grantRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         // Assert
-        bool res2 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res2 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertFalse(res2);
     }
 
     function testCannotRevokeFromNonAdmin() public {
         // Arrange - Assert
-        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res1 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertFalse(res1);
 
         // Arrange
         vm.prank(WHITELIST_OPERATOR_ADDRESS);
-        ruleWhitelist.grantRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        ruleWhitelist.grantRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         // Arrange - Assert
-        bool res2 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res2 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertEq(res2, true);
 
         // Act
@@ -107,10 +107,10 @@ contract RuleWhitelistAccessControlOZ is Test, HelperContract, AccessControl {
                 DEFAULT_ADMIN_ROLE
             )
         );
-        ruleWhitelist.revokeRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        ruleWhitelist.revokeRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
 
         // Assert
-        bool res3 = ruleWhitelist.hasRole(ADDRESS_LIST_ROLE, ADDRESS1);
+        bool res3 = ruleWhitelist.hasRole(ADDRESS_LIST_ADD_ROLE, ADDRESS1);
         assertEq(res3, true);
     }
 }
diff --git a/test/RuleWhitelist/RuleWhitelistDeployment.t.sol b/test/RuleWhitelist/RuleWhitelistDeployment.t.sol
@@ -33,7 +33,12 @@ contract RuleWhitelistTest is Test, HelperContract {
 
         // assert
         resBool = ruleWhitelist.hasRole(
-            ADDRESS_LIST_ROLE,
+            ADDRESS_LIST_ADD_ROLE,
+            WHITELIST_OPERATOR_ADDRESS
+        );
+        assertEq(resBool, true);
+        resBool = ruleWhitelist.hasRole(
+            ADDRESS_LIST_REMOVE_ROLE,
             WHITELIST_OPERATOR_ADDRESS
         );
         assertEq(resBool, true);