Implement UUPS proxy (ERC1822) (#2542)

Co-authored-by: Francisco Giordano <frangio.1@gmail.com>

contracts/mocks/StorageSlotMock.sol

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../utils/StorageSlot.sol";
+
+contract StorageSlotMock {
+    using StorageSlot for bytes32;
+    function setBoolean(bytes32 slot, bool value) public { slot.getBooleanSlot().value = value; }
+    function setAddress(bytes32 slot, address value) public { slot.getAddressSlot().value = value; }
+    function setBytes32(bytes32 slot, bytes32 value) public { slot.getBytes32Slot().value = value; }
+    function setUint256(bytes32 slot, uint256 value) public { slot.getUint256Slot().value = value; }
+    function getBoolean(bytes32 slot) public view returns (bool) { return slot.getBooleanSlot().value; }
+    function getAddress(bytes32 slot) public view returns (address) { return slot.getAddressSlot().value; }
+    function getBytes32(bytes32 slot) public view returns (bytes32) { return slot.getBytes32Slot().value; }
+    function getUint256(bytes32 slot) public view returns (uint256) { return slot.getUint256Slot().value; }
+}

contracts/mocks/UUPS/TestInProd.sol

@@ -0,0 +1,32 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../CountersImpl.sol";
+import "../../proxy/UUPS/UUPSUpgradeable.sol";
+
+contract UUPSUpgradeableMock is CountersImpl, UUPSUpgradeable {
+    // Not having any checks in this function is dangerous! Do not do this outside tests!
+    function _authorizeUpgrade(address) internal virtual override {}
+}
+
+contract UUPSUpgradeableUnsafeMock is UUPSUpgradeableMock {
+    function upgradeTo(address newImplementation) external virtual override {
+        ERC1967Upgrade._upgradeToAndCall(newImplementation, bytes(""), false);
+    }
+
+    function upgradeToAndCall(address newImplementation, bytes memory data) external payable virtual override {
+        ERC1967Upgrade._upgradeToAndCall(newImplementation, data, false);
+    }
+}
+
+contract UUPSUpgradeableBrokenMock is UUPSUpgradeableMock {
+    function upgradeTo(address) external virtual override {
+        // pass
+    }
+
+    function upgradeToAndCall(address, bytes memory) external payable virtual override {
+        // pass
+    }
+
+}

contracts/proxy/ERC1967/ERC1967Proxy.sol

@@ -3,18 +3,15 @@
 pragma solidity ^0.8.0;
 
 import "../Proxy.sol";
-import "../../utils/Address.sol";
+import "./ERC1967Upgrade.sol";
 
 /**
  * @dev This contract implements an upgradeable proxy. It is upgradeable because calls are delegated to an
  * implementation address that can be changed. This address is stored in storage in the location specified by
  * https://eips.ethereum.org/EIPS/eip-1967[EIP1967], so that it doesn't conflict with the storage layout of the
  * implementation behind the proxy.
- *
- * Upgradeability is only provided internally through {_upgradeTo}. For an externally upgradeable proxy see
- * {TransparentUpgradeableProxy}.
  */
-contract ERC1967Proxy is Proxy {
+contract ERC1967Proxy is Proxy, ERC1967Upgrade {
     /**
      * @dev Initializes the upgradeable proxy with an initial implementation specified by `_logic`.
      *
@@ -23,56 +20,13 @@ contract ERC1967Proxy is Proxy {
      */
     constructor(address _logic, bytes memory _data) payable {
         assert(_IMPLEMENTATION_SLOT == bytes32(uint256(keccak256("eip1967.proxy.implementation")) - 1));
-        _setImplementation(_logic);
-        if(_data.length > 0) {
-            Address.functionDelegateCall(_logic, _data);
-        }
+        _upgradeToAndCall(_logic, _data, false);
     }
 
-    /**
-     * @dev Emitted when the implementation is upgraded.
-     */
-    event Upgraded(address indexed implementation);
-
-    /**
-     * @dev Storage slot with the address of the current implementation.
-     * This is the keccak-256 hash of "eip1967.proxy.implementation" subtracted by 1, and is
-     * validated in the constructor.
-     */
-    bytes32 private constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
-
     /**
      * @dev Returns the current implementation address.
      */
     function _implementation() internal view virtual override returns (address impl) {
-        bytes32 slot = _IMPLEMENTATION_SLOT;
-        // solhint-disable-next-line no-inline-assembly
-        assembly {
-            impl := sload(slot)
-        }
-    }
-
-    /**
-     * @dev Upgrades the proxy to a new implementation.
-     *
-     * Emits an {Upgraded} event.
-     */
-    function _upgradeTo(address newImplementation) internal virtual {
-        _setImplementation(newImplementation);
-        emit Upgraded(newImplementation);
-    }
-
-    /**
-     * @dev Stores a new address in the EIP1967 implementation slot.
-     */
-    function _setImplementation(address newImplementation) private {
-        require(Address.isContract(newImplementation), "ERC1967Proxy: new implementation is not a contract");
-
-        bytes32 slot = _IMPLEMENTATION_SLOT;
-
-        // solhint-disable-next-line no-inline-assembly
-        assembly {
-            sstore(slot, newImplementation)
-        }
+        return ERC1967Storage._getImplementation();
     }
 }

contracts/proxy/ERC1967/ERC1967Storage.sol

@@ -0,0 +1,85 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../beacon/IBeacon.sol";
+import "../../utils/Address.sol";
+import "../../utils/StorageSlot.sol";
+
+/**
+ * @dev This abstract contract provides setters and getters for the different
+ * https://eips.ethereum.org/EIPS/eip-1967[EIP1967] storage slots.
+ */
+abstract contract ERC1967Storage {
+    /**
+     * @dev Storage slot with the address of the current implementation.
+     * This is the keccak-256 hash of "eip1967.proxy.implementation" subtracted by 1, and is
+     * validated in the constructor.
+     */
+    bytes32 internal constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
+
+    /**
+     * @dev Returns the current implementation address.
+     */
+    function _getImplementation() internal view returns (address) {
+        return StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value;
+    }
+
+    /**
+     * @dev Stores a new address in the EIP1967 implementation slot.
+     */
+    function _setImplementation(address newImplementation) internal {
+        require(Address.isContract(newImplementation), "ERC1967: new implementation is not a contract");
+        StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
+    }
+
+    /**
+     * @dev The storage slot of the UpgradeableBeacon contract which defines the implementation for this proxy.
+     * This is bytes32(uint256(keccak256('eip1967.proxy.beacon')) - 1)) and is validated in the constructor.
+     */
+    bytes32 internal constant _BEACON_SLOT = 0xa3f0ad74e5423aebfd80d3ef4346578335a9a72aeaee59ff6cb3582b35133d50;
+
+    /**
+     * @dev Returns the current beacon.
+     */
+    function _getBeacon() internal view returns (address) {
+        return StorageSlot.getAddressSlot(_BEACON_SLOT).value;
+    }
+
+    /**
+     * @dev Stores a new beacon in the EIP1967 beacon slot.
+     */
+    function _setBeacon(address newBeacon) internal {
+        require(
+            Address.isContract(newBeacon),
+            "ERC1967: new beacon is not a contract"
+        );
+        require(
+            Address.isContract(IBeacon(newBeacon).implementation()),
+            "ERC1967: beacon implementation is not a contract"
+        );
+        StorageSlot.getAddressSlot(_BEACON_SLOT).value = newBeacon;
+    }
+
+    /**
+     * @dev Storage slot with the admin of the contract.
+     * This is the keccak-256 hash of "eip1967.proxy.admin" subtracted by 1, and is
+     * validated in the constructor.
+     */
+    bytes32 internal constant _ADMIN_SLOT = 0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103;
+
+    /**
+     * @dev Returns the current admin.
+     */
+    function _getAdmin() internal view returns (address) {
+        return StorageSlot.getAddressSlot(_ADMIN_SLOT).value;
+    }
+
+    /**
+     * @dev Stores a new address in the EIP1967 admin slot.
+     */
+    function _setAdmin(address newAdmin) internal {
+        require(newAdmin != address(0), "ERC1967: new admin is the zero address");
+        StorageSlot.getAddressSlot(_ADMIN_SLOT).value = newAdmin;
+    }
+}

contracts/proxy/ERC1967/ERC1967Upgrade.sol

@@ -0,0 +1,113 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.3;
+
+import "./ERC1967Storage.sol";
+
+/**
+ * @dev This abstract contract provides event emitting update functions for
+ * https://eips.ethereum.org/EIPS/eip-1967[EIP1967] slots.
+ *
+ * @custom:oz-upgrades-unsafe-allow delegatecall
+ */
+abstract contract ERC1967Upgrade is ERC1967Storage {
+    // This is the keccak-256 hash of "eip1967.proxy.rollback" subtracted by 1
+    bytes32 private constant _ROLLBACK_SLOT = 0x4910fdfa16fed3260ed0e7147f7cc6da11a60208b5b9406d12a635614ffd9143;
+
+    /**
+     * @dev Emitted when the implementation is upgraded.
+     */
+    event Upgraded(address indexed implementation);
+
+    /**
+     * @dev Emitted when the beacon is upgraded.
+     */
+    event BeaconUpgraded(address indexed beacon);
+
+    /**
+     * @dev Emitted when the admin account has changed.
+     */
+    event AdminChanged(address previousAdmin, address newAdmin);
+
+    /**
+     * @dev Perform implementation upgrade
+     *
+     * Emits an {Upgraded} event.
+     */
+    function _upgradeTo(address newImplementation) internal {
+        _setImplementation(newImplementation);
+        emit Upgraded(newImplementation);
+    }
+
+    /**
+     * @dev Perform implementation upgrade with additional setup call.
+     *
+     * Emits an {Upgraded} event.
+     */
+    function _upgradeToAndCall(address newImplementation, bytes memory data, bool forceCall) internal {
+        _setImplementation(newImplementation);
+        emit Upgraded(newImplementation);
+        if (data.length > 0 || forceCall) {
+            Address.functionDelegateCall(newImplementation, data);
+        }
+    }
+
+    /**
+     * @dev Perform implementation upgrade with security checks for UUPS proxies, and additional setup call.
+     *
+     * Emits an {Upgraded} event.
+     */
+    function _upgradeToAndCallSecure(address newImplementation, bytes memory data, bool forceCall) internal {
+        address oldImplementation = _getImplementation();
+        // do inital upgrade
+        _setImplementation(newImplementation);
+        // do setup call
+        if (data.length > 0 || forceCall) {
+            Address.functionDelegateCall(newImplementation, data);
+        }
+        // check if nested in an upgrade check
+        StorageSlot.BooleanSlot storage rollbackTesting = StorageSlot.getBooleanSlot(_ROLLBACK_SLOT);
+        if (!rollbackTesting.value) {
+            // trigger upgrade check with flag set to true
+            rollbackTesting.value = true;
+            Address.functionDelegateCall(
+                newImplementation,
+                abi.encodeWithSignature(
+                    "upgradeTo(address)",
+                    oldImplementation
+                )
+            );
+            rollbackTesting.value = false;
+            // check upgrade was effective
+            require(oldImplementation == _getImplementation(), "ERC1967Upgrade: upgrade breaks further upgrades");
+            // reset upgrade
+            _setImplementation(newImplementation);
+            // emit event
+            emit Upgraded(newImplementation);
+        }
+    }
+
+    /**
+     * @dev Perform beacon upgrade with additional setup call. Note: This upgrades the address of the beacon, it does
+     * not upgrade the implementation contained in the beacon (see {UpgradeableBeacon-_setImplementation} for that).
+     *
+     * Emits a {BeaconUpgraded} event.
+     */
+    function _upgradeBeaconToAndCall(address newBeacon, bytes memory data, bool forceCall) internal {
+        _setBeacon(newBeacon);
+        emit BeaconUpgraded(newBeacon);
+        if (data.length > 0 || forceCall) {
+            Address.functionDelegateCall(IBeacon(newBeacon).implementation(), data);
+        }
+    }
+
+    /**
+     * @dev Changes the admin of the proxy.
+     *
+     * Emits an {AdminChanged} event.
+     */
+    function _changeAdmin(address newAdmin) internal {
+        emit AdminChanged(_getAdmin(), newAdmin);
+        _setAdmin(newAdmin);
+    }
+}

contracts/proxy/README.adoc

@@ -9,12 +9,29 @@ The abstract {Proxy} contract implements the core delegation functionality. If t
 
 {ERC1967Proxy} provides a simple, fully functioning, proxy. While this proxy is not by itself upgradeable, it includes an internal upgrade interface. For an upgrade interface exposed externally to an admin, we provide {TransparentUpgradeableProxy}. Both of these contracts use the storage slots specified in https://eips.ethereum.org/EIPS/eip-1967[EIP1967] to avoid clashes with the storage of the implementation contract behind the proxy.
 
-An alternative upgradeability mechanism is provided in <<Beacon>>. This pattern, popularized by Dharma, allows multiple proxies to be upgraded to a different implementation in a single transaction. In this pattern, the proxy contract doesn't hold the implementation address in storage like {UpgradeableProxy}, but the address of a {UpgradeableBeacon} contract, which is where the implementation address is actually stored and retrieved from. The `upgrade` operations that change the implementation contract address are then sent to the beacon instead of to the proxy contract, and all proxies that follow that beacon are automatically upgraded.
+An alternative upgradeability mechanism is provided in <<Beacon>>. This pattern, popularized by Dharma, allows multiple proxies to be upgraded to a different implementation in a single transaction. In this pattern, the proxy contract doesn't hold the implementation address in storage like {ERC1967Proxy}, but the address of a {UpgradeableBeacon} contract, which is where the implementation address is actually stored and retrieved from. The `upgrade` operations that change the implementation contract address are then sent to the beacon instead of to the proxy contract, and all proxies that follow that beacon are automatically upgraded.
 
 The {Clones} library provides a way to deploy minimal non-upgradeable proxies for cheap. This can be useful for applications that require deploying many instances of the same contract (for example one per user, or one per task). These instances are designed to be both cheap to deploy, and cheap to call. The drawback being that they are not upgradeable.
 
 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.
 
+== UUPS Design: Upgradeability as an Implementation Feature
+
+Upgradeable smart contracts rely on proxies to relay the calls in a way that is programmable. As discussed previously, we provide different proxy contracts that each come with a specific set of features. Other designs, not (yet) proposed as part of the OpenZeppelin products, also exist.
+
+The most simple, and common, design is known as Universal Upgradeable Proxy Standard (UUPS). This design is both lightweight and versatile and is proposed through the {ERC1967Proxy} and the {UUPSUpgradeable} contract.
+
+While {UUPSUpgradeable} uses the same interface as {TransparentUpgradeableProxy}, in the first case the upgrade is handled by the implementation, and can eventually be removed. {TransparentUpgradeableProxy}, on the other hand, includes the upgrade logic in the proxy. This means {TransparentUpgradeableProxy} is more expensive to deploy. Note that, since both proxies use the same storage slot for the implementation address, using a UUPS compliant implementation with a {TransparentUpgradeableProxy} might allow non-admins to perform upgrade operations.
+
+According to this design, the {ERC1967Proxy} is only capable of forwarding calls to an implementation contract. Unlike the more complex and expensive to deploy {TransparentUpgradeableProxy}, the {ERC1967Proxy} doesn't by itself provide any upgradeability mechanism. It is the role of the implementation to include, alongside the contract's logic, all the code necessary to update the implementation's address that is stored at a specific slot in the proxy's storage space. This is where the {UUPSUpgradeable} contract comes in. Inheriting from it (and overriding the {_authorizeUpgrade} function with the relevant access control mechanism) will turn your contract into a UUPS complaint implementation.
+
+By default, the upgrade mechanism included in {UUPSUpgradeable} contains a security mechanism that will prevent any upgrades to a non UUPS compliant implementation. This prevents upgrades to an implementation contract that wouldn't contain the necessary upgrade mechanism, as it would lock the proxy forever. This security mechanism can however be bypassed, either by:
+
+- 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 upgrade again to another implementation without the upgrade mechanism.
+
+When doing an upgrade, the second parameter of the {upgradeToAndCall} function allows for the atomic execution of an optional initialization/migration call.
+
 == Core
 
 {{Proxy}}
@@ -23,6 +40,14 @@ CAUTION: Using upgradeable proxies correctly and securely is a difficult task th
 
 {{ERC1967Proxy}}
 
+{{ERC1967Storage}}
+
+{{ERC1967Upgrade}}
+
+== UUPS
+
+{{UUPSUpgradeable}}
+
 == Transparent Proxy
 
 {{TransparentUpgradeableProxy}}

contracts/proxy/UUPS/UUPSUpgradeable.sol

@@ -0,0 +1,26 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "../ERC1967/ERC1967Upgrade.sol";
+
+/**
+ * @dev Base contract for building openzeppelin-upgrades compatible implementations for the {ERC1967Proxy}. It includes
+ * publicly available upgrade functions that are called by the plugin and by the secure upgrade mechanism to verify
+ * continuation of the upgradability.
+ *
+ * The {_authorizeUpgrade} function MUST be overridden to include access restriction to the upgrade mechanism.
+ */
+abstract contract UUPSUpgradeable is ERC1967Upgrade {
+    function upgradeTo(address newImplementation) external virtual {
+        _authorizeUpgrade(newImplementation);
+        _upgradeToAndCallSecure(newImplementation, bytes(""), false);
+    }
+
+    function upgradeToAndCall(address newImplementation, bytes memory data) external payable virtual {
+        _authorizeUpgrade(newImplementation);
+        _upgradeToAndCallSecure(newImplementation, data, true);
+    }
+
+    function _authorizeUpgrade(address newImplementation) internal virtual;
+}

contracts/proxy/beacon/BeaconProxy.sol

@@ -4,7 +4,7 @@ pragma solidity ^0.8.0;
 
 import "./IBeacon.sol";
 import "../Proxy.sol";
-import "../../utils/Address.sol";
+import "../ERC1967/ERC1967Upgrade.sol";
 
 /**
  * @dev This contract implements a proxy that gets the implementation address for each call from a {UpgradeableBeacon}.
@@ -14,13 +14,7 @@ import "../../utils/Address.sol";
  *
  * _Available since v3.4._
  */
-contract BeaconProxy is Proxy {
-    /**
-     * @dev The storage slot of the UpgradeableBeacon contract which defines the implementation for this proxy.
-     * This is bytes32(uint256(keccak256('eip1967.proxy.beacon')) - 1)) and is validated in the constructor.
-     */
-    bytes32 private constant _BEACON_SLOT = 0xa3f0ad74e5423aebfd80d3ef4346578335a9a72aeaee59ff6cb3582b35133d50;
-
+contract BeaconProxy is Proxy, ERC1967Upgrade {
     /**
      * @dev Initializes the proxy with `beacon`.
      *
@@ -34,29 +28,25 @@ contract BeaconProxy is Proxy {
      */
     constructor(address beacon, bytes memory data) payable {
         assert(_BEACON_SLOT == bytes32(uint256(keccak256("eip1967.proxy.beacon")) - 1));
-        _setBeacon(beacon, data);
+        _upgradeBeaconToAndCall(beacon, data, false);
     }
 
     /**
      * @dev Returns the current beacon address.
      */
-    function _beacon() internal view virtual returns (address beacon) {
-        bytes32 slot = _BEACON_SLOT;
-        // solhint-disable-next-line no-inline-assembly
-        assembly {
-            beacon := sload(slot)
-        }
+    function _beacon() internal view virtual returns (address) {
+        return _getBeacon();
     }
 
     /**
      * @dev Returns the current implementation address of the associated beacon.
      */
     function _implementation() internal view virtual override returns (address) {
-        return IBeacon(_beacon()).implementation();
+        return IBeacon(_getBeacon()).implementation();
     }
 
     /**
-     * @dev Changes the proxy to use a new beacon.
+     * @dev Changes the proxy to use a new beacon. Deprecated: see {_upgradeBeaconToAndCall}.
      *
      * If `data` is nonempty, it's used as data in a delegate call to the implementation returned by the beacon.
      *
@@ -66,23 +56,6 @@ contract BeaconProxy is Proxy {
      * - The implementation returned by `beacon` must be a contract.
      */
     function _setBeacon(address beacon, bytes memory data) internal virtual {
-        require(
-            Address.isContract(beacon),
-            "BeaconProxy: beacon is not a contract"
-        );
-        require(
-            Address.isContract(IBeacon(beacon).implementation()),
-            "BeaconProxy: beacon implementation is not a contract"
-        );
-        bytes32 slot = _BEACON_SLOT;
-
-        // solhint-disable-next-line no-inline-assembly
-        assembly {
-            sstore(slot, beacon)
-        }
-
-        if (data.length > 0) {
-            Address.functionDelegateCall(_implementation(), data, "BeaconProxy: function call failed");
-        }
+        _upgradeBeaconToAndCall(beacon, data, false);
     }
 }

contracts/proxy/transparent/TransparentUpgradeableProxy.sol

@@ -28,30 +28,18 @@ import "../ERC1967/ERC1967Proxy.sol";
 contract TransparentUpgradeableProxy is ERC1967Proxy {
     /**
      * @dev Initializes an upgradeable proxy managed by `_admin`, backed by the implementation at `_logic`, and
-     * optionally initialized with `_data` as explained in {UpgradeableProxy-constructor}.
+     * optionally initialized with `_data` as explained in {ERC1967Proxy-constructor}.
      */
     constructor(address _logic, address admin_, bytes memory _data) payable ERC1967Proxy(_logic, _data) {
         assert(_ADMIN_SLOT == bytes32(uint256(keccak256("eip1967.proxy.admin")) - 1));
-        _setAdmin(admin_);
+        _changeAdmin(admin_);
     }
 
-    /**
-     * @dev Emitted when the admin account has changed.
-     */
-    event AdminChanged(address previousAdmin, address newAdmin);
-
-    /**
-     * @dev Storage slot with the admin of the contract.
-     * This is the keccak-256 hash of "eip1967.proxy.admin" subtracted by 1, and is
-     * validated in the constructor.
-     */
-    bytes32 private constant _ADMIN_SLOT = 0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103;
-
     /**
      * @dev Modifier used internally that will delegate the call to the implementation unless the sender is the admin.
      */
     modifier ifAdmin() {
-        if (msg.sender == _admin()) {
+        if (msg.sender == _getAdmin()) {
             _;
         } else {
             _fallback();
@@ -68,7 +56,7 @@ contract TransparentUpgradeableProxy is ERC1967Proxy {
      * `0xb53127684a568b3173ae13b9f8a6016e243e63b6e8ee1178d6a717850b5d6103`
      */
     function admin() external ifAdmin returns (address admin_) {
-        admin_ = _admin();
+        admin_ = _getAdmin();
     }
 
     /**
@@ -92,9 +80,7 @@ contract TransparentUpgradeableProxy is ERC1967Proxy {
      * NOTE: Only the admin can call this function. See {ProxyAdmin-changeProxyAdmin}.
      */
     function changeAdmin(address newAdmin) external virtual ifAdmin {
-        require(newAdmin != address(0), "TransparentUpgradeableProxy: new admin is the zero address");
-        emit AdminChanged(_admin(), newAdmin);
-        _setAdmin(newAdmin);
+        _changeAdmin(newAdmin);
     }
 
     /**
@@ -102,8 +88,8 @@ contract TransparentUpgradeableProxy is ERC1967Proxy {
      *
      * NOTE: Only the admin can call this function. See {ProxyAdmin-upgrade}.
      */
-    function upgradeTo(address newImplementation) external virtual ifAdmin {
-        _upgradeTo(newImplementation);
+    function upgradeTo(address newImplementation) external ifAdmin {
+        _upgradeToAndCall(newImplementation, bytes(""), false);
     }
 
     /**
@@ -113,39 +99,22 @@ contract TransparentUpgradeableProxy is ERC1967Proxy {
      *
      * NOTE: Only the admin can call this function. See {ProxyAdmin-upgradeAndCall}.
      */
-    function upgradeToAndCall(address newImplementation, bytes calldata data) external payable virtual ifAdmin {
-        _upgradeTo(newImplementation);
-        Address.functionDelegateCall(newImplementation, data);
+    function upgradeToAndCall(address newImplementation, bytes calldata data) external payable ifAdmin {
+        _upgradeToAndCall(newImplementation, data, true);
     }
 
     /**
      * @dev Returns the current admin.
      */
-    function _admin() internal view virtual returns (address adm) {
-        bytes32 slot = _ADMIN_SLOT;
-        // solhint-disable-next-line no-inline-assembly
-        assembly {
-            adm := sload(slot)
-        }
-    }
-
-    /**
-     * @dev Stores a new address in the EIP1967 admin slot.
-     */
-    function _setAdmin(address newAdmin) private {
-        bytes32 slot = _ADMIN_SLOT;
-
-        // solhint-disable-next-line no-inline-assembly
-        assembly {
-            sstore(slot, newAdmin)
-        }
+    function _admin() internal view virtual returns (address) {
+        return _getAdmin();
     }
 
     /**
      * @dev Makes sure the admin cannot access the fallback function. See {Proxy-_beforeFallback}.
      */
     function _beforeFallback() internal virtual override {
-        require(msg.sender != _admin(), "TransparentUpgradeableProxy: admin cannot fallback to proxy target");
+        require(msg.sender != _getAdmin(), "TransparentUpgradeableProxy: admin cannot fallback to proxy target");
         super._beforeFallback();
     }
 }

contracts/proxy/utils/Initializable.sol

@@ -12,7 +12,7 @@ import "../../utils/Address.sol";
  * function so it can only be called once. The {initializer} modifier provided by this contract will have this effect.
  *
  * TIP: To avoid leaving the proxy in an uninitialized state, the initializer function should be called as early as
- * possible by providing the encoded function call as the `_data` argument to {UpgradeableProxy-constructor}.
+ * possible by providing the encoded function call as the `_data` argument to {ERC1967Proxy-constructor}.
  *
  * CAUTION: When used with inheritance, manual care must be taken to not invoke a parent initializer twice, or to ensure
  * that all initializers are idempotent. This is not verified automatically as constructors are by Solidity.

contracts/utils/README.adoc

@@ -96,4 +96,6 @@ Note that, in all cases, accounts simply _declare_ their interfaces, but they ar
 
 {{Strings}}
 
+{{StorageSlot}}
+
 {{Multicall}}

contracts/utils/StorageSlot.sol

@@ -0,0 +1,83 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+/**
+ * @dev Library for reading and writing primitive types to specific storage slots.
+ *
+ * Storage slots are often used to avoid storage conflict when dealing with upgradeable contracts.
+ * This library helps with reading and writing to such slots without the need for inline assembly.
+ *
+ * The functions in this library return Slot structs that contain a `value` member that can be used to read or write.
+ *
+ * Example usage to set ERC1967 implementation slot:
+ * ```
+ * contract ERC1967 {
+ *     bytes32 internal constant _IMPLEMENTATION_SLOT = 0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc;
+ *
+ *     function _getImplementation() internal view returns (address) {
+ *         return StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value;
+ *     }
+ *
+ *     function _setImplementation(address newImplementation) internal {
+ *         require(Address.isContract(newImplementation), "ERC1967: new implementation is not a contract");
+ *         StorageSlot.getAddressSlot(_IMPLEMENTATION_SLOT).value = newImplementation;
+ *     }
+ * }
+ * ```
+ *
+ * Available since v4.1.0 for `address`, `bool`, `bytes32`, and `uint256`.
+ */
+library StorageSlot {
+    struct AddressSlot {
+        address value;
+    }
+
+    struct BooleanSlot {
+        bool value;
+    }
+
+    struct Bytes32Slot {
+        bytes32 value;
+    }
+
+    struct Uint256Slot {
+        uint256 value;
+    }
+
+    /**
+     * @dev Returns an `AddressSlot` with member `value` located at `slot`.
+     */
+    function getAddressSlot(bytes32 slot) internal pure returns (AddressSlot storage r) {
+        assembly {
+            r.slot := slot
+        }
+    }
+
+    /**
+     * @dev Returns an `BooleanSlot` with member `value` located at `slot`.
+     */
+    function getBooleanSlot(bytes32 slot) internal pure returns (BooleanSlot storage r) {
+        assembly {
+            r.slot := slot
+        }
+    }
+
+    /**
+     * @dev Returns an `Bytes32Slot` with member `value` located at `slot`.
+     */
+    function getBytes32Slot(bytes32 slot) internal pure returns (Bytes32Slot storage r) {
+        assembly {
+            r.slot := slot
+        }
+    }
+
+    /**
+     * @dev Returns an `Uint256Slot` with member `value` located at `slot`.
+     */
+    function getUint256Slot(bytes32 slot) internal pure returns (Uint256Slot storage r) {
+        assembly {
+            r.slot := slot
+        }
+    }
+}

test/proxy/UUPS/UUPSUpgradeable.test.js

@@ -0,0 +1,72 @@
+const { expectEvent, expectRevert } = require('@openzeppelin/test-helpers');
+
+const ERC1967Proxy = artifacts.require('ERC1967Proxy');
+const UUPSUpgradeableMock = artifacts.require('UUPSUpgradeableMock');
+const UUPSUpgradeableUnsafeMock = artifacts.require('UUPSUpgradeableUnsafeMock');
+const UUPSUpgradeableBrokenMock = artifacts.require('UUPSUpgradeableBrokenMock');
+const CountersImpl = artifacts.require('CountersImpl');
+
+contract('UUPSUpgradeable', function (accounts) {
+  before(async function () {
+    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();
+  });
+
+  beforeEach(async function () {
+    const { address } = await ERC1967Proxy.new(this.implInitial.address, '0x');
+    this.instance = await UUPSUpgradeableMock.at(address);
+  });
+
+  it('upgrade to upgradeable implementation', async function () {
+    const { receipt } = await this.instance.upgradeTo(this.implUpgradeOk.address);
+    expect(receipt.logs.filter(({ event }) => event === 'Upgraded').length).to.be.equal(1);
+    expectEvent(receipt, 'Upgraded', { implementation: this.implUpgradeOk.address });
+  });
+
+  it('upgrade to upgradeable implementation with call', async function () {
+    expect(await this.instance.current()).to.be.bignumber.equal('0');
+
+    const { receipt } = await this.instance.upgradeToAndCall(
+      this.implUpgradeOk.address,
+      this.implUpgradeOk.contract.methods.increment().encodeABI(),
+    );
+    expect(receipt.logs.filter(({ event }) => event === 'Upgraded').length).to.be.equal(1);
+    expectEvent(receipt, 'Upgraded', { implementation: this.implUpgradeOk.address });
+
+    expect(await this.instance.current()).to.be.bignumber.equal('1');
+  });
+
+  it('upgrade to and unsafe upgradeable implementation', async function () {
+    const { receipt } = await this.instance.upgradeTo(this.implUpgradeUnsafe.address);
+    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',
+    );
+  });
+
+  it('reject proxy address as implementation', async function () {
+    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',
+    );
+  });
+});

test/proxy/beacon/BeaconProxy.test.js

@@ -25,7 +25,7 @@ contract('BeaconProxy', function (accounts) {
     it('non-contract beacon', async function () {
       await expectRevert(
         BeaconProxy.new(anotherAccount, '0x'),
-        'BeaconProxy: beacon is not a contract',
+        'ERC1967: new beacon is not a contract',
       );
     });
 
@@ -40,7 +40,7 @@ contract('BeaconProxy', function (accounts) {
       const beacon = await BadBeaconNotContract.new();
       await expectRevert(
         BeaconProxy.new(beacon.address, '0x'),
-        'BeaconProxy: beacon implementation is not a contract',
+        'ERC1967: beacon implementation is not a contract',
       );
     });
   });

test/proxy/transparent/TransparentUpgradeableProxy.behaviour.js

@@ -80,7 +80,7 @@ module.exports = function shouldBehaveLikeTransparentUpgradeableProxy (createPro
         it('reverts', async function () {
           await expectRevert(
             this.proxy.upgradeTo(ZERO_ADDRESS, { from }),
-            'ERC1967Proxy: new implementation is not a contract',
+            'ERC1967: new implementation is not a contract',
           );
         });
       });
@@ -304,7 +304,7 @@ module.exports = function shouldBehaveLikeTransparentUpgradeableProxy (createPro
       it('reverts', async function () {
         await expectRevert(
           this.proxy.changeAdmin(ZERO_ADDRESS, { from: proxyAdminAddress }),
-          'TransparentUpgradeableProxy: new admin is the zero address',
+          'ERC1967: new admin is the zero address',
         );
       });
     });

test/utils/StorageSlot.test.js

@@ -0,0 +1,110 @@
+const { constants, BN } = require('@openzeppelin/test-helpers');
+
+const { expect } = require('chai');
+
+const StorageSlotMock = artifacts.require('StorageSlotMock');
+
+const slot = web3.utils.keccak256('some.storage.slot');
+const otherSlot = web3.utils.keccak256('some.other.storage.slot');
+
+contract('StorageSlot', function (accounts) {
+  beforeEach(async function () {
+    this.store = await StorageSlotMock.new();
+  });
+
+  describe('boolean storage slot', function () {
+    beforeEach(async function () {
+      this.value = true;
+    });
+
+    it('set', async function () {
+      await this.store.setBoolean(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setBoolean(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.getBoolean(slot)).to.be.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.getBoolean(otherSlot)).to.be.equal(false);
+      });
+    });
+  });
+
+  describe('address storage slot', function () {
+    beforeEach(async function () {
+      this.value = accounts[1];
+    });
+
+    it('set', async function () {
+      await this.store.setAddress(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setAddress(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.getAddress(slot)).to.be.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.getAddress(otherSlot)).to.be.equal(constants.ZERO_ADDRESS);
+      });
+    });
+  });
+
+  describe('bytes32 storage slot', function () {
+    beforeEach(async function () {
+      this.value = web3.utils.keccak256('some byte32 value');
+    });
+
+    it('set', async function () {
+      await this.store.setBytes32(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setBytes32(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.getBytes32(slot)).to.be.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.getBytes32(otherSlot)).to.be.equal(constants.ZERO_BYTES32);
+      });
+    });
+  });
+
+  describe('uint256 storage slot', function () {
+    beforeEach(async function () {
+      this.value = new BN(1742);
+    });
+
+    it('set', async function () {
+      await this.store.setUint256(slot, this.value);
+    });
+
+    describe('get', function () {
+      beforeEach(async function () {
+        await this.store.setUint256(slot, this.value);
+      });
+
+      it('from right slot', async function () {
+        expect(await this.store.getUint256(slot)).to.be.bignumber.equal(this.value);
+      });
+
+      it('from other slot', async function () {
+        expect(await this.store.getUint256(otherSlot)).to.be.bignumber.equal('0');
+      });
+    });
+  });
+});