feat: add StakeVault and implements StakeManager interface into RewardsStreamerMP

This commit adds the following changes:
- Add the contract `StakeVault` so funds can be stored safely by the user. #14
- Added the interface `ITrustedCodehashAccess`  and contract `TrustedCodehashAccess` in the `src/access` directory, which implements the `ITrustedCodehashAccess` interface and provides functionality to set or update the trust status for a contract's codehash and implemented it on `RewardStreamerMP`. #15
- added the interface `IStakeManager` and implemented it on `RewardStreamerMP`  #13
These changes are necessary to enforce security measures and restrict access based on the codehash of the caller, and allow for better reuse of code between StakeManager and RewardStreamerMP.

src/IStakeManager.sol

@@ -0,0 +1,35 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.18;
+
+import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import { ITrustedCodehashAccess } from "./access/ITrustedCodehashAccess.sol";
+
+interface IStakeManager is ITrustedCodehashAccess {
+    error StakeManager__FundsLocked();
+    error StakeManager__InvalidLockTime();
+    error StakeManager__InsufficientFunds();
+    error StakeManager__StakeIsTooLow();
+
+    function stake(uint256 _amount, uint256 _seconds) external;
+    function unstake(uint256 _amount) external;
+    function lock(uint256 _secondsIncrease) external;
+    function leave() external returns (bool _leaveAccepted);
+    function acceptUpdate() external returns (address _migrated);
+
+    function potentialMP() external view returns (uint256);
+    function totalMP() external view returns (uint256);
+    function totalStaked() external view returns (uint256);
+    function totalSupply() external view returns (uint256 _totalSupply);
+    function totalSupplyMinted() external view returns (uint256 _totalSupply);
+    function pendingReward() external view returns (uint256);
+    function getStakedBalance(address _vault) external view returns (uint256 _balance);
+
+    function STAKE_TOKEN() external view returns (IERC20);
+    function REWARD_TOKEN() external view returns (IERC20);
+    function MIN_LOCKUP_PERIOD() external view returns (uint256);
+    function MAX_LOCKUP_PERIOD() external view returns (uint256);
+    function MP_APY() external view returns (uint256);
+    function MAX_BOOST() external view returns (uint256);
+
+    function calculateMP(uint256 _balance, uint256 _deltaTime) public pure returns (uint256);
+}

src/RewardsStreamerMP.sol

@@ -4,23 +4,22 @@ pragma solidity ^0.8.26;
 import { IERC20 } from "@openzeppelin/contracts/token/ERC20/IERC20.sol";
 import { ReentrancyGuard } from "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
 
-// Rewards Streamer with Multiplier Points
-contract RewardsStreamerMP is ReentrancyGuard {
-    error StakingManager__AmountCannotBeZero();
-    error StakingManager__TransferFailed();
-    error StakingManager__InsufficientBalance();
-    error StakingManager__InvalidLockingPeriod();
-    error StakingManager__CannotRestakeWithLockedFunds();
-    error StakingManager__TokensAreLocked();
+import { TrustedCodehashAccess } from "./access/TrustedCodehashAccess.sol";
+import { IStakeManager } from "./IStakeManager.sol";
 
-    IERC20 public immutable STAKING_TOKEN;
+// Rewards Streamer with Multiplier Points
+contract RewardsStreamerMP is IStakeManager, TrustedCodehashAccess, ReentrancyGuard {
+    error StakeManager__TransferFailed();
+    error StakeManager__CannotRestakeWithLockedFunds();
+
+    IERC20 public immutable STAKE_TOKEN;
     IERC20 public immutable REWARD_TOKEN;
 
     uint256 public constant SCALE_FACTOR = 1e18;
-    uint256 public constant MP_RATE_PER_YEAR = 1e18;
+    uint256 public constant MP_APY = 1e18;
 
-    uint256 public constant MIN_LOCKING_PERIOD = 90 days;
-    uint256 public constant MAX_LOCKING_PERIOD = 4 * 365 days;
+    uint256 public constant MIN_LOCKUP_PERIOD = 90 days;
+    uint256 public constant MAX_LOCKUP_PERIOD = 4 * 365 days;
     uint256 public constant MAX_MULTIPLIER = 4;
 
     uint256 public totalStaked;
@@ -42,18 +41,18 @@ contract RewardsStreamerMP is ReentrancyGuard {
     mapping(address account => UserInfo data) public users;
 
     constructor(address _stakingToken, address _rewardToken) {
-        STAKING_TOKEN = IERC20(_stakingToken);
+        STAKE_TOKEN = IERC20(_stakingToken);
         REWARD_TOKEN = IERC20(_rewardToken);
         lastMPUpdatedTime = block.timestamp;
     }
 
-    function stake(uint256 amount, uint256 lockPeriod) external nonReentrant {
-        if (amount == 0) {
-            revert StakingManager__AmountCannotBeZero();
+    function stake(uint256 _amount, uint256 _seconds) external onlyTrustedCodehash nonReentrant {
+        if (_amount == 0) {
+            revert StakeManager__StakeIsTooLow();
         }
 
-        if (lockPeriod != 0 && (lockPeriod < MIN_LOCKING_PERIOD || lockPeriod > MAX_LOCKING_PERIOD)) {
-            revert StakingManager__InvalidLockingPeriod();
+        if (_seconds != 0 && (_seconds < MIN_LOCKUP_PERIOD || _seconds > MAX_LOCKUP_PERIOD)) {
+            revert StakeManager__InvalidLockTime();
         }
 
         _updateGlobalState();
@@ -61,7 +60,7 @@ contract RewardsStreamerMP is ReentrancyGuard {
 
         UserInfo storage user = users[msg.sender];
         if (user.lockUntil != 0 && user.lockUntil > block.timestamp) {
-            revert StakingManager__CannotRestakeWithLockedFunds();
+            revert StakeManager__CannotRestakeWithLockedFunds();
         }
 
         uint256 userRewards = calculateUserRewards(msg.sender);
@@ -69,23 +68,23 @@ contract RewardsStreamerMP is ReentrancyGuard {
             distributeRewards(msg.sender, userRewards);
         }
 
-        bool success = STAKING_TOKEN.transferFrom(msg.sender, address(this), amount);
+        bool success = STAKE_TOKEN.transferFrom(msg.sender, address(this), _amount);
         if (!success) {
-            revert StakingManager__TransferFailed();
+            revert StakeManager__TransferFailed();
         }
 
-        user.stakedBalance += amount;
-        totalStaked += amount;
+        user.stakedBalance += _amount;
+        totalStaked += _amount;
 
-        uint256 initialMP = amount;
-        uint256 userPotentialMP = amount * MAX_MULTIPLIER;
+        uint256 initialMP = _amount;
+        uint256 userPotentialMP = _amount * MAX_MULTIPLIER;
 
-        if (lockPeriod != 0) {
-            uint256 lockMultiplier = (lockPeriod * MAX_MULTIPLIER * SCALE_FACTOR) / MAX_LOCKING_PERIOD;
+        if (_seconds != 0) {
+            uint256 lockMultiplier = (_seconds * MAX_MULTIPLIER * SCALE_FACTOR) / MAX_LOCKUP_PERIOD;
             lockMultiplier = lockMultiplier / SCALE_FACTOR;
-            initialMP += (amount * lockMultiplier);
-            userPotentialMP += (amount * lockMultiplier);
-            user.lockUntil = block.timestamp + lockPeriod;
+            initialMP += (_amount * lockMultiplier);
+            userPotentialMP += (_amount * lockMultiplier);
+            user.lockUntil = block.timestamp + _seconds;
         } else {
             user.lockUntil = 0;
         }
@@ -100,14 +99,14 @@ contract RewardsStreamerMP is ReentrancyGuard {
         user.lastMPUpdateTime = block.timestamp;
     }
 
-    function unstake(uint256 amount) external nonReentrant {
+    function unstake(uint256 _amount) external onlyTrustedCodehash nonReentrant {
         UserInfo storage user = users[msg.sender];
-        if (amount > user.stakedBalance) {
-            revert StakingManager__InsufficientBalance();
+        if (_amount > user.stakedBalance) {
+            revert StakeManager__InsufficientFunds();
         }
 
         if (block.timestamp < user.lockUntil) {
-            revert StakingManager__TokensAreLocked();
+            revert StakeManager__FundsLocked();
         }
 
         _updateGlobalState();
@@ -119,10 +118,10 @@ contract RewardsStreamerMP is ReentrancyGuard {
         }
 
         uint256 previousStakedBalance = user.stakedBalance;
-        user.stakedBalance -= amount;
-        totalStaked -= amount;
+        user.stakedBalance -= _amount;
+        totalStaked -= _amount;
 
-        uint256 amountRatio = (amount * SCALE_FACTOR) / previousStakedBalance;
+        uint256 amountRatio = (_amount * SCALE_FACTOR) / previousStakedBalance;
         uint256 mpToReduce = (user.userMP * amountRatio) / SCALE_FACTOR;
         uint256 potentialMPToReduce = (user.userPotentialMP * amountRatio) / SCALE_FACTOR;
 
@@ -131,14 +130,37 @@ contract RewardsStreamerMP is ReentrancyGuard {
         totalMP -= mpToReduce;
         potentialMP -= potentialMPToReduce;
 
-        bool success = STAKING_TOKEN.transfer(msg.sender, amount);
+        bool success = STAKE_TOKEN.transfer(msg.sender, _amount);
         if (!success) {
-            revert StakingManager__TransferFailed();
+            revert StakeManager__TransferFailed();
         }
 
         user.userRewardIndex = rewardIndex;
     }
 
+    function lock(uint256 _secondsIncrease) external onlyTrustedCodehash {
+        //TODO: increase lock time
+        revert("Not implemented");
+    }
+
+    function exit() external returns (bool _leaveAccepted) {
+        if (!isTrustedCodehash(msg.sender.codehash)) {
+            //case owner removed access from a class of StakeVault,. they might exit
+            delete user[msg.sender];
+            return true;
+        } else {
+            //TODO: handle other cases
+            //TODO: handle update/migration case
+            //TODO: handle emergency exit
+            revert("Not implemented");
+        }
+    }
+
+    function acceptUpdate() external onlyTrustedCodehash returns (address _migrated) {
+        //TODO: handle update/migration
+        revert("Not implemented");
+    }
+
     function _updateGlobalState() internal {
         updateGlobalMP();
         updateRewardIndex();
@@ -160,7 +182,7 @@ contract RewardsStreamerMP is ReentrancyGuard {
             return;
         }
 
-        uint256 accruedMP = (timeDiff * totalStaked * MP_RATE_PER_YEAR) / (365 days * SCALE_FACTOR);
+        uint256 accruedMP = calculateMP(totalStaked, timeDiff);
         if (accruedMP > potentialMP) {
             accruedMP = potentialMP;
         }
@@ -193,8 +215,8 @@ contract RewardsStreamerMP is ReentrancyGuard {
         }
     }
 
-    function updateUserMP(address userAddress) internal {
-        UserInfo storage user = users[userAddress];
+    function updateUserMP(address _vault) internal {
+        UserInfo storage user = users[_vault];
 
         if (user.userPotentialMP == 0 || user.stakedBalance == 0) {
             user.lastMPUpdateTime = block.timestamp;
@@ -206,7 +228,7 @@ contract RewardsStreamerMP is ReentrancyGuard {
             return;
         }
 
-        uint256 accruedMP = (timeDiff * user.stakedBalance * MP_RATE_PER_YEAR) / (365 days * SCALE_FACTOR);
+        uint256 accruedMP = calculateMP(user.stakedBalance, timeDiff);
 
         if (accruedMP > user.userPotentialMP) {
             accruedMP = user.userPotentialMP;
@@ -218,8 +240,8 @@ contract RewardsStreamerMP is ReentrancyGuard {
         user.lastMPUpdateTime = block.timestamp;
     }
 
-    function calculateUserRewards(address userAddress) public view returns (uint256) {
-        UserInfo storage user = users[userAddress];
+    function calculateUserRewards(address _vault) public view returns (uint256) {
+        UserInfo storage user = users[_vault];
         uint256 userWeight = user.stakedBalance + user.userMP;
         uint256 deltaRewardIndex = rewardIndex - user.userRewardIndex;
         return (userWeight * deltaRewardIndex) / SCALE_FACTOR;
@@ -236,19 +258,35 @@ contract RewardsStreamerMP is ReentrancyGuard {
 
         bool success = REWARD_TOKEN.transfer(to, amount);
         if (!success) {
-            revert StakingManager__TransferFailed();
+            revert StakeManager__TransferFailed();
         }
     }
 
-    function getStakedBalance(address userAddress) external view returns (uint256) {
-        return users[userAddress].stakedBalance;
+    function calculateMP(uint256 _balance, uint256 _deltaTime) public pure returns (uint256) {
+        return (_deltaTime * _balance * MP_APY) / (365 days * SCALE_FACTOR);
     }
 
-    function getPendingRewards(address userAddress) external view returns (uint256) {
-        return calculateUserRewards(userAddress);
+    function getStakedBalance(address _vault) external view returns (uint256 _balance) {
+        return users[_vault].stakedBalance;
     }
 
-    function getUserInfo(address userAddress) external view returns (UserInfo memory) {
-        return users[userAddress];
+    function getPendingRewards(address _vault) external view returns (uint256) {
+        return calculateUserRewards(_vault);
+    }
+
+    function getUserInfo(address _vault) external view returns (UserInfo memory) {
+        return users[_vault];
+    }
+
+    function totalSupplyMinted() external view returns (uint256 _totalSupply) {
+        return totalStaked + totalMP;
+    }
+
+    function totalSupply() external view returns (uint256 _totalSupply) {
+        return totalStaked + totalMP + potentialMP;
+    }
+
+    function pendingReward() external view returns (uint256 _pendingReward) {
+        return STAKE_TOKEN().balanceOf(address(this)) - accountedRewards;
     }
 }

src/StakeVault.sol

@@ -0,0 +1,229 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.18;
+
+import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { IERC20 } from "@openzeppelin/contracts/token/IERC20/IERC20.sol";
+import { IStakeManager } from "./IStakeManager.sol";
+
+/**
+ * @title StakeVault
+ * @author Ricardo Guilherme Schmidt <ricardo3@status.im>
+ * @notice A contract to secure user stakes and manage staking with IStakeManager.
+ * @dev This contract is owned by the user and allows staking, unstaking, and withdrawing tokens.
+ */
+contract StakeVault is Ownable {
+    error StakeVault__NoEnoughAvailableBalance();
+    error StakeVault__InvalidDestinationAddress();
+    error StakeVault__UpdateNotAvailable();
+    error StakeVault__StakingFailed();
+    error StakeVault__UnstakingFailed();
+
+    //STAKE_TOKEN must be kept as an immutable, otherwise, IStakeManager would accept StakeVaults with any token
+    //if is needed that STAKE_TOKEN to be a variable, IStakeManager should be changed to check codehash and
+    //StakeVault(msg.sender).STAKE_TOKEN()
+    IERC20 public immutable STAKE_TOKEN;
+    IStakeManager private stakeManager;
+
+    /**
+     * @dev Emitted when tokens are staked.
+     * @param from The address from which tokens are transferred.
+     * @param to The address receiving the staked tokens (this contract).
+     * @param amount The amount of tokens staked.
+     * @param time The time period for which tokens are staked.
+     */
+    event Staked(address indexed from, address indexed to, uint256 amount, uint256 time);
+
+    modifier validDestination(address _destination) {
+        if (_destination == address(0)) {
+            revert StakeVault__InvalidDestinationAddress();
+        }
+        _;
+    }
+
+    /**
+     * @notice Initializes the contract with the owner, staked token, and stake manager.
+     * @param _owner The address of the owner.
+     * @param _stakedToken The IERC20 token to be staked.
+     * @param _stakeManager The address of the IStakeManager contract.
+     */
+    constructor(address _owner, IStakeManager _stakeManager) {
+        _transferOwnership(_owner);
+        STAKE_TOKEN = _stakeManager.STAKE_TOKEN();
+        stakeManager = _stakeManager;
+    }
+
+    /**
+     * @notice Stake tokens for a specified time.
+     * @param _amount The amount of tokens to stake.
+     * @param _seconds The time period to stake for.
+     */
+    function stake(uint256 _amount, uint256 _seconds) external onlyOwner {
+        _stake(_amount, _seconds, msg.sender);
+    }
+
+    /**
+     * @notice Stake tokens from a specified address for a specified time.
+     * @param _amount The amount of tokens to stake.
+     * @param _seconds The time period to stake for.
+     * @param _from The address from which tokens will be transferred.
+     */
+    function stake(uint256 _amount, uint256 _seconds, address _from) external onlyOwner {
+        _stake(_amount, _seconds, _from);
+    }
+
+    /**
+     * @notice Extends the lock time of the stake.
+     * @param _secondsIncrease The additional time to lock the stake.
+     */
+    function lock(uint256 _secondsIncrease) external onlyOwner {
+        stakeManager.lock(_secondsIncrease);
+    }
+
+    /**
+     * @notice Unstake a specified amount of tokens and send to the owner.
+     * @param _amount The amount of tokens to unstake.
+     */
+    function unstake(uint256 _amount) external onlyOwner {
+        _unstake(_amount, msg.sender);
+    }
+
+    /**
+     * @notice Unstake a specified amount of tokens and send to a destination address.
+     * @param _amount The amount of tokens to unstake.
+     * @param _destination The address to receive the unstaked tokens.
+     */
+    function unstake(uint256 _amount, address _destination) external onlyOwner validDestination(_destination) {
+        _unstake(_amount, _destination);
+    }
+
+    /**
+     * @notice Withdraw tokens from the contract.
+     * @param _token The IERC20 token to withdraw.
+     * @param _amount The amount of tokens to withdraw.
+     */
+    function withdraw(IERC20 _token, uint256 _amount) external onlyOwner {
+        _withdraw(_token, _amount, msg.sender);
+    }
+
+    /**
+     * @notice Withdraw tokens from the contract to a destination address.
+     * @param _token The IERC20 token to withdraw.
+     * @param _amount The amount of tokens to withdraw.
+     * @param _destination The address to receive the tokens.
+     */
+    function withdraw(
+        IERC20 _token,
+        uint256 _amount,
+        address _destination
+    )
+        external
+        onlyOwner
+        validDestination(_destination)
+    {
+        _withdraw(_token, _amount, _destination);
+    }
+
+    /**
+     * @notice Withdraw Ether from the contract to the owner address.
+     * @param _amount The amount of Ether to withdraw.
+     */
+    function withdraw(uint256 _amount) external onlyOwner {
+        _withdraw(_amount, payable(msg.sender));
+    }
+
+    /**
+     * @notice Withdraw Ether from the contract to a destination address.
+     * @param _amount The amount of Ether to withdraw.
+     * @param _destination The address to receive the Ether.
+     */
+    function withdraw(
+        uint256 _amount,
+        address payable _destination
+    )
+        external
+        onlyOwner
+        validDestination(_destination)
+    {
+        _withdraw(_amount, _destination);
+    }
+
+    /**
+     * @notice Leave staking contract and withdraw all tokens to the owner, in case StakeManager have breached contract.
+     */
+    function exit() external onlyOwner {
+        _exit(msg.sender);
+    }
+
+    /**
+     * @notice Reject update, exit staking contract and withdraw all tokens to a destination address.
+     * @param _destination The address to receive the tokens.
+     */
+    function exit(address _destination) external onlyOwner validDestination(_destination) {
+        _exit(_destination);
+    }
+
+    /**
+     * @notice Opt-in update to a new IStakeManager contract.
+     * @dev Updates the stakeManager to the migrated contract.
+     */
+    function acceptUpdate() external onlyOwner {
+        IStakeManager migrated = stakeManager.acceptUpdate();
+        if (address(migrated) == address(0)) revert StakeVault__UpdateNotAvailable();
+        stakeManager = migrated;
+    }
+
+    /**
+     * @notice Returns the available amount of a token that can be withdrawn.
+     * @param _token The IERC20 token to check.
+     * @return The amount of token available for withdrawal.
+     */
+    function availableWithdraw(IERC20 _token) external view returns (uint256) {
+        if (_token == STAKE_TOKEN) {
+            return STAKE_TOKEN.balanceOf(address(this)) - amountStaked();
+        }
+        return _token.balanceOf(address(this));
+    }
+
+    function _stake(uint256 _amount, uint256 _seconds, address _source) internal {
+        bool success = STAKE_TOKEN.transferFrom(_source, address(this), _amount);
+        if (!success) {
+            revert StakeVault__StakingFailed();
+        }
+
+        stakeManager.stake(_amount, _seconds);
+
+        emit Staked(_source, address(this), _amount, _seconds);
+    }
+
+    function _unstake(uint256 _amount, address _destination) internal {
+        stakeManager.unstake(_amount);
+        bool success = STAKE_TOKEN.transfer(_destination, _amount);
+        if (!success) {
+            revert StakeVault__UnstakingFailed();
+        }
+    }
+
+    function _exit(address _destination) internal {
+        if (IStakeManager.isTrustedCodehash(this.codehash)) {
+            revert StakeVault__LeaveNotAvailable();
+        }
+        stakeManager.exit();
+        STAKE_TOKEN.transferFrom(address(this), _destination, STAKE_TOKEN.balanceOf(address(this)));
+    }
+
+    function _withdraw(IERC20 _token, uint256 _amount, address _destination) internal {
+        if (_token == STAKE_TOKEN && STAKE_TOKEN.balanceOf(address(this)) - amountStaked() < _amount) {
+            revert StakeVault__NoEnoughAvailableBalance();
+        }
+        _token.transfer(_destination, _amount);
+    }
+
+    function _withdraw(uint256 _amount, address payable _destination) internal {
+        _destination.transfer(_amount);
+    }
+
+    function amountStaked() public view returns (uint256) {
+        return stakeManager.getStakedBalance(address(this));
+    }
+}

src/access/ITrustedCodehashAccess.sol

@@ -0,0 +1,29 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.18;
+
+/**
+ * @title TrustedCodehashAccess
+ * @author Ricardo Guilherme Schmidt <ricardo3@status.im>
+ * @notice Ensures that only specific contract bytecode hashes are trusted to
+ *         interact with the functions using the `onlyTrustedCodehash` modifier.
+ */
+interface ITrustedCodehashAccess {
+    error TrustedCodehashAccess__UnauthorizedCodehash();
+
+    event TrustedCodehashUpdated(bytes32 indexed codehash, bool trusted);
+
+    /**
+     * @notice Allows the owner to set or update the trust status for a contract's codehash.
+     * @dev Emits the `TrustedCodehashUpdated` event whenever a codehash is updated.
+     * @param _codehash The bytecode hash of the contract.
+     * @param _trusted Boolean flag to designate the contract as trusted or not.
+     */
+    function setTrustedCodehash(bytes32 _codehash, bool _trusted) external;
+
+    /**
+     * @notice Checks if a contract's codehash is trusted to interact with protected functions.
+     * @param _codehash The bytecode hash of the contract.
+     * @return bool True if the codehash is trusted, false otherwise.
+     */
+    function isTrustedCodehash(bytes32 _codehash) external view returns (bool);
+}

src/access/TrustedCodehashAccess.sol

@@ -0,0 +1,47 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.18;
+
+import { Ownable } from "@openzeppelin/contracts/access/Ownable.sol";
+import { ITrustedCodehashAccess } from "./ITrustedCodehashAccess.sol";
+/**
+ * @title TrustedCodehashAccess
+ * @author Ricardo Guilherme Schmidt <ricardo3@status.im>
+ * @notice Ensures that only specific contract bytecode hashes are trusted to
+ *         interact with the functions using the `onlyTrustedCodehash` modifier.
+ */
+
+contract TrustedCodehashAccess is ITrustedCodehashAccess, Ownable {
+    mapping(bytes32 codehash => bool permission) private trustedCodehashes;
+
+    /**
+     * @notice Restricts access based on the codehash of the caller.
+     *         Only contracts with trusted codehashes can execute functions using this modifier.
+     */
+    modifier onlyTrustedCodehash() {
+        bytes32 codehash = msg.sender.codehash;
+        if (!trustedCodehashes[codehash]) {
+            revert TrustedCodehashAccess__UnauthorizedCodehash();
+        }
+        _;
+    }
+
+    /**
+     * @notice Allows the owner to set or update the trust status for a contract's codehash.
+     * @dev Emits the `TrustedCodehashUpdated` event whenever a codehash is updated.
+     * @param _codehash The bytecode hash of the contract.
+     * @param _trusted Boolean flag to designate the contract as trusted or not.
+     */
+    function setTrustedCodehash(bytes32 _codehash, bool _trusted) external onlyOwner {
+        trustedCodehashes[_codehash] = _trusted;
+        emit TrustedCodehashUpdated(_codehash, _trusted);
+    }
+
+    /**
+     * @notice Checks if a contract's codehash is trusted to interact with protected functions.
+     * @param _codehash The bytecode hash of the contract.
+     * @return bool True if the codehash is trusted, false otherwise.
+     */
+    function isTrustedCodehash(bytes32 _codehash) external view returns (bool) {
+        return trustedCodehashes[_codehash];
+    }
+}