diff --git a/src/contracts/AbstractARM.sol b/src/contracts/AbstractARM.sol index e41a6fd..8132261 100644 --- a/src/contracts/AbstractARM.sol +++ b/src/contracts/AbstractARM.sol @@ -159,4 +159,9 @@ abstract contract AbstractARM is OwnableOperable { function _transferAsset(address asset, address to, uint256 amount) internal virtual { IERC20(asset).transfer(to, amount); } + + /// @dev Hook to transfer assets into the ARM contract + function _transferAssetFrom(address asset, address from, address to, uint256 amount) internal virtual { + IERC20(asset).transferFrom(from, to, amount); + } } diff --git a/src/contracts/FixedPriceARM.sol b/src/contracts/FixedPriceARM.sol index 1d47cc5..c6587aa 100644 --- a/src/contracts/FixedPriceARM.sol +++ b/src/contracts/FixedPriceARM.sol @@ -54,7 +54,7 @@ abstract contract FixedPriceARM is AbstractARM { amountOut = amountIn * price / PRICE_SCALE; // Transfer the input tokens from the caller to this ARM contract - inToken.transferFrom(msg.sender, address(this), amountIn); + _transferAssetFrom(address(inToken), msg.sender, address(this), amountIn); // Transfer the output tokens to the recipient _transferAsset(address(outToken), to, amountOut); @@ -78,7 +78,7 @@ abstract contract FixedPriceARM is AbstractARM { amountIn = ((amountOut * PRICE_SCALE) / price) + 1; // +1 to always round in our favor // Transfer the input tokens from the caller to this ARM contract - inToken.transferFrom(msg.sender, address(this), amountIn); + _transferAssetFrom(address(inToken), msg.sender, address(this), amountIn); // Transfer the output tokens to the recipient _transferAsset(address(outToken), to, amountOut); diff --git a/src/contracts/LidoFixedPriceMultiLpARM.sol b/src/contracts/LidoFixedPriceMultiLpARM.sol index d0f2cd7..460a162 100644 --- a/src/contracts/LidoFixedPriceMultiLpARM.sol +++ b/src/contracts/LidoFixedPriceMultiLpARM.sol @@ -71,6 +71,14 @@ contract LidoFixedPriceMultiLpARM is MultiLP._transferAsset(asset, to, transferAmount); } + /// @dev Funds the ARM's withdrawal queue when swaps send WETH to the ARM + function _transferAssetFrom(address asset, address from, address to, uint256 amount) + internal + override(AbstractARM, MultiLP) + { + MultiLP._transferAssetFrom(asset, from, to, amount); + } + function _externalWithdrawQueue() internal view override(MultiLP, LidoLiquidityManager) returns (uint256) { return LidoLiquidityManager._externalWithdrawQueue(); } diff --git a/src/contracts/MultiLP.sol b/src/contracts/MultiLP.sol index 74048be..5ed6759 100644 --- a/src/contracts/MultiLP.sol +++ b/src/contracts/MultiLP.sol @@ -240,11 +240,20 @@ abstract contract MultiLP is AbstractARM, ERC20Upgradeable { /// @dev Ensure any liquidity assets reserved for the withdrawal queue are not used /// in swaps that send liquidity assets out of the ARM function _transferAsset(address asset, address to, uint256 amount) internal virtual override { - require(asset == liquidityAsset && amount <= _liquidityAvailable(), "ARM: Insufficient liquidity"); + if (asset == liquidityAsset) { + require(amount <= _liquidityAvailable(), "ARM: Insufficient liquidity"); + } IERC20(asset).transfer(to, amount); } + /// @dev Funds the ARM's withdrawal queue when swaps send liquidity assets to the ARM + function _transferAssetFrom(address asset, address from, address to, uint256 amount) internal virtual override { + IERC20(asset).transferFrom(from, to, amount); + + _addWithdrawalQueueLiquidity(); + } + /// @notice The total amount of assets in the ARM and external withdrawal queue, /// less the liquidity assets reserved for the withdrawal queue function totalAssets() public view virtual returns (uint256 assets) { diff --git a/test/fork/LidoFixedPriceMultiLpARM/SwapExactTokensForTokens.t.sol b/test/fork/LidoFixedPriceMultiLpARM/SwapExactTokensForTokens.t.sol index c2fd61e..c920476 100644 --- a/test/fork/LidoFixedPriceMultiLpARM/SwapExactTokensForTokens.t.sol +++ b/test/fork/LidoFixedPriceMultiLpARM/SwapExactTokensForTokens.t.sol @@ -116,7 +116,7 @@ contract Fork_Concrete_LidoFixedPriceMultiLpARM_SwapExactTokensForTokens_Test is initialBalance = weth.balanceOf(address(lidoFixedPriceMulltiLpARM)); deal(address(steth), address(this), initialBalance * 2); - vm.expectRevert(); // Lido error + vm.expectRevert("ARM: Insufficient liquidity"); lidoFixedPriceMulltiLpARM.swapExactTokensForTokens( steth, // inToken weth, // outToken