[lock/unlock]: support multi-stage finalization in rebalancer contract

contracts/gas-snapshots/rebalancer.gas-snapshot

@@ -15,24 +15,26 @@ OCR3Base_transmit:testTransmit2SignersSuccess_gas() (gas: 56892)
 OCR3Base_transmit:testUnAuthorizedTransmitterReverts() (gas: 28595)
 OCR3Base_transmit:testUnauthorizedSignerReverts() (gas: 44736)
 OCR3Base_transmit:testWrongNumberOfSignaturesReverts() (gas: 25655)
-OptimismL1BridgeAdapter_finalizeWithdrawERC20:testFinalizeWithdrawERC20Reverts() (gas: 11929)
-OptimismL1BridgeAdapter_finalizeWithdrawERC20:testfinalizeWithdrawERC20FinalizeSuccess() (gas: 16860)
-OptimismL1BridgeAdapter_finalizeWithdrawERC20:testfinalizeWithdrawERC20proveWithdrawalSuccess() (gas: 20650)
-Rebalancer__report:test_EmptyReportReverts() (gas: 11181)
-Rebalancer_addLiquidity:test_addLiquiditySuccess() (gas: 284776)
-Rebalancer_rebalanceLiquidity:test_InsufficientLiquidityReverts() (gas: 19579)
-Rebalancer_rebalanceLiquidity:test_InvalidRemoteChainReverts() (gas: 197579)
-Rebalancer_rebalanceLiquidity:test_rebalanceBetweenPoolsSuccess() (gas: 8207096)
-Rebalancer_rebalanceLiquidity:test_rebalanceBetweenPoolsSuccess_AlreadyFinalized() (gas: 8101676)
-Rebalancer_rebalanceLiquidity:test_rebalanceLiquiditySuccess() (gas: 384924)
+OptimismL1BridgeAdapter_finalizeWithdrawERC20:testFinalizeWithdrawERC20Reverts() (gas: 12932)
+OptimismL1BridgeAdapter_finalizeWithdrawERC20:testfinalizeWithdrawERC20FinalizeSuccess() (gas: 16972)
+OptimismL1BridgeAdapter_finalizeWithdrawERC20:testfinalizeWithdrawERC20proveWithdrawalSuccess() (gas: 20758)
+Rebalancer__report:test_EmptyReportReverts() (gas: 11183)
+Rebalancer_addLiquidity:test_addLiquiditySuccess() (gas: 284740)
+Rebalancer_rebalanceLiquidity:test_InsufficientLiquidityReverts() (gas: 19585)
+Rebalancer_rebalanceLiquidity:test_InvalidRemoteChainReverts() (gas: 197560)
+Rebalancer_rebalanceLiquidity:test_rebalanceBetweenPoolsSuccess() (gas: 8634573)
+Rebalancer_rebalanceLiquidity:test_rebalanceBetweenPoolsSuccess_AlreadyFinalized() (gas: 8499105)
+Rebalancer_rebalanceLiquidity:test_rebalanceBetweenPools_MultiStageFinalization() (gas: 8494313)
+Rebalancer_rebalanceLiquidity:test_rebalanceBetweenPools_NativeRewrap() (gas: 8412119)
+Rebalancer_rebalanceLiquidity:test_rebalanceLiquiditySuccess() (gas: 384928)
 Rebalancer_removeLiquidity:test_InsufficientLiquidityReverts() (gas: 191702)
 Rebalancer_removeLiquidity:test_OnlyOwnerReverts() (gas: 10923)
 Rebalancer_removeLiquidity:test_removeLiquiditySuccess() (gas: 243268)
-Rebalancer_setCrossChainRebalancer:test_OnlyOwnerReverts() (gas: 16997)
-Rebalancer_setCrossChainRebalancer:test_ZeroAddressReverts() (gas: 21624)
-Rebalancer_setCrossChainRebalancer:test_ZeroChainSelectorReverts() (gas: 13099)
-Rebalancer_setCrossChainRebalancer:test_setCrossChainRebalancerSuccess() (gas: 162244)
-Rebalancer_setLocalLiquidityContainer:test_OnlyOwnerReverts() (gas: 10986)
-Rebalancer_setLocalLiquidityContainer:test_setLocalLiquidityContainerSuccess() (gas: 3277943)
+Rebalancer_setCrossChainRebalancer:test_OnlyOwnerReverts() (gas: 16982)
+Rebalancer_setCrossChainRebalancer:test_ZeroAddressReverts() (gas: 21630)
+Rebalancer_setCrossChainRebalancer:test_ZeroChainSelectorReverts() (gas: 13105)
+Rebalancer_setCrossChainRebalancer:test_setCrossChainRebalancerSuccess() (gas: 162288)
+Rebalancer_setLocalLiquidityContainer:test_OnlyOwnerReverts() (gas: 11008)
+Rebalancer_setLocalLiquidityContainer:test_setLocalLiquidityContainerSuccess() (gas: 3277934)
 Rebalancer_setMinimumLiquidity:test_OnlyOwnerReverts() (gas: 10925)
-Rebalancer_setMinimumLiquidity:test_setMinimumLiquiditySuccess() (gas: 36422)
+Rebalancer_setMinimumLiquidity:test_setMinimumLiquiditySuccess() (gas: 36389)

contracts/src/v0.8/rebalancer/Rebalancer.sol

@@ -4,15 +4,15 @@ pragma solidity 0.8.19;
 import {IBridgeAdapter} from "./interfaces/IBridge.sol";
 import {IRebalancer} from "./interfaces/IRebalancer.sol";
 import {ILiquidityContainer} from "./interfaces/ILiquidityContainer.sol";
-
+import {IWrappedNative} from "../ccip/interfaces/IWrappedNative.sol";
 import {OCR3Base} from "./ocr/OCR3Base.sol";
 
 import {IERC20} from "../vendor/openzeppelin-solidity/v4.8.3/contracts/token/ERC20/IERC20.sol";
 import {SafeERC20} from "../vendor/openzeppelin-solidity/v4.8.3/contracts/token/ERC20/utils/SafeERC20.sol";
 
 /// @notice Rebalancer for a single token over multiple chains.
 /// @dev This contract is designed to be used with the LockReleaseTokenPool contract but
-/// isn't constraint to it. It can be used with any contract that implements the ILiquidityContainer
+/// isn't constrained to it. It can be used with any contract that implements the ILiquidityContainer
 /// interface.
 /// @dev The OCR3 DON should only be able to transfer funds to other pre-approved contracts
 /// on other chains. Under no circumstances should it be able to transfer funds to arbitrary
@@ -30,6 +30,11 @@ contract Rebalancer is IRebalancer, OCR3Base {
   error InsufficientLiquidity(uint256 requested, uint256 available);
   error EmptyReport();
 
+  event FinalizationStepCompleted(
+    uint64 indexed ocrSeqNum,
+    uint64 indexed remoteChainSelector,
+    bytes bridgeSpecificData
+  );
   event LiquidityTransferred(
     uint64 indexed ocrSeqNum,
     uint64 indexed fromChainSelector,
@@ -165,9 +170,10 @@ contract Rebalancer is IRebalancer, OCR3Base {
   function receiveLiquidity(
     uint64 remoteChainSelector,
     uint256 amount,
+    bool shouldWrapNative,
     bytes calldata bridgeSpecificPayload
   ) external onlyOwner {
-    _receiveLiquidity(remoteChainSelector, amount, bridgeSpecificPayload, type(uint64).max);
+    _receiveLiquidity(remoteChainSelector, amount, bridgeSpecificPayload, shouldWrapNative, type(uint64).max);
   }
 
   /// @notice Transfers liquidity to another chain.
@@ -217,6 +223,7 @@ contract Rebalancer is IRebalancer, OCR3Base {
     uint64 remoteChainSelector,
     uint256 amount,
     bytes memory bridgeSpecificPayload,
+    bool shouldWrapNative,
     uint64 ocrSeqNum
   ) internal {
     // check if the remote chain is supported
@@ -232,8 +239,20 @@ contract Rebalancer is IRebalancer, OCR3Base {
         address(this), // localReceiver: this contract
         bridgeSpecificPayload
       )
-    {
-      // successfully finalized the withdrawal
+    returns (bool fundsAvailable) {
+      if (fundsAvailable) {
+        // finalization was successful and we can inject the liquidity into the container.
+        // approve and liquidity container should transferFrom.
+        _injectLiquidity(amount, ocrSeqNum, remoteChainSelector, bridgeSpecificPayload, shouldWrapNative);
+      } else {
+        // a finalization step was completed, but funds are not available.
+        // hence, we cannot inject any liquidity yet.
+        emit FinalizationStepCompleted(ocrSeqNum, remoteChainSelector, bridgeSpecificPayload);
+      }
+
+      // return here on the happy path.
+      // sad path is when finalizeWithdrawERC20 reverts, which is handled after the catch block.
+      return;
     } catch (bytes memory lowLevelData) {
       // failed to finalize the withdrawal.
       // this could mean that the withdrawal was already finalized
@@ -242,8 +261,29 @@ contract Rebalancer is IRebalancer, OCR3Base {
       emit FinalizationFailed(ocrSeqNum, remoteChainSelector, bridgeSpecificPayload, lowLevelData);
     }
 
-    // inject liquidity into the liquidity container
-    // approve and liquidity container should transferFrom
+    // if we reach this point, the finalization failed.
+    // since we don't have enough information to know why it failed,
+    // we assume that it failed because the withdrawal was already finalized,
+    // and that the funds are available.
+    _injectLiquidity(amount, ocrSeqNum, remoteChainSelector, bridgeSpecificPayload, shouldWrapNative);
+  }
+
+  function _injectLiquidity(
+    uint256 amount,
+    uint64 ocrSeqNum,
+    uint64 remoteChainSelector,
+    bytes memory bridgeSpecificPayload,
+    bool shouldWrapNative
+  ) private {
+    // We trust the DON or the owner (the only two actors who can end up calling this function)
+    // to correctly set the shouldWrapNative flag.
+    // Some bridges only bridge native and not wrapped native.
+    // In such a case we need to re-wrap the native in order to inject it into the liquidity container.
+    // TODO: escape hatch in case of bug?
+    if (shouldWrapNative) {
+      _wrapNative(amount);
+    }
+
     i_localToken.safeIncreaseAllowance(address(s_localLiquidityContainer), amount);
     s_localLiquidityContainer.provideLiquidity(amount);
 
@@ -258,6 +298,11 @@ contract Rebalancer is IRebalancer, OCR3Base {
     );
   }
 
+  function _wrapNative(uint256 amount) private {
+    IWrappedNative weth = IWrappedNative(address(i_localToken));
+    weth.deposit{value: amount}();
+  }
+
   function _report(bytes calldata report, uint64 ocrSeqNum) internal override {
     IRebalancer.LiquidityInstructions memory instructions = abi.decode(report, (IRebalancer.LiquidityInstructions));
 
@@ -285,6 +330,7 @@ contract Rebalancer is IRebalancer, OCR3Base {
         instructions.receiveLiquidityParams[i].remoteChainSelector,
         instructions.receiveLiquidityParams[i].amount,
         instructions.receiveLiquidityParams[i].bridgeData,
+        instructions.receiveLiquidityParams[i].shouldWrapNative,
         ocrSeqNum
       );
     }

contracts/src/v0.8/rebalancer/bridge-adapters/ArbitrumL1BridgeAdapter.sol

@@ -58,7 +58,7 @@ contract ArbitrumL1BridgeAdapter is IBridgeAdapter {
   }
 
   /// @dev these are parameters provided by the caller of the sendERC20 function
-  /// @dev these must be determined offchain.
+  /// and must be determined offchain.
   struct SendERC20Params {
     uint256 gasLimit;
     uint256 maxSubmissionCost;
@@ -110,12 +110,10 @@ contract ArbitrumL1BridgeAdapter is IBridgeAdapter {
     return inboxSequenceNumber;
   }
 
-  /// @dev This function is so that we can easily abi-encode the arbitrum-specific
-  /// @dev payload for the sendERC20 function.
+  /// @dev This function is so that we can easily abi-encode the arbitrum-specific payload for the sendERC20 function.
   function exposeSendERC20Params(SendERC20Params memory params) public pure {}
 
-  /// @dev fees have to be determined offchain for arbitrum
-  /// @dev therefore revert here to discourage usage
+  /// @dev fees have to be determined offchain for arbitrum, therefore revert here to discourage usage.
   function getBridgeFeeInNative() public pure override returns (uint256) {
     revert Unimplemented();
   }
@@ -141,16 +139,19 @@ contract ArbitrumL1BridgeAdapter is IBridgeAdapter {
     bytes data;
   }
 
-  /// @dev This function is so that we can easily abi-encode the arbitrum-specific
-  /// @dev payload for the finalizeWithdrawERC20 function.
+  /// @dev This function is so that we can easily abi-encode the arbitrum-specific payload for the finalizeWithdrawERC20 function.
   function exposeArbitrumFinalizationPayload(ArbitrumFinalizationPayload memory payload) public pure {}
 
   /// @notice Finalize an L2 -> L1 transfer.
+  /// Arbitrum finalizations are single-step, so we always return true.
+  /// Calls to this function will revert in two cases, 1) if the finalization payload is wrong,
+  /// i.e incorrect merkle proof, or index and 2) if the withdrawal was already finalized.
+  /// @return true iff the finalization does not revert.
   function finalizeWithdrawERC20(
     address /* remoteSender */,
     address /* localReceiver */,
     bytes calldata arbitrumFinalizationPayload
-  ) external {
+  ) external override returns (bool) {
     ArbitrumFinalizationPayload memory payload = abi.decode(arbitrumFinalizationPayload, (ArbitrumFinalizationPayload));
     i_l1Outbox.executeTransaction(
       payload.proof,
@@ -163,8 +164,11 @@ contract ArbitrumL1BridgeAdapter is IBridgeAdapter {
       payload.value,
       payload.data
     );
+    return true;
   }
 
+  /// @notice Convenience function to get the L2 token address from the L1 token address.
+  /// @return The L2 token address for the given L1 token address.
   function getL2Token(address l1Token) external view returns (address) {
     return i_l1GatewayRouter.calculateL2TokenAddress(l1Token);
   }

contracts/src/v0.8/rebalancer/bridge-adapters/ArbitrumL2BridgeAdapter.sol

@@ -58,11 +58,14 @@ contract ArbitrumL2BridgeAdapter is IBridgeAdapter {
   }
 
   /// @notice No-op since L1 -> L2 transfers do not need finalization.
+  /// @return true always.
   function finalizeWithdrawERC20(
     address /* remoteSender */,
     address /* localReceiver */,
     bytes calldata /* bridgeSpecificPayload */
-  ) external {}
+  ) external override returns (bool) {
+    return true;
+  }
 
   /// @notice There are no fees to bridge back to L1
   function getBridgeFeeInNative() external pure returns (uint256) {

contracts/src/v0.8/rebalancer/bridge-adapters/OptimismL1BridgeAdapter.sol

@@ -131,18 +131,21 @@ contract OptimismL1BridgeAdapter is IBridgeAdapter {
   /// @notice Prove or finalize an ERC20 withdrawal from L2.
   /// The action to take is specified in the payload. See the docstring of FinalizeWithdrawERC20Payload for more details.
   /// @param data The payload for the action. This is an abi.encode'd FinalizeWithdrawERC20Payload with the appropriate data.
+  /// @return true iff finalization is successful, and false for proving a withdrawal. If either of these fail,
+  /// the call to this function will revert.
   function finalizeWithdrawERC20(
     address /* remoteSender */,
     address /* localReceiver */,
     bytes calldata data
-  ) external override {
+  ) external override returns (bool) {
     // decode the data into FinalizeWithdrawERC20Payload first and extract the action.
     FinalizeWithdrawERC20Payload memory payload = abi.decode(data, (FinalizeWithdrawERC20Payload));
     if (payload.action == FinalizationAction.ProveWithdrawal) {
       // The action being ProveWithdrawal indicates that this is a withdrawal proof payload.
       // Decode the data into OptimismProveWithdrawalPayload and call the proveWithdrawal function.
       OptimismProveWithdrawalPayload memory provePayload = abi.decode(payload.data, (OptimismProveWithdrawalPayload));
       _proveWithdrawal(provePayload);
+      return false;
     } else if (payload.action == FinalizationAction.FinalizeWithdrawal) {
       // decode the data into OptimismFinalizationPayload and call the finalizeWithdrawal function.
       OptimismFinalizationPayload memory finalizePayload = abi.decode(payload.data, (OptimismFinalizationPayload));
@@ -151,6 +154,7 @@ contract OptimismL1BridgeAdapter is IBridgeAdapter {
       // However, we can't do that from within this adapter because it doesn't actually have the ether.
       // So its up to the caller to rectify this by re-wrapping the ether.
       _finalizeWithdrawal(finalizePayload);
+      return true;
     } else {
       revert InvalidFinalizationAction();
     }

contracts/src/v0.8/rebalancer/bridge-adapters/OptimismL2BridgeAdapter.sol

@@ -88,11 +88,14 @@ contract OptimismL2BridgeAdapter is IBridgeAdapter {
   }
 
   /// @notice No-op since L1 -> L2 transfers do not need finalization.
+  /// @return true always.
   function finalizeWithdrawERC20(
     address /* remoteSender */,
     address /* localReceiver */,
     bytes calldata /* bridgeSpecificPayload */
-  ) external override {}
+  ) external override returns (bool) {
+    return true;
+  }
 
   /// @notice There are no fees to bridge back to L1
   function getBridgeFeeInNative() external pure returns (uint256) {

contracts/src/v0.8/rebalancer/interfaces/IBridge.sol

@@ -9,9 +9,9 @@ interface IBridgeAdapter {
   error MsgShouldNotContainValue(uint256 value);
 
   /// @notice Send the specified amount of the local token cross-chain to the remote chain.
-  /// @notice The tokens on the remote chain will then be sourced from the remoteToken address.
-  /// @notice The amount to be sent must be approved by the caller beforehand on the localToken contract.
-  /// @notice The caller must provide the bridging fee in native currency, i.e msg.value.
+  /// The tokens on the remote chain will then be sourced from the remoteToken address.
+  /// The amount to be sent must be approved by the caller beforehand on the localToken contract.
+  /// The caller must provide the bridging fee in native currency, i.e msg.value.
   /// @param localToken The address of the local ERC-20 token.
   /// @param remoteToken The address of the remote ERC-20 token.
   /// @param recipient The address of the recipient on the remote chain.
@@ -26,17 +26,24 @@ interface IBridgeAdapter {
   ) external payable returns (bytes memory);
 
   /// @notice Get the bridging fee in native currency. This fee must be provided upon sending tokens via
-  /// @notice the sendERC20 function.
+  /// the sendERC20 function.
   /// @return The bridging fee in native currency.
   function getBridgeFeeInNative() external view returns (uint256);
 
   /// @notice Finalize the withdrawal of a cross-chain transfer.
+  /// Not all implementations will finalize a transfer in a single call to this function.
+  /// Optimism, for example, requires a two-step process to finalize a transfer. The first
+  /// step requires proving the withdrawal that occurred on L2 on L1. The second step is then
+  /// the finalization, whereby funds become available to the recipient. So, in that particular
+  /// scenario, `false` is returned from `finalizeWithdrawERC20` when the first step is completed,
+  /// and `true` is returned when the second step is completed.
   /// @param remoteSender The address of the sender on the remote chain.
   /// @param localReceiver The address of the receiver on the local chain.
   /// @param bridgeSpecificPayload The payload of the cross-chain transfer, bridge-specific, i.e a proof of some kind.
+  /// @return true iff the funds are available, false otherwise.
   function finalizeWithdrawERC20(
     address remoteSender,
     address localReceiver,
     bytes calldata bridgeSpecificPayload
-  ) external;
+  ) external returns (bool);
 }

contracts/src/v0.8/rebalancer/interfaces/IRebalancer.sol

@@ -4,21 +4,29 @@ pragma solidity ^0.8.0;
 import {IBridgeAdapter} from "./IBridge.sol";
 
 interface IRebalancer {
+  /// @notice Parameters for sending liquidity to a remote chain.
+  /// @param amount The amount of tokens to be sent to the remote chain.
+  /// @param nativeBridgeFee The amount of native that should be sent by the rebalancer in the sendERC20 call.
+  ///        Used to pay for the bridge fees.
+  /// @param remoteChainSelector The selector of the remote chain.
+  /// @param bridgeData The bridge data that should be passed to the sendERC20 call.
   struct SendLiquidityParams {
-    /// @notice The amount of tokens to be sent to the remote chain.
     uint256 amount;
-    /// @notice The amount of native that should be sent by the rebalancer in the sendERC20 call.
-    /// @notice This is used to pay for the bridge fees.
     uint256 nativeBridgeFee;
-    /// @notice The selector of the remote chain.
     uint64 remoteChainSelector;
-    /// @notice The bridge data that should be passed to the sendERC20 call.
     bytes bridgeData;
   }
 
+  /// @notice Parameters for receiving liquidity from a remote chain.
+  /// @param amount The amount of tokens to be received from the remote chain.
+  /// @param remoteChainSelector The selector of the remote chain.
+  /// @param bridgeData The bridge data that should be passed to the finalizeWithdrawERC20 call.
+  /// @param shouldWrapNative Whether the received native token should be wrapped into wrapped native.
+  ///        This is needed for when the bridge being used doesn't bridge wrapped native but native directly.
   struct ReceiveLiquidityParams {
     uint256 amount;
     uint64 remoteChainSelector;
+    bool shouldWrapNative;
     bytes bridgeData;
   }
 
@@ -36,6 +44,7 @@ interface IRebalancer {
   }
 
   /// @notice Returns the current liquidity in the liquidity container.
+  /// @return currentLiquidity The current liquidity in the liquidity container.
   function getLiquidity() external view returns (uint256 currentLiquidity);
 
   function getAllCrossChainRebalancers() external view returns (CrossChainRebalancerArgs[] memory);