More router changes to make break 1:1 relation between on/off ramp an…

…d chain selector

contracts/src/v0.8/ccip/Router.sol

@@ -14,7 +14,7 @@ import {Internal} from "./libraries/Internal.sol";
 import {CallWithExactGas} from "./libraries/CallWithExactGas.sol";
 import {OwnerIsCreator} from "./../shared/access/OwnerIsCreator.sol";
 
-import {EnumerableMap} from "../vendor/openzeppelin-solidity/v4.8.0/contracts/utils/structs/EnumerableMap.sol";
+import {EnumerableSet} from "../vendor/openzeppelin-solidity/v4.8.0/contracts/utils/structs/EnumerableSet.sol";
 import {SafeERC20} from "../vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/utils/SafeERC20.sol";
 import {IERC20} from "../vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/IERC20.sol";
 
@@ -23,27 +23,21 @@ import {IERC20} from "../vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC
 /// @dev This contract is used as a router for both on-ramps and off-ramps
 contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
   using SafeERC20 for IERC20;
-  using EnumerableMap for EnumerableMap.AddressToUintMap;
+  using EnumerableSet for EnumerableSet.AddressSet;
 
   error FailedToSendValue();
   error InvalidRecipientAddress(address to);
-  error OffRampMismatch();
   error BadARMSignal();
 
-  event OnRampSet(uint64 indexed destChainSelector, address onRamp);
-  event OffRampAdded(uint64 indexed sourceChainSelector, address offRamp);
-  event OffRampRemoved(uint64 indexed sourceChainSelector, address offRamp);
-  event MessageExecuted(bytes32 messageId, uint64 sourceChainSelector, address offRamp, bytes32 calldataHash);
-
   struct OnRamp {
     uint64 destChainSelector;
     address onRamp;
   }
 
-  struct OffRamp {
-    uint64 sourceChainSelector;
-    address offRamp;
-  }
+  event OnRampSet(uint64 indexed destChainSelector, address onRamp);
+  event OffRampAdded(address offRamp);
+  event OffRampRemoved(address offRamp);
+  event MessageExecuted(bytes32 messageId, uint64 sourceChainSelector, address offRamp, bytes32 calldataHash);
 
   // solhint-disable-next-line chainlink-solidity/all-caps-constant-storage-variables
   string public constant override typeAndVersion = "Router 1.0.0";
@@ -63,7 +57,7 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
   // Mapping of offRamps to source chain ids
   // Can be multiple offRamps enabled at a time for a given sourceChainSelector,
   // for example during an no downtime upgrade while v1 messages are being flushed.
-  EnumerableMap.AddressToUintMap private s_offRamps;
+  EnumerableSet.AddressSet private s_offRamps;
 
   constructor(address wrappedNative, address armProxy) {
     // Zero address indicates unsupported auto-wrapping, therefore, unsupported
@@ -87,15 +81,15 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
     }
     address onRamp = s_onRamps[destinationChainSelector];
     if (onRamp == address(0)) revert UnsupportedDestinationChain(destinationChainSelector);
-    return IEVM2AnyOnRamp(onRamp).getFee(message);
+    return IEVM2AnyOnRamp(onRamp).getFee(destinationChainSelector, message);
   }
 
   /// @inheritdoc IRouterClient
   function getSupportedTokens(uint64 chainSelector) external view returns (address[] memory) {
     if (!isChainSupported(chainSelector)) {
       return new address[](0);
     }
-    return IEVM2AnyOnRamp(s_onRamps[uint256(chainSelector)]).getSupportedTokens();
+    return IEVM2AnyOnRamp(s_onRamps[uint256(chainSelector)]).getSupportedTokens(chainSelector);
   }
 
   /// @inheritdoc IRouterClient
@@ -117,7 +111,7 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
       // as part of the native fee coin payment.
       message.feeToken = s_wrappedNative;
       // We rely on getFee to validate that the feeToken is whitelisted.
-      feeTokenAmount = IEVM2AnyOnRamp(onRamp).getFee(message);
+      feeTokenAmount = IEVM2AnyOnRamp(onRamp).getFee(destinationChainSelector, message);
       // Ensure sufficient native.
       if (msg.value < feeTokenAmount) revert InsufficientFeeTokenAmount();
       // Wrap and send native payment.
@@ -128,7 +122,7 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
     } else {
       if (msg.value > 0) revert InvalidMsgValue();
       // We rely on getFee to validate that the feeToken is whitelisted.
-      feeTokenAmount = IEVM2AnyOnRamp(onRamp).getFee(message);
+      feeTokenAmount = IEVM2AnyOnRamp(onRamp).getFee(destinationChainSelector, message);
       IERC20(message.feeToken).safeTransferFrom(msg.sender, onRamp, feeTokenAmount);
     }
 
@@ -143,7 +137,7 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
       );
     }
 
-    return IEVM2AnyOnRamp(onRamp).forwardFromRouter(message, feeTokenAmount, msg.sender, destinationChainSelector);
+    return IEVM2AnyOnRamp(onRamp).forwardFromRouter(destinationChainSelector, message, feeTokenAmount, msg.sender);
   }
 
   // ================================================================
@@ -160,7 +154,7 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
   )
     external
     override
-    onlyOffRamp(message.sourceChainSelector)
+    onlyOffRamp()
     whenHealthy
     returns (bool success, bytes memory retData)
   {
@@ -210,28 +204,22 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
   }
 
   /// @notice Return a full list of configured offRamps.
-  function getOffRamps() external view returns (OffRamp[] memory) {
-    OffRamp[] memory offRamps = new OffRamp[](s_offRamps.length());
-    for (uint256 i = 0; i < offRamps.length; ++i) {
-      (address offRamp, uint256 sourceChainSelector) = s_offRamps.at(i);
-      offRamps[i] = OffRamp({sourceChainSelector: uint64(sourceChainSelector), offRamp: offRamp});
-    }
-    return offRamps;
+  function getOffRamps() external view returns (address[] memory) {
+    return s_offRamps.values();
   }
 
   /// @notice Returns true if the given address is a permissioned offRamp
   /// and sourceChainSelector if so.
-  function isOffRamp(address offRamp) external view returns (bool, uint64) {
-    (bool exists, uint256 sourceChainSelector) = s_offRamps.tryGet(offRamp);
-    return (exists, uint64(sourceChainSelector));
+  function isOffRamp(address offRamp) external view returns (bool) {
+    return s_offRamps.contains(offRamp);
   }
 
   /// @notice applyRampUpdates applies a set of ramp changes which provides
   /// the ability to add new chains and upgrade ramps.
   function applyRampUpdates(
     OnRamp[] calldata onRampUpdates,
-    OffRamp[] calldata offRampRemoves,
-    OffRamp[] calldata offRampAdds
+    address[] calldata offRampRemoves,
+    address[] calldata offRampAdds
   ) external onlyOwner {
     // Apply egress updates.
     // We permit zero address as way to disable egress.
@@ -243,21 +231,13 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
     // Apply ingress updates.
     // We permit an empty list as a way to disable ingress.
     for (uint256 i = 0; i < offRampRemoves.length; ++i) {
-      uint64 rampSelector = offRampRemoves[i].sourceChainSelector;
-      address rampAddress = offRampRemoves[i].offRamp;
-
-      if (s_offRamps.get(rampAddress) != uint256(rampSelector)) revert OffRampMismatch();
-
-      if (s_offRamps.remove(rampAddress)) {
-        emit OffRampRemoved(rampSelector, rampAddress);
+      if (s_offRamps.remove(offRampRemoves[i])) {
+        emit OffRampRemoved(offRampRemoves[i]);
       }
     }
     for (uint256 i = 0; i < offRampAdds.length; ++i) {
-      uint64 rampSelector = offRampAdds[i].sourceChainSelector;
-      address rampAddress = offRampAdds[i].offRamp;
-
-      if (s_offRamps.set(rampAddress, rampSelector)) {
-        emit OffRampAdded(rampSelector, rampAddress);
+      if (s_offRamps.add(offRampAdds[i])) {
+        emit OffRampAdded(offRampAdds[i]);
       }
     }
   }
@@ -283,10 +263,8 @@ contract Router is IRouter, IRouterClient, ITypeAndVersion, OwnerIsCreator {
   // ================================================================
 
   /// @notice only lets permissioned offRamps execute
-  /// @dev We additionally restrict offRamps to specific source chains for defense in depth.
-  modifier onlyOffRamp(uint64 expectedSourceChainSelector) {
-    (bool exists, uint256 sourceChainSelector) = s_offRamps.tryGet(msg.sender);
-    if (!exists || expectedSourceChainSelector != uint64(sourceChainSelector)) revert OnlyOffRamp();
+  modifier onlyOffRamp() {
+    if (!s_offRamps.contains(msg.sender)) revert OnlyOffRamp();
     _;
   }
 

contracts/src/v0.8/ccip/interfaces/IEVM2AnyOnRamp.sol

@@ -8,21 +8,9 @@ import {Internal} from "../libraries/Internal.sol";
 
 import {IERC20} from "../../vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/IERC20.sol";
 
-interface IEVM2AnyOnRamp {
-  /// @notice Get the fee for a given ccip message
-  /// @param message The message to calculate the cost for
-  /// @return fee The calculated fee
-  function getFee(Client.EVM2AnyMessage calldata message) external view returns (uint256 fee);
-
-  /// @notice Get the pool for a specific token
-  /// @param sourceToken The source chain token to get the pool for
-  /// @return pool Token pool
-  function getPoolBySourceToken(IERC20 sourceToken) external view returns (IPool);
-
-  /// @notice Gets a list of all supported source chain tokens.
-  /// @return tokens The addresses of all tokens that this onRamp supports for sending.
-  function getSupportedTokens() external view returns (address[] memory tokens);
+import {IEVM2AnyOnRampRouter} from "./IEVM2AnyOnRampRouter.sol";
 
+interface IEVM2AnyOnRamp is IEVM2AnyOnRampRouter {
   /// @notice Gets the next sequence number to be used in the onRamp
   /// @return the next sequence number to be used
   function getExpectedNextSequenceNumber() external view returns (uint64);
@@ -36,18 +24,4 @@ interface IEVM2AnyOnRamp {
   /// @param removes The tokens and pools to be removed
   /// @param adds The tokens and pools to be added.
   function applyPoolUpdates(Internal.PoolUpdate[] memory removes, Internal.PoolUpdate[] memory adds) external;
-
-  /// @notice Send a message to the remote chain
-  /// @dev only callable by the Router
-  /// @dev approve() must have already been called on the token using the this ramp address as the spender.
-  /// @dev if the contract is paused, this function will revert.
-  /// @param message Message struct to send
-  /// @param originalSender The original initiator of the CCIP request
-  /// @param destChainSelector The destination chain selector
-  function forwardFromRouter(
-    Client.EVM2AnyMessage memory message,
-    uint256 feeTokenAmount,
-    address originalSender,
-    uint64 destChainSelector
-  ) external returns (bytes32);
 }

contracts/src/v0.8/ccip/interfaces/IEVM2AnyOnRampRouter.sol

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.0;
+
+import {IPool} from "./pools/IPool.sol";
+
+import {Client} from "../libraries/Client.sol";
+
+import {IERC20} from "../../vendor/openzeppelin-solidity/v4.8.0/contracts/token/ERC20/IERC20.sol";
+
+interface IEVM2AnyOnRampRouter {
+  /// @notice Get the fee for a given ccip message
+  /// @param destChainSelector The destination chain selector
+  /// @param message The message to calculate the cost for
+  /// @return fee The calculated fee
+  function getFee(uint64 destChainSelector, Client.EVM2AnyMessage calldata message) external view returns (uint256 fee);
+
+  /// @notice Get the pool for a specific token
+  /// @param sourceToken The source chain token to get the pool for
+  /// @return pool Token pool
+  function getPoolBySourceToken(IERC20 sourceToken) external view returns (IPool);
+
+  /// @notice Gets a list of all supported source chain tokens.
+  /// @param destChainSelector The destination chain selector
+  /// @return tokens The addresses of all tokens that this onRamp supports the given destination chain
+  function getSupportedTokens(uint64 destChainSelector) external view returns (address[] memory tokens);
+
+  /// @notice Send a message to the remote chain
+  /// @dev only callable by the Router
+  /// @dev approve() must have already been called on the token using the this ramp address as the spender.
+  /// @dev if the contract is paused, this function will revert.
+  /// @param destChainSelector The destination chain selector
+  /// @param message Message struct to send
+  /// @param feeTokenAmount Amount of fee tokens for payment
+  /// @param originalSender The original initiator of the CCIP request
+  function forwardFromRouter(
+    uint64 destChainSelector,
+    Client.EVM2AnyMessage memory message,
+    uint256 feeTokenAmount,
+    address originalSender
+  ) external returns (bytes32);
+}

contracts/src/v0.8/ccip/onRamp/EVM2EVMOnRamp.sol

@@ -6,6 +6,7 @@ import {IPool} from "../interfaces/pools/IPool.sol";
 import {IARM} from "../interfaces/IARM.sol";
 import {IPriceRegistry} from "../interfaces/IPriceRegistry.sol";
 import {IEVM2AnyOnRamp} from "../interfaces/IEVM2AnyOnRamp.sol";
+import {IEVM2AnyOnRampRouter} from "../interfaces/IEVM2AnyOnRampRouter.sol";
 import {ILinkAvailable} from "../interfaces/automation/ILinkAvailable.sol";
 
 import {AggregateRateLimiter} from "../AggregateRateLimiter.sol";
@@ -253,12 +254,12 @@ contract EVM2EVMOnRamp is IEVM2AnyOnRamp, ILinkAvailable, AggregateRateLimiter,
     return uint64(senderNonce);
   }
 
-  /// @inheritdoc IEVM2AnyOnRamp
+  /// @inheritdoc IEVM2AnyOnRampRouter
   function forwardFromRouter(
+    uint64 destChainSelector,
     Client.EVM2AnyMessage calldata message,
     uint256 feeTokenAmount,
-    address originalSender,
-    uint64 destChainSelector
+    address originalSender
   ) external whenHealthy returns (bytes32) {
     // Validate message sender is set and allowed. Not validated in `getFee` since it is not user-driven.
     if (originalSender == address(0)) revert RouterMustSetOriginalSender();
@@ -436,16 +437,16 @@ contract EVM2EVMOnRamp is IEVM2AnyOnRamp, ILinkAvailable, AggregateRateLimiter,
   // │                      Tokens and pools                        │
   // ================================================================
 
-  /// @inheritdoc IEVM2AnyOnRamp
-  function getSupportedTokens() external view returns (address[] memory) {
+  /// @inheritdoc IEVM2AnyOnRampRouter
+  function getSupportedTokens(uint64 /*destChainSelector*/) external view returns (address[] memory) {
     address[] memory sourceTokens = new address[](s_poolsBySourceToken.length());
     for (uint256 i = 0; i < sourceTokens.length; ++i) {
       (sourceTokens[i], ) = s_poolsBySourceToken.at(i);
     }
     return sourceTokens;
   }
 
-  /// @inheritdoc IEVM2AnyOnRamp
+  /// @inheritdoc IEVM2AnyOnRampRouter
   function getPoolBySourceToken(IERC20 sourceToken) public view returns (IPool) {
     if (!s_poolsBySourceToken.contains(address(sourceToken))) revert UnsupportedToken(sourceToken);
     return IPool(s_poolsBySourceToken.get(address(sourceToken)));
@@ -492,11 +493,14 @@ contract EVM2EVMOnRamp is IEVM2AnyOnRamp, ILinkAvailable, AggregateRateLimiter,
   // │                             Fees                             │
   // ================================================================
 
-  /// @inheritdoc IEVM2AnyOnRamp
+  /// @inheritdoc IEVM2AnyOnRampRouter
   /// @dev getFee MUST revert if the feeToken is not listed in the fee token config, as the router assumes it does.
+  /// @param destChainSelector The destination chain selector.
   /// @param message The message to get quote for.
   /// @return feeTokenAmount The amount of fee token needed for the fee, in smallest denomination of the fee token.
-  function getFee(Client.EVM2AnyMessage calldata message) external view returns (uint256 feeTokenAmount) {
+  function getFee(uint64 destChainSelector, Client.EVM2AnyMessage calldata message) external view returns (uint256 feeTokenAmount) {
+    if (destChainSelector != i_destChainSelector) revert InvalidChainSelector(destChainSelector);
+
     Client.EVMExtraArgsV1 memory extraArgs = _fromBytes(message.extraArgs);
     // Validate the message with various checks
     _validateMessage(message.data.length, extraArgs.gasLimit, message.tokenAmounts.length);

contracts/src/v0.8/ccip/pools/ThirdPartyBurnMintTokenPool.sol

@@ -33,8 +33,7 @@ contract ThirdPartyBurnMintTokenPool is BurnMintTokenPool {
       // If the offRamp is being added do an additional check if the offRamp is
       // permission by the router. If not, we revert because we tried to add an
       // invalid offRamp.
-      (bool exists, ) = Router(s_router).isOffRamp(offRamps[i].ramp);
-      if (!exists) revert InvalidOffRamp(offRamps[i].ramp);
+      if (!Router(s_router).isOffRamp(offRamps[i].ramp)) revert InvalidOffRamp(offRamps[i].ramp);
     }
     _applyRampUpdates(onRamps, offRamps);
   }