Raunak/dispatcher agnostic fee implementation

contracts/base/GeneralMiddleware.sol

@@ -45,9 +45,9 @@
      * one MW stack.
      * - be 1 << N, where N is a non-negative integer, and not in conflict with other MWs.
      */
     uint256 public MW_ID;
 
     event UCHPacketSent(address source, bytes32 destination);
     /**
      * @param _middleware The middleware contract address this contract sends packets to and receives packets from.
      */
@@ -56,28 +56,37 @@
         MW_ID = mwId;
     }
 
     function sendUniversalPacket(
         bytes32 channelId,
         bytes32 destPortAddr,
         bytes calldata appData,
         uint64 timeoutTimestamp
-    ) external override {
+    ) external override returns (uint64 sequence) {
         emit UCHPacketSent(msg.sender, destPortAddr);
-        _sendPacket(channelId, IbcUtils.toBytes32(msg.sender), destPortAddr, 0, appData, timeoutTimestamp);
+        return _sendPacket(channelId, IbcUtils.toBytes32(msg.sender), destPortAddr, 0, appData, timeoutTimestamp);
     }
 
+    function sendUniversalPacketWithFee(
+        bytes32 channelId,
+        bytes32 destPortAddr,
+        bytes calldata appData,
+        uint64 timeoutTimestamp,
+        uint256[2] calldata gasLimits,
+        uint256[2] calldata gasPrices
+    ) external payable override returns (uint64 sequence) {}
+
     function sendMWPacket(
         bytes32 channelId,
         bytes32 srcPortAddr,
         bytes32 destPortAddr,
         uint256 srcMwIds,
         bytes calldata appData,
         uint64 timeoutTimestamp
-    ) external override {
-        _sendPacket(channelId, srcPortAddr, destPortAddr, srcMwIds, appData, timeoutTimestamp);
+    ) external override returns (uint64 sequence) {
+        return _sendPacket(channelId, srcPortAddr, destPortAddr, srcMwIds, appData, timeoutTimestamp);
     }
 
     function onRecvMWPacket(
         bytes32 channelId,
         UniversalPacket calldata ucPacket,
         // address srcPortAddr,
@@ -191,7 +200,7 @@
         uint256 srcMwIds,
         bytes calldata appData,
         uint64 timeoutTimestamp
-    ) internal virtual {
+    ) internal virtual returns (uint64 sequence) {
         // extra MW custom logic here to process packet, eg. emit MW events, mutate state, etc.
         // implementer can emit custom data fields suitable for their use cases.
         // Here we use MW_ID as the custom MW data field.
@@ -200,7 +209,7 @@
         );
 
         // send packet to next MW
-        IbcMwPacketSender(mw).sendMWPacket(
+        return IbcMwPacketSender(mw).sendMWPacket(
             channelId, srcPortAddr, destPortAddr, srcMwIds | MW_ID, appData, timeoutTimestamp
         );
     }

contracts/core/Dispatcher.sol

@@ -28,6 +28,7 @@ import {ReentrancyGuard} from "@openzeppelin/contracts/security/ReentrancyGuard.
 import {Channel, ChannelEnd, ChannelOrder, IbcPacket, ChannelState, AckPacket, Ibc} from "../libs/Ibc.sol";
 import {IBCErrors} from "../libs/IbcErrors.sol";
 import {IbcUtils} from "../libs/IbcUtils.sol";
+import {IFeeVault} from "../interfaces/IFeeVault.sol";
 
 /**
  * @title Dispatcher
@@ -64,6 +65,7 @@ contract Dispatcher is OwnableUpgradeable, UUPSUpgradeable, ReentrancyGuard, IDi
     ILightClient _UNUSED; // From previous dispatcher version
     mapping(bytes32 => string) private _channelIdToConnection;
     mapping(string => ILightClient) private _connectionToLightClient;
+    IFeeVault public feeVault;
 
     constructor() {
         _disableInitializers();
@@ -75,13 +77,17 @@ contract Dispatcher is OwnableUpgradeable, UUPSUpgradeable, ReentrancyGuard, IDi
      * @dev This method should be called only once during contract deployment.
      * @dev For contract upgarades, which need to reinitialize the contract, use the reinitializer modifier.
      */
-    function initialize(string memory initPortPrefix) public virtual initializer nonReentrant {
+    function initialize(string memory initPortPrefix, IFeeVault _feeVault) public virtual initializer nonReentrant {
         if (bytes(initPortPrefix).length == 0) {
             revert IBCErrors.invalidPortPrefix();
         }
+        if (address(_feeVault) == address(0)) {
+            revert IBCErrors.invalidAddress();
+        }
         __Ownable_init();
         portPrefix = initPortPrefix;
         portPrefixLen = uint32(bytes(initPortPrefix).length);
+        feeVault = _feeVault;
     }
 
     /**
@@ -473,8 +479,14 @@ contract Dispatcher is OwnableUpgradeable, UUPSUpgradeable, ReentrancyGuard, IDi
      * @param packet The packet data to send.
      * @param timeoutTimestamp The timestamp, in seconds after the unix epoch, after which the packet times out if it
      * has not been received.
+     * @return sequence The sequence number of the packet, starting from 1 and incremented with each sendPacket. This
+     * sequence number is used to link with depositing packet fees in the fee vault
      */
-    function sendPacket(bytes32 channelId, bytes calldata packet, uint64 timeoutTimestamp) external nonReentrant {
+    function sendPacket(bytes32 channelId, bytes calldata packet, uint64 timeoutTimestamp)
+        external
+        nonReentrant
+        returns (uint64 sequence)
+    {
         // ensure port owns channel
         if (_portChannelMap[msg.sender][channelId].counterpartyChannelId == bytes32(0)) {
             revert IBCErrors.channelNotOwnedBySender();
@@ -484,7 +496,7 @@ contract Dispatcher is OwnableUpgradeable, UUPSUpgradeable, ReentrancyGuard, IDi
         }
 
         // current packet sequence
-        uint64 sequence = _nextSequenceSend[msg.sender][channelId];
+        sequence = _nextSequenceSend[msg.sender][channelId];
         if (sequence == 0) {
             revert IBCErrors.invalidPacketSequence();
         }

contracts/core/FeeVault.sol

@@ -0,0 +1,91 @@
+// SPDX-License-Identifier: Apache-2.0
+/*
+ * Copyright 2024, Polymer Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+pragma solidity 0.8.15;
+
+import {Ownable} from "@openzeppelin/contracts/access/Ownable.sol";
+import {IFeeVault} from "../interfaces/IFeeVault.sol";
+import {ReentrancyGuard} from "@openzeppelin/contracts/security/ReentrancyGuard.sol";
+import {ChannelOrder} from "../libs/Ibc.sol";
+
+contract FeeVault is Ownable, ReentrancyGuard, IFeeVault {
+    /**
+     * @notice Deposits the send packet fee for a given channel and sequence that is used for relaying recieve and
+     * acknowledge steps of a packet handhsake after a dapp has called the sendPacket on dispatcher.
+     * @dev This function calculates the required fee based on provided gas limits and gas prices,
+     *      and reverts if the sent value does not match the calculated fee.
+     *      The first entry in `gasLimits` and `gasPrices` arrays corresponds to `recvPacket` fees,
+     *      and the second entry corresponds to `ackPacket` fees.
+     * @param channelId The identifier of the channel.
+     * @param sequence The sequence number for the packet, returned from the dispatcher sendPacket call.
+     * @param gasLimits An array containing two gas limit values:
+     *                  - gasLimits[0] for `recvPacket` fees
+     *                  - gasLimits[1] for `ackPacket` fees.
+     * @param gasPrices An array containing two gas price values:
+     *                  - gasPrices[0] for `recvPacket` fees, for the dest chain
+     *                  - gasPrices[1] for `ackPacket` fees, for the src chain
+     */
+    function depositSendPacketFee(
+        bytes32 channelId,
+        uint64 sequence,
+        uint256[2] calldata gasLimits,
+        uint256[2] calldata gasPrices
+    ) external payable nonReentrant {
+        uint256 fee = gasLimits[0] * gasPrices[0] + gasLimits[1] * gasPrices[1];
+        if ((fee) != msg.value) {
+            revert IncorrectFeeSent(fee, msg.value);
+        }
+        emit SendPacketFeeDeposited(channelId, sequence, gasLimits, gasPrices);
+    }
+
+    /**
+     * @notice Deposits the fee for a channel handshake, to pay a relayer for relaying the channelOpenTry,
+     * channelOpenConfirm, and channelOpenAck steps after a dapp has called channelOpenInit
+     * @dev The fee amount that needs to be sent for Polymer to relay the whole channel handshake can be queried on the
+     * web2 layer.
+     * @param src The address of the sender, should be the address in the localportId.
+     * @param version The version string of the channel, the same argument as that sent in the
+     * dispatcher.channelOpenInit call
+     * @param ordering The ordering of the channel, the same argument as that sent in the dispatcher.channelOpenInit
+     * call
+     * @param connectionHops An array of connection hops, the same argument as that sent in the
+     * dispatcher.channelOpenInit call
+     * @param counterpartyPortId The counterparty port identifier, the same argument as that sent in the
+     * dispatcher.channelOpenInit call
+     */
+    function depositOpenChannelFee(
+        address src,
+        string memory version,
+        ChannelOrder ordering,
+        string[] calldata connectionHops,
+        string calldata counterpartyPortId
+    ) external payable nonReentrant {
+        if (msg.value == 0) {
+            revert NoFeeSent();
+        }
+        emit OpenChannelFeeDeposited(src, version, ordering, connectionHops, counterpartyPortId, msg.value);
+    }
+
+    /**
+     * @notice Withdraws all collected fees to the contract owner's address.
+     * @dev Transfers the entire balance of this contract to the owner.
+     * @dev Anyone can call this, but it will always only be sent to the owner.
+     */
+    function withdrawFeesToOwner() external {
+        payable(owner()).transfer(address(this).balance);
+    }
+}

contracts/core/UniversalChannelHandler.sol

@@ -19,10 +19,11 @@ pragma solidity 0.8.15;
 
 import {IbcDispatcher} from "../interfaces/IbcDispatcher.sol";
 import {IbcUniversalChannelMW, IbcUniversalPacketReceiver} from "../interfaces/IbcMiddleware.sol";
-import {IbcReceiverBaseUpgradeable} from "../interfaces/IbcReceiverUpgradeable.sol";
+import {IbcReceiverBaseUpgradeable} from "../implementation_templates/IbcReceiverUpgradeable.sol";
 import {ChannelOrder, IbcPacket, AckPacket, UniversalPacket} from "../libs/Ibc.sol";
 import {IbcUtils} from "../libs/IbcUtils.sol";
 import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
+import {FeeSender} from "../implementation_templates/FeeSender.sol";
 
 /**
  * @title Universal Channel Handler
@@ -31,7 +32,7 @@ import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeab
  * channel handshake to establish a channel.
  * @dev This contract can integrate directly with dapps, or a middleware stack for packet routing.
  */
-contract UniversalChannelHandler is IbcReceiverBaseUpgradeable, UUPSUpgradeable, IbcUniversalChannelMW {
+contract UniversalChannelHandler is IbcReceiverBaseUpgradeable, FeeSender, UUPSUpgradeable, IbcUniversalChannelMW {
     bytes32 private _UNUSED; // Storage placeholder to ensure upgrade from this version is backwards compatible
 
     string public constant VERSION = "1.0";
@@ -82,12 +83,38 @@ contract UniversalChannelHandler is IbcReceiverBaseUpgradeable, UUPSUpgradeable,
         bytes32 destPortAddr,
         bytes calldata appData,
         uint64 timeoutTimestamp
-    ) external {
+    ) external returns (uint64 sequence) {
         bytes memory packetData = IbcUtils.toUniversalPacketBytes(
             UniversalPacket(IbcUtils.toBytes32(msg.sender), MW_ID, destPortAddr, appData)
         );
         emit UCHPacketSent(msg.sender, destPortAddr);
-        dispatcher.sendPacket(channelId, packetData, timeoutTimestamp);
+        sequence = dispatcher.sendPacket(channelId, packetData, timeoutTimestamp);
+    }
+
+    /**
+     * @notice Sends a universal packet over an IBC channel
+     * @param channelId The channel ID through which the packet is sent on the dispatcher
+     * @param destPortAddr The destination port address
+     * @param appData The packet data to be sent
+     * @param timeoutTimestamp of when the packet can timeout
+     */
+    function sendUniversalPacketWithFee(
+        bytes32 channelId,
+        bytes32 destPortAddr,
+        bytes calldata appData,
+        uint64 timeoutTimestamp,
+        uint256[2] calldata gasLimits,
+        uint256[2] calldata gasPrices
+    ) external payable returns (uint64 sequence) {
+        // Cache dispatcher for gas savings
+        IbcDispatcher _dispatcher = dispatcher;
+
+        bytes memory packetData = IbcUtils.toUniversalPacketBytes(
+            UniversalPacket(IbcUtils.toBytes32(msg.sender), MW_ID, destPortAddr, appData)
+        );
+        emit UCHPacketSent(msg.sender, destPortAddr);
+        sequence = _dispatcher.sendPacket(channelId, packetData, timeoutTimestamp);
+        _depositSendPacketFee(dispatcher, channelId, sequence, gasLimits, gasPrices);
     }
 
     /**

contracts/examples/Earth.sol

@@ -20,6 +20,7 @@ pragma solidity ^0.8.9;
 import {UniversalPacket, AckPacket} from "../libs/Ibc.sol";
 import {IbcUtils} from "../libs/IbcUtils.sol";
 import {IbcUniversalPacketReceiverBase, IbcUniversalPacketSender} from "../interfaces/IbcMiddleware.sol";
+import {IUniversalChannelHandler} from "../interfaces/IUniversalChannelHandler.sol";
 
 /**
  * @title Earth
@@ -48,12 +49,33 @@ contract Earth is IbcUniversalPacketReceiverBase {
 
     constructor(address _middleware) IbcUniversalPacketReceiverBase(_middleware) {}
 
+    /**
+     * @notice Send a packet to a destination chain. without a fee
+     * @notice this is useful for self-relaying apckets which don't rely on polymer to fund.
+     * @param destPortAddr The destination chain's port address.
+     * @param channelId The channel id to send the packet on.
+     * @param message The message to send.
+     * @param timeoutTimestamp The timeout timestamp for the packet.
+     */
     function greet(address destPortAddr, bytes32 channelId, bytes calldata message, uint64 timeoutTimestamp) external {
         IbcUniversalPacketSender(mw).sendUniversalPacket(
             channelId, IbcUtils.toBytes32(destPortAddr), message, timeoutTimestamp
         );
     }
 
+    function greetWithFee(
+        address destPortAddr,
+        bytes32 channelId,
+        bytes calldata message,
+        uint64 timeoutTimestamp,
+        uint256[2] memory gasLimits,
+        uint256[2] memory gasPrices
+    ) external payable returns (uint64 sequence) {
+        return IUniversalChannelHandler(mw).sendUniversalPacketWithFee{value: msg.value}(
+            channelId, IbcUtils.toBytes32(destPortAddr), message, timeoutTimestamp, gasLimits, gasPrices
+        );
+    }
+
     function onRecvUniversalPacket(bytes32 channelId, UniversalPacket calldata packet)
         external
         onlyIbcMw

contracts/examples/Mars.sol

@@ -20,14 +20,15 @@ pragma solidity ^0.8.9;
 import {AckPacket, ChannelOrder} from "../libs/Ibc.sol";
 import {IbcReceiverBase, IbcReceiver, IbcPacket} from "../interfaces/IbcReceiver.sol";
 import {IbcDispatcher} from "../interfaces/IbcDispatcher.sol";
+import {FeeSender} from "../implementation_templates/FeeSender.sol";
 
 /**
  * @title Mars
  * @notice Mars is a simple IBC receiver contract that receives packets and sends acks.
  * @dev This contract is used for only testing IBC functionality and as an example for dapp developers on how to
  * integrate with the vibc protocol.
  */
-contract Mars is IbcReceiverBase, IbcReceiver {
+contract Mars is IbcReceiverBase, IbcReceiver, FeeSender {
     // received packet as chain B
     IbcPacket[] public recvedPackets;
     // received ack packet as chain A
@@ -50,6 +51,18 @@ contract Mars is IbcReceiverBase, IbcReceiver {
         dispatcher.channelOpenInit(version, ordering, feeEnabled, connectionHops, counterpartyPortId);
     }
 
+    function triggerChannelInitWithFee(
+        string calldata version,
+        ChannelOrder ordering,
+        bool feeEnabled,
+        string[] calldata connectionHops,
+        string calldata counterpartyPortId
+    ) external payable onlyOwner {
+        IbcDispatcher _dispatcher = dispatcher; // cache for gas savings to avoid 2 SLOADS
+        _dispatcher.channelOpenInit(version, ordering, feeEnabled, connectionHops, counterpartyPortId);
+        _depositOpenChannelFee(_dispatcher, version, ordering, connectionHops, counterpartyPortId);
+    }
+
     function onRecvPacket(IbcPacket memory packet)
         external
         virtual
@@ -92,13 +105,36 @@ contract Mars is IbcReceiverBase, IbcReceiver {
     }
 
     /**
-     * @dev Sends a packet with a greeting message over a specified channel.
+     * @dev Sends a packet with a greeting message over a specified channel, without depositing any sendPacket relaying
+     * fees.
+     * @notice Use greetWithFee for sending packets with fees.
      * @param message The greeting message to be sent.
      * @param channelId The ID of the channel to send the packet to.
      * @param timeoutTimestamp The timestamp at which the packet will expire if not received.
+     * @dev This method also returns sequence from the dispatcher for easy testing
      */
-    function greet(string calldata message, bytes32 channelId, uint64 timeoutTimestamp) external {
-        dispatcher.sendPacket(channelId, bytes(message), timeoutTimestamp);
+    function greet(string calldata message, bytes32 channelId, uint64 timeoutTimestamp)
+        external
+        returns (uint64 sequence)
+    {
+        sequence = dispatcher.sendPacket(channelId, bytes(message), timeoutTimestamp);
+    }
+
+    /**
+     * @dev Sends a packet with a greeting message over a specified channel, and deposits a fee for relaying the packet
+     * @param message The greeting message to be sent.
+     * @param channelId The ID of the channel to send the packet to.
+     * @param timeoutTimestamp The timestamp at which the packet will expire if not received.
+     */
+    function greetWithFee(
+        string calldata message,
+        bytes32 channelId,
+        uint64 timeoutTimestamp,
+        uint256[2] calldata gasLimits,
+        uint256[2] calldata gasPrices
+    ) external payable returns (uint64 sequence) {
+        sequence = dispatcher.sendPacket(channelId, bytes(message), timeoutTimestamp);
+        _depositSendPacketFee(dispatcher, channelId, sequence, gasLimits, gasPrices);
     }
 
     function onChanOpenInit(ChannelOrder, string[] calldata, string calldata, string calldata version)