Adds test for agent execute transact.

Snowfork · Aug 22, 2023 · 2577727 · 2577727
1 parent 6b3b09d
commit 2577727
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 11 deletions.
diff --git a/contracts/src/Gateway.sol b/contracts/src/Gateway.sol
@@ -8,7 +8,7 @@ import {Verification} from "./Verification.sol";
 import {Assets} from "./Assets.sol";
 import {AgentExecutor} from "./AgentExecutor.sol";
 import {Agent} from "./Agent.sol";
-import {Channel, InboundMessage, OperatingMode, ParaID, Config, Command} from "./Types.sol";
+import {Channel, InboundMessage, OperatingMode, ParaID, Config, Command, AgentExecuteCommand} from "./Types.sol";
 import {IGateway} from "./interfaces/IGateway.sol";
 import {IInitializable} from "./interfaces/IInitializable.sol";
 import {ERC1967} from "./utils/ERC1967.sol";
@@ -133,6 +133,9 @@ contract Gateway is IGateway, IInitializable {
         // Otherwise malicious relayers can break the bridge by allowing the message handlers below to run out gas and fail silently.
         // In this scenario case, the channel's state would have been updated to accept the message (by virtue of the nonce increment), yet the actual message
         // dispatch would have failed
+
+        uint256 gas = computeGas(message.command, message.params);
+
         if (gasleft() < DISPATCH_GAS + BUFFER_GAS) {
             revert NotEnoughGas();
         }
@@ -141,37 +144,37 @@ contract Gateway is IGateway, IInitializable {
 
         // Dispatch message to a handler
         if (message.command == Command.AgentExecute) {
-            try Gateway(this).agentExecute{gas: DISPATCH_GAS}(message.params) {}
+            try Gateway(this).agentExecute{gas: gas}(message.params) {}
             catch {
                 success = false;
             }
         } else if (message.command == Command.CreateAgent) {
-            try Gateway(this).createAgent{gas: DISPATCH_GAS}(message.params) {}
+            try Gateway(this).createAgent{gas: gas}(message.params) {}
             catch {
                 success = false;
             }
         } else if (message.command == Command.CreateChannel) {
-            try Gateway(this).createChannel{gas: DISPATCH_GAS}(message.params) {}
+            try Gateway(this).createChannel{gas: gas}(message.params) {}
             catch {
                 success = false;
             }
         } else if (message.command == Command.UpdateChannel) {
-            try Gateway(this).updateChannel{gas: DISPATCH_GAS}(message.params) {}
+            try Gateway(this).updateChannel{gas: gas}(message.params) {}
             catch {
                 success = false;
             }
         } else if (message.command == Command.SetOperatingMode) {
-            try Gateway(this).setOperatingMode{gas: DISPATCH_GAS}(message.params) {}
+            try Gateway(this).setOperatingMode{gas: gas}(message.params) {}
             catch {
                 success = false;
             }
         } else if (message.command == Command.TransferNativeFromAgent) {
-            try Gateway(this).transferNativeFromAgent{gas: DISPATCH_GAS}(message.params) {}
+            try Gateway(this).transferNativeFromAgent{gas: gas}(message.params) {}
             catch {
                 success = false;
             }
         } else if (message.command == Command.Upgrade) {
-            try Gateway(this).upgrade{gas: DISPATCH_GAS}(message.params) {}
+            try Gateway(this).upgrade{gas: gas}(message.params) {}
             catch {
                 success = false;
             }
@@ -464,6 +467,22 @@ contract Gateway is IGateway, IInitializable {
         }
     }
 
+    function computeGas(Command command, bytes memory params) internal view returns (uint256) {
+        if (command != Command.AgentExecute) {
+            return DISPATCH_GAS;
+        }
+
+        (AgentExecuteCommand agentExecuteCommand, bytes memory innerParams) = abi.decode(params, (AgentExecuteCommand, bytes));
+
+        if (agentExecuteCommand != AgentExecuteCommand.Transact) {
+            return DISPATCH_GAS;
+        }
+
+        (, , uint256 dynamicGas) = abi.decode(innerParams, (address, bytes, uint256));
+
+        return DISPATCH_GAS + dynamicGas;
+    }
+
     // Submit an outbound message to Polkadot
     function _submitOutbound(ParaID dest, bytes memory payload, uint256 extraFee) internal {
         Channel storage channel = _ensureChannel(dest);

diff --git a/contracts/test/Gateway.t.sol b/contracts/test/Gateway.t.sol
@@ -25,6 +25,7 @@ import {NativeTransferFailed} from "../src/utils/SafeTransfer.sol";
 import {AgentExecuteCommand, InboundMessage, OperatingMode, ParaID, Config, Command} from "../src/Types.sol";
 
 import {WETH9} from "canonical-weth/WETH9.sol";
+import "./mocks/HelloWorld.sol";
 
 contract GatewayTest is Test {
     event InboundMessageDispatched(ParaID indexed origin, uint64 nonce, bool result);
@@ -37,6 +38,7 @@ contract GatewayTest is Test {
     event AgentCreated(bytes32 agentID, address agent);
     event ChannelCreated(ParaID indexed paraID);
     event ChannelUpdated(ParaID indexed paraID);
+    event SaidHello(string indexed message);
 
     event Upgraded(address indexed implementation);
 
@@ -255,7 +257,7 @@ contract GatewayTest is Test {
      * Handlers
      */
 
-    function testAgentExecution() public {
+    function testAgentExecutionTransferToken() public {
         token.transfer(address(assetHubAgent), 1);
 
         Gateway.AgentExecuteParams memory params = Gateway.AgentExecuteParams({
@@ -266,6 +268,23 @@ contract GatewayTest is Test {
         GatewayMock(address(gateway)).agentExecutePublic(abi.encode(params));
     }
 
+    function testAgentExecutionTransact() public {
+        HelloWorld helloWorld = new HelloWorld();
+
+        bytes memory payload = abi.encodeWithSignature("sayHello(string)", "Clara");
+
+        Gateway.AgentExecuteParams memory params = Gateway.AgentExecuteParams({
+            agentID: assetHubAgentID,
+            payload: abi.encode(AgentExecuteCommand.Transact, abi.encode(address(helloWorld), payload, 100))
+        });
+
+        // Expect the HelloWorld contract to emit `SaidHello`
+        vm.expectEmit({checkTopic1: true, checkTopic2: false, checkTopic3: false, checkData: false});
+        emit SaidHello("Hello there, Clara");
+
+        GatewayMock(address(gateway)).agentExecutePublic(abi.encode(params));
+    }
+
     function testAgentExecutionBadOrigin() public {
         Gateway.AgentExecuteParams memory params = Gateway.AgentExecuteParams({
             agentID: bytes32(0),

diff --git a/contracts/test/mocks/HelloWorld.sol b/contracts/test/mocks/HelloWorld.sol
@@ -3,10 +3,10 @@
 pragma solidity ^0.8.0;
 
 contract HelloWorld {
-    event SaidHello(string message);
+    event SaidHello(string indexed message);
 
     function sayHello(string memory _text) public {
-        string memory fullMessage = string(abi.encodePacked("Hello world ", _text));
+        string memory fullMessage = string(abi.encodePacked("Hello there, ", _text));
         emit SaidHello(fullMessage);
     }
 }