feat(ControllerToken): Adding validation in ControllerToken transfer/…

…from/andCall
monerium · May 16, 2024 · 1c8d364 · 1c8d364
1 parent b5e2473
commit 1c8d364
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 16 deletions.
diff --git a/src/ControllerToken.sol b/src/ControllerToken.sol
@@ -47,7 +47,7 @@ contract ControllerToken is Token {
 
     modifier onlyFrontend() {
         require(
-          isFrontend(msg.sender),
+            isFrontend(msg.sender),
             "ControllerToken: caller is not the frontend"
         );
         _;
@@ -74,6 +74,10 @@ contract ControllerToken is Token {
         address to,
         uint256 amount
     ) external onlyFrontend returns (bool) {
+        require(
+            validator.validate(caller, to, amount),
+            "Transfer not validated"
+        );
         _transfer(caller, to, amount);
         return true;
     }
@@ -84,6 +88,7 @@ contract ControllerToken is Token {
         address to,
         uint256 amount
     ) external onlyFrontend returns (bool) {
+        require(validator.validate(from, to, amount), "Transfer not validated");
         _spendAllowance(from, caller, amount);
         _transfer(from, to, amount);
         return true;
@@ -104,13 +109,17 @@ contract ControllerToken is Token {
         uint256 amount,
         bytes calldata data
     ) external onlyFrontend returns (bool) {
-      _transfer(caller, to, amount);
-      IERC677Recipient recipient = IERC677Recipient(to);
-      require(
-        recipient.onTokenTransfer(caller, amount, data),
-        "token handler returns false"
-      );
-      return true;
+        require(
+            validator.validate(caller, to, amount),
+            "Transfer not validated"
+        );
+        _transfer(caller, to, amount);
+        IERC677Recipient recipient = IERC677Recipient(to);
+        require(
+            recipient.onTokenTransfer(caller, amount, data),
+            "token handler returns false"
+        );
+        return true;
     }
 
     function mintTo_withCaller(
@@ -126,14 +135,14 @@ contract ControllerToken is Token {
 
     function burnFrom_withCaller(
         address caller,
-        address ,//from
-        uint256 ,//amount
-        bytes32 ,//h
-        uint8 ,//v
-        bytes32 ,//r
+        address, //from
+        uint256, //amount
+        bytes32, //h
+        uint8, //v
+        bytes32, //r
         bytes32 //s
     ) external view onlyFrontend onlySystemAccount(caller) returns (bool) {
-      revert("deprecated");
+        revert("deprecated");
     }
 
     function recover_withCaller(
@@ -145,7 +154,7 @@ contract ControllerToken is Token {
         bytes32, //r
         bytes32 //s
     ) external view onlyFrontend onlySystemAccount(caller) returns (uint256) {
-      revert("deprecated");
+        revert("deprecated");
     }
 
     function claimOwnership() external onlyFrontend {

diff --git a/src/Token.sol b/src/Token.sol
@@ -25,7 +25,7 @@ contract Token is
     SystemRoleUpgradeable
 {
     // Subsequent contract versions must retain this variable to avoid storage conflicts with the proxy.
-    IValidator private validator;
+    IValidator internal validator;
     using SignatureChecker for address;
 
     /**

diff --git a/test/ControllerToken.t.sol b/test/ControllerToken.t.sol
@@ -98,6 +98,45 @@ contract ControllerTokenTest is Test {
         assertEq(frontend.balanceOf(user2), 2e18);
     }
 
+    function test_shouldNotTransferIfBlacklisted() public {
+        // Add user2 to blacklist
+        vm.prank(admin);
+        validator.ban(user1);
+        assertTrue(validator.isBan(user1));
+
+        vm.prank(user1);
+        vm.expectRevert("Transfer not validated");
+        frontend.transfer(user2, 1);
+    }
+
+    function test_from_banned_user_should_not_transferFrom() public {
+        // Add user1 to blacklist
+        vm.prank(admin);
+        validator.ban(user1);
+        assertTrue(validator.isBan(user1));
+
+        vm.prank(user1);
+        frontend.approve(user2, 1e18);
+
+        vm.prank(user2);
+        vm.expectRevert("Transfer not validated");
+        frontend.transferFrom(user1, user2, 1);
+    }
+
+    function test_from_banned_user_should_not_transferFrom() public {
+        // Add user1 to blacklist
+        vm.prank(admin);
+        validator.ban(user1);
+        assertTrue(validator.isBan(user1));
+
+        vm.prank(user1);
+        frontend.approve(user2, 1e18);
+
+        vm.prank(user2);
+        vm.expectRevert("Transfer not validated");
+        frontend.transferFrom(user1, user2, 1);
+    }
+
     function testFail_shouldNotTransferIfNotFromFrontend() public {
         vm.prank(user1);
         token.transfer_withCaller(user1, user2, 1e18);