Fee for Gateway.sendToken should cover all costs

contracts/src/AgentExecutor.sol

@@ -17,8 +17,7 @@ contract AgentExecutor {
     /// @dev Execute a message which originated from the Polkadot side of the bridge. In other terms,
     /// the `data` parameter is constructed by the BridgeHub parachain.
     ///
-    /// NOTE: There are no authorization checks here. Since this contract is only used for its code.
-    function execute(address, bytes memory data) external {
+    function execute(bytes memory data) external {
         (AgentExecuteCommand command, bytes memory params) = abi.decode(data, (AgentExecuteCommand, bytes));
         if (command == AgentExecuteCommand.TransferToken) {
             (address token, address recipient, uint128 amount) = abi.decode(params, (address, address, uint128));
@@ -29,7 +28,6 @@ contract AgentExecutor {
     /// @dev Transfer ether to `recipient`. Unlike `_transferToken` This logic is not nested within `execute`,
     /// as the gateway needs to control an agent's ether balance directly.
     ///
-    /// NOTE: There are no authorization checks here. Since this contract is only used for its code.
     function transferNative(address payable recipient, uint256 amount) external {
         recipient.safeNativeTransfer(amount);
     }

contracts/src/Assets.sol

@@ -24,11 +24,11 @@ library Assets {
     error Unsupported();
 
     // This library requires state which must be initialized in the gateway's storage.
-    function initialize(uint256 registerTokenFee, uint256 sendTokenFee) external {
+    function initialize(uint256 _registerTokenFee, uint256 _sendTokenFee) external {
         AssetsStorage.Layout storage $ = AssetsStorage.layout();
 
-        $.registerTokenFee = registerTokenFee;
-        $.sendTokenFee = sendTokenFee;
+        $.registerTokenFee = _registerTokenFee;
+        $.sendTokenFee = _sendTokenFee;
     }
 
     /// @dev transfer tokens from the sender to the specified
@@ -44,6 +44,16 @@ library Assets {
         IERC20(token).safeTransferFrom(sender, assetHubAgent, amount);
     }
 
+    function sendTokenFee(ParaID assetHubParaID, ParaID destinationChain) external view returns (uint256) {
+        AssetsStorage.Layout storage $ = AssetsStorage.layout();
+        if (assetHubParaID == destinationChain) {
+            return $.sendTokenFee;
+        }
+        // If the final destination chain is not AssetHub, then the fee needs to additionally
+        // include the cost of executing an XCM on the final destination parachain.
+        return 2 * $.sendTokenFee;
+    }
+
     function sendToken(
         ParaID assetHubParaID,
         address assetHubAgent,
@@ -61,8 +71,10 @@ library Assets {
             if (destinationAddress.isAddress32()) {
                 payload = SubstrateTypes.SendTokenToAssetHubAddress32(token, destinationAddress.asAddress32(), amount);
             } else {
+                // AssetHub does not support 20-byte account IDs
                 revert Unsupported();
             }
+            extraFee = $.sendTokenFee;
         } else {
             if (destinationAddress.isAddress32()) {
                 payload = SubstrateTypes.SendTokenToAddress32(
@@ -75,12 +87,19 @@ library Assets {
             } else {
                 revert Unsupported();
             }
+            // If the final destination chain is not AssetHub, then the fee needs to additionally
+            // include the cost of executing an XCM on the final destination parachain.
+            extraFee = 2 * $.sendTokenFee;
         }
-        extraFee = $.sendTokenFee;
 
         emit IGateway.TokenSent(sender, token, destinationChain, destinationAddress, amount);
     }
 
+    function registerTokenFee() external view returns (uint256) {
+        AssetsStorage.Layout storage $ = AssetsStorage.layout();
+        return $.registerTokenFee;
+    }
+
     /// @dev Enqueues a create native token message to substrate.
     /// @param token The ERC20 token address.
     function registerToken(address token) external returns (bytes memory payload, uint256 extraFee) {

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, ChannelID, InboundMessage, OperatingMode, ParaID, Command, MultiAddress} from "./Types.sol";
+import {Channel, ChannelID, InboundMessage, OperatingMode, ParaID, Command, MultiAddress, Fee} from "./Types.sol";
 import {IGateway} from "./interfaces/IGateway.sol";
 import {IInitializable} from "./interfaces/IInitializable.sol";
 import {ERC1967} from "./utils/ERC1967.sol";
@@ -238,11 +238,6 @@ contract Gateway is IGateway, IInitializable {
         return ERC1967.load();
     }
 
-    function tokenTransferFees() external view returns (uint256, uint256) {
-        AssetsStorage.Layout storage $ = AssetsStorage.layout();
-        return ($.registerTokenFee, $.sendTokenFee);
-    }
-
     /**
      * Handlers
      */
@@ -262,7 +257,7 @@ contract Gateway is IGateway, IInitializable {
             revert InvalidAgentExecutionPayload();
         }
 
-        bytes memory call = abi.encodeCall(AgentExecutor.execute, (address(this), params.payload));
+        bytes memory call = abi.encodeCall(AgentExecutor.execute, params.payload);
 
         (bool success, bytes memory returndata) = Agent(payable(agent)).invoke(AGENT_EXECUTOR, call);
         if (!success) {
@@ -444,13 +439,25 @@ contract Gateway is IGateway, IInitializable {
      * Assets
      */
 
+    // Total fee for registering a token
+    function registerTokenFee() external view returns (Fee memory) {
+        Channel storage channel = _ensureChannel(ASSET_HUB_PARA_ID.into());
+        return Fee({bridge: channel.outboundFee, xcm: Assets.registerTokenFee()});
+    }
+
     // Register a token on AssetHub
     function registerToken(address token) external payable {
         (bytes memory payload, uint256 extraFee) = Assets.registerToken(token);
 
         _submitOutbound(ASSET_HUB_PARA_ID, payload, extraFee);
     }
 
+    // Total fee for sending a token
+    function sendTokenFee(address, ParaID destinationChain) external view returns (Fee memory) {
+        Channel storage channel = _ensureChannel(ASSET_HUB_PARA_ID.into());
+        return Fee({bridge: channel.outboundFee, xcm: Assets.sendTokenFee(ASSET_HUB_PARA_ID, destinationChain)});
+    }
+
     // Transfer ERC20 tokens to a Polkadot parachain
     function sendToken(address token, ParaID destinationChain, MultiAddress calldata destinationAddress, uint128 amount)
         external

contracts/src/Types.sol

@@ -87,3 +87,15 @@ enum Command {
 }
 
 enum AgentExecuteCommand {TransferToken}
+
+
+using {total} for Fee global;
+
+struct Fee {
+    uint256 bridge;
+    uint256 xcm;
+}
+
+function total(Fee memory fee) pure returns (uint256) {
+    return fee.bridge + fee.xcm;
+}

contracts/src/interfaces/IGateway.sol

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: 2023 Snowfork <[email protected]>
 pragma solidity 0.8.22;
 
-import {OperatingMode, InboundMessage, ParaID, ChannelID, MultiAddress} from "../Types.sol";
+import {OperatingMode, InboundMessage, ParaID, ChannelID, MultiAddress, Fee} from "../Types.sol";
 import {Verification} from "../Verification.sol";
 
 interface IGateway {
@@ -75,13 +75,20 @@ interface IGateway {
     /// @dev Emitted when a command is sent to register a new wrapped token on AssetHub
     event TokenRegistrationSent(address token);
 
-    // @dev Fees in Ether for registering and sending tokens respectively
-    function tokenTransferFees() external view returns (uint256, uint256);
+    /// @dev Fee schedule in Ether for registering a token, covering
+    /// 1. Delivery costs to BridgeHub
+    /// 2. XCM Execution costs on AssetHub
+    function registerTokenFee() external view returns (Fee memory);
 
     /// @dev Send a message to the AssetHub parachain to register a new fungible asset
     ///      in the `ForeignAssets` pallet.
     function registerToken(address token) external payable;
 
+    /// @dev Fees in Ether for sending a token
+    /// 1. Delivery costs to BridgeHub
+    /// 2. XCM execution costs on destinationChain
+    function sendTokenFee(address token, ParaID destinationChain) external view returns (Fee memory);
+
     /// @dev Send ERC20 tokens to parachain `destinationChain` and deposit into account `destinationAddress`
     function sendToken(address token, ParaID destinationChain, MultiAddress calldata destinationAddress, uint128 amount)
         external

contracts/test/Gateway.t.sol

@@ -10,6 +10,7 @@ import {BeefyClient} from "../src/BeefyClient.sol";
 import {IGateway} from "../src/interfaces/IGateway.sol";
 import {IInitializable} from "../src/interfaces/IInitializable.sol";
 import {Gateway} from "../src/Gateway.sol";
+import {Fee} from "../src/Types.sol";
 import {GatewayMock, GatewayV2} from "./mocks/GatewayMock.sol";
 
 import {GatewayProxy} from "../src/GatewayProxy.sol";
@@ -276,14 +277,16 @@ contract GatewayTest is Test {
         address user = makeAddr("user");
         deal(address(token), user, 1);
 
+        Fee memory fee = IGateway(address(gateway)).sendTokenFee(address(token), ParaID.wrap(0));
+
         // Let gateway lock up to 1 tokens
         hoax(user);
         token.approve(address(gateway), 1);
 
-        hoax(user, 2 ether);
-        IGateway(address(gateway)).sendToken{value: 2 ether}(address(token), ParaID.wrap(0), recipientAddress32, 1);
+        hoax(user, fee.total());
+        IGateway(address(gateway)).sendToken{value: fee.total()}(address(token), ParaID.wrap(0), recipientAddress32, 1);
 
-        assertEq(user.balance, 0 ether);
+        assertEq(user.balance, 0);
     }
 
     // User doesn't have enough funds to send message
@@ -581,14 +584,16 @@ contract GatewayTest is Test {
         // Multilocation for recipient
         ParaID destPara = ParaID.wrap(2043);
 
+        Fee memory fee = IGateway(address(gateway)).sendTokenFee(address(token), destPara);
+
         vm.expectEmit(true, true, false, true);
         emit IGateway.TokenSent(address(this), address(token), destPara, recipientAddress32, 1);
 
         // Expect the gateway to emit `OutboundMessageAccepted`
         vm.expectEmit(true, false, false, false);
         emit IGateway.OutboundMessageAccepted(assetHubParaID.into(), 1, messageID, bytes(""));
 
-        IGateway(address(gateway)).sendToken{value: 2 ether}(address(token), destPara, recipientAddress32, 1);
+        IGateway(address(gateway)).sendToken{value: fee.total()}(address(token), destPara, recipientAddress32, 1);
     }
 
     function testSendTokenAddress32ToAssetHub() public {
@@ -598,14 +603,16 @@ contract GatewayTest is Test {
         // Multilocation for recipient
         ParaID destPara = assetHubParaID;
 
+        Fee memory fee = IGateway(address(gateway)).sendTokenFee(address(token), destPara);
+
         vm.expectEmit(true, true, false, true);
         emit IGateway.TokenSent(address(this), address(token), destPara, recipientAddress32, 1);
 
         // Expect the gateway to emit `OutboundMessageAccepted`
         vm.expectEmit(true, false, false, false);
         emit IGateway.OutboundMessageAccepted(assetHubParaID.into(), 1, messageID, bytes(""));
 
-        IGateway(address(gateway)).sendToken{value: 2 ether}(address(token), destPara, recipientAddress32, 1);
+        IGateway(address(gateway)).sendToken{value: fee.total()}(address(token), destPara, recipientAddress32, 1);
     }
 
     function testSendTokenAddress20() public {
@@ -615,14 +622,16 @@ contract GatewayTest is Test {
         // Multilocation for recipient
         ParaID destPara = ParaID.wrap(2043);
 
+        Fee memory fee = IGateway(address(gateway)).sendTokenFee(address(token), destPara);
+
         vm.expectEmit(true, true, false, true);
         emit IGateway.TokenSent(address(this), address(token), destPara, recipientAddress20, 1);
 
         // Expect the gateway to emit `OutboundMessageAccepted`
         vm.expectEmit(true, false, false, false);
         emit IGateway.OutboundMessageAccepted(assetHubParaID.into(), 1, messageID, bytes(""));
 
-        IGateway(address(gateway)).sendToken{value: 2 ether}(address(token), destPara, recipientAddress20, 1);
+        IGateway(address(gateway)).sendToken{value: fee.total()}(address(token), destPara, recipientAddress20, 1);
     }
 
     function testSendTokenAddress20FailsInvalidDestination() public {
@@ -760,15 +769,13 @@ contract GatewayTest is Test {
     }
 
     function testSetTokenFees() public {
-        (uint256 register, uint256 send) = IGateway(address(gateway)).tokenTransferFees();
-        assertEq(register, 1 ether);
-        assertEq(send, 1 ether);
+        Fee memory fee = IGateway(address(gateway)).registerTokenFee();
+        assertEq(fee.xcm, 1 ether);
         GatewayMock(address(gateway)).setTokenTransferFeesPublic(
-            abi.encode(Gateway.SetTokenTransferFeesParams({register: 1, send: 1}))
+            abi.encode(Gateway.SetTokenTransferFeesParams({register: 2, send: 1}))
         );
-        (register, send) = IGateway(address(gateway)).tokenTransferFees();
-        assertEq(register, 1);
-        assertEq(send, 1);
+        fee = IGateway(address(gateway)).registerTokenFee();
+        assertEq(fee.xcm, 2);
     }
 
     bytes32 public expectChannelIDBytes = bytes32(0xc173fac324158e77fb5840738a1a541f633cbec8884c6a601c567d2b376a0539);

contracts/test/mocks/GatewayUpgradeMock.sol

@@ -2,7 +2,7 @@
 // SPDX-FileCopyrightText: 2023 Snowfork <[email protected]>
 pragma solidity 0.8.22;
 
-import {Channel, InboundMessage, OperatingMode, ParaID, Command, ChannelID, MultiAddress} from "../../src/Types.sol";
+import {Channel, InboundMessage, OperatingMode, ParaID, Command, ChannelID, MultiAddress, Fee} from "../../src/Types.sol";
 import {IGateway} from "../../src/interfaces/IGateway.sol";
 import {IInitializable} from "../../src/interfaces/IInitializable.sol";
 import {Verification} from "../../src/Verification.sol";
@@ -32,17 +32,22 @@ contract GatewayUpgradeMock is IGateway, IInitializable {
         return address(0);
     }
 
-    function tokenTransferFees() external pure returns (uint256, uint256) {
-        return (1, 1);
-    }
-
     function implementation() external pure returns (address) {
         return address(0);
     }
 
     function submitInbound(InboundMessage calldata, bytes32[] calldata, Verification.Proof calldata) external {}
 
+    function registerTokenFee() external pure returns (Fee memory) {
+        return Fee(1, 1);
+    }
+
     function registerToken(address) external payable {}
+
+    function sendTokenFee(address, ParaID) external pure returns (Fee memory) {
+        return Fee(1, 1);
+    }
+
     function sendToken(address, ParaID, MultiAddress calldata, uint128) external payable {}
 
     event Initialized(uint256 d0, uint256 d1);