Simplify UUPSUpgradeable along the lines of ERC1822 (#3021)

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

CHANGELOG.md

@@ -17,10 +17,13 @@
  * `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))
 
-### 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/interfaces/draft-IERC1822.sol

@@ -0,0 +1,20 @@
+// 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 IERC1822Proxiable {
+    /**
+     * @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/UUPSLegacy.sol

@@ -0,0 +1,58 @@
+// SPDX-License-Identifier: MIT
+
+pragma solidity ^0.8.0;
+
+import "./UUPSUpgradeableMock.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 UUPSUpgradeableLegacyMock is UUPSUpgradeableMock {
+    // 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(Address.isContract(newImplementation), "ERC1967: new implementation is not a contract");
+        StorageSlot.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) {
+            Address.functionDelegateCall(newImplementation, data);
+        }
+
+        // Perform rollback test if not already in progress
+        StorageSlot.BooleanSlot storage rollbackTesting = StorageSlot.getBooleanSlot(_ROLLBACK_SLOT);
+        if (!rollbackTesting.value) {
+            // Trigger rollback using upgradeTo from the new implementation
+            rollbackTesting.value = true;
+            Address.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);
+    }
+}

contracts/mocks/UUPS/UUPSUpgradeableMock.sol

@@ -19,13 +19,3 @@ contract UUPSUpgradeableUnsafeMock is UUPSUpgradeableMock {
         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/ERC1967Upgrade.sol

@@ -4,6 +4,7 @@
 pragma solidity ^0.8.2;
 
 import "../beacon/IBeacon.sol";
+import "../../interfaces/draft-IERC1822.sol";
 import "../../utils/Address.sol";
 import "../../utils/StorageSlot.sol";
 
@@ -77,33 +78,23 @@ abstract contract ERC1967Upgrade {
      *
      * 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) {
-            Address.functionDelegateCall(newImplementation, data);
-        }
-
-        // Perform rollback test if not already in progress
-        StorageSlot.BooleanSlot storage rollbackTesting = StorageSlot.getBooleanSlot(_ROLLBACK_SLOT);
-        if (!rollbackTesting.value) {
-            // Trigger rollback using upgradeTo from the new implementation
-            rollbackTesting.value = true;
-            Address.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 (StorageSlot.getBooleanSlot(_ROLLBACK_SLOT).value) {
+            _setImplementation(newImplementation);
+        } else {
+            try IERC1822Proxiable(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-IERC1822.sol";
 import "../ERC1967/ERC1967Upgrade.sol";
 
 /**
@@ -17,7 +18,7 @@ import "../ERC1967/ERC1967Upgrade.sol";
  *
  * _Available since v4.1._
  */
-abstract contract UUPSUpgradeable is ERC1967Upgrade {
+abstract contract UUPSUpgradeable is IERC1822Proxiable, ERC1967Upgrade {
     /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment
     address private immutable __self = address(this);
 
@@ -34,6 +35,27 @@ abstract contract UUPSUpgradeable is ERC1967Upgrade {
         _;
     }
 
+    /**
+     * @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`.
      *
@@ -43,7 +65,7 @@ abstract contract UUPSUpgradeable is ERC1967Upgrade {
      */
     function upgradeTo(address newImplementation) external virtual onlyProxy {
         _authorizeUpgrade(newImplementation);
-        _upgradeToAndCallSecure(newImplementation, new bytes(0), false);
+        _upgradeToAndCallUUPS(newImplementation, new bytes(0), false);
     }
 
     /**
@@ -56,7 +78,7 @@ abstract contract UUPSUpgradeable is ERC1967Upgrade {
      */
     function upgradeToAndCall(address newImplementation, bytes memory data) external payable virtual onlyProxy {
         _authorizeUpgrade(newImplementation);
-        _upgradeToAndCallSecure(newImplementation, data, true);
+        _upgradeToAndCallUUPS(newImplementation, data, true);
     }
 
     /**

test/helpers/erc1967.js

@@ -0,0 +1,24 @@
+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

@@ -1,16 +1,10 @@
-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

@@ -1,6 +1,5 @@
-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

@@ -1,6 +1,6 @@
 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

@@ -1,9 +1,11 @@
 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);
+  });
 });