feat(protocol-contracts): make OperatorStaking UUPS upgradeable (#1562)

Author: melanciani

protocol-contracts/staking/contracts/OperatorStaking.sol

@@ -2,9 +2,9 @@
 
 pragma solidity ^0.8.27;
 
+import {UUPSUpgradeable} from "@openzeppelin/contracts/proxy/utils/UUPSUpgradeable.sol";
 import {IERC20} from "@openzeppelin/contracts/interfaces/IERC20.sol";
-import {ERC20} from "@openzeppelin/contracts/token/ERC20/ERC20.sol";
-import {ERC1363} from "@openzeppelin/contracts/token/ERC20/extensions/ERC1363.sol";
+import {ERC1363Upgradeable} from "@openzeppelin/contracts-upgradeable/token/ERC20/extensions/ERC1363Upgradeable.sol";
 import {ERC4626, IERC4626} from "@openzeppelin/contracts/token/ERC20/extensions/ERC4626.sol";
 import {SafeERC20} from "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
 import {Math} from "@openzeppelin/contracts/utils/math/Math.sol";
@@ -27,17 +27,24 @@ import {ProtocolStaking} from "./ProtocolStaking.sol";
  * may decrease due to slashing. These losses are symmetrically passed to delegators on the `OperatorStaking` level.
  * Slashing must first decrease the `ProtocolStaking` balance of this contract before affecting pending withdrawals.
  */
-contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
+contract OperatorStaking is ERC1363Upgradeable, ReentrancyGuardTransient, UUPSUpgradeable {
     using Math for uint256;
     using Checkpoints for Checkpoints.Trace208;
 
-    ProtocolStaking private immutable _protocolStaking;
-    IERC20 private immutable _asset;
-    address private _rewarder;
-    uint256 private _totalSharesInRedemption;
-    mapping(address => uint256) private _sharesReleased;
-    mapping(address => Checkpoints.Trace208) private _redeemRequests;
-    mapping(address => mapping(address => bool)) private _operator;
+    /// @custom:storage-location erc7201:fhevm_protocol.storage.OperatorStaking
+    struct OperatorStakingStorage {
+        ProtocolStaking _protocolStaking;
+        IERC20 _asset;
+        address _rewarder;
+        uint256 _totalSharesInRedemption;
+        mapping(address => uint256) _sharesReleased;
+        mapping(address => Checkpoints.Trace208) _redeemRequests;
+        mapping(address => mapping(address => bool)) _operator;
+    }
+
+    // keccak256(abi.encode(uint256(keccak256("fhevm_protocol.storage.OperatorStaking")) - 1)) & ~bytes32(uint256(0xff))
+    bytes32 private constant OPERATOR_STAKING_STORAGE_LOCATION =
+        0x7fc851282090a0d8832502c48739eac98a0856539351f17cb5d5950c860fd200;
 
     /// @dev Emitted when an operator is set or unset for a controller.
     event OperatorSet(address indexed controller, address indexed operator, bool approved);
@@ -72,6 +79,11 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
         _;
     }
 
+    /// @custom:oz-upgrades-unsafe-allow constructor
+    constructor() {
+        _disableInitializers();
+    }
+
     /**
      * @notice Initializes the OperatorStaking contract.
      * @param name The name of the ERC20 token.
@@ -81,24 +93,34 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
      * @param initialMaxFeeBasisPoints_ The initial maximum fee basis points for the OperatorRewarder contract.
      * @param initialFeeBasisPoints_ The initial fee basis points for the OperatorRewarder contract.
      */
-    constructor(
+    function initialize(
         string memory name,
         string memory symbol,
         ProtocolStaking protocolStaking_,
         address beneficiary_,
         uint16 initialMaxFeeBasisPoints_,
         uint16 initialFeeBasisPoints_
-    ) ERC20(name, symbol) {
-        _asset = IERC20(protocolStaking_.stakingToken());
-        _protocolStaking = protocolStaking_;
+    ) public initializer {
+        __ERC20_init(name, symbol);
+
+        OperatorStakingStorage storage $ = _getOperatorStakingStorage();
+
+        $._asset = IERC20(protocolStaking_.stakingToken());
+        $._protocolStaking = protocolStaking_;
 
         IERC20(asset()).approve(address(protocolStaking_), type(uint256).max);
 
         address rewarder_ = address(
-            new OperatorRewarder(beneficiary_, protocolStaking_, this, initialMaxFeeBasisPoints_, initialFeeBasisPoints_)
+            new OperatorRewarder(
+                beneficiary_,
+                protocolStaking_,
+                this,
+                initialMaxFeeBasisPoints_,
+                initialFeeBasisPoints_
+            )
         );
         protocolStaking_.setRewardsRecipient(rewarder_);
-        _rewarder = rewarder_;
+        $._rewarder = rewarder_;
 
         emit RewarderSet(address(0), rewarder_);
     }
@@ -133,19 +155,21 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
         }
         _burn(ownerRedeem, shares);
 
+        OperatorStakingStorage storage $ = _getOperatorStakingStorage();
+
         uint256 newTotalSharesInRedemption = totalSharesInRedemption() + shares;
-        _totalSharesInRedemption = newTotalSharesInRedemption;
+        $._totalSharesInRedemption = newTotalSharesInRedemption;
 
         ProtocolStaking protocolStaking_ = protocolStaking();
         int256 assetsToWithdraw = SafeCast.toInt256(previewRedeem(newTotalSharesInRedemption)) -
             SafeCast.toInt256(
                 IERC20(asset()).balanceOf(address(this)) + protocolStaking_.awaitingRelease(address(this))
             );
 
-        (, uint48 lastReleaseTime, uint208 controllerSharesRedeemed) = _redeemRequests[controller].latestCheckpoint();
+        (, uint48 lastReleaseTime, uint208 controllerSharesRedeemed) = $._redeemRequests[controller].latestCheckpoint();
         uint48 releaseTime = protocolStaking_.unstake(SafeCast.toUint256(SignedMath.max(assetsToWithdraw, 0)));
         assert(releaseTime >= lastReleaseTime); // should never happen
-        _redeemRequests[controller].push(releaseTime, controllerSharesRedeemed + shares);
+        $._redeemRequests[controller].push(releaseTime, controllerSharesRedeemed + shares);
 
         emit RedeemRequest(controller, ownerRedeem, 0, msg.sender, shares, releaseTime);
     }
@@ -174,8 +198,9 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
         uint256 assets = previewRedeem(shares);
 
         if (assets > 0) {
-            _totalSharesInRedemption -= shares;
-            _sharesReleased[controller] += shares;
+            OperatorStakingStorage storage $ = _getOperatorStakingStorage();
+            $._totalSharesInRedemption -= shares;
+            $._sharesReleased[controller] += shares;
             _doTransferOut(receiver, assets);
 
             emit IERC4626.Withdraw(msg.sender, receiver, controller, assets, shares);
@@ -200,15 +225,15 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
     }
 
     /**
-     * @dev Set a new rewarder contract.
+     * @dev Set a new rewarder contract. Only callable by the owner.
      * @param newRewarder The new rewarder contract address. This contract must not be the same as the current
      * and must have code.
      */
     function setRewarder(address newRewarder) public virtual onlyOwner {
         address oldRewarder = rewarder();
         require(newRewarder != oldRewarder && newRewarder.code.length > 0, InvalidRewarder(newRewarder));
         OperatorRewarder(oldRewarder).shutdown();
-        _rewarder = newRewarder;
+        _getOperatorStakingStorage()._rewarder = newRewarder;
         protocolStaking().setRewardsRecipient(newRewarder);
 
         emit RewarderSet(oldRewarder, newRewarder);
@@ -220,7 +245,7 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
      * @param approved True to approve, false to revoke.
      */
     function setOperator(address operator, bool approved) public virtual {
-        _operator[msg.sender][operator] = approved;
+        _getOperatorStakingStorage()._operator[msg.sender][operator] = approved;
 
         emit OperatorSet(msg.sender, operator, approved);
     }
@@ -238,23 +263,23 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
      * @return The asset address.
      */
     function asset() public view virtual returns (address) {
-        return address(_asset);
+        return address(_getOperatorStakingStorage()._asset);
     }
 
     /**
      * @notice Returns the ProtocolStaking contract address.
      * @return The ProtocolStaking contract address.
      */
     function protocolStaking() public view virtual returns (ProtocolStaking) {
-        return _protocolStaking;
+        return _getOperatorStakingStorage()._protocolStaking;
     }
 
     /**
      * @notice Returns the rewarder contract address.
      * @return The rewarder contract address.
      */
     function rewarder() public view virtual returns (address) {
-        return _rewarder;
+        return _getOperatorStakingStorage()._rewarder;
     }
 
     /**
@@ -275,7 +300,8 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
      * @return Amount of shares pending redeem.
      */
     function pendingRedeemRequest(uint256, address controller) public view virtual returns (uint256) {
-        return _redeemRequests[controller].latest() - _redeemRequests[controller].upperLookup(Time.timestamp());
+        OperatorStakingStorage storage $ = _getOperatorStakingStorage();
+        return $._redeemRequests[controller].latest() - $._redeemRequests[controller].upperLookup(Time.timestamp());
     }
 
     /**
@@ -284,15 +310,16 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
      * @return Amount of claimable shares.
      */
     function claimableRedeemRequest(uint256, address controller) public view virtual returns (uint256) {
-        return _redeemRequests[controller].upperLookup(Time.timestamp()) - _sharesReleased[controller];
+        OperatorStakingStorage storage $ = _getOperatorStakingStorage();
+        return $._redeemRequests[controller].upperLookup(Time.timestamp()) - $._sharesReleased[controller];
     }
 
     /**
      * @notice Returns the total shares in redemption.
      * @return The total shares in redemption.
      */
     function totalSharesInRedemption() public view virtual returns (uint256) {
-        return _totalSharesInRedemption;
+        return _getOperatorStakingStorage()._totalSharesInRedemption;
     }
 
     /**
@@ -337,7 +364,7 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
      * @return True if operator is approved, false otherwise.
      */
     function isOperator(address controller, address operator) public view virtual returns (bool) {
-        return _operator[controller][operator];
+        return _getOperatorStakingStorage()._operator[controller][operator];
     }
 
     function _doTransferOut(address to, uint256 amount) internal {
@@ -371,6 +398,8 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
         emit IERC4626.Deposit(caller, receiver, assets, shares);
     }
 
+    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}
+
     function _convertToShares(uint256 assets, Math.Rounding rounding) internal view virtual returns (uint256) {
         // Shares in redemption have not yet received assets, so we need to account for them in the conversion.
         return
@@ -394,4 +423,10 @@ contract OperatorStaking is ERC1363, ReentrancyGuardTransient {
     function _decimalsOffset() internal view virtual returns (uint8) {
         return 0;
     }
+
+    function _getOperatorStakingStorage() private pure returns (OperatorStakingStorage storage $) {
+        assembly {
+            $.slot := OPERATOR_STAKING_STORAGE_LOCATION
+        }
+    }
 }

protocol-contracts/staking/contracts/ProtocolStaking.sol

@@ -28,7 +28,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
     using SafeERC20 for IERC20;
     using Math for uint256;
 
-    /// @custom:storage-location erc7201:zama.storage.ProtocolStaking
+    /// @custom:storage-location erc7201:fhevm_protocol.storage.ProtocolStaking
     struct ProtocolStakingStorage {
         // Stake - general
         address _stakingToken;
@@ -48,10 +48,9 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         int256 _totalVirtualPaid;
     }
 
-    // keccak256(abi.encode(uint256(keccak256("zama.storage.ProtocolStaking")) - 1)) & ~bytes32(uint256(0xff))
+    // keccak256(abi.encode(uint256(keccak256("fhevm_protocol.storage.ProtocolStaking")) - 1)) & ~bytes32(uint256(0xff))
     bytes32 private constant PROTOCOL_STAKING_STORAGE_LOCATION =
-        0x6867237db38693700f305f18dff1dbf600e282237f7d452b4c792e6b019c6b00;
-    bytes32 private constant UPGRADER_ROLE = keccak256("UPGRADER_ROLE");
+        0xd955b2342c0487c5e5b5f50f5620ec67dcb16d94462ba5d080d7b7472b67b900;
     bytes32 private constant MANAGER_ROLE = keccak256("MANAGER_ROLE");
     bytes32 private constant ELIGIBLE_ACCOUNT_ROLE = keccak256("ELIGIBLE_ACCOUNT_ROLE");
 
@@ -89,13 +88,11 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         string memory version,
         address stakingToken_,
         address governor,
-        address upgrader,
         address manager,
         uint48 initialUnstakeCooldownPeriod,
         uint256 initialRewardRate
     ) public initializer {
         __AccessControlDefaultAdminRules_init(0, governor);
-        _grantRole(UPGRADER_ROLE, upgrader);
         _grantRole(MANAGER_ROLE, manager);
         _setRoleAdmin(ELIGIBLE_ACCOUNT_ROLE, MANAGER_ROLE);
         __ERC20_init(name, symbol);
@@ -370,7 +367,7 @@ contract ProtocolStaking is AccessControlDefaultAdminRulesUpgradeable, ERC20Vote
         super._update(from, to, value);
     }
 
-    function _authorizeUpgrade(address newImplementation) internal override onlyRole(UPGRADER_ROLE) {}
+    function _authorizeUpgrade(address newImplementation) internal override onlyRole(DEFAULT_ADMIN_ROLE) {}
 
     function _historicalReward() internal view returns (uint256) {
         ProtocolStakingStorage storage $ = _getProtocolStakingStorage();

protocol-contracts/staking/contracts/mocks/ProtocolStakingSlashingMock.sol

@@ -13,7 +13,7 @@ contract ProtocolStakingSlashingMock is ProtocolStaking {
     using SafeERC20 for IERC20;
 
     bytes32 private constant PROTOCOL_STAKING_STORAGE_LOCATION =
-        0x6867237db38693700f305f18dff1dbf600e282237f7d452b4c792e6b019c6b00;
+        0xd955b2342c0487c5e5b5f50f5620ec67dcb16d94462ba5d080d7b7472b67b900;
 
     mapping(address => uint256) private _slashedAmount;
 

protocol-contracts/staking/tasks/deployment/operatorStaking.ts

@@ -8,7 +8,12 @@ export const OPERATOR_REWARDER_CONTRACT_NAME = 'OperatorRewarder';
 
 // Get the name of the operator staking contract to save in the deployments
 export function getOperatorStakingName(tokenName: string): string {
-  return tokenName + '_Staking';
+  return tokenName + '_Staking_Proxy';
+}
+
+// Get the name of the implementation contract to save in the deployments
+export function getOperatorStakingImplName(tokenName: string): string {
+  return tokenName + '_Staking_Impl';
 }
 
 // Get the name of the operator rewarder contract to save in the deployments
@@ -26,43 +31,47 @@ async function deployOperatorStaking(
   initialFeeBasisPoints: number,
   hre: HardhatRuntimeEnvironment,
 ) {
-  const { getNamedAccounts, ethers, deployments, network } = hre;
+  const { getNamedAccounts, ethers, deployments, network, upgrades } = hre;
   const { save, getArtifact } = deployments;
 
   // Get the deployer account
   const { deployer } = await getNamedAccounts();
   const deployerSigner = await ethers.getSigner(deployer);
 
-  // Get the contract factory and deploy the operator staking and rewarder contracts
-  const operatorStakingFactory = await ethers.getContractFactory(OPERATOR_STAKING_CONTRACT_NAME, deployerSigner);
-  const operatorStaking = await operatorStakingFactory.deploy(
-    tokenName,
-    symbol,
-    protocolStakingAddress,
-    beneficiaryAddress,
-    initialMaxFeeBasisPoints,
-    initialFeeBasisPoints,
+  // Get the contract factory and deploy the proxy + the implementation + the rewarder contract
+  const protocolStakingFactory = await ethers.getContractFactory(OPERATOR_STAKING_CONTRACT_NAME, deployerSigner);
+  const proxy = await upgrades.deployProxy(
+    protocolStakingFactory,
+    [tokenName, symbol, protocolStakingAddress, beneficiaryAddress, initialMaxFeeBasisPoints, initialFeeBasisPoints],
+    { kind: 'uups', initializer: 'initialize' },
   );
-  await operatorStaking.waitForDeployment();
+  await proxy.waitForDeployment();
 
-  // Get the operator staking and rewarder addresses
-  const operatorStakingAddress = await operatorStaking.getAddress();
-  const operatorRewarderAddress = await operatorStaking.rewarder();
+  // Get the operator staking proxy and rewarder addresses
+  const operatorStakingProxyAddress = await proxy.getAddress();
+  const operatorRewarderAddress = await proxy.rewarder();
 
   console.log(
     [
       `✅ Deployed ${tokenName} OperatorStaking:`,
-      `  - Operator staking address:  ${operatorStakingAddress}`,
+      `  - Operator staking proxy address:  ${operatorStakingProxyAddress}`,
       `  - Operator rewarder address: ${operatorRewarderAddress}`,
       `  - Deployed by deployer account: ${deployer}`,
       `  - Network: ${network.name}`,
       '',
     ].join('\n'),
   );
 
-  // Save the OperatorStaking and OperatorRewarder contract artifacts
+  // Save the OperatorStaking proxy and implementation contract artifacts
   const operatorStakingArtifact = await getArtifact(OPERATOR_STAKING_CONTRACT_NAME);
-  await save(getOperatorStakingName(tokenName), { address: operatorStakingAddress, abi: operatorStakingArtifact.abi });
+  const implAddress = await upgrades.erc1967.getImplementationAddress(operatorStakingProxyAddress);
+  await save(getOperatorStakingName(tokenName), {
+    address: operatorStakingProxyAddress,
+    abi: operatorStakingArtifact.abi,
+  });
+  await save(getOperatorStakingImplName(tokenName), { address: implAddress, abi: operatorStakingArtifact.abi });
+
+  // Save the OperatorRewarder contract artifacts
   const operatorRewarderArtifact = await getArtifact(OPERATOR_REWARDER_CONTRACT_NAME);
   await save(getOperatorRewarderName(tokenName), {
     address: operatorRewarderAddress,