How To Implement a UUPSUpgradeable Flash Loan Vault With Foundry

Making An Upgradeable FlashLoan Vault

Building off of our last tutorial, Creating a FlashLoan Vault we will now be able to learn how to develop and deploy Upgradeable contracts from foundry using the openzeppelin-upgrades library.

In our last version we built a Flashloan vault with a fixed fee for loans of any size. However, what happens when we realize that this pricing model makes little sense and decide instead to charge a percentage of the loan size as the fee?

With our previous version we would be forced to deploy an entirely new lending protocol and instruct all of our users, both depositors and borrowers to target a new contract and migrate all of their funds. This is particularly problematic given that our borrowers have developed and deployed specially built IERC3156FlashBorrower contracts that are designed to target our lending facility. This puts a major overhead on our users and customers for such a trivial change.

The solution to this issue is to deploy our flash lending protocol as an upgradeable contract. This way we can make changes to the protocol without having to redeploy the entire contract and migrate all of our users. All previous users can continue to target the same address but they will be interacting with our new V2 logic.

In this post we will show how easily we can reimplement our initial FlashLoan vault design as an upgradeable contract and then make a change to the fee structure and, again easily, deploy the version 2 to our users without the need for a migration or downtime. If you would like to follow along please see the previous post for the initial implementation of the FlashLoanVault.

Installation

The first thing we will need to do is install the necessary tools. One being the openzeppelin-foundry-upgrades library to facilitate our management of the upgradeable contracts via foundry. The second being the openzeppelin-contracts-upgradeable library which contains the upgradeable versions of the OpenZeppelin standard contracts we will be inheriting in our smart contract.

forge install OpenZeppelin/openzeppelin-foundry-upgrades --no-commit
forge install OpenZeppelin/openzeppelin-contracts-upgradeable@v5.0.1 --no-commit

Once we have these installed we can add the following to our remapping configurations in foundry.toml. We will also need to add the extra_output configuration to include the storageLayout output. These configs will help us manage the build artifacts necessary for our upgradeable deployment.

foundry.toml

build_info = true
extra_output = ["storageLayout"]
remappings = [
    '@openzeppelin/contracts/=lib/openzeppelin-contracts-upgradeable/lib/openzeppelin-contracts/contracts/',
    '@openzeppelin/contracts-upgradeable/=lib/openzeppelin-contracts-upgradeable/contracts/',
    '@solady/contracts/=lib/solady/src/',
]

(Re)Implementation

Now our reimplementation of the FlashLoan vault will only require some minor changes. Starting with the imports, we will need to import the upgradeable versions of the OpenZeppelin contracts we will be inheriting. For this example, we only need to import the ERC4626Upgradeable contract to replace our previous ERC4626 import. Then, we can add imports for the UUPSUpgradeable and Ownable2StepUpgradeable contracts which will allow us to manage the upgradeability of our contract as well as require special permissions from the owner to perform the upgrade.

// FlashLoanVault.sol
import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
import { ERC4626Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/extensions/ERC4626Upgradeable.sol";
import { Ownable2StepUpgradeable } from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.sol";

Once we have the upgrades imported, we can change the inheritance structure of our contract to match it.

contract FlashLoanVault is ERC4626Upgradeable, IERC3156FlashLender, UUPSUpgradeable, Ownable2StepUpgradeable {}

Now, one thing in particular to note about upgradeable contracts is that they do not have a constructor function, instead we will use an initialize function to set the initial state of our contract as well as it's inherited contracts.

    // constructor(address _depositToken) ERC4626(IERC20(_depositToken)) ERC20("Shares Vault", "SV") {}
    function initialize(
        address _depositToken
    ) external initializer {
        __Ownable2Step_init();
        __Ownable_init(msg.sender);
        __ERC4626_init(IERC20(_depositToken));
        __ERC20_init("Shares Vault", "SV");
    }

After we delete the constructor, we can implement our intialize function to set the initial state of our base contracts. Note the initializer modifier which is used to ensure that the initialize function is only called once.

The final change we need to make to our implementation is adding an _authorizeUpgrade function which will be used to ensure that only the owner of the contract can perform the upgrade. In order that we may achieve this, we simply add the onlyOwner modifier to the function.

    function _authorizeUpgrade(address) internal override onlyOwner {}

Deploying a UUPSUpgradable Contract in Foundry Tests

Now that we have the upgradeable reimplementation of our FlashLoanVault complete we can move on to testing the upgradeability. The first thing we will need to change in our FlashLoanVault.t.sol test file is the deployment setup of our contract.

import {Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";

We import the openzeppelin-foundry-upgrades/Upgrades.sol library which will assist us in deploying our new contracts.

After that we modify our FlashLoan Vault deployment in the setup function to use the deployUUPSProxy helper. Once our proxy is deployed we will declare our vault variable used in testing to be an instance of FlashLoanVault at the proxy address.

        proxy = Upgrades.deployUUPSProxy("FlashLoanVault.sol",
            abi.encodeCall(FlashLoanVault.initialize, (address(usdc)))
        );
        vault = FlashLoanVault(proxy);
        //vault = new FlashLoanVault(address(usdc));

Now, we can run all of our previously written tests and see that they all pass and the expected logic is the same as before.

Note: When we build and run this new implementation we will have to add the --ffi flag to enable the openzeppelin foundry upgrades library to be used in our tests.

forge clean && forge build && forge test -vv --ffi

Making an Upgrade

At this point, we have our upgradeable FlashLoanVault with exactly the same functional logic as before. Including the fixed fee of 0.1 USDC

Now, we can test changing the logic of the fee structure in a V2.

All we need to do to implement the upgrade is make a copy of our FlashLoanVault.sol file and name it FlashLoanVaultV2.sol. After this, we simply change the name of the contract to indicate it's the second version and add a reference tag to the previous implementation so that the openzeppelin upgrades library knows how to manage the upgrade. The reference is simply added as a comment above the contract declaration detailing the contract to be upgraded.

FlashLoanVaultV2.sol

/// @custom:oz-upgrades-from FlashLoanVault
contract FlashLoanVaultV2 is ERC4626Upgradeable, IERC3156FlashLender, UUPSUpgradeable, Ownable2StepUpgradeable {}

With the version 2 now setup, we can upgrade the logic of the flashFee function to change the fee structure of our flash loan protocol.

    function flashFee(address token, uint256 amount) public view override returns (uint256) {
        return (amount * 100) / 10000;
    }

Here, we have simply changed the formula to calculate the fee to be 100 basis points or 1% of the loan amount.

Deploying an Upgrade

Now that we have our V2 implementation complete, we can deploy it in our tests to see the procedure and test that the new functionality is working properly.

All we do is add a new test in our FlashLoanVault.t.sol file called test_upgrade_fee().

Deploying the upgrade is actually rather simple utilizing the Upgrades library we simply call the Upgrade.upgradeProxy function and pass the proxy address and the new implementation file "FlashLoanVaultV2.sol" as arguments.

FlashLoanVault.t.sol

    function test_upgrade_fee() public {
        Upgrades.upgradeProxy(
            proxy,
            "FlashLoanVaultV2.sol",
            ""
        );
        vm.startPrank(user1);
        usdc.transfer(address(borrower), 0.01 ether);
        vault.flashLoan(borrower, address(usdc), 1 ether, "");

        // Check we paid the fee 1% of 1 ether = 0.01 ether
        assertEq(usdc.balanceOf(address(vault)), 10.01 ether); 
        assertEq(usdc.balanceOf(address(user1)), 9.99 ether);
        assertEq(usdc.balanceOf(address(borrower)), 0);

        usdc.transfer(address(borrower), 0.1 ether);
        vault.flashLoan(borrower, address(usdc), 10 ether, "");

        // no we paid the fee 1% of 10 ether = 0.1 ether
        assertEq(usdc.balanceOf(address(vault)), 10.11 ether); 
        assertEq(usdc.balanceOf(address(user1)), 9.89 ether);
        assertEq(usdc.balanceOf(address(borrower)), 0);
    }

Once we have the upgrade deployed, we can execute our tests using user1 to make two flashloans of 1 USDC and 10 USDC, each time checking that the fee was paid correctly and in the amount of 1% of the loan size.

Conclusion

In this post we have seen how to reimplement our FlashLoanVault contract as a UUPSupgradeable contract to allow for future changes to the protocol logic without requiring action from our user or any downtime to the protocol.

We have also seen how to utilize openzeppelin's foundry-upgrades library to deploy and test the contracts as well as how to upgrade a live contract to a version 2 implementation.

Full Code

FlashLoanVault.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity 0.8.20;
import { UUPSUpgradeable } from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
import { ERC4626Upgradeable } from "@openzeppelin/contracts-upgradeable/token/ERC20/extensions/ERC4626Upgradeable.sol";
import { Ownable2StepUpgradeable } from "@openzeppelin/contracts-upgradeable/access/Ownable2StepUpgradeable.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 ERC4626Upgradeable, IERC3156FlashLender, UUPSUpgradeable, Ownable2StepUpgradeable {
    using SafeTransferLib for IERC20;
    // constructor(address _depositToken) ERC4626(IERC20(_depositToken)) ERC20("Shares Vault", "SV") {
    // }
    function initialize(
        address _depositToken
    ) external initializer {
        __Ownable2Step_init();
        __Ownable_init(msg.sender);
        __ERC4626_init(IERC20(_depositToken));
        __ERC20_init("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;
    }

    function _authorizeUpgrade(address) internal override onlyOwner {}

}