lol, tests pass on first try (#3)

Co-authored-by: alcueca <alberto@yield.is>

src/aave/AaveWrapper.sol

@@ -0,0 +1,94 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+// Thanks to ultrasecr.eth
+pragma solidity ^0.8.0;
+
+// contract AaveFlashLoanProvider is BaseFlashLoanProvider, FlashLoanReceiverBase {
+//     using Math for *;
+//     using SafeERC20 for IERC20;
+//     using ReserveConfiguration for DataTypes.ReserveConfigurationMap;
+//     using TransferLib for *;
+// 
+//     FlashLoanProvider public constant override id = FlashLoanProvider.Aave;
+// 
+//     AaveMoneyMarket public immutable moneyMarket;
+// 
+//     constructor(AaveMoneyMarket _moneyMarket) FlashLoanReceiverBase(_moneyMarket.provider()) {
+//         moneyMarket = _moneyMarket;
+//     }
+// 
+//     function calculateFlashLoanFeeAmount(IERC20Metadata asset, uint256 amount)
+//         external
+//         view
+//         returns (uint256 feeAmount)
+//     {
+//         DataTypes.ReserveData memory reserve = POOL.getReserveData(address(asset));
+//         DataTypes.ReserveConfigurationMap memory configuration = reserve.configuration;
+// 
+//         if (
+//             !configuration.getPaused() && configuration.getActive() && configuration.getFlashLoanEnabled()
+//                 && amount < asset.balanceOf(reserve.aTokenAddress)
+//         ) feeAmount = amount.mulDiv(fee(), 1e18, Math.Rounding.Up);
+//         else feeAmount = type(uint256).max;
+//     }
+// 
+//     function fee() public view override returns (uint256) {
+//         return POOL.FLASHLOAN_PREMIUM_TOTAL() * 0.0001e18;
+//     }
+// 
+//     function flashLoan(
+//         IERC20 asset,
+//         uint256 amount,
+//         address onBehalfOf,
+//         bytes calldata params,
+//         bool flashBorrow,
+//         function(IERC20, uint256, bytes memory, address) external returns (bytes memory) callback
+//     ) external override returns (bytes memory result) {
+//         tmpResult = "";
+//         MetaParams memory metaParams = MetaParams({params: params, flashBorrow: flashBorrow, callback: callback});
+// 
+//         if (flashBorrow) {
+//             moneyMarket.delegateIfNecessary(asset, onBehalfOf, address(this));
+//             // console.log("AAVE: flashBorrow %s", amount);
+//         }
+// 
+//         POOL.flashLoan({
+//             receiverAddress: address(this),
+//             assets: toArray(address(asset)),
+//             amounts: toArray(amount),
+//             interestRateModes: toArray(flashBorrow ? 2 : 0),
+//             onBehalfOf: flashBorrow ? onBehalfOf : address(this),
+//             params: abi.encode(metaParams),
+//             referralCode: 0
+//         });
+// 
+//         result = tmpResult;
+//     }
+// 
+//     function executeOperation(
+//         address[] calldata assets,
+//         uint256[] calldata amounts,
+//         uint256[] calldata fees,
+//         address initiator,
+//         bytes calldata params
+//     ) external override returns (bool) {
+//         require(msg.sender == address(POOL), "not pool");
+//         require(initiator == address(this), "AaveFlashLoanProvider: not initiator");
+// 
+//         MetaParams memory metaParams = abi.decode(params, (MetaParams));
+// 
+//         if (!metaParams.flashBorrow) IERC20(assets[0]).approveIfNecessary(address(POOL));
+// 
+//         IERC20(assets[0]).safeTransfer(metaParams.callback.address, amounts[0]);
+// 
+//         bytes memory result = metaParams.callback(
+//             IERC20(assets[0]),
+//             amounts[0] + (metaParams.flashBorrow ? 0 : fees[0]),
+//             metaParams.params,
+//             metaParams.flashBorrow ? address(0) : address(this)
+//         );
+// 
+//         if (result.length > 0) tmpResult = result;
+// 
+//         return true;
+//     }
+// }

src/balancer/BalancerWrapper.sol

@@ -0,0 +1,103 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+// Thanks to ultrasecr.eth
+pragma solidity ^0.8.0;
+
+import { IProtocolFeesCollector } from "./interfaces/IProtocolFeesCollector.sol";
+import { IFlashLoanRecipient } from "./interfaces/IFlashLoanRecipient.sol";
+import { IFlashLoaner } from "./interfaces/IFlashLoaner.sol";
+
+import { TransferHelper } from "../utils/TransferHelper.sol";
+import { FunctionCodec } from "../utils/FunctionCodec.sol";
+import { Arrays } from "../utils/Arrays.sol";
+
+import { IERC3156PPFlashLender } from "lib/erc3156pp/src/interfaces/IERC3156PPFlashLender.sol";
+import { FixedPointMathLib } from "lib/solmate/src/utils/FixedPointMathLib.sol";
+import { IERC20 } from "lib/erc3156pp/src/interfaces/IERC20.sol";
+
+
+contract BalancerWrapper is IFlashLoanRecipient, IERC3156PPFlashLender {
+    using TransferHelper for IERC20;
+    using FunctionCodec for function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory);
+    using FunctionCodec for bytes24;
+    using Arrays for uint256;
+    using Arrays for address;
+    using FixedPointMathLib for uint256;
+
+    bytes32 private flashLoanDataHash;
+    bytes internal _callbackResult;
+
+    IFlashLoaner public immutable balancer;
+
+    constructor(IFlashLoaner _balancer) {
+        balancer = _balancer;
+    }
+
+    function flashFee(IERC20 asset, uint256 amount)
+        external
+        view
+        returns (uint256 fee)
+    {
+        if (amount >= asset.balanceOf(address(balancer))) fee = type(uint256).max;
+        else fee = amount.mulWadUp(balancer.getProtocolFeesCollector().getFlashLoanFeePercentage());
+    }
+
+    /// @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) {
+        bytes memory data = abi.encode(msg.sender, loanReceiver, asset, amount, callback.encodeFunction(), initiatorData);
+
+        flashLoanDataHash = keccak256(data);
+        balancer.flashLoan(this, address(asset).toArray(), amount.toArray(), data);
+
+        bytes memory result = _callbackResult;
+        delete _callbackResult; // TODO: Confirm that this deletes the storage variable
+        return result;
+    }
+
+    function receiveFlashLoan(
+        address[] memory assets,
+        uint256[] memory amounts,
+        uint256[] memory fees,
+        bytes memory data
+    ) external override {
+        require(msg.sender == address(balancer), "not balancer");
+        require(keccak256(data) == flashLoanDataHash, "params hash mismatch");
+        delete flashLoanDataHash;
+
+        // decode data
+        (address initiator, address loanReceiver, IERC20 asset, uint256 amount, bytes24 encodedCallback, bytes memory initiatorData) = abi
+            .decode(data, (address, address, IERC20, uint256, bytes24, bytes));
+
+        function(address, address, IERC20, uint256, uint256, bytes memory) external returns (bytes memory) callback = encodedCallback.decodeFunction();
+
+        // send the borrowed amount to the loan receiver
+        asset.safeTransfer(address(loanReceiver), amount);
+
+        // call the callback and tell the calback receiver to pay to the balancer contract
+        // the callback result is kept in a storage variable to be retrieved later in this tx
+        _callbackResult = callback(initiator, msg.sender, IERC20(assets[0]), amounts[0], fees[0], initiatorData); // TODO: Skip the storage write if result.length == 0
+    }
+}

src/balancer/interfaces/IFlashLoanRecipient.sol

@@ -0,0 +1,21 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.0;
+
+
+interface IFlashLoanRecipient {
+    /**
+     * @dev When `flashLoan` is called on the Vault, it invokes the `receiveFlashLoan` hook on the recipient.
+     *
+     * At the time of the call, the Vault will have transferred `amounts` for `tokens` to the recipient. Before this
+     * call returns, the recipient must have transferred `amounts` plus `feeAmounts` for each token back to the
+     * Vault, or else the entire flash loan will revert.
+     *
+     * `userData` is the same value passed in the `IVault.flashLoan` call.
+     */
+    function receiveFlashLoan(
+        address[] memory tokens,
+        uint256[] memory amounts,
+        uint256[] memory feeAmounts,
+        bytes memory userData
+    ) external;
+}

src/balancer/interfaces/IFlashLoaner.sol

@@ -0,0 +1,15 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.0;
+import { IFlashLoanRecipient } from "./IFlashLoanRecipient.sol";
+import { IProtocolFeesCollector } from "./IProtocolFeesCollector.sol";
+
+interface IFlashLoaner {
+    function flashLoan(
+        IFlashLoanRecipient recipient,
+        address[] memory tokens,
+        uint256[] memory amounts,
+        bytes memory userData
+    ) external;
+
+    function getProtocolFeesCollector() external view returns (IProtocolFeesCollector);
+}

src/balancer/interfaces/IProtocolFeesCollector.sol

@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+pragma solidity ^0.8.0;
+
+interface IProtocolFeesCollector {
+    function getFlashLoanFeePercentage() external view returns (uint256);
+}

src/utils/Arrays.sol

@@ -0,0 +1,16 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+// Thanks to ultrasecr.eth
+pragma solidity ^0.8.0;
+
+
+library Arrays {
+    function toArray(uint256 n) external pure returns (uint256[] memory arr) {
+        arr = new uint[](1);
+        arr[0] = n;
+    }
+
+    function toArray(address a) external pure returns (address[] memory arr) {
+        arr = new address[](1);
+        arr[0] = a;
+    }
+}

test/BalancerWrapper.t.sol

@@ -0,0 +1,63 @@
+// SPDX-License-Identifier: UNLICENSED
+pragma solidity >=0.8.19 <0.9.0;
+
+import { PRBTest } from "@prb/test/PRBTest.sol";
+import { console2 } from "forge-std/console2.sol";
+import { StdCheats } from "forge-std/StdCheats.sol";
+
+import { IFlashLoaner } from "../src/balancer/interfaces/IFlashLoaner.sol";
+import { FlashBorrower } from "../src/test/FlashBorrower.sol";
+import { IERC20, BalancerWrapper } from "../src/balancer/BalancerWrapper.sol";
+
+
+/// @dev If this is your first time with Forge, read this tutorial in the Foundry Book:
+/// https://book.getfoundry.sh/forge/writing-tests
+contract BalancerWrapperTest is PRBTest, StdCheats {
+    BalancerWrapper internal wrapper;
+    FlashBorrower internal borrower;
+    IERC20 internal dai;
+    IFlashLoaner internal balancer;
+
+    /// @dev A function invoked before each test case is run.
+    function setUp() public virtual {
+        // Revert if there is no API key.
+        string memory alchemyApiKey = vm.envOr("API_KEY_ALCHEMY", string(""));
+        if (bytes(alchemyApiKey).length == 0) {
+            revert("API_KEY_ALCHEMY variable missing");
+        }
+
+        vm.createSelectFork({ urlOrAlias: "mainnet", blockNumber: 16_428_000 });
+        balancer = IFlashLoaner(0xBA12222222228d8Ba445958a75a0704d566BF2C8);
+        dai = IERC20(0x6B175474E89094C44Da98b954EedeAC495271d0F);
+
+        wrapper = new BalancerWrapper(balancer);
+        borrower = new FlashBorrower(wrapper);
+        deal(address(dai), address(this), 1e18); // For fees
+    }
+
+    /// @dev Basic test. Run it with `forge test -vvv` to see the console log.
+    function test_flashFee() external {
+        console2.log("test_flashFee");
+        assertEq(wrapper.flashFee(dai, 1e18), 0, "Fee not zero");
+        assertEq(wrapper.flashFee(dai, type(uint256).max), type(uint256).max, "Fee not max");
+    }
+
+    function test_flashLoan() external {
+        console2.log("test_flashLoan");
+        uint256 loan = 1e18;
+        uint256 fee = wrapper.flashFee(dai, loan);
+        dai.transfer(address(borrower), fee);
+        bytes memory result = borrower.flashBorrow(dai, loan);
+
+        // Test the return values
+        (bytes32 callbackReturn) = abi.decode(result, (bytes32));
+        assertEq(uint256(callbackReturn), uint256(borrower.ERC3156PP_CALLBACK_SUCCESS()), "Callback failed");
+
+        // Test the borrower state
+        assertEq(borrower.flashInitiator(), address(borrower));
+        assertEq(address(borrower.flashAsset()), address(dai));
+        assertEq(borrower.flashAmount(), loan);
+        assertEq(borrower.flashBalance(), loan + fee); // The amount we transferred to pay for fees, plus the amount we borrowed
+        assertEq(borrower.flashFee(), fee);
+    }
+}