From ff3a7b09ea7563df4b88682a066a1a5fcfeca165 Mon Sep 17 00:00:00 2001
From: Ricardo Guilherme Schmidt <3esmit@gmail.com>
Date: Fri, 23 Jun 2023 15:03:52 -0300
Subject: [PATCH] Only vaults that Re previosly allowed

---
 contracts/StakeManager.sol | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/contracts/StakeManager.sol b/contracts/StakeManager.sol
index c122ba5..355a3c4 100644
--- a/contracts/StakeManager.sol
+++ b/contracts/StakeManager.sol
@@ -4,7 +4,7 @@ pragma solidity 0.8.19;
 
 import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
 
-contract StakeManager is ERC20 {
+contract StakeManager is Controlled {
 
     struct Account {
         uint256 lockUntil;
@@ -27,12 +27,18 @@ contract StakeManager is ERC20 {
 
     mapping (address => Account) accounts;
     mapping (uint256 => Epoch) epoch;
+    mapping (bytes32 => bool) isVault;
 
     ERC20 stakedToken;
     uint256 currentEpoch;
     uint256 pendingReward;
     uint256 public totalSupply;
 
+    modifier onlyVault {
+        require(isVault[msg.sender.codehash], "Not a vault")
+        _;
+    }
+
     constructor() {
         epoch[0].startTime = now();
     }
@@ -135,6 +141,14 @@ contract StakeManager is ERC20 {
         pendingReward -= userReward;
         stakedToken.transfer(_vault, userReward);
     }
+    
+    /**
+     * @notice Enables a contract class to interact with staking functions
+     * @param _codehash bytecode hash of contract
+     */
+    function setVault(bytes32 _codehash) external onlyController {
+        isVault[_codehash] = true;  
+    }
 
     function checkMaxMultiplier(uint256 _increasedMultiplier, uint256 _currentMp) private view returns(uint256 _maxToIncrease) {
         uint256 newMp = _increasedMultiplier + _currentMp;