Implementing a Flash Loan Vault with ERC3156 and ERC4626

Flashloan Vault

If you've followed our previous article on implementing Vault with ERC4626, then you are already aware of how easy it is to implement a token vault with a profit sharing distribution mechanism.

Building on this knowledge we will implement a flashloan lending facility into our vault allowing our shareholding depositors to earn additional yield from the flashloan fees.

Implementing flashloan functionality is standardized by ERC3156 which defines the interface for for both the flashloan lender and the flashloan borrower.

We will again use the OpenZeppelin library as scaffold to implement the flashloan functionality.

Getting Started

FlashLoanVault.sol

import { ERC4626 } from "@openzeppelin/contracts/token/ERC20/extensions/ERC4626.sol";
import { ERC20, IERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import {IERC3156FlashBorrower} from "@openzeppelin/contracts/interfaces/IERC3156FlashBorrower.sol";
import {IERC3156FlashLender} from "@openzeppelin/contracts/interfaces/IERC3156FlashLender.sol";
import { SafeTransferLib } from "solady/utils/SafeTransferLib.sol";
error FlashLoanVault__WrongToken();
error FlashLoanVault__RepayFailed();
error FlashLoanVault__MaxLoanExceeded();

The first thing we need to do is import the necessary libraries and interfaces. We'll need the ERC4626 contract in order to implement the tokenized vault shares. Additionally we need to import both IERC3156FlashBorrower and IERC3156FlashLender interfaces from the OpenZeppelin library.

In addition to inheriting, we are also defining some custom errors that we will be using later on.

We'll also be using the SafeTransferLib from solady in order to confidently handle token transfers. To install solady simply run:

forge install vectorized/solady

and in your foundry.toml add a remapping the import with:

remappings = [
   '@solady/contracts/=lib/solady/src/',
]

save your remappings with:

forge remappings > remappings.txt

Once we have that, we can define our FlashLoanVault contract by inheriting the token vault and the flashlender interface in addition to declaring our usage of the SafeTransferLib.

contract FlashLoanVault is ERC4626, IERC3156FlashLender {
    using SafeTransferLib for IERC20;
    constructor(address _depositToken) ERC4626(IERC20(_depositToken)) ERC20("Shares Vault", "SV") {

    }
}

Now, if you've ready our previous article about implementing a vault with ERC4626, you'll notice that we wont be implementing the shareProfits function in this contract. Our flashlending protocol will be able to receive its fees directly and they will automatically and immediately be available to our shareholder.

In order to implement an ERC3156 compliant flash lender we need to implement three simple functions: maxFlashLoan(address token), flashFee(address token, uint256 amount) and flashLoan(IERC3156FlashBorrower receiver, address token, uint256 amount, bytes calldata data).

Firstly, we are only going to allow flashloans of the deposit token that is being pooled in our vault. So we are going to enforce this in the maxFlashLoan function by reverting if the token requested is not the deposit token. Otherwise, the max flashloan allowed will be the total balance of the deposit tokens in the vault.

    function maxFlashLoan(address token) public view override returns (uint256) {
        if (token != address(asset())) revert FlashLoanVault__WrongToken();
        return IERC20(token).balanceOf(address(this));
    }

Next, we need to implement the flashFee function which will return the fee that will be charged for our flashloan. Here, we will be using a flat rate (rather expensive) fee of 0.1 ether, regardless of the amount of the loan.

    function flashFee(address token, uint256 amount) public view override returns (uint256) {
        return 0.1 ether;
    }

Finally we can implement the flashLoan function which will be called when a flashloan is initiated. The caller must pass a compliant IERC3156FlashBorrower as the loan receiver, note: only a smart contract can receive a flashloan since it must contain the appropriate logic to pay the loan back within the same transaction. Next, our borrower must specify the token and amount of the loan, and any additional data that the borrower may need to process the loan in the receiver contract, this data is simply passed along to the receiver contract after the loan is distributed.

    function flashLoan(
        IERC3156FlashBorrower receiver,
        address token,
        uint256 amount,
        bytes calldata data
    ) external override returns (bool)
    {
        if (amount > maxFlashLoan(token)) revert FlashLoanVault__MaxLoanExceeded();

        uint256 fee = flashFee(token, amount);
        uint256 totalDebt = amount + fee;
        SafeTransferLib.safeTransfer(address(asset()), address(receiver), amount);
        if (receiver.onFlashLoan(msg.sender, token, amount, fee, data) != keccak256("IERC3156FlashBorrower.onFlashLoan")) revert FlashLoanVault__RepayFailed();

        SafeTransferLib.safeTransferFrom(address(asset()), address(receiver), address(this), totalDebt);
        return true;
    }

Here we are simply checking that the flashloan amount requested is within the maximum allowed, then we calculate the fee and cache the amount we expect to be paid back, loan + fee, in the totalDebt variable.

After, that we can transfer the loan to the receiver. Once the receiver has the loan our lender contract will call the receivers onFlashLoan() function this is where the receiver can implement any logic desired utilizing the loan funds as long as it does two things: 1. Results in the funds being transferred back to the lender at the end of the transaction and 2. Returns the keccak256 hash of the IERC3156FlashBorrower.onFlashLoan function signature. If the receiver fails to do either of these things the transaction will revert and the loan will not be completed.

Once, we've called the onFlashLoan function we can call safetTransferFrom on the receiver for the amount of the loan + fee, and the flashloan is repaid.

Note: since we've implemented the repayment using safeTransferFrom any receiver contract wishing to borrow from our vault must execute an approve function granting the vault an allowance to transfer the total debt amount by the end of the transaction.

Secondly, we've used the SafeTransferLib here to handle the token transfers, the solady SafeTransferLib is an optimized gas efficient implementation that allows us to be sure the whole transaction will revert if the loan repayment fails.

Testing

Now that we have implemented our flashloan lender vault, we can write some tests to have a look at how it works and how to interact with it.

We'll start with the basic foundry test contract. We'll also import the FlashLoanVault contract and the IERC3156FlashBorrower interface as well as the ERC20 contracts. All of which we will need for testing. Then, we can setup some test users and declare the custom errors we defined in the FlashLoanVault contract, which we will also want to test.

FlashLoanVaultTest.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.20;

import {Test, console2} from "forge-std/Test.sol";
import { FlashLoanVault } from "../src/FlashLoanVault.sol";
import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import { IERC3156FlashBorrower } from "@openzeppelin/contracts/interfaces/IERC3156FlashBorrower.sol";

contract FlashLoanVaultTest is Test {
    address OwnerWallet;
    address user1;
    address user2;
    address user3;
    FlashLoanVault vault;

    error FlashLoanVault__WrongToken();
    error FlashLoanVault__RepayFailed();
    error FlashLoanVault__MaxLoanExceeded();

    function setUp() public {
        OwnerWallet = address(69);

        user1 = address(420);
        user2 = address(666);
        user3 = address(777);
    }

}

Once we have the basics, we notice that we will need some extra setup in order to test all our functionality. The first thing we'll need is a deposit token to use in our vault. So we will create a MockUSDC token similar to what we did in our ERC4626 vault article. At the bottom of our test file we can just define a simple mock token contract.

FlashLoanVaultTest.t.sol

// FlashLoanVaultTest.t.sol
contract MockUSDC is ERC20 {
    constructor() ERC20("USDC", "USDC") {
    }
    function mint(address to, uint256 amount) public {
        _mint(to, amount);
    }
}
`

Additionally, we will need a Mock Borrower to test our flashloans. The ERC3156FlashBorrower interface only needs to implement the onFlashLoan function in order to be compliant. Our mock borrower will take the lender address as a constructor argument to ensure that it knows to where to repay the loan. Note: a this only for testing purposes, a real flash borrower would need to implement further checks to ensure security.

contract MockFlashLoanReceiver is IERC3156FlashBorrower, Test{
    address public lender;

    constructor(address _lender) {
        lender = _lender;
    }
    //
    function onFlashLoan(address initiator, address token, uint256 amount, uint256 fee, bytes calldata data) external override returns (bytes32) {
        // check we have received the loan
        assertEq(IERC20(token).balanceOf(address(this)), amount+fee);
        // Hypothetically do something with the loan here
        // approve the repayment 
        IERC20(token).approve(msg.sender, amount + fee);
        // return this signature so the lender knows we are a compliant borrower
        return keccak256("IERC3156FlashBorrower.onFlashLoan");
    }
}

Normally, this is where the borrower will implement its own custom logic to make use of the loan funds. In our case we will simply want to test that we have received the appropriate funds and have enough to back the loan with an assertion. For this reason, we have inherited the Test base contract from from foundry in order to use the assertEq function. After that we approved the repayment of funds and returned the appropriate selector hash to signal to the lender that we are a compliant borrower.

Now that we have the appropriate mocks, we can return to our test setup function and create our flash loan protocol.

FlashLoanVaultTest.t.sol

// FlashLoanVaultTest.t.sol
    function setUp() public {
        OwnerWallet = address(69);

        user1 = address(420);
        user2 = address(666);
        user3 = address(777);

        usdc = new MockUSDC();
        usdc.mint(user1, 10 ether);
        usdc.mint(user2, 10 ether);
        usdc.mint(OwnerWallet, 10 ether);
        vault = new FlashLoanVault(address(usdc));

        vm.startPrank(OwnerWallet);
        usdc.approve(address(vault), 10 ether); 
        vault.deposit(7 ether, OwnerWallet);
        vm.stopPrank();

        vm.startPrank(user2);
        usdc.approve(address(vault), 10 ether);
        vault.deposit(3 ether, user2);
        vm.stopPrank();

        vm.startPrank(user1);
        borrower = new MockFlashLoanReceiver(address(vault));
        vm.stopPrank();
    }

All we've done here is setup our user wallets with some USDC and then created a vault and deposited some funds into it. We now have a vault with 10 USDC and two users with shares equivalent to 1 USDC each. We have also deployed a borrower contract targeting our vault as the lender.

We are ready to test the flashloan with the test_flashloan() function.

    function test_flashLoan() public {
        vm.startPrank(user1);
        usdc.transfer(address(borrower), 0.1 ether);
        vault.flashLoan(borrower, address(usdc), 10, "");
        assertEq(usdc.balanceOf(address(vault)), 10.1 ether);
        assertEq(usdc.balanceOf(address(borrower)), 0);
    }

In this test, we've had our test user1 transfer 0.1 USDC to the borrower contract to use as the loan fee. Then we executed a flashloan on the fault. Our borrower contract will test that it has received the loan by calling assertEq mid transaction in the onFlashLoan function. After the flashloan call completes, our test checks that the vault balance has increased by 0.1 USDC and maintained its initial balance of 10 USDC as well as checking that the borrower contract is now empty. All of which are indicative of a successful flashloan.

Testing the Shareholder's Profits

Now that we have a basic flashloan working, we can run a function that executes several flashloans thus accumulating profits to the vault and check that the shareholder profits are being distributed correctly.

    function test_ten_flashloans() public {
        vm.startPrank(user1);
        for (uint i = 0; i < 10; i++) {
            usdc.transfer(address(borrower), 0.1 ether);
            vault.flashLoan(borrower, address(usdc), 1, "");
            assertEq(usdc.balanceOf(address(borrower)), 0);
        }
        vm.stopPrank();
        assertEq(usdc.balanceOf(address(vault)), 11 ether);
        assertApproxEqAbs(vault.previewRedeem(7 ether), 7.7 ether, 1);
        assertApproxEqAbs(vault.previewRedeem(3 ether), 3.3 ether, 1);
    }

In this test, we've use a simple loop to do 10 flashloans which should result in 1 USDC of profit to the vault. We then check that the vault has accumulated 11 USDC and that the shareholder profits are being distributed correctly according to their holding. Since, the vault has profited 10% on its 10 USDC of liquidity, the shareholders should be able to redeem 10% more than their initial deposit for each share.

Note here that we are using assertApproxEqAbs to check the shareholder redeem value are within 1 unit of what we expect. This is due to the rounding precision effects of integer division which was discussed in the shareholder vault article.

Testing Errors

Finally, we can test that our custom errors are being thrown when they should be. Namely, when a flashloan is requested for a token that is not the vault deposit token and when a flashloan is requested that exceeds the amount of liquidity available in the vault.

    function test_flashloan_revert_wrong_token() public {
        vm.startPrank(user1);
        vm.expectRevert(abi.encodeWithSelector(FlashLoanVault__WrongToken.selector));
        vault.flashLoan(borrower, address(vault), 10, "");
    }
    function test_revertMaxLoanExceeded() public {
        vm.startPrank(user1);
        vm.expectRevert(abi.encodeWithSelector(FlashLoanVault__MaxLoanExceeded.selector));
        vault.flashLoan(borrower, address(usdc), 100 ether, "");
    }

To implement these checks, we utilize vm.expectRevert and compare the returned error data to the selector of the custom error we declared earlier.

Testing an Attack

Finally, we can test that our flashloan protocol is secure by attempting to attack it. We can do this by creating a malicious borrower contract that does not repay the loan. We will first need to define an attacker contract that attempts to take the loan and transfer it directly to the attacker. at the bottom of our test file we can define the attacker contract.

contract FlashLoanAttacker is IERC3156FlashBorrower{
    constructor(){}
    function onFlashLoan(address initiator, address token, uint256 amount, uint256 fee, bytes calldata data) external override returns (bytes32) {
        // try to send the loan to ourselves
        IERC20(token).transfer(initiator, amount + fee); 
        return keccak256("IERC3156FlashBorrower.onFlashLoan");
    }
}

Finally, in our test we can have the attacker call for a flashloan and expect it to revert.

    function test_revert_loan_repay_failed() public {
        vm.startPrank(user1);
        FlashLoanAttacker attacker = new FlashLoanAttacker();
        usdc.transfer(address(attacker), 0.1 ether);
        vm.expectRevert();
        vault.flashLoan(attacker, address(usdc), 10, "");
    }

In this test, we've used a general vm.expectRevert() since we are expecting the SafeTransferLib to revert when it fails to safeTrasferFrom the attacker contract in the appropriate amount to pay back the loan and fee.

Conclusion

Here we have seen the relative ease with which a flashloan protocol can be implemented utilizing the ERC3156 interface to provide standardized flashloan functionality and the ERC4626 interface to provide a simple profit distributing vault mechanism for depositors to provide pooled capital and receive a yield.

Full Code

FlashLoanVault.sol

import { ERC4626 } from "@openzeppelin/contracts/token/ERC20/extensions/ERC4626.sol";
import {IERC3156FlashBorrower} from "@openzeppelin/contracts/interfaces/IERC3156FlashBorrower.sol";
import {IERC3156FlashLender} from "@openzeppelin/contracts/interfaces/IERC3156FlashLender.sol";
import { SafeTransferLib } from "solady/utils/SafeTransferLib.sol";
import { ERC20, IERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
error FlashLoanVault__WrongToken();
error FlashLoanVault__RepayFailed();
error FlashLoanVault__MaxLoanExceeded();
contract FlashLoanVault is ERC4626, IERC3156FlashLender {
    using SafeTransferLib for IERC20;
    constructor(address _depositToken) ERC4626(IERC20(_depositToken)) ERC20("Shares Vault", "SV") {

    }
    function maxFlashLoan(address token) public view override returns (uint256) {
        if (token != address(asset())) revert FlashLoanVault__WrongToken();
        return IERC20(token).balanceOf(address(this));
    }

    function flashFee(address token, uint256 amount) public view override returns (uint256) {
        return 0.1 ether;
    }

    function flashLoan(
        IERC3156FlashBorrower receiver,
        address token,
        uint256 amount,
        bytes calldata data
    ) external override returns (bool)
    {
        if (amount > maxFlashLoan(token)) revert FlashLoanVault__MaxLoanExceeded();

        uint256 fee = flashFee(token, amount);
        uint256 totalDebt = amount + fee;
        SafeTransferLib.safeTransfer(address(asset()), address(receiver), amount);
        if (receiver.onFlashLoan(msg.sender, token, amount, fee, data) != keccak256("IERC3156FlashBorrower.onFlashLoan")) revert FlashLoanVault__RepayFailed();

        SafeTransferLib.safeTransferFrom(address(asset()), address(receiver), address(this), totalDebt);
        return true;
    }


}

FlashLoanVaultTest.t.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.20;

import {Test, console2} from "forge-std/Test.sol";
import { FlashLoanVault } from "../src/FlashLoanVault.sol";
import { ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import { IERC3156FlashBorrower } from "@openzeppelin/contracts/interfaces/IERC3156FlashBorrower.sol";

contract FlashLoanVaultTest is Test {
    address OwnerWallet;
    address user1;
    address user2;
    address user3;
    FlashLoanVault vault;
    MockFlashLoanReceiver borrower;
    MockUSDC usdc;

    error FlashLoanVault__WrongToken();
    error FlashLoanVault__RepayFailed();
    error FlashLoanVault__MaxLoanExceeded();

    function setUp() public {
        OwnerWallet = address(69);

        user1 = address(420);
        user2 = address(666);
        user3 = address(777);

        usdc = new MockUSDC();
        usdc.mint(user1, 10 ether);
        usdc.mint(user2, 10 ether);
        usdc.mint(OwnerWallet, 10 ether);
        vault = new FlashLoanVault(address(usdc));

        vm.startPrank(OwnerWallet);
        usdc.approve(address(vault), 10 ether); 
        vault.deposit(7 ether, OwnerWallet);
        vm.stopPrank();

        vm.startPrank(user2);
        usdc.approve(address(vault), 10 ether);
        vault.deposit(3 ether, user2);
        vm.stopPrank();

        vm.startPrank(user1);
        borrower = new MockFlashLoanReceiver(address(vault));
        vm.stopPrank();
    }

    function test_flashLoan() public {
        vm.startPrank(user1);
        usdc.transfer(address(borrower), 0.1 ether);
        vault.flashLoan(borrower, address(usdc), 10, "");
        assertEq(usdc.balanceOf(address(vault)), 10.1 ether);
        assertEq(usdc.balanceOf(address(borrower)), 0);
    }
    function test_ten_flashloans() public {
        vm.startPrank(user1);
        for (uint i = 0; i < 10; i++) {
            usdc.transfer(address(borrower), 0.1 ether);
            vault.flashLoan(borrower, address(usdc), 1, "");
            assertEq(usdc.balanceOf(address(borrower)), 0);
        }
        vm.stopPrank();
        assertEq(usdc.balanceOf(address(vault)), 11 ether);
        assertApproxEqAbs(vault.previewRedeem(7 ether), 7.7 ether, 1);
        assertApproxEqAbs(vault.previewRedeem(3 ether), 3.3 ether, 1);
    }

    function test_revert_loan_repay_failed() public {
        vm.startPrank(user1);
        FlashLoanAttacker attacker = new FlashLoanAttacker();
        usdc.transfer(address(attacker), 0.1 ether);
        vm.expectRevert();
        vault.flashLoan(attacker, address(usdc), 10, "");
    }
}

contract MockFlashLoanReceiver is IERC3156FlashBorrower, Test{
    address public lender;

    constructor(address _lender) {
        lender = _lender;
    }
    //
    function onFlashLoan(address initiator, address token, uint256 amount, uint256 fee, bytes calldata data) external override returns (bytes32) {
        // check we have received the loan
        assertEq(IERC20(token).balanceOf(address(this)), amount+fee);
        // Hypothetically do something with the loan
        // approve the repayment 
        IERC20(token).approve(msg.sender, amount + fee);
        return keccak256("IERC3156FlashBorrower.onFlashLoan");
    }
}

contract FlashLoanAttacker is IERC3156FlashBorrower{
    constructor(){}
    function onFlashLoan(address initiator, address token, uint256 amount, uint256 fee, bytes calldata data) external override returns (bytes32) {
        // try to send the loan to ourselves
        IERC20(token).transfer(initiator, amount + fee); 
        return keccak256("IERC3156FlashBorrower.onFlashLoan");
    }
}

contract MockUSDC is ERC20 {
    constructor() ERC20("USDC", "USDC") {
    }
    function mint(address to, uint256 amount) public {
        _mint(to, amount);
    }
}