Disable token operation through transact

Snowfork · Apr 24, 2024 · b89cc6c · b89cc6c
1 parent 1fb852d
commit b89cc6c
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 9 deletions.
diff --git a/contracts/src/Gateway.sol b/contracts/src/Gateway.sol
@@ -281,8 +281,11 @@ contract Gateway is IGateway, IInitializable, IUpgradable {
             abi.decode(params.payload, (AgentExecuteCommand, bytes));
         if (command == AgentExecuteCommand.Transact) {
             (address target,,) = abi.decode(commandParams, (address, bytes, uint64));
-            //Todo: Add registered ERC20 to the blacklist
-            if (target == address(this) || target == agent || target == AGENT_EXECUTOR) {
+            // blacklist to check with
+            if (
+                target == address(this) || target == agent || target == AGENT_EXECUTOR
+                    || Assets.isTokenRegistered(target)
+            ) {
                 revert NoPermission();
             }
         }

diff --git a/contracts/test/Gateway.t.sol b/contracts/test/Gateway.t.sol
@@ -52,6 +52,7 @@ import {
 import {WETH9} from "canonical-weth/WETH9.sol";
 import {UD60x18, ud60x18, convert} from "prb/math/src/UD60x18.sol";
 import {HelloWorld} from "./mocks/HelloWorld.sol";
+import {IERC20} from "../src/interfaces/IERC20.sol";
 
 contract GatewayTest is Test {
     ParaID public bridgeHubParaID = ParaID.wrap(1001);
@@ -100,16 +101,12 @@ contract GatewayTest is Test {
     HelloWorld helloWorld;
 
     event SaidHello(string indexed message);
+    event Transfer(address indexed src, address indexed dst, uint256 wad);
 
     function setUp() public {
         AgentExecutor executor = new AgentExecutor();
         gatewayLogic = new MockGateway(
-            address(0),
-            address(executor),
-            bridgeHubParaID,
-            bridgeHubAgentID,
-            foreignTokenDecimals,
-            maxDestinationFee
+            address(0), address(executor), bridgeHubParaID, bridgeHubAgentID, foreignTokenDecimals, maxDestinationFee
         );
         Gateway.Config memory config = Gateway.Config({
             mode: OperatingMode.Normal,
@@ -928,4 +925,27 @@ contract GatewayTest is Test {
 
         MockGateway(address(gateway)).agentExecutePublic(abi.encode(params));
     }
+
+    function testAgentExecutionTransactToTransferERC20() public {
+        testRegisterToken();
+
+        token.transfer(address(assetHubAgent), 200);
+
+        uint256 balanceBefore = IERC20(address(token)).balanceOf(account1);
+
+        uint256 amount = 50;
+        bytes memory payload = abi.encodeCall(IERC20.transfer, (account1, amount));
+
+        AgentExecuteParams memory params = AgentExecuteParams({
+            agentID: assetHubAgentID,
+            payload: abi.encode(AgentExecuteCommand.Transact, abi.encode(address(token), payload, 100000))
+        });
+
+        vm.expectRevert(Gateway.NoPermission.selector);
+
+        MockGateway(address(gateway)).agentExecutePublic(abi.encode(params));
+
+        uint256 balanceAfter = IERC20(address(token)).balanceOf(account1);
+        assertEq(balanceBefore, balanceAfter);
+    }
 }
diff --git a/contracts/test/mocks/MockGatewayV2.sol b/contracts/test/mocks/MockGatewayV2.sol
@@ -19,7 +19,7 @@ library AdditionalStorage {
 }
 
 // Used to test upgrades.
-contract MockGatewayV2 is IInitializable  {
+contract MockGatewayV2 is IInitializable {
     // Reinitialize gateway with some additional storage fields
     function initialize(bytes memory params) external {
         AdditionalStorage.Layout storage $ = AdditionalStorage.layout();