Transpile debaffad

CHANGELOG.md

@@ -17,10 +17,14 @@
  * `ERC20`: do not update allowance on `transferFrom` when allowance is `type(uint256).max`. ([#3085](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/#3085))
  * `ERC777`: do not update allowance on `transferFrom` when allowance is `type(uint256).max`. ([#3085](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/#3085))
  * `SignedMath`: a new signed version of the Math library with `max`, `min`,  and `average`. ([#2686](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/2686))
+ * `ERC1967Upgrade`: Refactor the secure upgrade to use `ERC1822` instead of the previous rollback mechanism. This reduces code complexity and attack surface with similar security guarantees. ([#3021](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3021))
+ * `UUPSUpgradeable`: Add `ERC1822` compliance to support the updated secure upgrade mechanism. ([#3021](https://github.com/OpenZeppelin/openzeppelin-contracts/pull/3021))
+ * Some more functions have been made virtual to customize them via overrides. In many cases this will not imply that other functions in the contract will automatically adapt to the overridden definitions. People who wish to override should consult the source code to understand the impact and if they need to override any additional functions to achieve the desired behavior.
 
-### Breaking change
+### Breaking changes
 
-Solidity pragma in `utils/Address.sol` is increased from `^0.8.0` to `^0.8.1`. This is required by the `account.code.length` syntax that replaces inline assembly. This may require users to bump their compiler version from `0.8.0` to `0.8.1` or later. Note that other parts of the code already include stricter requirements.
+* `ERC1967Upgrade`: The function `_upgradeToAndCallSecure` was renamed to `_upgradeToAndCallUUPS`, along with the change in security mechanism described above.
+* `Address`: The Solidity pragma is increased from `^0.8.0` to `^0.8.1`. This is required by the `account.code.length` syntax that replaces inline assembly. This may require users to bump their compiler version from `0.8.0` to `0.8.1` or later. Note that other parts of the code already include stricter requirements.
 
 ## 4.4.2 (2022-01-11)
 

contracts/access/AccessControlEnumerableUpgradeable.sol

@@ -44,7 +44,7 @@ abstract contract AccessControlEnumerableUpgradeable is Initializable, IAccessCo
      * https://forum.openzeppelin.com/t/iterating-over-elements-on-enumerableset-in-openzeppelin-contracts/2296[forum post]
      * for more information.
      */
-    function getRoleMember(bytes32 role, uint256 index) public view override returns (address) {
+    function getRoleMember(bytes32 role, uint256 index) public view virtual override returns (address) {
         return _roleMembers[role].at(index);
     }
 
@@ -52,7 +52,7 @@ abstract contract AccessControlEnumerableUpgradeable is Initializable, IAccessCo
      * @dev Returns the number of accounts that have `role`. Can be used
      * together with {getRoleMember} to enumerate all bearers of a role.
      */
-    function getRoleMemberCount(bytes32 role) public view override returns (uint256) {
+    function getRoleMemberCount(bytes32 role) public view virtual override returns (uint256) {
         return _roleMembers[role].length();
     }
 

contracts/access/AccessControlUpgradeable.sol

@@ -90,7 +90,7 @@ abstract contract AccessControlUpgradeable is Initializable, ContextUpgradeable,
     /**
      * @dev Returns `true` if `account` has been granted `role`.
      */
-    function hasRole(bytes32 role, address account) public view override returns (bool) {
+    function hasRole(bytes32 role, address account) public view virtual override returns (bool) {
         return _roles[role].members[account];
     }
 
@@ -101,7 +101,7 @@ abstract contract AccessControlUpgradeable is Initializable, ContextUpgradeable,
      *
      *  /^AccessControl: account (0x[0-9a-f]{40}) is missing role (0x[0-9a-f]{64})$/
      */
-    function _checkRole(bytes32 role, address account) internal view {
+    function _checkRole(bytes32 role, address account) internal view virtual {
         if (!hasRole(role, account)) {
             revert(
                 string(
@@ -122,7 +122,7 @@ abstract contract AccessControlUpgradeable is Initializable, ContextUpgradeable,
      *
      * To change a role's admin, use {_setRoleAdmin}.
      */
-    function getRoleAdmin(bytes32 role) public view override returns (bytes32) {
+    function getRoleAdmin(bytes32 role) public view virtual override returns (bytes32) {
         return _roles[role].adminRole;
     }
 

contracts/governance/GovernorUpgradeable.sol

@@ -379,7 +379,7 @@ abstract contract GovernorUpgradeable is Initializable, ContextUpgradeable, ERC1
         address target,
         uint256 value,
         bytes calldata data
-    ) external onlyGovernance {
+    ) external virtual onlyGovernance {
         AddressUpgradeable.functionCallWithValue(target, data, value);
     }
 

contracts/interfaces/draft-IERC1822Upgradeable.sol

+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts v4.x.0 (proxy/ERC1822/IProxiable.sol)
+
+pragma solidity ^0.8.0;
+
+/**
+ * @dev ERC1822: Universal Upgradeable Proxy Standard (UUPS) documents a method for upgradeability through a simplified
+ * proxy whose upgrades are fully controlled by the current implementation.
+ */
+interface IERC1822ProxiableUpgradeable {
+    /**
+     * @dev Returns the storage slot that the proxiable contract assumes is being used to store the implementation
+     * address.
+     *
+     * IMPORTANT: A proxy pointing at a proxiable contract should not be considered proxiable itself, because this risks
+     * bricking a proxy that upgrades to it, by delegating to itself until out of gas. Thus it is critical that this
+     * function revert if invoked through a proxy.
+     */
+    function proxiableUUID() external view returns (bytes32);
+}

contracts/mocks/UUPS/UUPSLegacyUpgradeable.sol

+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "./UUPSUpgradeableMockUpgradeable.sol";
+import "../../proxy/utils/Initializable.sol";
+
+// This contract implements the pre-4.5 UUPS upgrade function with a rollback test.
+// It's used to test that newer UUPS contracts are considered valid upgrades by older UUPS contracts.
+contract UUPSUpgradeableLegacyMockUpgradeable is Initializable, UUPSUpgradeableMockUpgradeable {
+    function __UUPSUpgradeableLegacyMock_init() internal onlyInitializing {
+        __CountersImpl_init_unchained();
+        __ERC1967Upgrade_init_unchained();
+        __UUPSUpgradeable_init_unchained();
+        __UUPSUpgradeableMock_init_unchained();
+        __UUPSUpgradeableLegacyMock_init_unchained();
+    }
+
+    function __UUPSUpgradeableLegacyMock_init_unchained() internal onlyInitializing {
+    }
+    // Inlined from ERC1967Upgrade
+    bytes32 private constant _ROLLBACK_SLOT = 0x4910fdfa16fed3260ed0e7147f7cc6da11a60208b5b9406d12a635614ffd9143;
+
+    // ERC1967Upgrade._setImplementation is private so we reproduce it here.
+    // An extra underscore prevents a name clash error.
+    function __setImplementation(address newImplementation) private {
+        require(AddressUpgradeable.isContract(newImplementation), "ERC1967: new implementation is not a contract");
+        StorageSlotUpgradeable.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
+    }
+
+    function _upgradeToAndCallSecureLegacyV1(
+        address newImplementation,
+        bytes memory data,
+        bool forceCall
+    ) internal {
+        address oldImplementation = _getImplementation();
+
+        // Initial upgrade and setup call
+        __setImplementation(newImplementation);
+        if (data.length > 0 || forceCall) {
+            __functionDelegateCall(newImplementation, data);
+        }
+
+        // Perform rollback test if not already in progress
+        StorageSlotUpgradeable.BooleanSlot storage rollbackTesting = StorageSlotUpgradeable.getBooleanSlot(_ROLLBACK_SLOT);
+        if (!rollbackTesting.value) {
+            // Trigger rollback using upgradeTo from the new implementation
+            rollbackTesting.value = true;
+            __functionDelegateCall(
+                newImplementation,
+                abi.encodeWithSignature("upgradeTo(address)", oldImplementation)
+            );
+            rollbackTesting.value = false;
+            // Check rollback was effective
+            require(oldImplementation == _getImplementation(), "ERC1967Upgrade: upgrade breaks further upgrades");
+            // Finally reset to the new implementation and log the upgrade
+            _upgradeTo(newImplementation);
+        }
+    }
+
+    // hooking into the old mechanism
+    function upgradeTo(address newImplementation) external virtual override {
+        _upgradeToAndCallSecureLegacyV1(newImplementation, bytes(""), false);
+    }
+
+    function upgradeToAndCall(address newImplementation, bytes memory data) external payable virtual override {
+        _upgradeToAndCallSecureLegacyV1(newImplementation, data, false);
+    }
+
+    // ERC1967Upgrade._functionDelegateCall is private so we reproduce it here.
+    // An extra underscore prevents a name clash error.
+    function __functionDelegateCall(address target, bytes memory data) private returns (bytes memory) {
+        require(AddressUpgradeable.isContract(target), "Address: delegate call to non-contract");
+
+        // solhint-disable-next-line avoid-low-level-calls
+        (bool success, bytes memory returndata) = target.delegatecall(data);
+        return AddressUpgradeable.verifyCallResult(success, returndata, "Address: low-level delegate call failed");
+    }
+    uint256[50] private __gap;
+}

contracts/mocks/UUPS/TestInProdUpgradeable.sol → contracts/mocks/UUPS/UUPSUpgradeableMockUpgradeable.sol

@@ -41,24 +41,3 @@ contract UUPSUpgradeableUnsafeMockUpgradeable is Initializable, UUPSUpgradeableM
     }
     uint256[50] private __gap;
 }
-
-contract UUPSUpgradeableBrokenMockUpgradeable is Initializable, UUPSUpgradeableMockUpgradeable {
-    function __UUPSUpgradeableBrokenMock_init() internal onlyInitializing {
-        __CountersImpl_init_unchained();
-        __ERC1967Upgrade_init_unchained();
-        __UUPSUpgradeable_init_unchained();
-        __UUPSUpgradeableMock_init_unchained();
-        __UUPSUpgradeableBrokenMock_init_unchained();
-    }
-
-    function __UUPSUpgradeableBrokenMock_init_unchained() internal onlyInitializing {
-    }
-    function upgradeTo(address) external virtual override {
-        // pass
-    }
-
-    function upgradeToAndCall(address, bytes memory) external payable virtual override {
-        // pass
-    }
-    uint256[50] private __gap;
-}

contracts/mocks/WithInit.sol

@@ -661,25 +661,25 @@ contract CountersImplUpgradeableWithInit is CountersImplUpgradeable {
         __CountersImpl_init();
     }
 }
-import "./UUPS/TestInProdUpgradeable.sol";
+import "./UUPS/UUPSUpgradeableMockUpgradeable.sol";
 
 contract UUPSUpgradeableMockUpgradeableWithInit is UUPSUpgradeableMockUpgradeable {
     constructor() payable initializer {
         __UUPSUpgradeableMock_init();
     }
 }
-import "./UUPS/TestInProdUpgradeable.sol";
+import "./UUPS/UUPSUpgradeableMockUpgradeable.sol";
 
 contract UUPSUpgradeableUnsafeMockUpgradeableWithInit is UUPSUpgradeableUnsafeMockUpgradeable {
     constructor() payable initializer {
         __UUPSUpgradeableUnsafeMock_init();
     }
 }
-import "./UUPS/TestInProdUpgradeable.sol";
+import "./UUPS/UUPSLegacyUpgradeable.sol";
 
-contract UUPSUpgradeableBrokenMockUpgradeableWithInit is UUPSUpgradeableBrokenMockUpgradeable {
+contract UUPSUpgradeableLegacyMockUpgradeableWithInit is UUPSUpgradeableLegacyMockUpgradeable {
     constructor() payable initializer {
-        __UUPSUpgradeableBrokenMock_init();
+        __UUPSUpgradeableLegacyMock_init();
     }
 }
 import "./ReentrancyAttackUpgradeable.sol";

contracts/proxy/ERC1967/ERC1967UpgradeUpgradeable.sol

@@ -4,6 +4,7 @@
 pragma solidity ^0.8.2;
 
 import "../beacon/IBeaconUpgradeable.sol";
+import "../../interfaces/draft-IERC1822Upgradeable.sol";
 import "../../utils/AddressUpgradeable.sol";
 import "../../utils/StorageSlotUpgradeable.sol";
 import "../utils/Initializable.sol";
@@ -84,33 +85,23 @@ abstract contract ERC1967UpgradeUpgradeable is Initializable {
      *
      * Emits an {Upgraded} event.
      */
-    function _upgradeToAndCallSecure(
+    function _upgradeToAndCallUUPS(
         address newImplementation,
         bytes memory data,
         bool forceCall
     ) internal {
-        address oldImplementation = _getImplementation();
-
-        // Initial upgrade and setup call
-        _setImplementation(newImplementation);
-        if (data.length > 0 || forceCall) {
-            _functionDelegateCall(newImplementation, data);
-        }
-
-        // Perform rollback test if not already in progress
-        StorageSlotUpgradeable.BooleanSlot storage rollbackTesting = StorageSlotUpgradeable.getBooleanSlot(_ROLLBACK_SLOT);
-        if (!rollbackTesting.value) {
-            // Trigger rollback using upgradeTo from the new implementation
-            rollbackTesting.value = true;
-            _functionDelegateCall(
-                newImplementation,
-                abi.encodeWithSignature("upgradeTo(address)", oldImplementation)
-            );
-            rollbackTesting.value = false;
-            // Check rollback was effective
-            require(oldImplementation == _getImplementation(), "ERC1967Upgrade: upgrade breaks further upgrades");
-            // Finally reset to the new implementation and log the upgrade
-            _upgradeTo(newImplementation);
+        // Upgrades from old implementations will perform a rollback test. This test requires the new
+        // implementation to upgrade back to the old, non-ERC1822 compliant, implementation. Removing
+        // this special case will break upgrade paths from old UUPS implementation to new ones.
+        if (StorageSlotUpgradeable.getBooleanSlot(_ROLLBACK_SLOT).value) {
+            _setImplementation(newImplementation);
+        } else {
+            try IERC1822ProxiableUpgradeable(newImplementation).proxiableUUID() returns (bytes32 slot) {
+                require(slot == _IMPLEMENTATION_SLOT, "ERC1967Upgrade: unsupported proxiableUUID");
+            } catch {
+                revert("ERC1967Upgrade: new implementation is not UUPS");
+            }
+            _upgradeToAndCall(newImplementation, data, forceCall);
         }
     }
 

contracts/proxy/README.adoc

@@ -17,14 +17,14 @@ In order to avoid clashes with the storage variables of the implementation contr
 There are two alternative ways to add upgradeability to an ERC1967 proxy. Their differences are explained below in <<transparent-vs-uups>>.
 
 - {TransparentUpgradeableProxy}: A proxy with a built in admin and upgrade interface.
-- {UUPSUpgradeable}: An upgradeability mechanism to be included in the implementation for an ERC1967 proxy.
+- {UUPSUpgradeable}: An upgradeability mechanism to be included in the implementation contract.
 
 CAUTION: Using upgradeable proxies correctly and securely is a difficult task that requires deep knowledge of the proxy pattern, Solidity, and the EVM. Unless you want a lot of low level control, we recommend using the xref:upgrades-plugins::index.adoc[OpenZeppelin Upgrades Plugins] for Truffle and Hardhat.
 
 A different family of proxies are beacon proxies. This pattern, popularized by Dharma, allows multiple proxies to be upgraded to a different implementation in a single transaction.
 
 - {BeaconProxy}: A proxy that retreives its implementation from a beacon contract.
-- {UpgradeableBeacon}: A beacon contract that can be upgraded.
+- {UpgradeableBeacon}: A beacon contract with a built in admin that can upgrade the {BeaconProxy} pointing to it.
 
 In this pattern, the proxy contract doesn't hold the implementation address in storage like an ERC1967 proxy, instead the address is stored in a separate beacon contract. The `upgrade` operations that are sent to the beacon instead of to the proxy contract, and all proxies that follow that beacon are automatically upgraded.
 
@@ -48,6 +48,8 @@ By default, the upgrade functionality included in {UUPSUpgradeable} contains a s
 - Adding a flag mechanism in the implementation that will disable the upgrade function when triggered.
 - Upgrading to an implementation that features an upgrade mechanism without the additional security check, and then upgrading again to another implementation without the upgrade mechanism.
 
+The current implementation of this security mechanism uses https://eips.ethereum.org/EIPS/eip-1822[EIP1822] to detect the storage slot used by the implementation. A previous implementation, now deprecated, relied on a rollback check. It is possible to upgrade from a contract using the old mechanism to a new one. The inverse is however not possible, as old implementations (before version 4.5) did not include the `ERC1822` interface.
+
 == Core
 
 {{Proxy}}

contracts/proxy/utils/UUPSUpgradeable.sol

@@ -3,6 +3,7 @@
 
 pragma solidity ^0.8.0;
 
+import "../../interfaces/draft-IERC1822Upgradeable.sol";
 import "../ERC1967/ERC1967UpgradeUpgradeable.sol";
 import "./Initializable.sol";
 
@@ -18,7 +19,7 @@ import "./Initializable.sol";
  *
  * _Available since v4.1._
  */
-abstract contract UUPSUpgradeable is Initializable, ERC1967UpgradeUpgradeable {
+abstract contract UUPSUpgradeable is Initializable, IERC1822ProxiableUpgradeable, ERC1967UpgradeUpgradeable {
     function __UUPSUpgradeable_init() internal onlyInitializing {
         __ERC1967Upgrade_init_unchained();
         __UUPSUpgradeable_init_unchained();
@@ -43,6 +44,27 @@ abstract contract UUPSUpgradeable is Initializable, ERC1967UpgradeUpgradeable {
     }
 
     /**
+     * @dev Check that the execution is not being performed through a delegate call. This allows a function to be
+     * callable on the implementing contract but not through proxies.
+     */
+    modifier notDelegated() {
+        require(address(this) == __self, "UUPSUpgradeable: must not be called through delegatecall");
+        _;
+    }
+
+    /**
+     * @dev Implementation of the ERC1822 {proxiableUUID} function. This returns the storage slot used by the
+     * implementation. It is used to validate that the this implementation remains valid after an upgrade.
+     *
+     * IMPORTANT: A proxy pointing at a proxiable contract should not be considered proxiable itself, because this risks
+     * bricking a proxy that upgrades to it, by delegating to itself until out of gas. Thus it is critical that this
+     * function revert if invoked through a proxy. This is guaranteed by the `notDelegated` modifier.
+     */
+    function proxiableUUID() external view virtual override notDelegated returns (bytes32) {
+        return _IMPLEMENTATION_SLOT;
+    }
+
+    /**
      * @dev Upgrade the implementation of the proxy to `newImplementation`.
      *
      * Calls {_authorizeUpgrade}.
@@ -51,7 +73,7 @@ abstract contract UUPSUpgradeable is Initializable, ERC1967UpgradeUpgradeable {
      */
     function upgradeTo(address newImplementation) external virtual onlyProxy {
         _authorizeUpgrade(newImplementation);
-        _upgradeToAndCallSecure(newImplementation, new bytes(0), false);
+        _upgradeToAndCallUUPS(newImplementation, new bytes(0), false);
     }
 
     /**
@@ -64,7 +86,7 @@ abstract contract UUPSUpgradeable is Initializable, ERC1967UpgradeUpgradeable {
      */
     function upgradeToAndCall(address newImplementation, bytes memory data) external payable virtual onlyProxy {
         _authorizeUpgrade(newImplementation);
-        _upgradeToAndCallSecure(newImplementation, data, true);
+        _upgradeToAndCallUUPS(newImplementation, data, true);
     }
 
     /**

contracts/token/ERC20/extensions/ERC20FlashMintUpgradeable.sol

@@ -31,7 +31,7 @@ abstract contract ERC20FlashMintUpgradeable is Initializable, ERC20Upgradeable,
      * @param token The address of the token that is requested.
      * @return The amont of token that can be loaned.
      */
-    function maxFlashLoan(address token) public view override returns (uint256) {
+    function maxFlashLoan(address token) public view virtual override returns (uint256) {
         return token == address(this) ? type(uint256).max - ERC20Upgradeable.totalSupply() : 0;
     }
 
@@ -70,6 +70,7 @@ abstract contract ERC20FlashMintUpgradeable is Initializable, ERC20Upgradeable,
         uint256 amount,
         bytes calldata data
     ) public virtual override returns (bool) {
+        require(amount <= maxFlashLoan(token), "ERC20FlashMint: amount exceeds maxFlashLoan");
         uint256 fee = flashFee(token, amount);
         _mint(address(receiver), amount);
         require(

contracts/token/ERC20/extensions/ERC20VotesCompUpgradeable.sol

@@ -29,14 +29,14 @@ abstract contract ERC20VotesCompUpgradeable is Initializable, ERC20VotesUpgradea
     /**
      * @dev Comp version of the {getVotes} accessor, with `uint96` return type.
      */
-    function getCurrentVotes(address account) external view returns (uint96) {
+    function getCurrentVotes(address account) external view virtual returns (uint96) {
         return SafeCastUpgradeable.toUint96(getVotes(account));
     }
 
     /**
      * @dev Comp version of the {getPastVotes} accessor, with `uint96` return type.
      */
-    function getPriorVotes(address account, uint256 blockNumber) external view returns (uint96) {
+    function getPriorVotes(address account, uint256 blockNumber) external view virtual returns (uint96) {
         return SafeCastUpgradeable.toUint96(getPastVotes(account, blockNumber));
     }
 

contracts/token/ERC20/extensions/ERC20VotesUpgradeable.sol

@@ -64,7 +64,7 @@ abstract contract ERC20VotesUpgradeable is Initializable, IVotesUpgradeable, ERC
     /**
      * @dev Gets the current votes balance for `account`
      */
-    function getVotes(address account) public view override returns (uint256) {
+    function getVotes(address account) public view virtual override returns (uint256) {
         uint256 pos = _checkpoints[account].length;
         return pos == 0 ? 0 : _checkpoints[account][pos - 1].votes;
     }
@@ -76,7 +76,7 @@ abstract contract ERC20VotesUpgradeable is Initializable, IVotesUpgradeable, ERC
      *
      * - `blockNumber` must have been already mined
      */
-    function getPastVotes(address account, uint256 blockNumber) public view override returns (uint256) {
+    function getPastVotes(address account, uint256 blockNumber) public view virtual override returns (uint256) {
         require(blockNumber < block.number, "ERC20Votes: block not yet mined");
         return _checkpointsLookup(_checkpoints[account], blockNumber);
     }
@@ -89,7 +89,7 @@ abstract contract ERC20VotesUpgradeable is Initializable, IVotesUpgradeable, ERC
      *
      * - `blockNumber` must have been already mined
      */
-    function getPastTotalSupply(uint256 blockNumber) public view override returns (uint256) {
+    function getPastTotalSupply(uint256 blockNumber) public view virtual override returns (uint256) {
         require(blockNumber < block.number, "ERC20Votes: block not yet mined");
         return _checkpointsLookup(_totalSupplyCheckpoints, blockNumber);
     }

contracts/token/ERC777/ERC777Upgradeable.sol

@@ -477,7 +477,7 @@ contract ERC777Upgradeable is Initializable, ContextUpgradeable, IERC777Upgradea
         address holder,
         address spender,
         uint256 value
-    ) internal {
+    ) internal virtual {
         require(holder != address(0), "ERC777: approve from the zero address");
         require(spender != address(0), "ERC777: approve to the zero address");
 

contracts/utils/MulticallUpgradeable.sol

@@ -21,7 +21,7 @@ abstract contract MulticallUpgradeable is Initializable {
     /**
      * @dev Receives and executes a batch of function calls on this contract.
      */
-    function multicall(bytes[] calldata data) external returns (bytes[] memory results) {
+    function multicall(bytes[] calldata data) external virtual returns (bytes[] memory results) {
         results = new bytes[](data.length);
         for (uint256 i = 0; i < data.length; i++) {
             results[i] = _functionDelegateCall(address(this), data[i]);

test/helpers/erc1967.js

+const ImplementationLabel = 'eip1967.proxy.implementation';
+const AdminLabel = 'eip1967.proxy.admin';
+const BeaconLabel = 'eip1967.proxy.beacon';
+
+function labelToSlot (label) {
+  return '0x' + web3.utils.toBN(web3.utils.keccak256(label)).subn(1).toString(16);
+}
+
+function getSlot (address, slot) {
+  return web3.eth.getStorageAt(
+    web3.utils.isAddress(address) ? address : address.address,
+    web3.utils.isHex(slot) ? slot : labelToSlot(slot),
+  );
+}
+
+module.exports = {
+  ImplementationLabel,
+  AdminLabel,
+  BeaconLabel,
+  ImplementationSlot: labelToSlot(ImplementationLabel),
+  AdminSlot: labelToSlot(AdminLabel),
+  BeaconSlot: labelToSlot(BeaconLabel),
+  getSlot,
+};

test/proxy/Proxy.behaviour.js

-const { BN, expectRevert } = require('@openzeppelin/test-helpers');
-const ethereumjsUtil = require('ethereumjs-util');
+const { expectRevert } = require('@openzeppelin/test-helpers');
+const { getSlot, ImplementationSlot } = require('../helpers/erc1967');
 
 const { expect } = require('chai');
 
 const DummyImplementation = artifacts.require('DummyImplementation');
 
-const IMPLEMENTATION_LABEL = 'eip1967.proxy.implementation';
-
-function toChecksumAddress (address) {
-  return ethereumjsUtil.toChecksumAddress('0x' + address.replace(/^0x/, '').padStart(40, '0'));
-}
-
 module.exports = function shouldBehaveLikeProxy (createProxy, proxyAdminAddress, proxyCreator) {
   it('cannot be initialized with a non-contract address', async function () {
     const nonContractAddress = proxyCreator;
@@ -28,9 +22,9 @@ module.exports = function shouldBehaveLikeProxy (createProxy, proxyAdminAddress,
 
   const assertProxyInitialization = function ({ value, balance }) {
     it('sets the implementation address', async function () {
-      const slot = '0x' + new BN(ethereumjsUtil.keccak256(Buffer.from(IMPLEMENTATION_LABEL))).subn(1).toString(16);
-      const implementation = toChecksumAddress((await web3.eth.getStorageAt(this.proxy, slot)).substr(-40));
-      expect(implementation).to.be.equal(this.implementation);
+      const implementationSlot = await getSlot(this.proxy, ImplementationSlot);
+      const implementationAddress = web3.utils.toChecksumAddress(implementationSlot.substr(-40));
+      expect(implementationAddress).to.be.equal(this.implementation);
     });
 
     it('initializes the proxy', async function () {

test/proxy/beacon/BeaconProxy.test.js

-const { BN, expectRevert } = require('@openzeppelin/test-helpers');
-const ethereumjsUtil = require('ethereumjs-util');
-const { keccak256 } = ethereumjsUtil;
+const { expectRevert } = require('@openzeppelin/test-helpers');
+const { getSlot, BeaconSlot } = require('../../helpers/erc1967');
 
 const { expect } = require('chai');
 
@@ -11,13 +10,6 @@ const DummyImplementationV2 = artifacts.require('DummyImplementationV2');
 const BadBeaconNoImpl = artifacts.require('BadBeaconNoImpl');
 const BadBeaconNotContract = artifacts.require('BadBeaconNotContract');
 
-function toChecksumAddress (address) {
-  return ethereumjsUtil.toChecksumAddress('0x' + address.replace(/^0x/, '').padStart(40, '0').substr(-40));
-}
-
-const BEACON_LABEL = 'eip1967.proxy.beacon';
-const BEACON_SLOT = '0x' + new BN(keccak256(Buffer.from(BEACON_LABEL))).subn(1).toString(16);
-
 contract('BeaconProxy', function (accounts) {
   const [anotherAccount] = accounts;
 
@@ -53,7 +45,8 @@ contract('BeaconProxy', function (accounts) {
   describe('initialization', function () {
     before(function () {
       this.assertInitialized = async ({ value, balance }) => {
-        const beaconAddress = toChecksumAddress(await web3.eth.getStorageAt(this.proxy.address, BEACON_SLOT));
+        const beaconSlot = await getSlot(this.proxy, BeaconSlot);
+        const beaconAddress = web3.utils.toChecksumAddress(beaconSlot.substr(-40));
         expect(beaconAddress).to.equal(this.beacon.address);
 
         const dummy = new DummyImplementation(this.proxy.address);

test/proxy/transparent/TransparentUpgradeableProxy.behaviour.js

 const { BN, expectRevert, expectEvent, constants } = require('@openzeppelin/test-helpers');
 const { ZERO_ADDRESS } = constants;
-const ethereumjsUtil = require('ethereumjs-util');
+const { getSlot, ImplementationSlot, AdminSlot } = require('../../helpers/erc1967');
 
 const { expect } = require('chai');
 
@@ -16,13 +16,6 @@ const InitializableMock = artifacts.require('InitializableMock');
 const DummyImplementation = artifacts.require('DummyImplementation');
 const ClashingImplementation = artifacts.require('ClashingImplementation');
 
-const IMPLEMENTATION_LABEL = 'eip1967.proxy.implementation';
-const ADMIN_LABEL = 'eip1967.proxy.admin';
-
-function toChecksumAddress (address) {
-  return ethereumjsUtil.toChecksumAddress('0x' + address.replace(/^0x/, '').padStart(40, '0').substr(-40));
-}
-
 module.exports = function shouldBehaveLikeTransparentUpgradeableProxy (createProxy, accounts) {
   const [proxyAdminAddress, proxyAdminOwner, anotherAccount] = accounts;
 
@@ -312,15 +305,15 @@ module.exports = function shouldBehaveLikeTransparentUpgradeableProxy (createPro
 
   describe('storage', function () {
     it('should store the implementation address in specified location', async function () {
-      const slot = '0x' + new BN(ethereumjsUtil.keccak256(Buffer.from(IMPLEMENTATION_LABEL))).subn(1).toString(16);
-      const implementation = toChecksumAddress(await web3.eth.getStorageAt(this.proxyAddress, slot));
-      expect(implementation).to.be.equal(this.implementationV0);
+      const implementationSlot = await getSlot(this.proxy, ImplementationSlot);
+      const implementationAddress = web3.utils.toChecksumAddress(implementationSlot.substr(-40));
+      expect(implementationAddress).to.be.equal(this.implementationV0);
     });
 
     it('should store the admin proxy in specified location', async function () {
-      const slot = '0x' + new BN(ethereumjsUtil.keccak256(Buffer.from(ADMIN_LABEL))).subn(1).toString(16);
-      const proxyAdmin = toChecksumAddress(await web3.eth.getStorageAt(this.proxyAddress, slot));
-      expect(proxyAdmin).to.be.equal(proxyAdminAddress);
+      const proxyAdminSlot = await getSlot(this.proxy, AdminSlot);
+      const proxyAdminAddress = web3.utils.toChecksumAddress(proxyAdminSlot.substr(-40));
+      expect(proxyAdminAddress).to.be.equal(proxyAdminAddress);
     });
   });
 

test/proxy/utils/UUPSUpgradeable.test.js

 const { expectEvent, expectRevert } = require('@openzeppelin/test-helpers');
+const { web3 } = require('@openzeppelin/test-helpers/src/setup');
+const { getSlot, ImplementationSlot } = require('../../helpers/erc1967');
 
 const ERC1967Proxy = artifacts.require('ERC1967Proxy');
 const UUPSUpgradeableMock = artifacts.require('UUPSUpgradeableMock');
 const UUPSUpgradeableUnsafeMock = artifacts.require('UUPSUpgradeableUnsafeMock');
-const UUPSUpgradeableBrokenMock = artifacts.require('UUPSUpgradeableBrokenMock');
+const UUPSUpgradeableLegacyMock = artifacts.require('UUPSUpgradeableLegacyMock');
 const CountersImpl = artifacts.require('CountersImpl');
 
 contract('UUPSUpgradeable', function (accounts) {
@@ -11,7 +13,6 @@ contract('UUPSUpgradeable', function (accounts) {
     this.implInitial = await UUPSUpgradeableMock.new();
     this.implUpgradeOk = await UUPSUpgradeableMock.new();
     this.implUpgradeUnsafe = await UUPSUpgradeableUnsafeMock.new();
-    this.implUpgradeBroken = await UUPSUpgradeableBrokenMock.new();
     this.implUpgradeNonUUPS = await CountersImpl.new();
   });
 
@@ -44,18 +45,11 @@ contract('UUPSUpgradeable', function (accounts) {
     expectEvent(receipt, 'Upgraded', { implementation: this.implUpgradeUnsafe.address });
   });
 
-  it('reject upgrade to broken upgradeable implementation', async function () {
-    await expectRevert(
-      this.instance.upgradeTo(this.implUpgradeBroken.address),
-      'ERC1967Upgrade: upgrade breaks further upgrades',
-    );
-  });
-
   // delegate to a non existing upgradeTo function causes a low level revert
   it('reject upgrade to non uups implementation', async function () {
     await expectRevert(
       this.instance.upgradeTo(this.implUpgradeNonUUPS.address),
-      'Address: low-level delegate call failed',
+      'ERC1967Upgrade: new implementation is not UUPS',
     );
   });
 
@@ -63,10 +57,29 @@ contract('UUPSUpgradeable', function (accounts) {
     const { address } = await ERC1967Proxy.new(this.implInitial.address, '0x');
     const otherInstance = await UUPSUpgradeableMock.at(address);
 
-    // infinite loop reverts when a nested call is out-of-gas
     await expectRevert(
       this.instance.upgradeTo(otherInstance.address),
-      'Address: low-level delegate call failed',
+      'ERC1967Upgrade: new implementation is not UUPS',
     );
   });
+
+  it('can upgrade from legacy implementations', async function () {
+    const legacyImpl = await UUPSUpgradeableLegacyMock.new();
+    const legacyInstance = await ERC1967Proxy.new(legacyImpl.address, '0x')
+      .then(({ address }) => UUPSUpgradeableLegacyMock.at(address));
+
+    const receipt = await legacyInstance.upgradeTo(this.implInitial.address);
+
+    const UpgradedEvents = receipt.logs.filter(({ address, event }) =>
+      address === legacyInstance.address &&
+      event === 'Upgraded',
+    );
+    expect(UpgradedEvents.length).to.be.equal(1);
+
+    expectEvent(receipt, 'Upgraded', { implementation: this.implInitial.address });
+
+    const implementationSlot = await getSlot(legacyInstance, ImplementationSlot);
+    const implementationAddress = web3.utils.toChecksumAddress(implementationSlot.substr(-40));
+    expect(implementationAddress).to.be.equal(this.implInitial.address);
+  });
 });