Refactor token source contracts

contracts/src/CrossChainApplications/examples/NativeTokenBridge/ERC20TokenSource.sol

@@ -5,16 +5,13 @@
 
 pragma solidity 0.8.18;
 
-import {ReentrancyGuard} from "@openzeppelin/contracts/security/ReentrancyGuard.sol";
-import {IWarpMessenger} from "@subnet-evm-contracts/interfaces/IWarpMessenger.sol";
 import {IERC20TokenSource} from "./IERC20TokenSource.sol";
-import {ITokenSource} from "./ITokenSource.sol";
+import {TokenSource} from "./TokenSource.sol";
 import {
     ITeleporterMessenger,
     TeleporterMessageInput,
     TeleporterFeeInfo
 } from "@teleporter/ITeleporterMessenger.sol";
-import {TeleporterOwnerUpgradeable} from "@teleporter/upgrades/TeleporterOwnerUpgradeable.sol";
 import {SafeERC20TransferFrom} from "@teleporter/SafeERC20TransferFrom.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
@@ -24,46 +21,22 @@ import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol
  * DO NOT USE THIS CODE IN PRODUCTION.
  */
 
-contract ERC20TokenSource is
-    TeleporterOwnerUpgradeable,
-    IERC20TokenSource,
-    ITokenSource,
-    ReentrancyGuard
-{
-    // Designated Blackhole Address for this contract. Tokens are sent here to be "burned" when
-    // a SourceAction.Burn message is received from the destination chain.
-    address public constant BURN_ADDRESS = 0x0100000000000000000000000000000000010203;
-    uint256 public constant MINT_NATIVE_TOKENS_REQUIRED_GAS = 100_000;
-    // Used to keep track of tokens burned through transactions on the destination chain. They can
-    // be reported to this contract to burn an equivalent number of tokens on this chain.
-    uint256 public destinationBurnedTotal;
-    bytes32 public immutable destinationBlockchainID;
-    address public immutable nativeTokenDestinationAddress;
+/**
+ * @dev Implementation of the {TokenSource} abstract contract.
+ *
+ * This contracts implements {TokenSource} and uses a specified ERC20 token as the currency.
+ */
+contract ERC20TokenSource is IERC20TokenSource, TokenSource {
     address public immutable erc20ContractAddress;
 
     constructor(
         address teleporterRegistryAddress,
         bytes32 destinationBlockchainID_,
         address nativeTokenDestinationAddress_,
         address erc20ContractAddress_
-    ) TeleporterOwnerUpgradeable(teleporterRegistryAddress) {
-        require(
-            destinationBlockchainID_ != bytes32(0),
-            "ERC20TokenSource: zero destination blockchain ID"
-        );
-        require(
-            destinationBlockchainID_
-                != IWarpMessenger(0x0200000000000000000000000000000000000005).getBlockchainID(),
-            "ERC20TokenSource: cannot bridge with same blockchain"
-        );
-        destinationBlockchainID = destinationBlockchainID_;
-
-        require(
-            nativeTokenDestinationAddress_ != address(0),
-            "ERC20TokenSource: zero destination contract address"
-        );
-        nativeTokenDestinationAddress = nativeTokenDestinationAddress_;
-
+    )
+        TokenSource(teleporterRegistryAddress, destinationBlockchainID_, nativeTokenDestinationAddress_)
+    {
         require(
             erc20ContractAddress_ != address(0), "ERC20TokenSource: zero ERC20 contract address"
         );
@@ -122,46 +95,9 @@ contract ERC20TokenSource is
     }
 
     /**
-     * @dev See {TeleporterUpgradeable-receiveTeleporterMessage}.
-     *
-     * Receives a Teleporter message and routes to the appropriate internal function call.
+     * @dev See {TokenSource-_unlockTokens}
      */
-    function _receiveTeleporterMessage(
-        bytes32 senderBlockchainID,
-        address originSenderAddress,
-        bytes memory message
-    ) internal override {
-        // Only allow messages from the destination chain.
-        require(
-            senderBlockchainID == destinationBlockchainID,
-            "ERC20TokenSource: invalid destination chain"
-        );
-
-        // Only allow the partner contract to send messages.
-        require(
-            originSenderAddress == nativeTokenDestinationAddress,
-            "ERC20TokenSource: unauthorized sender"
-        );
-
-        // Decode the payload to recover the action and corresponding function parameters
-        (SourceAction action, bytes memory actionData) = abi.decode(message, (SourceAction, bytes));
-
-        // Route to the appropriate function.
-        if (action == SourceAction.Unlock) {
-            (address recipient, uint256 amount) = abi.decode(actionData, (address, uint256));
-            _unlockTokens(recipient, amount);
-        } else if (action == SourceAction.Burn) {
-            uint256 newBurnTotal = abi.decode(actionData, (uint256));
-            _handleBurnTokens(newBurnTotal);
-        } else {
-            revert("ERC20TokenSource: invalid action");
-        }
-    }
-
-    /**
-     * @dev Unlocks tokens to recipient.
-     */
-    function _unlockTokens(address recipient, uint256 amount) private {
+    function _unlockTokens(address recipient, uint256 amount) internal override {
         require(recipient != address(0), "ERC20TokenSource: zero recipient address");
 
         // Transfer to recipient
@@ -170,21 +106,21 @@ contract ERC20TokenSource is
     }
 
     /**
-     * @dev Sends tokens to BURNED_TX_FEES_ADDRESS.
+     * @dev See {TokenSource-_handleBurnTokens}
      */
-    function _burnTokens(uint256 amount) private {
-        emit BurnTokens(amount);
-        SafeERC20.safeTransfer(IERC20(erc20ContractAddress), BURN_ADDRESS, amount);
-    }
-
-    /**
-     * @dev Update destinationBurnedTotal sent from destination chain
-     */
-    function _handleBurnTokens(uint256 newBurnTotal) private {
+    function _handleBurnTokens(uint256 newBurnTotal) internal override {
         if (newBurnTotal > destinationBurnedTotal) {
             uint256 difference = newBurnTotal - destinationBurnedTotal;
             _burnTokens(difference);
             destinationBurnedTotal = newBurnTotal;
         }
     }
+
+    /**
+     * @dev See {TokenSource-_burnTokens}
+     */
+    function _burnTokens(uint256 amount) private {
+        emit BurnTokens(amount);
+        SafeERC20.safeTransfer(IERC20(erc20ContractAddress), BURN_ADDRESS, amount);
+    }
 }

contracts/src/CrossChainApplications/examples/NativeTokenBridge/NativeTokenDestination.sol

@@ -29,6 +29,12 @@ import {IAllowList} from "@subnet-evm-contracts/interfaces/IAllowList.sol";
  * DO NOT USE THIS CODE IN PRODUCTION.
  */
 
+/**
+ * @dev Implementation of the {INativeTokenDestination} interface.
+ *
+ * This contract pairs with exactly one `TokenSource` contract on the source chain.
+ * It mints and burns native tokens on the destination chain corresponding to locks and unlocks on the source chain.
+ */
 contract NativeTokenDestination is
     TeleporterOwnerUpgradeable,
     INativeTokenDestination,
@@ -188,13 +194,14 @@ contract NativeTokenDestination is
      * Receives a Teleporter message.
      */
     function _receiveTeleporterMessage(
-        bytes32 senderBlockchainID,
+        bytes32 sourceBlockchainID_,
         address originSenderAddress,
         bytes memory message
     ) internal override {
         // Only allow messages from the source chain.
         require(
-            senderBlockchainID == sourceBlockchainID, "NativeTokenDestination: invalid source chain"
+            sourceBlockchainID_ == sourceBlockchainID,
+            "NativeTokenDestination: invalid source chain"
         );
 
         // Only allow the partner contract to send messages.

contracts/src/CrossChainApplications/examples/NativeTokenBridge/NativeTokenSource.sol

@@ -5,17 +5,14 @@
 
 pragma solidity 0.8.18;
 
-import {ReentrancyGuard} from "@openzeppelin/contracts/security/ReentrancyGuard.sol";
 import {Address} from "@openzeppelin/contracts/utils/Address.sol";
-import {IWarpMessenger} from "@subnet-evm-contracts/interfaces/IWarpMessenger.sol";
 import {INativeTokenSource} from "./INativeTokenSource.sol";
-import {ITokenSource} from "./ITokenSource.sol";
+import {TokenSource} from "./TokenSource.sol";
 import {
     ITeleporterMessenger,
     TeleporterFeeInfo,
     TeleporterMessageInput
 } from "@teleporter/ITeleporterMessenger.sol";
-import {TeleporterOwnerUpgradeable} from "@teleporter/upgrades/TeleporterOwnerUpgradeable.sol";
 import {SafeERC20TransferFrom} from "@teleporter/SafeERC20TransferFrom.sol";
 import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
@@ -25,44 +22,19 @@ import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol
  * DO NOT USE THIS CODE IN PRODUCTION.
  */
 
-contract NativeTokenSource is
-    TeleporterOwnerUpgradeable,
-    INativeTokenSource,
-    ITokenSource,
-    ReentrancyGuard
-{
-    // Designated Blackhole Address for this contract. Tokens are sent here to be "burned" when
-    // a SourceAction.Burn message is received from the destination chain.
-    address public constant BURN_ADDRESS = 0x0100000000000000000000000000000000010203;
-    uint256 public constant MINT_NATIVE_TOKENS_REQUIRED_GAS = 100_000;
-    // Used to keep track of tokens burned through transactions on the destination chain. They can
-    // be reported to this contract to burn an equivalent number of tokens on this chain.
-    uint256 public destinationBurnedTotal;
-    bytes32 public immutable destinationBlockchainID;
-    address public immutable nativeTokenDestinationAddress;
-
+/**
+ * @dev Implementation of the {TokenSource} abstract contract.
+ *
+ * This contracts implements {TokenSource} and uses native tokens as the currency.
+ */
+contract NativeTokenSource is INativeTokenSource, TokenSource {
     constructor(
         address teleporterRegistryAddress,
         bytes32 destinationBlockchainID_,
         address nativeTokenDestinationAddress_
-    ) TeleporterOwnerUpgradeable(teleporterRegistryAddress) {
-        require(
-            destinationBlockchainID_ != bytes32(0),
-            "NativeTokenSource: zero destination blockchain ID"
-        );
-        require(
-            destinationBlockchainID_
-                != IWarpMessenger(0x0200000000000000000000000000000000000005).getBlockchainID(),
-            "NativeTokenSource: cannot bridge with same blockchain"
-        );
-        destinationBlockchainID = destinationBlockchainID_;
-
-        require(
-            nativeTokenDestinationAddress_ != address(0),
-            "NativeTokenSource: zero destination contract address"
-        );
-        nativeTokenDestinationAddress = nativeTokenDestinationAddress_;
-    }
+    )
+        TokenSource(teleporterRegistryAddress, destinationBlockchainID_, nativeTokenDestinationAddress_)
+    {}
 
     /**
      * @dev See {INativeTokenSource-transferToDestination}.
@@ -110,46 +82,9 @@ contract NativeTokenSource is
     }
 
     /**
-     * @dev See {TeleporterUpgradeable-receiveTeleporterMessage}.
-     *
-     * Receives a Teleporter message and routes to the appropriate internal function call.
-     */
-    function _receiveTeleporterMessage(
-        bytes32 senderBlockchainID,
-        address originSenderAddress,
-        bytes memory message
-    ) internal override {
-        // Only allow messages from the destination chain.
-        require(
-            senderBlockchainID == destinationBlockchainID,
-            "NativeTokenSource: invalid destination chain"
-        );
-
-        // Only allow the partner contract to send messages.
-        require(
-            originSenderAddress == nativeTokenDestinationAddress,
-            "NativeTokenSource: unauthorized sender"
-        );
-
-        // Decode the payload to recover the action and corresponding function parameters
-        (SourceAction action, bytes memory actionData) = abi.decode(message, (SourceAction, bytes));
-
-        // Route to the appropriate function.
-        if (action == SourceAction.Unlock) {
-            (address recipient, uint256 amount) = abi.decode(actionData, (address, uint256));
-            _unlockTokens(recipient, amount);
-        } else if (action == SourceAction.Burn) {
-            uint256 newBurnTotal = abi.decode(actionData, (uint256));
-            _handleBurnTokens(newBurnTotal);
-        } else {
-            revert("NativeTokenSource: invalid action");
-        }
-    }
-
-    /**
-     * @dev Unlocks tokens to recipient.
+     * @dev See {TokenSource-_unlockTokens}
      */
-    function _unlockTokens(address recipient, uint256 amount) private {
+    function _unlockTokens(address recipient, uint256 amount) internal override {
         require(recipient != address(0), "NativeTokenSource: zero recipient address");
         require(address(this).balance >= amount, "NativeTokenSource: insufficient collateral");
 
@@ -159,24 +94,21 @@ contract NativeTokenSource is
     }
 
     /**
-     * @dev Sends tokens to BURNED_TX_FEES_ADDRESS.
-     */
-    function _burnTokens(uint256 amount) private {
-        emit BurnTokens(amount);
-        Address.sendValue(payable(BURN_ADDRESS), amount);
-    }
-
-    /**
-     * @dev Update destinationBurnedTotal sent from destination chain
-     * If the new burned total is less than the highest known burned total, this transaction is a no-op.
-     * The burned total on the destination will only ever increase, but new totals may be relayed to this
-     * chain out of order.
+     * @dev See {TokenSource-_handleBurnTokens}
      */
-    function _handleBurnTokens(uint256 newBurnTotal) private {
+    function _handleBurnTokens(uint256 newBurnTotal) internal override {
         if (newBurnTotal > destinationBurnedTotal) {
             uint256 difference = newBurnTotal - destinationBurnedTotal;
             _burnTokens(difference);
             destinationBurnedTotal = newBurnTotal;
         }
     }
+
+    /**
+     * @dev See {TokenSource-_burnTokens}
+     */
+    function _burnTokens(uint256 amount) private {
+        emit BurnTokens(amount);
+        Address.sendValue(payable(BURN_ADDRESS), amount);
+    }
 }