diff --git a/contracts/BorrowController.sol b/contracts/BorrowController.sol index 979b3ac..c9c701a 100644 --- a/contracts/BorrowController.sol +++ b/contracts/BorrowController.sol @@ -8,6 +8,7 @@ import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol"; import {IBorrowController} from "./interfaces/IBorrowController.sol"; import {ILoanPositionManager} from "./interfaces/ILoanPositionManager.sol"; import {ISubstitute} from "./interfaces/ISubstitute.sol"; +import {SubstituteLibrary} from "./libraries/Substitute.sol"; import {IPositionLocker} from "./interfaces/IPositionLocker.sol"; import {LoanPosition} from "./libraries/LoanPosition.sol"; import {Coupon} from "./libraries/Coupon.sol"; @@ -18,6 +19,7 @@ import {ERC20PermitParams, PermitSignature, PermitParamsLibrary} from "./librari contract BorrowController is IBorrowController, Controller, IPositionLocker { using PermitParamsLibrary for *; using EpochLibrary for Epoch; + using SubstituteLibrary for ISubstitute; ILoanPositionManager private immutable _loanPositionManager; address private immutable _router; @@ -88,7 +90,7 @@ contract BorrowController is IBorrowController, Controller, IPositionLocker { ); if (collateralDelta > 0) { - _ensureBalance(position.collateralToken, user, uint256(collateralDelta)); + ISubstitute(position.collateralToken).ensureBalance(user, uint256(collateralDelta)); IERC20(position.collateralToken).approve(address(_loanPositionManager), uint256(collateralDelta)); _loanPositionManager.depositToken(position.collateralToken, uint256(collateralDelta)); } @@ -113,7 +115,7 @@ contract BorrowController is IBorrowController, Controller, IPositionLocker { Epoch expiredWith, SwapParams calldata swapParams, ERC20PermitParams calldata collateralPermitParams - ) external payable nonReentrant wrapETH returns (uint256 positionId) { + ) external payable nonReentrant wrapAndRefundETH returns (uint256 positionId) { collateralPermitParams.tryPermit(_getUnderlyingToken(collateralToken), msg.sender, address(this)); bytes memory lockData = abi.encode(collateralAmount, debtAmount, expiredWith, maxPayInterest); @@ -121,8 +123,8 @@ contract BorrowController is IBorrowController, Controller, IPositionLocker { bytes memory result = _loanPositionManager.lock(lockData); positionId = abi.decode(result, (uint256)); - _burnAllSubstitute(collateralToken, msg.sender); - _burnAllSubstitute(debtToken, msg.sender); + ISubstitute(collateralToken).burnAll(msg.sender); + ISubstitute(debtToken).burnAll(msg.sender); _loanPositionManager.transferFrom(address(this), msg.sender, positionId); } @@ -136,7 +138,7 @@ contract BorrowController is IBorrowController, Controller, IPositionLocker { PermitSignature calldata positionPermitParams, ERC20PermitParams calldata collateralPermitParams, ERC20PermitParams calldata debtPermitParams - ) external payable nonReentrant wrapETH onlyPositionOwner(positionId) { + ) external payable nonReentrant wrapAndRefundETH onlyPositionOwner(positionId) { positionPermitParams.tryPermit(_loanPositionManager, positionId, address(this)); LoanPosition memory position = _loanPositionManager.getPosition(positionId); collateralPermitParams.tryPermit(_getUnderlyingToken(position.collateralToken), msg.sender, address(this)); @@ -148,8 +150,8 @@ contract BorrowController is IBorrowController, Controller, IPositionLocker { _loanPositionManager.lock(_encodeAdjustData(positionId, position, interestThreshold, swapParams)); - _burnAllSubstitute(position.collateralToken, msg.sender); - _burnAllSubstitute(position.debtToken, msg.sender); + ISubstitute(position.collateralToken).burnAll(msg.sender); + ISubstitute(position.debtToken).burnAll(msg.sender); } function _swap(address inSubstitute, address outSubstitute, uint256 inAmount, bytes memory swapParams) diff --git a/contracts/DepositController.sol b/contracts/DepositController.sol index 4d271a3..8f71b9c 100644 --- a/contracts/DepositController.sol +++ b/contracts/DepositController.sol @@ -12,12 +12,15 @@ import {BondPosition} from "./libraries/BondPosition.sol"; import {Epoch, EpochLibrary} from "./libraries/Epoch.sol"; import {CouponKey} from "./libraries/CouponKey.sol"; import {Coupon} from "./libraries/Coupon.sol"; +import {ISubstitute} from "./interfaces/ISubstitute.sol"; +import {SubstituteLibrary} from "./libraries/Substitute.sol"; import {Controller} from "./libraries/Controller.sol"; import {ERC20PermitParams, PermitSignature, PermitParamsLibrary} from "./libraries/PermitParams.sol"; contract DepositController is IDepositController, Controller, IPositionLocker { using PermitParamsLibrary for *; using EpochLibrary for Epoch; + using SubstituteLibrary for ISubstitute; IBondPositionManager private immutable _bondPositionManager; @@ -87,13 +90,13 @@ contract DepositController is IDepositController, Controller, IPositionLocker { Epoch expiredWith, int256 minEarnInterest, ERC20PermitParams calldata tokenPermitParams - ) external payable nonReentrant wrapETH returns (uint256 positionId) { + ) external payable nonReentrant wrapAndRefundETH returns (uint256 positionId) { tokenPermitParams.tryPermit(_getUnderlyingToken(asset), msg.sender, address(this)); bytes memory lockData = abi.encode(amount, expiredWith, -minEarnInterest); bytes memory result = _bondPositionManager.lock(abi.encode(0, msg.sender, abi.encode(asset, lockData))); positionId = abi.decode(result, (uint256)); - _burnAllSubstitute(asset, msg.sender); + ISubstitute(asset).burnAll(msg.sender); _bondPositionManager.transferFrom(address(this), msg.sender, positionId); } @@ -105,7 +108,7 @@ contract DepositController is IDepositController, Controller, IPositionLocker { int256 interestThreshold, ERC20PermitParams calldata tokenPermitParams, PermitSignature calldata positionPermitParams - ) external payable nonReentrant wrapETH onlyPositionOwner(positionId) { + ) external payable nonReentrant wrapAndRefundETH onlyPositionOwner(positionId) { positionPermitParams.tryPermit(_bondPositionManager, positionId, address(this)); BondPosition memory position = _bondPositionManager.getPosition(positionId); tokenPermitParams.tryPermit(position.asset, msg.sender, address(this)); @@ -113,6 +116,6 @@ contract DepositController is IDepositController, Controller, IPositionLocker { bytes memory lockData = abi.encode(amount, expiredWith, interestThreshold); _bondPositionManager.lock(abi.encode(positionId, msg.sender, lockData)); - _burnAllSubstitute(position.asset, msg.sender); + ISubstitute(position.asset).burnAll(msg.sender); } } diff --git a/contracts/SimpleBondController.sol b/contracts/SimpleBondController.sol index ff4448b..bb2fd82 100644 --- a/contracts/SimpleBondController.sol +++ b/contracts/SimpleBondController.sol @@ -46,9 +46,18 @@ contract SimpleBondController is IPositionLocker, ERC1155Holder, ISimpleBondCont _couponManager.setApprovalForAll(address(_couponWrapper), true); } - modifier wrapETH() { - if (address(this).balance > 0) _weth.deposit{value: address(this).balance}(); + modifier wrapAndRefundETH() { + bool hasMsgValue = address(this).balance > 0; + if (hasMsgValue) _weth.deposit{value: address(this).balance}(); _; + if (hasMsgValue) { + uint256 leftBalance = _weth.balanceOf(address(this)); + if (leftBalance > 0) { + _weth.withdraw(leftBalance); + (bool success,) = msg.sender.call{value: leftBalance}(""); + require(success); + } + } } function positionLockAcquired(bytes calldata data) external returns (bytes memory result) { @@ -62,7 +71,7 @@ contract SimpleBondController is IPositionLocker, ERC1155Holder, ISimpleBondCont (Coupon[] memory couponsToMint, Coupon[] memory couponsToBurn, int256 amountDelta) = _bondPositionManager.adjustPosition(tokenId, amount, expiredWith); if (amountDelta > 0) { - ISubstitute(asset).ensureThisBalance(user, uint256(amountDelta)); + ISubstitute(asset).ensureBalance(user, uint256(amountDelta)); IERC20(asset).approve(address(_bondPositionManager), uint256(amountDelta)); _bondPositionManager.depositToken(address(asset), uint256(amountDelta)); } else if (amountDelta < 0) { @@ -108,7 +117,7 @@ contract SimpleBondController is IPositionLocker, ERC1155Holder, ISimpleBondCont uint256 amount, Epoch expiredWith, bool wrapCoupons - ) internal wrapETH returns (uint256 positionId) { + ) internal wrapAndRefundETH returns (uint256 positionId) { address underlyingToken = ISubstitute(asset).underlyingToken(); permitParams.tryPermit(underlyingToken, msg.sender, address(this)); bytes memory result = @@ -147,7 +156,7 @@ contract SimpleBondController is IPositionLocker, ERC1155Holder, ISimpleBondCont uint256 amount, Epoch expiredWith, bool wrapCoupons - ) internal wrapETH { + ) internal wrapAndRefundETH { positionPermitParams.tryPermit(_bondPositionManager, tokenId, address(this)); couponPermitParams.tryPermit(_couponManager, msg.sender, address(this), true); address asset = _bondPositionManager.getPosition(tokenId).asset; diff --git a/contracts/libraries/Controller.sol b/contracts/libraries/Controller.sol index 9808e0b..d3783f7 100644 --- a/contracts/libraries/Controller.sol +++ b/contracts/libraries/Controller.sol @@ -24,6 +24,7 @@ import {IERC721Permit} from "../interfaces/IERC721Permit.sol"; import {ISubstitute} from "../interfaces/ISubstitute.sol"; import {IController} from "../interfaces/IController.sol"; import {ReentrancyGuard} from "./ReentrancyGuard.sol"; +import {SubstituteLibrary} from "./Substitute.sol"; import {Epoch} from "./Epoch.sol"; @@ -38,6 +39,7 @@ abstract contract Controller is using SafeERC20 for IERC20; using CouponKeyLibrary for CouponKey; using CouponLibrary for Coupon; + using SubstituteLibrary for ISubstitute; IWrapped1155Factory internal immutable _wrapped1155Factory; CloberMarketFactory internal immutable _cloberMarketFactory; @@ -55,9 +57,18 @@ abstract contract Controller is _weth = IWETH9(weth); } - modifier wrapETH() { - if (address(this).balance > 0) _weth.deposit{value: address(this).balance}(); + modifier wrapAndRefundETH() { + bool hasMsgValue = address(this).balance > 0; + if (hasMsgValue) _weth.deposit{value: address(this).balance}(); _; + if (hasMsgValue) { + uint256 leftBalance = _weth.balanceOf(address(this)); + if (leftBalance > 0) { + _weth.withdraw(leftBalance); + (bool success,) = msg.sender.call{value: leftBalance}(""); + require(success); + } + } } function _executeCouponTrade( @@ -97,7 +108,7 @@ abstract contract Controller is market.marketOrder(address(this), 0, 0, lastCoupon.amount, 2, data); } else { if (remainingInterest < 0) revert ControllerSlippage(); - _ensureBalance(token, user, amountToPay); + ISubstitute(token).ensureBalance(user, amountToPay); } } @@ -143,29 +154,6 @@ abstract contract Controller is return ISubstitute(substitute).underlyingToken(); } - function _burnAllSubstitute(address substitute, address to) internal { - uint256 leftAmount = IERC20(substitute).balanceOf(address(this)); - if (leftAmount == 0) return; - ISubstitute(substitute).burn(leftAmount, to); - } - - function _ensureBalance(address token, address user, uint256 amount) internal { - // TODO: consider to use SubstituteLibrary - address underlyingToken = ISubstitute(token).underlyingToken(); - uint256 thisBalance = IERC20(token).balanceOf(address(this)); - uint256 underlyingBalance = IERC20(underlyingToken).balanceOf(address(this)); - if (amount > thisBalance + underlyingBalance) { - unchecked { - IERC20(underlyingToken).safeTransferFrom(user, address(this), amount - thisBalance - underlyingBalance); - underlyingBalance = amount - thisBalance; - } - } - if (underlyingBalance > 0) { - IERC20(underlyingToken).approve(token, underlyingBalance); - ISubstitute(token).mint(underlyingBalance, address(this)); - } - } - function _wrapCoupons(Coupon[] memory coupons) internal { // wrap 1155 to 20 bytes memory metadata = Wrapped1155MetadataBuilder.buildWrapped1155BatchMetadata(coupons); diff --git a/contracts/libraries/Substitute.sol b/contracts/libraries/Substitute.sol index 1f5023c..46d1daf 100644 --- a/contracts/libraries/Substitute.sol +++ b/contracts/libraries/Substitute.sol @@ -10,21 +10,27 @@ import {ISubstitute} from "../interfaces/ISubstitute.sol"; library SubstituteLibrary { using SafeERC20 for IERC20; - function ensureThisBalance(ISubstitute substitute, address payer, uint256 amount) internal { + function ensureBalance(ISubstitute substitute, address payer, uint256 amount) internal { uint256 balance = IERC20(address(substitute)).balanceOf(address(this)); if (balance >= amount) { return; } - unchecked { - amount -= balance; - } - address underlyingToken = substitute.underlyingToken(); uint256 underlyingBalance = IERC20(underlyingToken).balanceOf(address(this)); - if (underlyingBalance < amount) { - IERC20(underlyingToken).safeTransferFrom(payer, address(this), amount - underlyingBalance); + unchecked { + amount -= balance; + if (underlyingBalance < amount) { + IERC20(underlyingToken).safeTransferFrom(payer, address(this), amount - underlyingBalance); + } } IERC20(underlyingToken).approve(address(substitute), amount); substitute.mint(amount, address(this)); } + + function burnAll(ISubstitute substitute, address to) internal { + uint256 leftAmount = IERC20(address(substitute)).balanceOf(address(this)); + if (leftAmount > 0) { + ISubstitute(substitute).burn(leftAmount, to); + } + } } diff --git a/test/foundry/unit/libraries/Substitute.sol b/test/foundry/unit/libraries/Substitute.t.sol similarity index 97% rename from test/foundry/unit/libraries/Substitute.sol rename to test/foundry/unit/libraries/Substitute.t.sol index 9c65f00..91bd605 100644 --- a/test/foundry/unit/libraries/Substitute.sol +++ b/test/foundry/unit/libraries/Substitute.t.sol @@ -25,7 +25,7 @@ contract SubstituteLibraryUnitTest is Test { IERC20(weth).approve(address(waweth), type(uint256).max); } - function testEnsureThisBalance() public { + function testEnsureBalance() public { _testEnsureThisBalance( Input({substitute: 1 ether, underlying: 1 ether, ensure: 0 ether}), Expected({thisSubstitute: 0, thisUnderlying: 0, wrapperSubstitute: 0, wrapperUnderlying: 0}) @@ -98,6 +98,6 @@ contract SubstituteLibraryUnitTest is Test { contract SubstituteLibraryWrapper { function ensureThisBalance(address substitute, address payer, uint256 amount) external { - SubstituteLibrary.ensureThisBalance(ISubstitute(substitute), payer, amount); + SubstituteLibrary.ensureBalance(ISubstitute(substitute), payer, amount); } }