upgrades: support EIP1967 for Proxy Storage Slots

Author: abarmat

contracts/DisputeManager.sol

@@ -143,6 +143,9 @@ contract DisputeManager is Governed {
         uint256 tokens
     );
 
+    /**
+     * @dev Check if the caller is the arbitrator.
+     */
     modifier onlyArbitrator {
         require(msg.sender == arbitrator, "Caller is not the Arbitrator");
         _;
@@ -165,6 +168,7 @@ contract DisputeManager is Governed {
         uint256 _fishermanRewardPercentage,
         uint256 _slashingPercentage
     ) public {
+        Governed._initialize(msg.sender);
         arbitrator = _arbitrator;
         token = IGraphToken(_token);
         staking = IStaking(_staking);

contracts/GNS.sol

@@ -179,6 +179,7 @@ contract GNS is Governed, BancorFormula {
         address _curation,
         address _token
     ) public {
+        Governed._initialize(msg.sender);
         erc1056Registry = IEthereumDIDRegistry(_didRegistry);
         curation = ICuration(_curation);
         token = IGraphToken(_token);

contracts/curation/Curation.sol

@@ -4,6 +4,7 @@ import "@openzeppelin/contracts/math/SafeMath.sol";
 
 import "../governance/Governed.sol";
 import "../upgrades/GraphProxy.sol";
+import "../upgrades/GraphUpgradeable.sol";
 
 import "./CurationStorage.sol";
 import "./ICuration.sol";
@@ -20,7 +21,7 @@ import "./ICuration.sol";
  * Holders can burn GST tokens using this contract to get GRT tokens back according to the
  * bonding curve.
  */
-contract Curation is CurationV1Storage, ICuration, Governed {
+contract Curation is CurationV1Storage, GraphUpgradeable, ICuration, Governed {
     using SafeMath for uint256;
 
     // 100% in parts per million
@@ -61,23 +62,18 @@ contract Curation is CurationV1Storage, ICuration, Governed {
      */
     event Collected(bytes32 indexed subgraphDeploymentID, uint256 tokens);
 
-    /**
-     * @dev Check if the caller is the governor or initializing the implementation.
-     */
-    modifier onlyGovernorOrInit {
-        require(msg.sender == governor || msg.sender == implementation, "Only Governor can call");
-        _;
-    }
-
     /**
      * @dev Initialize this contract.
      */
     function initialize(
+        address _governor,
         address _token,
         uint32 _defaultReserveRatio,
         uint256 _minimumCurationStake
-    ) external onlyGovernorOrInit {
+    ) external onlyImpl {
+        Governed._initialize(_governor);
         BancorFormula._initialize();
+
         token = IGraphToken(_token);
         defaultReserveRatio = _defaultReserveRatio;
         minimumCurationStake = _minimumCurationStake;
@@ -96,21 +92,24 @@ contract Curation is CurationV1Storage, ICuration, Governed {
         uint32 _defaultReserveRatio,
         uint256 _minimumCurationStake
     ) external {
-        require(msg.sender == _proxy.governor(), "Only proxy governor can upgrade");
-
         // Accept to be the implementation for this proxy
-        _proxy.acceptImplementation();
+        _acceptUpgrade(_proxy);
 
         // Initialization
-        Curation(address(_proxy)).initialize(_token, _defaultReserveRatio, _minimumCurationStake);
+        Curation(address(_proxy)).initialize(
+            _proxy.admin(), // default governor is proxy admin
+            _token,
+            _defaultReserveRatio,
+            _minimumCurationStake
+        );
     }
 
     /**
      * @dev Set the staking contract used for fees distribution.
      * @notice Update the staking contract to `_staking`
      * @param _staking Address of the staking contract
      */
-    function setStaking(address _staking) external override onlyGovernorOrInit {
+    function setStaking(address _staking) external override onlyGovernor {
         staking = IStaking(_staking);
         emit ParameterUpdated("staking");
     }
@@ -120,11 +119,7 @@ contract Curation is CurationV1Storage, ICuration, Governed {
      * @notice Update the default reserver ratio to `_defaultReserveRatio`
      * @param _defaultReserveRatio Reserve ratio (in PPM)
      */
-    function setDefaultReserveRatio(uint32 _defaultReserveRatio)
-        external
-        override
-        onlyGovernorOrInit
-    {
+    function setDefaultReserveRatio(uint32 _defaultReserveRatio) external override onlyGovernor {
         // Reserve Ratio must be within 0% to 100% (exclusive, in PPM)
         require(_defaultReserveRatio > 0, "Default reserve ratio must be > 0");
         require(
@@ -141,11 +136,7 @@ contract Curation is CurationV1Storage, ICuration, Governed {
      * @notice Update the minimum stake amount to `_minimumCurationStake`
      * @param _minimumCurationStake Minimum amount of tokens required stake
      */
-    function setMinimumCurationStake(uint256 _minimumCurationStake)
-        external
-        override
-        onlyGovernorOrInit
-    {
+    function setMinimumCurationStake(uint256 _minimumCurationStake) external override onlyGovernor {
         require(_minimumCurationStake > 0, "Minimum curation stake cannot be 0");
         minimumCurationStake = _minimumCurationStake;
         emit ParameterUpdated("minimumCurationStake");
@@ -155,7 +146,7 @@ contract Curation is CurationV1Storage, ICuration, Governed {
      * @dev Set the fee percentage to charge when a curator withdraws stake.
      * @param _percentage Percentage fee charged when withdrawing stake
      */
-    function setWithdrawalFeePercentage(uint32 _percentage) external override onlyGovernorOrInit {
+    function setWithdrawalFeePercentage(uint32 _percentage) external override onlyGovernor {
         // Must be within 0% to 100% (inclusive)
         require(
             _percentage <= MAX_PPM,

contracts/curation/CurationStorage.sol

@@ -3,11 +3,10 @@ pragma solidity ^0.6.4;
 import "../bancor/BancorFormula.sol";
 import "../staking/IStaking.sol";
 import "../token/IGraphToken.sol";
-import "../upgrades/GraphProxyStorage.sol";
 
 import "./GraphSignalToken.sol";
 
-contract CurationV1Storage is GraphProxyStorage, BancorFormula {
+contract CurationV1Storage is BancorFormula {
     // -- Pool --
 
     struct CurationPool {

contracts/epochs/EpochManager.sol

@@ -4,6 +4,7 @@ import "@openzeppelin/contracts/math/SafeMath.sol";
 
 import "../governance/Governed.sol";
 import "../upgrades/GraphProxy.sol";
+import "../upgrades/GraphUpgradeable.sol";
 
 import "./EpochManagerStorage.sol";
 import "./IEpochManager.sol";
@@ -12,28 +13,22 @@ import "./IEpochManager.sol";
  * @title EpochManager contract
  * @dev Tracks epochs based on its block duration to sync contracts in the protocol.
  */
-contract EpochManager is EpochManagerV1Storage, IEpochManager, Governed {
+contract EpochManager is EpochManagerV1Storage, GraphUpgradeable, IEpochManager, Governed {
     using SafeMath for uint256;
 
     // -- Events --
 
     event EpochRun(uint256 indexed epoch, address caller);
     event EpochLengthUpdate(uint256 indexed epoch, uint256 epochLength);
 
-    /**
-     * @dev Check if the caller is the governor or initializing the implementation.
-     */
-    modifier onlyGovernorOrInit {
-        require(msg.sender == governor || msg.sender == implementation, "Only Governor can call");
-        _;
-    }
-
     /**
      * @dev Initialize this contract.
      */
-    function initialize(uint256 _epochLength) external onlyGovernorOrInit {
+    function initialize(address _governor, uint256 _epochLength) external onlyImpl {
         require(_epochLength > 0, "Epoch length cannot be 0");
 
+        Governed._initialize(_governor);
+
         lastLengthUpdateEpoch = 0;
         lastLengthUpdateBlock = blockNum();
         epochLength = _epochLength;
@@ -47,13 +42,11 @@ contract EpochManager is EpochManagerV1Storage, IEpochManager, Governed {
      * @param _epochLength Epoch length in blocks
      */
     function acceptProxy(GraphProxy _proxy, uint256 _epochLength) external {
-        require(msg.sender == _proxy.governor(), "Only proxy governor can upgrade");
-
         // Accept to be the implementation for this proxy
-        _proxy.acceptImplementation();
+        _acceptUpgrade(_proxy);
 
         // Initialization
-        EpochManager(address(_proxy)).initialize(_epochLength);
+        EpochManager(address(_proxy)).initialize(_proxy.admin(), _epochLength);
     }
 
     /**

contracts/epochs/EpochManagerStorage.sol

@@ -1,8 +1,6 @@
 pragma solidity ^0.6.4;
 
-import "../upgrades/GraphProxyStorage.sol";
-
-contract EpochManagerV1Storage is GraphProxyStorage {
+contract EpochManagerV1Storage {
     // -- State --
 
     // Epoch length in blocks

contracts/governance/Governed.sol

@@ -1,12 +1,15 @@
 pragma solidity ^0.6.4;
 
-import "./GovernedStorage.sol";
-
 /**
  * @title Graph Governance contract
  * @dev All contracts that will be owned by a Governor entity should extend this contract.
  */
-contract Governed is GovernedStorage {
+contract Governed {
+    // -- State --
+
+    address public governor;
+    address public pendingGovernor;
+
     // -- Events --
 
     event NewPendingOwnership(address indexed from, address indexed to);
@@ -22,10 +25,10 @@ contract Governed is GovernedStorage {
     }
 
     /**
-     * @dev Constructor. Set the governor to the contract caller.
+     * @dev Initialize the governor to the contract caller.
      */
-    constructor() public {
-        governor = msg.sender;
+    function _initialize(address _governor) internal {
+        governor = _governor;
     }
 
     /**

contracts/governance/GovernedStorage.sol

@@ -8,6 +8,7 @@ import "../epochs/IEpochManager.sol";
 import "../governance/Governed.sol";
 import "../token/IGraphToken.sol";
 import "../upgrades/GraphProxy.sol";
+import "../upgrades/GraphUpgradeable.sol";
 
 import "./IStaking.sol";
 import "./StakingStorage.sol";
@@ -17,7 +18,7 @@ import "./libs/Stakes.sol";
 /**
  * @title Staking contract
  */
-contract Staking is StakingV1Storage, IStaking, Governed {
+contract Staking is StakingV1Storage, GraphUpgradeable, IStaking, Governed {
     using SafeMath for uint256;
     using Stakes for Stakes.Indexer;
     using Rebates for Rebates.Pool;
@@ -168,14 +169,6 @@ contract Staking is StakingV1Storage, IStaking, Governed {
      */
     event SetOperator(address indexed indexer, address operator, bool allowed);
 
-    /**
-     * @dev Check if the caller is the governor or initializing the implementation.
-     */
-    modifier onlyGovernorOrInit {
-        require(msg.sender == governor || msg.sender == implementation, "Only Governor can call");
-        _;
-    }
-
     /**
      * @dev Check if the caller is the slasher.
      */
@@ -199,11 +192,15 @@ contract Staking is StakingV1Storage, IStaking, Governed {
     }
 
     /**
-     * @dev Staking Contract Constructor.
-     * @param _token Address of the Graph Protocol token
-     * @param _epochManager Address of the EpochManager contract
+     * @dev Initialize this contract.
      */
-    function initialize(address _token, address _epochManager) external onlyGovernorOrInit {
+    function initialize(
+        address _governor,
+        address _token,
+        address _epochManager
+    ) external onlyImpl {
+        Governed._initialize(_governor);
+
         token = IGraphToken(_token);
         epochManager = IEpochManager(_epochManager);
     }
@@ -219,13 +216,11 @@ contract Staking is StakingV1Storage, IStaking, Governed {
         address _token,
         address _epochManager
     ) external {
-        require(msg.sender == _proxy.governor(), "Only proxy governor can upgrade");
-
         // Accept to be the implementation for this proxy
-        _proxy.acceptImplementation();
+        _acceptUpgrade(_proxy);
 
         // Initialization
-        Staking(address(_proxy)).initialize(_token, _epochManager);
+        Staking(address(_proxy)).initialize(_proxy.admin(), _token, _epochManager);
     }
 
     /**

contracts/staking/Staking.sol

@@ -4,13 +4,12 @@ import "../curation/ICuration.sol";
 import "../epochs/IEpochManager.sol";
 import "../staking/IStaking.sol";
 import "../token/IGraphToken.sol";
-import "../upgrades/GraphProxyStorage.sol";
 
 import "./IStaking.sol";
 import "./libs/Rebates.sol";
 import "./libs/Stakes.sol";
 
-contract StakingV1Storage is GraphProxyStorage {
+contract StakingV1Storage {
     // -- Staking --
 
     // Time in blocks to unstake