Make master vault upgradable

contracts/tokenbridge/libraries/vault/IMasterVault.sol

@@ -2,5 +2,5 @@
 pragma solidity ^0.8.0;
 
 interface IMasterVault {
-    function setSubVault(address subVault) external;
+    function setSubVault(address subVault, uint256 minSubVaultExchRateWad) external;
 }

contracts/tokenbridge/libraries/vault/IMasterVaultFactory.sol

@@ -9,5 +9,5 @@ interface IMasterVaultFactory {
     function deployVault(address token) external returns (address vault);
     function calculateVaultAddress(address token) external view returns (address);
     function getVault(address token) external returns (address);
-    function setSubVault(address masterVault, address subVault) external;
+    function setSubVault(address masterVault, address subVault, uint256 minSubVaultExchRateWad) external;
 }

contracts/tokenbridge/libraries/vault/MasterVault.sol

@@ -1,16 +1,20 @@
 // SPDX-License-Identifier: Apache-2.0
 pragma solidity ^0.8.0;
 
-import {ERC4626} from "@openzeppelin/contracts/token/ERC20/extensions/ERC4626.sol";
-import { IERC20, ERC20 } from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
-import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import {ERC4626Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC20/extensions/ERC4626Upgradeable.sol";
+import {ERC20Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC20/ERC20Upgradeable.sol";
+import {IERC20} from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import {IERC20Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC20/IERC20Upgradeable.sol";
+import {IERC4626} from "@openzeppelin/contracts/interfaces/IERC4626.sol";
+import {OwnableUpgradeable} from "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
+import {Initializable} from "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
 import {IERC20Metadata} from "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
-import {Math} from "@openzeppelin/contracts/utils/math/Math.sol";
+import {MathUpgradeable} from "@openzeppelin/contracts-upgradeable/utils/math/MathUpgradeable.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 
-contract MasterVault is ERC4626, Ownable {
+contract MasterVault is Initializable, ERC4626Upgradeable, OwnableUpgradeable {
     using SafeERC20 for IERC20;
-    using Math for uint256;
+    using MathUpgradeable for uint256;
 
     error TooFewSharesReceived();
     error TooManySharesBurned();
@@ -26,16 +30,21 @@
     error NewSubVaultExchangeRateTooLow();
     error BeneficiaryNotSet();
     error PerformanceFeeDisabled();
+    error NoSharesRedeemed();
+    error NoSubvaultShares();
+    error NoSharesBurned();
+    error InvalidAsset();
+    error InvalidOwner();
 
     // todo: avoid inflation, rounding, other common 4626 vulns
     // we may need a minimum asset or master share amount when setting subvaults (bc of exchange rate calc)
-    ERC4626 public subVault;
+    IERC4626 public subVault;
 
     // how many subVault shares one MV2 share can be redeemed for
     // initially 1 to 1
     // constant per subvault
     // changes when subvault is set
-    uint256 public subVaultExchRateWad = 1e18;
+    uint256 public subVaultExchRateWad;
 
     // note: the performance fee can be avoided if the underlying strategy can be sandwiched (eg ETH to wstETH dex swap)
     // maybe a simpler and more robust implementation would be for the owner to adjust the subVaultExchRateWad directly
@@ -49,16 +58,27 @@
     event PerformanceFeeToggled(bool enabled);
     event BeneficiaryUpdated(address indexed oldBeneficiary, address indexed newBeneficiary);
 
-    constructor(IERC20 _asset, string memory _name, string memory _symbol) ERC20(_name, _symbol) ERC4626(_asset) Ownable() {}
+    function vaultInit(IERC20 _asset, string memory _name, string memory _symbol, address _owner) external initializer {
+        if (address(_asset) == address(0)) revert InvalidAsset();
+        if (_owner == address(0)) revert InvalidOwner();
+
+        __ERC20_init(_name, _symbol);
+        __ERC4626_init(IERC20Upgradeable(address(_asset)));
+        __Ownable_init();
+        _transferOwnership(_owner);
+
+        subVaultExchRateWad = 1e18;
+    }
+
 
     function deposit(uint256 assets, address receiver, uint256 minSharesMinted) public returns (uint256) {
-        uint256 shares = super.deposit(assets, receiver);
+        uint256 shares = deposit(assets, receiver);
         if (shares < minSharesMinted) revert TooFewSharesReceived();
         return shares;
     }
 
     function withdraw(uint256 assets, address receiver, address _owner, uint256 maxSharesBurned) public returns (uint256) {
-        uint256 shares = super.withdraw(assets, receiver, _owner);
+        uint256 shares = withdraw(assets, receiver, _owner);
         if (shares > maxSharesBurned) revert TooManySharesBurned();
         return shares;
     }
@@ -78,7 +98,7 @@
     /// @notice Set a subvault. Can only be called if there is not already a subvault set.
     /// @param  _subVault The subvault to set. Must be an ERC4626 vault with the same asset as this MasterVault.
     /// @param  minSubVaultExchRateWad Minimum acceptable ratio (times 1e18) of new subvault shares to outstanding MasterVault shares after deposit.
-    function setSubVault(ERC4626 _subVault, uint256 minSubVaultExchRateWad) external onlyOwner {
+    function setSubVault(IERC4626 _subVault, uint256 minSubVaultExchRateWad) external onlyOwner {
         if (address(subVault) != address(0)) revert SubVaultAlreadySet();
         _setSubVault(_subVault, minSubVaultExchRateWad);
     }
@@ -89,56 +109,62 @@
         _revokeSubVault(minAssetExchRateWad);
     }
 
-    function _setSubVault(ERC4626 _subVault, uint256 minSubVaultExchRateWad) internal {
+    function _setSubVault(IERC4626 _subVault, uint256 minSubVaultExchRateWad) internal {
         if (address(_subVault) == address(0)) revert SubVaultCannotBeZeroAddress();
         if (totalSupply() == 0) revert MustHaveSupplyBeforeSettingSubVault();
         if (address(_subVault.asset()) != address(asset())) revert SubVaultAssetMismatch();
 
+        uint256 _totalAssets = totalAssets();
+        uint256 _totalSupply = totalSupply();
+
+        subVault = _subVault;
+
         IERC20(asset()).safeApprove(address(_subVault), type(uint256).max);
-        uint256 subShares = _subVault.deposit(totalAssets(), address(this));
+        uint256 subShares = _subVault.deposit(_totalAssets, address(this));
 
-        uint256 _subVaultExchRateWad = subShares.mulDiv(1e18, totalSupply(), Math.Rounding.Down);
+        uint256 _subVaultExchRateWad = subShares.mulDiv(1e18, _totalSupply, MathUpgradeable.Rounding.Down);
         if (_subVaultExchRateWad < minSubVaultExchRateWad) revert SubVaultExchangeRateTooLow();
         subVaultExchRateWad = _subVaultExchRateWad;
 
-        subVault = _subVault;
-
         emit SubvaultChanged(address(0), address(_subVault));
     }
 
     function _revokeSubVault(uint256 minAssetExchRateWad) internal {
-        ERC4626 oldSubVault = subVault;
+        IERC4626 oldSubVault = subVault;
         if (address(oldSubVault) == address(0)) revert NoExistingSubVault();
 
         uint256 _totalSupply = totalSupply();
-        uint256 assetReceived = oldSubVault.withdraw(oldSubVault.maxWithdraw(address(this)), address(this), address(this));
-        uint256 effectiveAssetExchRateWad = assetReceived.mulDiv(1e18, _totalSupply, Math.Rounding.Down);
-        if (effectiveAssetExchRateWad < minAssetExchRateWad) revert TooFewAssetsReceived();
+        uint256 maxWithdrawAmount = oldSubVault.maxWithdraw(address(this));
 
-        IERC20(asset()).safeApprove(address(oldSubVault), 0);
-        subVault = ERC4626(address(0));
+        subVault = IERC4626(address(0));
         subVaultExchRateWad = 1e18;
 
+        uint256 assetReceived = oldSubVault.withdraw(maxWithdrawAmount, address(this), address(this));
+        IERC20(asset()).safeApprove(address(oldSubVault), 0);
+
+        uint256 effectiveAssetExchRateWad = assetReceived.mulDiv(1e18, _totalSupply, MathUpgradeable.Rounding.Down);
+        if (effectiveAssetExchRateWad < minAssetExchRateWad) revert TooFewAssetsReceived();
+
         emit SubvaultChanged(address(oldSubVault), address(0));
     }
 
     /// @notice Switches to a new subvault or revokes current subvault if newSubVault is zero address
     /// @param newSubVault The new subvault to switch to, or zero address to revoke current subvault
     /// @param minAssetExchRateWad Minimum acceptable ratio (times 1e18) of assets received from old subvault to outstanding MasterVault shares
     /// @param minNewSubVaultExchRateWad Minimum acceptable ratio (times 1e18) of new subvault shares to outstanding MasterVault shares after deposit
-    function switchSubVault(ERC4626 newSubVault, uint256 minAssetExchRateWad, uint256 minNewSubVaultExchRateWad) external onlyOwner {
+    function switchSubVault(IERC4626 newSubVault, uint256 minAssetExchRateWad, uint256 minNewSubVaultExchRateWad) external onlyOwner {
         _revokeSubVault(minAssetExchRateWad);
 
         if (address(newSubVault) != address(0)) {
             _setSubVault(newSubVault, minNewSubVaultExchRateWad);
         }
     }
 
-    function masterSharesToSubShares(uint256 masterShares, Math.Rounding rounding) public view returns (uint256) {
+    function masterSharesToSubShares(uint256 masterShares, MathUpgradeable.Rounding rounding) public view returns (uint256) {
         return masterShares.mulDiv(subVaultExchRateWad, 1e18, rounding);
     }
 
-    function subSharesToMasterShares(uint256 subShares, Math.Rounding rounding) public view returns (uint256) {
+    function subSharesToMasterShares(uint256 subShares, MathUpgradeable.Rounding rounding) public view returns (uint256) {
         return subShares.mulDiv(1e18, subVaultExchRateWad, rounding);
     }
 
@@ -165,17 +191,18 @@
 
         uint256 totalProfits = totalProfit();
         if (totalProfits > 0) {
-            ERC4626 _subVault = subVault;
+            IERC4626 _subVault = subVault;
             if (address(_subVault) != address(0)) {
-                _subVault.withdraw(totalProfits, address(this), address(this));
+                uint256 sharesRedeemed = _subVault.withdraw(totalProfits, address(this), address(this));
+                if (sharesRedeemed == 0) revert NoSharesRedeemed();
             }
             IERC20(asset()).safeTransfer(beneficiary, totalProfits);
         }
     }
 
     /** @dev See {IERC4626-totalAssets}. */
     function totalAssets() public view virtual override returns (uint256) {
-        ERC4626 _subVault = subVault;
+        IERC4626 _subVault = subVault;
         if (address(_subVault) == address(0)) {
             return super.totalAssets();
         }
@@ -196,7 +223,7 @@
         if (subShares == type(uint256).max) {
             return type(uint256).max;
         }
-        return subSharesToMasterShares(subShares, Math.Rounding.Down);
+        return subSharesToMasterShares(subShares, MathUpgradeable.Rounding.Down);
     }
 
     /**
@@ -205,25 +232,25 @@
      * Will revert if assets > 0, totalSupply > 0 and totalAssets = 0. That corresponds to a case where any asset
      * would represent an infinite amount of shares.
      */
-    function _convertToShares(uint256 assets, Math.Rounding rounding) internal view virtual override returns (uint256 shares) {
-        ERC4626 _subVault = subVault;
+    function _convertToShares(uint256 assets, MathUpgradeable.Rounding rounding) internal view virtual override returns (uint256 shares) {
+        IERC4626 _subVault = subVault;
         if (address(_subVault) == address(0)) {
             return super._convertToShares(assets, rounding);
         }
-        uint256 subShares = rounding == Math.Rounding.Up ? _subVault.previewWithdraw(assets) : _subVault.previewDeposit(assets);
+        uint256 subShares = rounding == MathUpgradeable.Rounding.Up ? _subVault.previewWithdraw(assets) : _subVault.previewDeposit(assets);
         return subSharesToMasterShares(subShares, rounding);
     }
 
     /**
      * @dev Internal conversion function (from shares to assets) with support for rounding direction.
      */
-    function _convertToAssets(uint256 shares, Math.Rounding rounding) internal view virtual override returns (uint256 assets) {
-        ERC4626 _subVault = subVault;
+    function _convertToAssets(uint256 shares, MathUpgradeable.Rounding rounding) internal view virtual override returns (uint256 assets) {
+        IERC4626 _subVault = subVault;
         if (address(_subVault) == address(0)) {
             return super._convertToAssets(shares, rounding);
         }
         uint256 subShares = masterSharesToSubShares(shares, rounding);
-        return rounding == Math.Rounding.Up ? _subVault.previewMint(subShares) : _subVault.previewRedeem(subShares);
+        return rounding == MathUpgradeable.Rounding.Up ? _subVault.previewMint(subShares) : _subVault.previewRedeem(subShares);
     }
 
     function totalProfit() public view returns (uint256) {
@@ -241,10 +268,12 @@
         uint256 shares
     ) internal virtual override {
         super._deposit(caller, receiver, assets, shares);
+
         totalPrincipal += assets;
-        ERC4626 _subVault = subVault;
+        IERC4626 _subVault = subVault;
         if (address(_subVault) != address(0)) {
-            _subVault.deposit(assets, address(this));
+            uint256 subShares = _subVault.deposit(assets, address(this));
+            if (subShares == 0) revert NoSubvaultShares();
         }
     }
 
@@ -260,9 +289,10 @@
     ) internal virtual override {
         totalPrincipal -= assets;
 
-        ERC4626 _subVault = subVault;
+        IERC4626 _subVault = subVault;
         if (address(_subVault) != address(0)) {
-            _subVault.withdraw(assets, address(this), address(this));
+            uint256 sharesBurned = _subVault.withdraw(assets, address(this), address(this));
+            if (sharesBurned == 0) revert NoSharesBurned();
         }
 
         super._withdraw(caller, receiver, _owner, assets, shares);

contracts/tokenbridge/libraries/vault/MasterVaultFactory.sol

@@ -5,49 +5,58 @@ pragma solidity ^0.8.0;
 import "@openzeppelin/contracts/utils/Create2.sol";
 import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
 import "@openzeppelin/contracts/token/ERC20/extensions/IERC20Metadata.sol";
+import "@openzeppelin/contracts/proxy/beacon/UpgradeableBeacon.sol";
+import "../ClonableBeaconProxy.sol";
 import "./IMasterVault.sol";
 import "./IMasterVaultFactory.sol";
 import "./MasterVault.sol";
 
 contract MasterVaultFactory is IMasterVaultFactory, OwnableUpgradeable {
-
     error ZeroAddress();
+    error BeaconNotDeployed();
+
+    UpgradeableBeacon public beacon;
+    BeaconProxyFactory public beaconProxyFactory;
 
     function initialize(address _owner) public initializer {
         _transferOwnership(_owner);
+
+        MasterVault masterVaultImplementation = new MasterVault();
+        beacon = new UpgradeableBeacon(address(masterVaultImplementation));
+        beaconProxyFactory = new BeaconProxyFactory();
+        beaconProxyFactory.initialize(address(beacon));
+        beacon.transferOwnership(_owner);
     }
 
     function deployVault(address token) public returns (address vault) {
         if (token == address(0)) {
             revert ZeroAddress();
         }
+        if (
+            address(beaconProxyFactory) == address(0) && beaconProxyFactory.beacon() == address(0)
+        ) {
+            revert BeaconNotDeployed();
+        }
+
+        bytes32 userSalt = _getUserSalt(token);
+        vault = beaconProxyFactory.createProxy(userSalt);
 
         IERC20Metadata tokenMetadata = IERC20Metadata(token);
         string memory name = string(abi.encodePacked("Master ", tokenMetadata.name()));
         string memory symbol = string(abi.encodePacked("m", tokenMetadata.symbol()));
 
-        bytes memory bytecode = abi.encodePacked(
-            type(MasterVault).creationCode,
-            abi.encode(token, name, symbol)
-        );
-
-        vault = Create2.deploy(0, bytes32(0), bytecode);
+        MasterVault(vault).vaultInit(IERC20(token), name, symbol, address(this));
 
         emit VaultDeployed(token, vault);
     }
 
-    function calculateVaultAddress(address token) public view returns (address) {
-        IERC20Metadata tokenMetadata = IERC20Metadata(token);
-        string memory name = string(abi.encodePacked("Master ", tokenMetadata.name()));
-        string memory symbol = string(abi.encodePacked("m", tokenMetadata.symbol()));
+    function _getUserSalt(address token) internal pure returns (bytes32) {
+        return keccak256(abi.encode(token));
+    }
 
-        bytes32 bytecodeHash = keccak256(
-            abi.encodePacked(
-                type(MasterVault).creationCode,
-                abi.encode(token, name, symbol)
-            )
-        );
-        return Create2.computeAddress(bytes32(0), bytecodeHash);
+    function calculateVaultAddress(address token) public view returns (address) {
+        bytes32 userSalt = _getUserSalt(token);
+        return beaconProxyFactory.calculateExpectedAddress(address(this), userSalt);
     }
 
     function getVault(address token) external returns (address) {
@@ -61,9 +70,10 @@ contract MasterVaultFactory is IMasterVaultFactory, OwnableUpgradeable {
     // todo: consider a method to enable bridge owner to transfer specific master vault ownership to new address
     function setSubVault(
         address masterVault,
-        address subVault
+        address subVault,
+        uint256 minSubVaultExchRateWad
     ) external onlyOwner {
-        IMasterVault(masterVault).setSubVault(subVault);
+        IMasterVault(masterVault).setSubVault(subVault, minSubVaultExchRateWad);
         emit SubVaultSet(masterVault, subVault);
     }
-}
+}