Make ERC3156Wrapper more flexible (#44)

alcueca · May 6, 2024 · d55c154 · d55c154
1 parent d61ab3e
commit d55c154
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 7 deletions.
diff --git a/.forge-snapshots/MakerDAO.snap b/.forge-snapshots/MakerDAO.snap
@@ -1 +1 @@
-313172
+313302
diff --git a/.forge-snapshots/MorphoPendle.snap b/.forge-snapshots/MorphoPendle.snap
@@ -1 +1 @@
-558536
+558528
diff --git a/src/erc3156/ERC3156Wrapper.sol b/src/erc3156/ERC3156Wrapper.sol
@@ -29,13 +29,13 @@ contract ERC3156Wrapper is BaseWrapper, IERC3156FlashBorrower {
 
     /// @inheritdoc IERC7399
     function maxFlashLoan(address asset) external view returns (uint256) {
-        IERC3156FlashLender lender = lenders[asset];
+        IERC3156FlashLender lender = _lender(asset);
         return address(lender) != address(0) ? _maxFlashLoan(lender, asset) : 0;
     }
 
     /// @inheritdoc IERC7399
     function flashFee(address asset, uint256 amount) external view returns (uint256) {
-        IERC3156FlashLender lender = lenders[asset];
+        IERC3156FlashLender lender = _lender(asset);
         require(address(lender) != address(0), "Unsupported currency");
         return amount >= _maxFlashLoan(lender, asset) ? type(uint256).max : _flashFee(lender, asset, amount);
     }
@@ -52,15 +52,15 @@ contract ERC3156Wrapper is BaseWrapper, IERC3156FlashBorrower {
         returns (bytes32)
     {
         require(erc3156initiator == address(this), "External loan initiator");
-        require(msg.sender == address(lenders[asset]), "Unknown lender");
+        require(msg.sender == address(_lender(asset)), "Unknown lender");
 
         _bridgeToCallback(asset, amount, fee, params);
 
         return CALLBACK_SUCCESS;
     }
 
     function _flashLoan(address asset, uint256 amount, bytes memory data) internal override {
-        IERC3156FlashLender lender = lenders[asset];
+        IERC3156FlashLender lender = _lender(asset);
         require(address(lender) != address(0), "Unsupported currency");
 
         // We get funds from an ERC3156 lender to serve the ERC7399 flash loan in our ERC3156 callback
@@ -74,4 +74,9 @@ contract ERC3156Wrapper is BaseWrapper, IERC3156FlashBorrower {
     function _flashFee(IERC3156FlashLender lender, address asset, uint256 amount) internal view returns (uint256) {
         return lender.flashFee(asset, amount);
     }
+
+    function _lender(address asset) internal view returns (IERC3156FlashLender lender) {
+        lender = lenders[asset];
+        if (address(lender) == address(0)) lender = IERC3156FlashLender(asset);
+    }
 }
diff --git a/test/ERC3156Wrapper.t.sol b/test/ERC3156Wrapper.t.sol
@@ -3,7 +3,6 @@ pragma solidity >=0.8.19 <0.9.0;
 
 import { Test } from "forge-std/Test.sol";
 import { console2 } from "forge-std/console2.sol";
-import { StdCheats } from "forge-std/StdCheats.sol";
 import { IERC3156FlashLender } from "lib/erc3156/contracts/interfaces/IERC3156FlashLender.sol";
 
 import { IERC20Metadata as IERC20 } from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";