Native Token Bridges uses Teleporter Registry

contracts/src/CrossChainApplications/NativeTokenBridge/ERC20TokenSource.sol

@@ -14,13 +14,13 @@ import {
     TeleporterMessageInput,
     TeleporterFeeInfo
 } from "../../Teleporter/ITeleporterMessenger.sol";
-import {ITeleporterReceiver} from "../../Teleporter/ITeleporterReceiver.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";
 
 contract ERC20TokenSource is
-    ITeleporterReceiver,
+    TeleporterOwnerUpgradeable,
     IERC20TokenSource,
     ITokenSource,
     ReentrancyGuard
@@ -37,21 +37,12 @@ contract ERC20TokenSource is
     address public immutable nativeTokenDestinationAddress;
     address public immutable erc20ContractAddress;
 
-    // Used for sending and receiving Teleporter messages.
-    ITeleporterMessenger public immutable teleporterMessenger;
-
     constructor(
-        address teleporterMessengerAddress,
+        address teleporterRegistryAddress,
         bytes32 destinationBlockchainID_,
         address nativeTokenDestinationAddress_,
         address erc20ContractAddress_
-    ) {
-        require(
-            teleporterMessengerAddress != address(0),
-            "ERC20TokenSource: zero TeleporterMessenger address"
-        );
-        teleporterMessenger = ITeleporterMessenger(teleporterMessengerAddress);
-
+    ) TeleporterOwnerUpgradeable(teleporterRegistryAddress) {
         require(
             destinationBlockchainID_ != bytes32(0),
             "ERC20TokenSource: zero destination blockchain ID"
@@ -75,48 +66,6 @@ contract ERC20TokenSource is
         erc20ContractAddress = erc20ContractAddress_;
     }
 
-    /**
-     * @dev See {ITeleporterReceiver-receiveTeleporterMessage}.
-     *
-     * Receives a Teleporter message and routes to the appropriate internal function call.
-     */
-    function receiveTeleporterMessage(
-        bytes32 senderBlockchainID,
-        address senderAddress,
-        bytes calldata message
-    ) external nonReentrant {
-        // Only allow the Teleporter messenger to deliver messages.
-        require(
-            msg.sender == address(teleporterMessenger),
-            "ERC20TokenSource: unauthorized TeleporterMessenger contract"
-        );
-
-        // Only allow messages from the destination chain.
-        require(
-            senderBlockchainID == destinationBlockchainID,
-            "ERC20TokenSource: invalid destination chain"
-        );
-
-        // Only allow the partner contract to send messages.
-        require(
-            senderAddress == 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 See {IERC20TokenSource-transferToDestination}.
      */
@@ -126,6 +75,8 @@ contract ERC20TokenSource is
         uint256 feeAmount,
         address[] calldata allowedRelayerAddresses
     ) external nonReentrant {
+        ITeleporterMessenger teleporterMessenger = _getTeleporterMessenger();
+
         // The recipient cannot be the zero address.
         require(recipient != address(0), "ERC20TokenSource: zero recipient address");
 
@@ -166,6 +117,42 @@ contract ERC20TokenSource is
         });
     }
 
+    /**
+     * @dev See {TeleporterUpgradeable-receiveTeleporterMessage}.
+     *
+     * Receives a Teleporter message and routes to the appropriate internal function call.
+     */
+    function _receiveTeleporterMessage(
+        bytes32 senderBlockchainID,
+        address senderAddress,
+        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(
+            senderAddress == 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.
      */

contracts/src/CrossChainApplications/NativeTokenBridge/NativeTokenDestination.sol

@@ -16,15 +16,19 @@ import {
     TeleporterFeeInfo,
     TeleporterMessageInput
 } from "../../Teleporter/ITeleporterMessenger.sol";
-import {ITeleporterReceiver} from "../../Teleporter/ITeleporterReceiver.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";
 // We need IAllowList as an indirect dependency in order to compile.
 // solhint-disable-next-line no-unused-import
 import {IAllowList} from "@subnet-evm-contracts/interfaces/IAllowList.sol";
 
-contract NativeTokenDestination is ITeleporterReceiver, INativeTokenDestination, ReentrancyGuard {
+contract NativeTokenDestination is
+    TeleporterOwnerUpgradeable,
+    INativeTokenDestination,
+    ReentrancyGuard
+{
     // The address where the burned transaction fees are credited.
     // Defined as BLACKHOLE_ADDRESS at
     // https://github.com/ava-labs/subnet-evm/blob/e23ab058d039ff9c8469c89b139d21d52c4bd283/constants/constants.go
@@ -48,21 +52,12 @@ contract NativeTokenDestination is ITeleporterReceiver, INativeTokenDestination,
     uint256 public currentReserveImbalance;
     uint256 public totalMinted;
 
-    // Used for sending and receiving Teleporter messages.
-    ITeleporterMessenger public immutable teleporterMessenger;
-
     constructor(
-        address teleporterMessengerAddress,
+        address teleporterRegistryAddress,
         bytes32 sourceBlockchainID_,
         address nativeTokenSourceAddress_,
         uint256 initialReserveImbalance_
-    ) {
-        require(
-            teleporterMessengerAddress != address(0),
-            "NativeTokenDestination: zero TeleporterMessenger address"
-        );
-        teleporterMessenger = ITeleporterMessenger(teleporterMessengerAddress);
-
+    ) TeleporterOwnerUpgradeable(teleporterRegistryAddress) {
         require(
             sourceBlockchainID_ != bytes32(0), "NativeTokenDestination: zero source blockchain ID"
         );
@@ -87,58 +82,6 @@ contract NativeTokenDestination is ITeleporterReceiver, INativeTokenDestination,
         currentReserveImbalance = initialReserveImbalance_;
     }
 
-    /**
-     * @dev See {ITeleporterReceiver-receiveTeleporterMessage}.
-     *
-     * Receives a Teleporter message.
-     */
-    function receiveTeleporterMessage(
-        bytes32 senderBlockchainID,
-        address senderAddress,
-        bytes calldata message
-    ) external nonReentrant {
-        // Only allow the Teleporter messenger to deliver messages.
-        require(
-            msg.sender == address(teleporterMessenger),
-            "NativeTokenDestination: unauthorized TeleporterMessenger contract"
-        );
-
-        // Only allow messages from the source chain.
-        require(
-            senderBlockchainID == sourceBlockchainID, "NativeTokenDestination: invalid source chain"
-        );
-
-        // Only allow the partner contract to send messages.
-        require(
-            senderAddress == nativeTokenSourceAddress, "NativeTokenDestination: unauthorized sender"
-        );
-
-        (address recipient, uint256 amount) = abi.decode(message, (address, uint256));
-        require(recipient != address(0), "NativeTokenDestination: zero recipient address");
-        require(amount != 0, "NativeTokenDestination: zero transfer value");
-
-        // If the contract has not yet been collateralized, we will deduct as many tokens
-        // as needed from the transfer as needed. If there are any excess tokens, they will
-        // be minted and sent to the recipient.
-        uint256 adjustedAmount = amount;
-        if (currentReserveImbalance > 0) {
-            if (amount > currentReserveImbalance) {
-                emit CollateralAdded({amount: currentReserveImbalance, remaining: 0});
-                adjustedAmount = amount - currentReserveImbalance;
-                currentReserveImbalance = 0;
-            } else {
-                currentReserveImbalance -= amount;
-                emit CollateralAdded({amount: amount, remaining: currentReserveImbalance});
-                return;
-            }
-        }
-
-        totalMinted += adjustedAmount;
-        emit NativeTokensMinted(recipient, adjustedAmount);
-        // Calls NativeMinter precompile through INativeMinter interface.
-        _nativeMinter.mintNativeCoin(recipient, adjustedAmount);
-    }
-
     /**
      * @dev See {INativeTokenDestination-transferToSource}.
      */
@@ -147,6 +90,7 @@ contract NativeTokenDestination is ITeleporterReceiver, INativeTokenDestination,
         TeleporterFeeInfo calldata feeInfo,
         address[] calldata allowedRelayerAddresses
     ) external payable nonReentrant {
+        ITeleporterMessenger teleporterMessenger = teleporterRegistry.getLatestTeleporter();
         // The recipient cannot be the zero address.
         require(recipient != address(0), "NativeTokenDestination: zero recipient address");
 
@@ -196,6 +140,8 @@ contract NativeTokenDestination is ITeleporterReceiver, INativeTokenDestination,
         TeleporterFeeInfo calldata feeInfo,
         address[] calldata allowedRelayerAddresses
     ) external {
+        ITeleporterMessenger teleporterMessenger = _getTeleporterMessenger();
+
         uint256 totalBurnedTxFees = address(BURNED_TX_FEES_ADDRESS).balance;
         uint256 messageID = teleporterMessenger.sendCrossChainMessage(
             TeleporterMessageInput({
@@ -231,4 +177,50 @@ contract NativeTokenDestination is ITeleporterReceiver, INativeTokenDestination,
 
         return created - burned;
     }
+
+    /**
+     * @dev See {TeleporterUpgradeable-receiveTeleporterMessage}.
+     *
+     * Receives a Teleporter message.
+     */
+    function _receiveTeleporterMessage(
+        bytes32 senderBlockchainID,
+        address senderAddress,
+        bytes memory message
+    ) internal override {
+        // Only allow messages from the source chain.
+        require(
+            senderBlockchainID == sourceBlockchainID, "NativeTokenDestination: invalid source chain"
+        );
+
+        // Only allow the partner contract to send messages.
+        require(
+            senderAddress == nativeTokenSourceAddress, "NativeTokenDestination: unauthorized sender"
+        );
+
+        (address recipient, uint256 amount) = abi.decode(message, (address, uint256));
+        require(recipient != address(0), "NativeTokenDestination: zero recipient address");
+        require(amount != 0, "NativeTokenDestination: zero transfer value");
+
+        // If the contract has not yet been collateralized, we will deduct as many tokens
+        // as needed from the transfer as needed. If there are any excess tokens, they will
+        // be minted and sent to the recipient.
+        uint256 adjustedAmount = amount;
+        if (currentReserveImbalance > 0) {
+            if (amount > currentReserveImbalance) {
+                emit CollateralAdded({amount: currentReserveImbalance, remaining: 0});
+                adjustedAmount = amount - currentReserveImbalance;
+                currentReserveImbalance = 0;
+            } else {
+                currentReserveImbalance -= amount;
+                emit CollateralAdded({amount: amount, remaining: currentReserveImbalance});
+                return;
+            }
+        }
+
+        totalMinted += adjustedAmount;
+        emit NativeTokensMinted(recipient, adjustedAmount);
+        // Calls NativeMinter precompile through INativeMinter interface.
+        _nativeMinter.mintNativeCoin(recipient, adjustedAmount);
+    }
 }