Update upgradeability util tests

with tests of implementation variable/slot getter functions

tests/test_upgradeability_util.py

@@ -1,14 +1,14 @@
 import os
-import json
-import re
 
 from solc_select import solc_select
-from deepdiff import DeepDiff
 
 from slither import Slither
-from slither.core.declarations import Function
-from slither.core.variables import StateVariable
-from slither.utils.upgradeability import compare
+from slither.core.expressions import Literal
+from slither.utils.upgradeability import (
+    compare,
+    get_proxy_implementation_var,
+    get_proxy_implementation_slot,
+)
 
 SLITHER_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 UPGRADE_TEST_ROOT = os.path.join(SLITHER_ROOT, "tests", "upgradeability-util")
@@ -18,7 +18,7 @@ UPGRADE_TEST_ROOT = os.path.join(SLITHER_ROOT, "tests", "upgradeability-util")
 def test_upgrades_compare() -> None:
     solc_select.switch_global_version("0.8.2", always_install=True)
 
-    sl = Slither(os.path.join(UPGRADE_TEST_ROOT, "TestUpgrades.sol"))
+    sl = Slither(os.path.join(UPGRADE_TEST_ROOT, "TestUpgrades-0.8.2.sol"))
     v1 = sl.get_contract_from_name("ContractV1")[0]
     v2 = sl.get_contract_from_name("ContractV2")[0]
     missing_vars, new_vars, tainted_vars, new_funcs, modified_funcs, tainted_funcs = compare(v1, v2)
@@ -26,11 +26,53 @@ def test_upgrades_compare() -> None:
     assert new_vars == [v2.get_state_variable_from_name("stateC")]
     assert tainted_vars == [
         v2.get_state_variable_from_name("stateB"),
-        v2.get_state_variable_from_name("bug")
+        v2.get_state_variable_from_name("bug"),
     ]
     assert new_funcs == [v2.get_function_from_signature("i()")]
     assert modified_funcs == [v2.get_function_from_signature("checkB()")]
     assert tainted_funcs == [
         v2.get_function_from_signature("g(uint256)"),
-        v2.get_function_from_signature("h()")
+        v2.get_function_from_signature("h()"),
     ]
+
+
+def test_upgrades_implementation_var() -> None:
+    solc_select.switch_global_version("0.8.2", always_install=True)
+    sl = Slither(os.path.join(UPGRADE_TEST_ROOT, "TestUpgrades-0.8.2.sol"))
+
+    erc_1967_proxy = sl.get_contract_from_name("ERC1967Proxy")[0]
+    storage_proxy = sl.get_contract_from_name("InheritedStorageProxy")[0]
+
+    target = get_proxy_implementation_var(erc_1967_proxy)
+    slot = get_proxy_implementation_slot(erc_1967_proxy)
+    assert target == erc_1967_proxy.get_state_variable_from_name("_IMPLEMENTATION_SLOT")
+    assert slot.slot == 0x360894A13BA1A3210667C828492DB98DCA3E2076CC3735A920A3CA505D382BBC
+    target = get_proxy_implementation_var(storage_proxy)
+    slot = get_proxy_implementation_slot(storage_proxy)
+    assert target == storage_proxy.get_state_variable_from_name("implementation")
+    assert slot.slot == 1
+
+    solc_select.switch_global_version("0.5.0", always_install=True)
+    sl = Slither(os.path.join(UPGRADE_TEST_ROOT, "TestUpgrades-0.5.0.sol"))
+
+    eip_1822_proxy = sl.get_contract_from_name("EIP1822Proxy")[0]
+    zos_proxy = sl.get_contract_from_name("ZosProxy")[0]
+    master_copy_proxy = sl.get_contract_from_name("MasterCopyProxy")[0]
+
+    target = get_proxy_implementation_var(eip_1822_proxy)
+    slot = get_proxy_implementation_slot(eip_1822_proxy)
+    assert target not in eip_1822_proxy.state_variables_ordered
+    assert target.name == "contractLogic" and isinstance(target.expression, Literal)
+    assert (
+        target.expression.value
+        == "0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7"
+    )
+    assert slot.slot == 0xC5F16F0FCC639FA48A6947836D9850F504798523BF8C9A3A87D5876CF622BCF7
+    target = get_proxy_implementation_var(zos_proxy)
+    slot = get_proxy_implementation_slot(zos_proxy)
+    assert target == zos_proxy.get_state_variable_from_name("IMPLEMENTATION_SLOT")
+    assert slot.slot == 0x7050C9E0F4CA769C69BD3A8EF740BC37934F8E2C036E5A723FD8EE048ED3F8C3
+    target = get_proxy_implementation_var(master_copy_proxy)
+    slot = get_proxy_implementation_slot(master_copy_proxy)
+    assert target == master_copy_proxy.get_state_variable_from_name("masterCopy")
+    assert slot.slot == 0

tests/upgradeability-util/TestUpgrades-0.5.0.sol

@@ -0,0 +1,5 @@
+pragma solidity ^0.5.0;
+
+import "./src/EIP1822Proxy.sol";
+import "./src/ZosProxy.sol";
+import "./src/MasterCopyProxy.sol";

tests/upgradeability-util/src/EIP1822Proxy.sol

@@ -0,0 +1,47 @@
+pragma solidity ^0.5.0;
+
+contract EIP1822Proxy {
+    // Code position in storage is keccak256("PROXIABLE") = "0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7"
+    constructor(bytes memory constructData, address contractLogic) public {
+        // save the code address
+        assembly { // solium-disable-line
+            sstore(0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7, contractLogic)
+        }
+        (bool success, bytes memory _ ) = contractLogic.delegatecall(constructData); // solium-disable-line
+        require(success, "Construction failed");
+    }
+
+    function() external payable {
+        assembly { // solium-disable-line
+            let contractLogic := sload(0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7)
+            calldatacopy(0x0, 0x0, calldatasize)
+            let success := delegatecall(sub(gas, 10000), contractLogic, 0x0, calldatasize, 0, 0)
+            let retSz := returndatasize
+            returndatacopy(0, 0, retSz)
+            switch success
+            case 0 {
+                revert(0, retSz)
+            }
+            default {
+                return(0, retSz)
+            }
+        }
+    }
+}
+
+contract EIP1822Proxiable {
+    // Code position in storage is keccak256("PROXIABLE") = "0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7"
+
+    function updateCodeAddress(address newAddress) internal {
+        require(
+            bytes32(0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7) == EIP1822Proxiable(newAddress).proxiableUUID(),
+            "Not compatible"
+        );
+        assembly { // solium-disable-line
+            sstore(0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7, newAddress)
+        }
+    }
+    function proxiableUUID() public pure returns (bytes32) {
+        return 0xc5f16f0fcc639fa48a6947836d9850f504798523bf8c9a3a87d5876cf622bcf7;
+    }
+}

tests/upgradeability-util/src/MasterCopyProxy.sol

@@ -0,0 +1,27 @@
+pragma solidity ^0.5.0;
+
+contract MasterCopyProxy {
+    address internal masterCopy;
+
+    constructor(address _masterCopy)
+        public
+    {
+        require(_masterCopy != address(0), "Invalid master copy address provided");
+        masterCopy = _masterCopy;
+    }
+
+    /// @dev Fallback function forwards all transactions and returns all received return data.
+    function ()
+        external
+        payable
+    {
+        // solium-disable-next-line security/no-inline-assembly
+        assembly {
+            calldatacopy(0, 0, calldatasize())
+            let success := delegatecall(gas, sload(0), 0, calldatasize(), 0, 0)
+            returndatacopy(0, 0, returndatasize())
+            if eq(success, 0) { revert(0, returndatasize()) }
+            return(0, returndatasize())
+        }
+    }
+}

tests/upgradeability-util/src/ZosProxy.sol

@@ -0,0 +1,67 @@
+pragma solidity ^0.5.0;
+
+contract ZosProxy {
+  function () payable external {
+    _fallback();
+  }
+
+  function _implementation() internal view returns (address);
+
+  function _delegate(address implementation) internal {
+    assembly {
+      calldatacopy(0, 0, calldatasize)
+      let result := delegatecall(gas, implementation, 0, calldatasize, 0, 0)
+      returndatacopy(0, 0, returndatasize)
+      switch result
+      case 0 { revert(0, returndatasize) }
+      default { return(0, returndatasize) }
+    }
+  }
+
+  function _willFallback() internal {
+  }
+
+  function _fallback() internal {
+    _willFallback();
+    _delegate(_implementation());
+  }
+}
+
+library AddressUtils {
+  function isContract(address addr) internal view returns (bool) {
+    uint256 size;
+    assembly { size := extcodesize(addr) }
+    return size > 0;
+  }
+}
+
+contract UpgradeabilityProxy is ZosProxy {
+  event Upgraded(address indexed implementation);
+
+  bytes32 private constant IMPLEMENTATION_SLOT = 0x7050c9e0f4ca769c69bd3a8ef740bc37934f8e2c036e5a723fd8ee048ed3f8c3;
+
+  constructor(address _implementation) public payable {
+    assert(IMPLEMENTATION_SLOT == keccak256("org.zeppelinos.proxy.implementation"));
+    _setImplementation(_implementation);
+  }
+
+  function _implementation() internal view returns (address impl) {
+    bytes32 slot = IMPLEMENTATION_SLOT;
+    assembly {
+      impl := sload(slot)
+    }
+  }
+
+  function _upgradeTo(address newImplementation) internal {
+    _setImplementation(newImplementation);
+    emit Upgraded(newImplementation);
+  }
+
+  function _setImplementation(address newImplementation) private {
+    require(AddressUtils.isContract(newImplementation), "Cannot set a proxy implementation to a non-contract address");
+    bytes32 slot = IMPLEMENTATION_SLOT;
+    assembly {
+      sstore(slot, newImplementation)
+    }
+  }
+}