WIP: Code tweaks

FlashLender.t.sol

@@ -6,7 +6,7 @@ import { console2 } from "forge-std/console2.sol";
 import { StdCheats } from "forge-std/StdCheats.sol";
 
 import { FlashLender } from "../src/FlashLender.sol";
-import { FlashBorrower } from "../src/test/FlashBorrower.sol";
+import { FlashBorrower } from "./FlashBorrower.sol";
 import { ERC20Mock } from "../src/test/ERC20Mock.sol";
 import { IERC20 } from "../src/interfaces/IERC20.sol";
 

script/Deploy.s.sol

@@ -1,13 +1,9 @@
 // SPDX-License-Identifier: UNLICENSED
 pragma solidity >=0.8.19 <=0.9.0;
 
-import { Foo } from "../src/Foo.sol";
-
 import { BaseScript } from "./Base.s.sol";
 
 /// @dev See the Solidity Scripting tutorial: https://book.getfoundry.sh/tutorials/solidity-scripting
 contract Deploy is BaseScript {
-    function run() public broadcast returns (Foo foo) {
-        foo = new Foo();
-    }
+    function run() public broadcast { }
 }

src/BaseWrapper.sol

@@ -0,0 +1,88 @@
+// SPDX-License-Identifier: MIT
+// Thanks to ultrasecr.eth
+pragma solidity ^0.8.0;
+
+import { IERC3156PPFlashLender } from "lib/erc3156pp/src/interfaces/IERC3156PPFlashLender.sol";
+import { IERC20 } from "lib/erc3156pp/src/interfaces/IERC20.sol";
+
+import { TransferHelper } from "./utils/TransferHelper.sol";
+
+abstract contract BaseWrapper is IERC3156PPFlashLender {
+    using TransferHelper for IERC20;
+
+    struct Data {
+        address loanReceiver;
+        address initiator;
+        function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory) callback;
+        bytes initiatorData;
+    }
+
+    bytes internal _callbackResult;
+
+    /// @inheritdoc IERC3156PPFlashLender
+    function flashLoan(
+        address loanReceiver,
+        IERC20 asset,
+        uint256 amount,
+        bytes calldata initiatorData,
+        function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory) callback
+    )
+        external
+        returns (bytes memory result)
+    {
+        Data memory data = Data({
+            loanReceiver: loanReceiver,
+            initiator: msg.sender,
+            callback: callback,
+            initiatorData: initiatorData
+        });
+
+        _flashLoan(asset, amount, abi.encode(data));
+
+        result = _callbackResult;
+        // Avoid storage write if not needed
+        if (result.length > 0) {
+            delete _callbackResult;
+        }
+        return result;
+    }
+
+    /// @dev Call the flashloan function in the child contract
+    function _flashLoan(IERC20 asset, uint256 amount, bytes memory params) internal virtual;
+
+    /// @dev Handle the common parts of bridging the callback
+    function bridgeToCallback(IERC20 asset, uint256 amount, uint256 fee, bytes memory params) internal {
+        Data memory data = abi.decode(params, (Data));
+        _transferAssets(asset, amount, data.loanReceiver);
+
+        // call the callback and tell the callback receiver to repay the loan to this contract
+        bytes memory result = data.callback(data.initiator, _repayTo(), IERC20(asset), amount, fee, data.initiatorData);
+
+        _approveRepayment(asset, amount, fee);
+
+        if (result.length > 0) {
+            // if there's any result, it is kept in a storage variable to be retrieved later in this tx
+            _callbackResult = result;
+        }
+    }
+
+    /// @dev Transfer the assets to the loan receiver.
+    /// Override it if the provider can send the funds directly
+    function _transferAssets(IERC20 asset, uint256 amount, address loanReceiver) internal virtual {
+        asset.safeTransfer(loanReceiver, amount);
+    }
+
+    /// @dev Approve the repayment of the loan to the provider if needed.
+    /// Override it if the provider can receive the funds directly and you want to avoid the if condition
+    function _approveRepayment(IERC20 asset, uint256 amount, uint256 fee) internal virtual {
+        if (_repayTo() == address(this)) {
+            asset.approve(msg.sender, amount + fee);
+        }
+    }
+
+    /// @dev Where should the end client send the funds to repay the loan
+    /// Override it if the provider can receive the funds directly
+    function _repayTo() internal view virtual returns (address) {
+        return address(this);
+    }
+}

src/aave/AaveWrapper.sol

@@ -1,6 +1,6 @@
-// SPDX-License-Identifier: GPL-3.0-or-later
+// SPDX-License-Identifier: MIT
 // Thanks to ultrasecr.eth
 pragma solidity ^0.8.0;
 
 import { IPool } from "./interfaces/IPool.sol";
 import { DataTypes } from "./interfaces/DataTypes.sol";
@@ -8,27 +8,17 @@
 import { IPoolAddressesProvider } from "./interfaces/IPoolAddressesProvider.sol";
 import { IFlashLoanSimpleReceiver } from "./interfaces/IFlashLoanSimpleReceiver.sol";
 
-import { FunctionCodec } from "../utils/FunctionCodec.sol";
-import { TransferHelper } from "../utils/TransferHelper.sol";
-
 import { IERC20 } from "lib/erc3156pp/src/interfaces/IERC20.sol";
 import { FixedPointMathLib } from "lib/solmate/src/utils/FixedPointMathLib.sol";
-import { IERC3156PPFlashLender } from "lib/erc3156pp/src/interfaces/IERC3156PPFlashLender.sol";
 
-import { console2 } from "forge-std/console2.sol";
+import { BaseWrapper } from "../BaseWrapper.sol";
 
-contract AaveWrapper is IERC3156PPFlashLender, IFlashLoanSimpleReceiver {
-    using TransferHelper for IERC20;
-    using
-    FunctionCodec
-    for function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory);
-    using FunctionCodec for bytes24;
+contract AaveWrapper is BaseWrapper, IFlashLoanSimpleReceiver {
     using FixedPointMathLib for uint256;
     using ReserveConfiguration for DataTypes.ReserveConfigurationMap;
 
-    bytes internal _callbackResult;
     IPoolAddressesProvider public immutable ADDRESSES_PROVIDER;
     IPool public POOL;
 
     constructor(IPoolAddressesProvider provider) {
         ADDRESSES_PROVIDER = provider;
@@ -39,6 +29,13 @@
         POOL = IPool(ADDRESSES_PROVIDER.getPool());
     }
 
+    /**
+     * @dev From ERC-3156. The fee to be charged for a given loan.
+     * @param asset The loan currency.
+     * @param amount The amount of assets lent.
+     * @return fee The amount of `asset` to be charged for the loan, on top of the returned principal.
+     * type(uint256).max if the loan is not possible.
+     */
     function flashFee(IERC20 asset, uint256 amount) external view returns (uint256 fee) {
         DataTypes.ReserveData memory reserve = POOL.getReserveData(address(asset));
         DataTypes.ReserveConfigurationMap memory configuration = reserve.configuration;
@@ -50,95 +47,32 @@
         else fee = type(uint256).max;
     }
 
-    /// @dev Use the aggregator to serve an ERC3156++ flash loan.
-    /// @dev Forward the callback to the callback receiver. The borrower only needs to trust the aggregator and its
-    /// governance, instead of the underlying lenders.
-    /// @param loanReceiver The address receiving the flash loan
-    /// @param asset The asset to be loaned
-    /// @param amount The amount to loaned
-    /// @param initiatorData The ABI encoded initiator data
-    /// @param callback The address and signature of the callback function
-    /// @return result ABI encoded result of the callback
-    function flashLoan(
-        address loanReceiver,
-        IERC20 asset,
-        uint256 amount,
-        bytes calldata initiatorData,
-        /// @dev callback.
-        /// This is a concatenation of (address, bytes4), where the address is the callback receiver, and the bytes4 is
-        /// the signature of callback function.
-        /// The arguments in the callback function are fixed.
-        /// If the callback receiver needs to know the loan receiver, it should be encoded by the initiator in `data`.
-        /// @param initiator The address that called this function
-        /// @param paymentReceiver The address that needs to receive the amount plus fee at the end of the callback
-        /// @param asset The asset to be loaned
-        /// @param amount The amount to loaned
-        /// @param fee The fee to be paid
-        /// @param data The ABI encoded data to be passed to the callback
-        /// @return result ABI encoded result of the callback
-        function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory) callback
-    )
-        external
-        returns (bytes memory result)
-    {
-        bytes memory data = abi.encode(msg.sender, loanReceiver, callback.encodeFunction(), initiatorData);
-
+    function _flashLoan(IERC20 asset, uint256 amount, bytes memory data) internal override {
         POOL.flashLoanSimple({
             receiverAddress: address(this),
             asset: address(asset),
             amount: amount,
             params: data,
             referralCode: 0
         });
-
-        result = _callbackResult;
-        // Avoid storage write if not needed
-        if (result.length > 0) {
-            delete _callbackResult;
-        }
-        return result;
     }
 
+    /// @inheritdoc IFlashLoanSimpleReceiver
     function executeOperation(
         address asset,
         uint256 amount,
         uint256 fee,
-        address aaveInitiator,
-        bytes calldata data
+        address initiator,
+        bytes calldata params
     )
         external
         override
         returns (bool)
     {
-        console2.log("executeOperation");
-        require(msg.sender == address(POOL), "not pool");
-        require(aaveInitiator == address(this), "AaveFlashLoanProvider: not initiator");
-
-        address initiator;
-        bytes memory initiatorData;
-        function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory) callback;
-        {
-            address loanReceiver;
-            bytes24 encodedCallback;
-
-            // decode data
-            console2.log("abi decoding...");
-            (initiator, loanReceiver, encodedCallback, initiatorData) =
-                abi.decode(data, (address, address, bytes24, bytes));
-            console2.log("callback decoding...");
-            callback = encodedCallback.decodeFunction();
-
-            IERC20(asset).approve(address(POOL), amount + fee);
-            IERC20(asset).safeTransfer(loanReceiver, amount);
-        } // release loanReceiver and encodedCallback from the stack
-
-        // call the callback and tell the calback receiver to repay the loan to this contract
-        bytes memory result = callback(initiator, address(this), IERC20(asset), amount, fee, initiatorData);
+        require(msg.sender == address(POOL), "AaveFlashLoanProvider: not pool");
+        require(initiator == address(this), "AaveFlashLoanProvider: not initiator");
 
-        if (result.length > 0) {
-            // if there's any result, it is kept in a storage variable to be retrieved later in this tx
-            _callbackResult = result;
-        }
+        bridgeToCallback(IERC20(asset), amount, fee, params);
 
         return true;
     }