feat: make master vault beacon upgradable

waelsy123 · waelsy123 · commit 7922536aa59c · 2025-10-01T15:45:10.000+02:00
diff --git a/contracts/tokenbridge/libraries/vault/IMasterVault.sol b/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;
 }
diff --git a/contracts/tokenbridge/libraries/vault/IMasterVaultFactory.sol b/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;
 }
diff --git a/contracts/tokenbridge/libraries/vault/MasterVault.sol b/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();
@@ -29,13 +33,13 @@ contract MasterVault is ERC4626, Ownable {
 
     // 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 +53,27 @@ contract MasterVault is ERC4626, Ownable {
     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 {
+        require(address(_asset) != address(0), "INVALID_ASSET");
+        require(_owner != address(0), "INVALID_OWNER");
+
+        __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 +93,7 @@ contract MasterVault is ERC4626, Ownable {
     /// @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,15 +104,15 @@ contract MasterVault is ERC4626, Ownable {
         _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();
 
         IERC20(asset()).safeApprove(address(_subVault), type(uint256).max);
         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;
 
@@ -107,16 +122,16 @@ contract MasterVault is ERC4626, Ownable {
     }
 
     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);
+        uint256 effectiveAssetExchRateWad = assetReceived.mulDiv(1e18, _totalSupply, MathUpgradeable.Rounding.Down);
         if (effectiveAssetExchRateWad < minAssetExchRateWad) revert TooFewAssetsReceived();
 
         IERC20(asset()).safeApprove(address(oldSubVault), 0);
-        subVault = ERC4626(address(0));
+        subVault = IERC4626(address(0));
         subVaultExchRateWad = 1e18;
 
         emit SubvaultChanged(address(oldSubVault), address(0));
@@ -126,19 +141,19 @@ contract MasterVault is ERC4626, Ownable {
     /// @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,7 +180,7 @@ contract MasterVault is ERC4626, Ownable {
 
         uint256 totalProfits = totalProfit();
         if (totalProfits > 0) {
-            ERC4626 _subVault = subVault;
+            IERC4626 _subVault = subVault;
             if (address(_subVault) != address(0)) {
                 _subVault.withdraw(totalProfits, address(this), address(this));
             }
@@ -175,7 +190,7 @@ contract MasterVault is ERC4626, Ownable {
 
     /** @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 +211,7 @@ contract MasterVault is ERC4626, Ownable {
         if (subShares == type(uint256).max) {
             return type(uint256).max;
         }
-        return subSharesToMasterShares(subShares, Math.Rounding.Down);
+        return subSharesToMasterShares(subShares, MathUpgradeable.Rounding.Down);
     }
 
     /**
@@ -205,25 +220,25 @@ contract MasterVault is ERC4626, Ownable {
      * 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,8 +256,9 @@ contract MasterVault is ERC4626, Ownable {
         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));
         }
@@ -260,7 +276,7 @@ contract MasterVault is ERC4626, Ownable {
     ) internal virtual override {
         totalPrincipal -= assets;
 
-        ERC4626 _subVault = subVault;
+        IERC4626 _subVault = subVault;
         if (address(_subVault) != address(0)) {
             _subVault.withdraw(assets, address(this), address(this));
         }
diff --git a/contracts/tokenbridge/libraries/vault/MasterVaultFactory.sol b/contracts/tokenbridge/libraries/vault/MasterVaultFactory.sol
@@ -5,49 +5,56 @@ 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(address(this));
     }
 
     function deployVault(address token) public returns (address vault) {
         if (token == address(0)) {
             revert ZeroAddress();
         }
+        if (address(beaconProxyFactory) == 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 +68,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);
     }
-}
+}
diff --git a/test-foundry/libraries/vault/MasterVault.t.sol b/test-foundry/libraries/vault/MasterVault.t.sol
@@ -6,10 +6,14 @@ import { MasterVault } from "../../../contracts/tokenbridge/libraries/vault/Mast
 import { TestERC20 } from "../../../contracts/tokenbridge/test/TestERC20.sol";
 import { MockSubVault } from "../../../contracts/tokenbridge/test/MockSubVault.sol";
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { UpgradeableBeacon } from "@openzeppelin/contracts/proxy/beacon/UpgradeableBeacon.sol";
+import { BeaconProxyFactory, ClonableBeaconProxy } from "../../../contracts/tokenbridge/libraries/ClonableBeaconProxy.sol";
 
 contract MasterVaultTest is Test {
     MasterVault public vault;
     TestERC20 public token;
+    UpgradeableBeacon public beacon;
+    BeaconProxyFactory public beaconProxyFactory;
 
     event SubvaultChanged(address indexed oldSubvault, address indexed newSubvault);
 
@@ -19,7 +23,18 @@ contract MasterVaultTest is Test {
 
     function setUp() public {
         token = new TestERC20();
-        vault = new MasterVault(IERC20(address(token)), name, symbol);
+
+        MasterVault implementation = new MasterVault();
+        beacon = new UpgradeableBeacon(address(implementation));
+
+        beaconProxyFactory = new BeaconProxyFactory();
+        beaconProxyFactory.initialize(address(beacon));
+
+        bytes32 salt = keccak256("test");
+        address proxyAddress = beaconProxyFactory.createProxy(salt);
+        vault = MasterVault(proxyAddress);
+
+        vault.vaultInit(IERC20(address(token)), name, symbol, address(this));
     }
 
     function test_initialize() public {

Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`	`pragma solidity ^0.8.0;`
`3`	`3`
`4`	`4`	`interface IMasterVault {`
`5`		`- function setSubVault(address subVault) external;`
	`5`	`+ function setSubVault(address subVault, uint256 minSubVaultExchRateWad) external;`
`6`	`6`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,5 +9,5 @@ interface IMasterVaultFactory {`
`9`	`9`	`function deployVault(address token) external returns (address vault);`
`10`	`10`	`function calculateVaultAddress(address token) external view returns (address);`
`11`	`11`	`function getVault(address token) external returns (address);`
`12`		`- function setSubVault(address masterVault, address subVault) external;`
	`12`	`+ function setSubVault(address masterVault, address subVault, uint256 minSubVaultExchRateWad) external;`
`13`	`13`	`}`