Move upgradeToAndCallUUPS to UUPSUpgradeable (#4356)

Co-authored-by: ernestognw <ernestognw@gmail.com>

.changeset/hip-beds-provide.md

@@ -0,0 +1,5 @@
+---
+'openzeppelin-solidity': major
+---
+
+Move the logic to validate ERC-1822 during an upgrade from `ERC1967Utils` to `UUPSUpgradeable`.

contracts/mocks/proxy/UUPSLegacy.sol

@@ -1,54 +0,0 @@
-// SPDX-License-Identifier: MIT
-
-pragma solidity ^0.8.19;
-
-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 ERC1967Utils
-    bytes32 private constant _ROLLBACK_SLOT = 0x4910fdfa16fed3260ed0e7147f7cc6da11a60208b5b9406d12a635614ffd9143;
-
-    // ERC1967Utils._setImplementation is private so we reproduce it here.
-    // An extra underscore prevents a name clash error.
-    function __setImplementation(address newImplementation) private {
-        require(newImplementation.code.length > 0, "ERC1967: new implementation is not a contract");
-        StorageSlot.getAddressSlot(ERC1967Utils.IMPLEMENTATION_SLOT).value = newImplementation;
-    }
-
-    function _upgradeToAndCallSecureLegacyV1(address newImplementation, bytes memory data, bool forceCall) internal {
-        address oldImplementation = ERC1967Utils.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.encodeCall(this.upgradeTo, (oldImplementation)));
-            rollbackTesting.value = false;
-            // Check rollback was effective
-            require(
-                oldImplementation == ERC1967Utils.getImplementation(),
-                "ERC1967Utils: upgrade breaks further upgrades"
-            );
-            // Finally reset to the new implementation and log the upgrade
-            ERC1967Utils.upgradeTo(newImplementation);
-        }
-    }
-
-    // hooking into the old mechanism
-    function upgradeTo(address newImplementation) public override {
-        _upgradeToAndCallSecureLegacyV1(newImplementation, bytes(""), false);
-    }
-
-    function upgradeToAndCall(address newImplementation, bytes memory data) public payable override {
-        _upgradeToAndCallSecureLegacyV1(newImplementation, data, false);
-    }
-}

contracts/mocks/proxy/UUPSUpgradeableMock.sol

@@ -30,3 +30,9 @@ contract UUPSUpgradeableUnsafeMock is UUPSUpgradeableMock {
         ERC1967Utils.upgradeToAndCall(newImplementation, data, false);
     }
 }
+
+contract UUPSUnsupportedProxiableUUID is UUPSUpgradeableMock {
+    function proxiableUUID() external pure override returns (bytes32) {
+        return keccak256("invalid UUID");
+    }
+}

contracts/proxy/ERC1967/ERC1967Utils.sol

@@ -4,8 +4,6 @@
 pragma solidity ^0.8.20;
 
 import "../beacon/IBeacon.sol";
-import "../../interfaces/IERC1967.sol";
-import "../../interfaces/draft-IERC1822.sol";
 import "../../utils/Address.sol";
 import "../../utils/StorageSlot.sol";
 
@@ -33,9 +31,6 @@ library ERC1967Utils {
      */
     event BeaconUpgraded(address indexed beacon);
 
-    // This is the keccak-256 hash of "eip1967.proxy.rollback" subtracted by 1
-    bytes32 private constant _ROLLBACK_SLOT = 0x4910fdfa16fed3260ed0e7147f7cc6da11a60208b5b9406d12a635614ffd9143;
-
     /**
      * @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
@@ -59,11 +54,6 @@ library ERC1967Utils {
      */
     error ERC1967InvalidBeacon(address beacon);
 
-    /**
-     * @dev The storage `slot` is unsupported as a UUID.
-     */
-    error ERC1967UnsupportedProxiableUUID(bytes32 slot);
-
     /**
      * @dev Returns the current implementation address.
      */
@@ -103,30 +93,6 @@ library ERC1967Utils {
         }
     }
 
-    /**
-     * @dev Perform implementation upgrade with security checks for UUPS proxies, and additional setup call.
-     *
-     * Emits an {IERC1967-Upgraded} event.
-     */
-    function upgradeToAndCallUUPS(address newImplementation, bytes memory data, bool forceCall) internal {
-        // 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) {
-                if (slot != IMPLEMENTATION_SLOT) {
-                    revert ERC1967UnsupportedProxiableUUID(slot);
-                }
-            } catch {
-                // The implementation is not UUPS
-                revert ERC1967InvalidImplementation(newImplementation);
-            }
-            upgradeToAndCall(newImplementation, data, forceCall);
-        }
-    }
-
     /**
      * @dev Storage slot with the admin of the contract.
      * This is the keccak-256 hash of "eip1967.proxy.admin" subtracted by 1, and is

contracts/proxy/transparent/TransparentUpgradeableProxy.sol

@@ -4,6 +4,7 @@
 pragma solidity ^0.8.19;
 
 import "../ERC1967/ERC1967Proxy.sol";
+import "../../interfaces/IERC1967.sol";
 
 /**
  * @dev Interface for {TransparentUpgradeableProxy}. In order to implement transparency, {TransparentUpgradeableProxy}

contracts/proxy/utils/UUPSUpgradeable.sol

@@ -27,6 +27,11 @@ abstract contract UUPSUpgradeable is IERC1822Proxiable {
      */
     error UUPSUnauthorizedCallContext();
 
+    /**
+     * @dev The storage `slot` is unsupported as a UUID.
+     */
+    error UUPSUnsupportedProxiableUUID(bytes32 slot);
+
     /**
      * @dev Check that the execution is being performed through a delegatecall call and that the execution context is
      * a proxy contract with an implementation (as defined in ERC1967) pointing to self. This should only be the case
@@ -35,12 +40,10 @@ abstract contract UUPSUpgradeable is IERC1822Proxiable {
      * fail.
      */
     modifier onlyProxy() {
-        if (address(this) == __self) {
-            // Must be called through delegatecall
-            revert UUPSUnauthorizedCallContext();
-        }
-        if (ERC1967Utils.getImplementation() != __self) {
-            // Must be called through an active proxy
+        if (
+            address(this) == __self || // Must be called through delegatecall
+            ERC1967Utils.getImplementation() != __self // Must be called through an active proxy
+        ) {
             revert UUPSUnauthorizedCallContext();
         }
         _;
@@ -81,7 +84,7 @@ abstract contract UUPSUpgradeable is IERC1822Proxiable {
      */
     function upgradeTo(address newImplementation) public virtual onlyProxy {
         _authorizeUpgrade(newImplementation);
-        ERC1967Utils.upgradeToAndCallUUPS(newImplementation, new bytes(0), false);
+        _upgradeToAndCallUUPS(newImplementation, new bytes(0), false);
     }
 
     /**
@@ -96,7 +99,7 @@ abstract contract UUPSUpgradeable is IERC1822Proxiable {
      */
     function upgradeToAndCall(address newImplementation, bytes memory data) public payable virtual onlyProxy {
         _authorizeUpgrade(newImplementation);
-        ERC1967Utils.upgradeToAndCallUUPS(newImplementation, data, true);
+        _upgradeToAndCallUUPS(newImplementation, data, true);
     }
 
     /**
@@ -110,4 +113,21 @@ abstract contract UUPSUpgradeable is IERC1822Proxiable {
      * ```
      */
     function _authorizeUpgrade(address newImplementation) internal virtual;
+
+    /**
+     * @dev Perform implementation upgrade with security checks for UUPS proxies, and additional setup call.
+     *
+     * Emits an {IERC1967-Upgraded} event.
+     */
+    function _upgradeToAndCallUUPS(address newImplementation, bytes memory data, bool forceCall) private {
+        try IERC1822Proxiable(newImplementation).proxiableUUID() returns (bytes32 slot) {
+            if (slot != ERC1967Utils.IMPLEMENTATION_SLOT) {
+                revert UUPSUnsupportedProxiableUUID(slot);
+            }
+            ERC1967Utils.upgradeToAndCall(newImplementation, data, forceCall);
+        } catch {
+            // The implementation is not UUPS
+            revert ERC1967Utils.ERC1967InvalidImplementation(newImplementation);
+        }
+    }
 }

test/proxy/utils/UUPSUpgradeable.test.js

@@ -1,13 +1,13 @@
 const { expectEvent } = require('@openzeppelin/test-helpers');
-const { web3 } = require('@openzeppelin/test-helpers/src/setup');
-const { getSlot, ImplementationSlot } = require('../../helpers/erc1967');
+const { getAddressInSlot, ImplementationSlot } = require('../../helpers/erc1967');
 const { expectRevertCustomError } = require('../../helpers/customError');
 
 const ERC1967Proxy = artifacts.require('ERC1967Proxy');
 const UUPSUpgradeableMock = artifacts.require('UUPSUpgradeableMock');
 const UUPSUpgradeableUnsafeMock = artifacts.require('UUPSUpgradeableUnsafeMock');
-const UUPSUpgradeableLegacyMock = artifacts.require('UUPSUpgradeableLegacyMock');
 const NonUpgradeableMock = artifacts.require('NonUpgradeableMock');
+const UUPSUnsupportedProxiableUUID = artifacts.require('UUPSUnsupportedProxiableUUID');
+const Address = artifacts.require('$Address');
 
 contract('UUPSUpgradeable', function () {
   before(async function () {
@@ -15,6 +15,8 @@ contract('UUPSUpgradeable', function () {
     this.implUpgradeOk = await UUPSUpgradeableMock.new();
     this.implUpgradeUnsafe = await UUPSUpgradeableUnsafeMock.new();
     this.implUpgradeNonUUPS = await NonUpgradeableMock.new();
+    this.implUnsupportedUUID = await UUPSUnsupportedProxiableUUID.new();
+    this.helper = await Address.new();
   });
 
   beforeEach(async function () {
@@ -26,6 +28,7 @@ contract('UUPSUpgradeable', 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 });
+    expect(await getAddressInSlot(this.instance, ImplementationSlot)).to.be.equal(this.implUpgradeOk.address);
   });
 
   it('upgrade to upgradeable implementation with call', async function () {
@@ -37,13 +40,64 @@ contract('UUPSUpgradeable', function () {
     );
     expect(receipt.logs.filter(({ event }) => event === 'Upgraded').length).to.be.equal(1);
     expectEvent(receipt, 'Upgraded', { implementation: this.implUpgradeOk.address });
+    expect(await getAddressInSlot(this.instance, ImplementationSlot)).to.be.equal(this.implUpgradeOk.address);
 
     expect(await this.instance.current()).to.be.bignumber.equal('1');
   });
 
+  it('calling upgradeTo on the implementation reverts', async function () {
+    await expectRevertCustomError(
+      this.implInitial.upgradeTo(this.implUpgradeOk.address),
+      'UUPSUnauthorizedCallContext',
+      [],
+    );
+  });
+
+  it('calling upgradeToAndCall on the implementation reverts', async function () {
+    await expectRevertCustomError(
+      this.implInitial.upgradeToAndCall(
+        this.implUpgradeOk.address,
+        this.implUpgradeOk.contract.methods.increment().encodeABI(),
+      ),
+      'UUPSUnauthorizedCallContext',
+      [],
+    );
+  });
+
+  it('calling upgradeTo from a contract that is not an ERC1967 proxy (with the right implementation) reverts', async function () {
+    await expectRevertCustomError(
+      this.helper.$functionDelegateCall(
+        this.implUpgradeOk.address,
+        this.implUpgradeOk.contract.methods.upgradeTo(this.implUpgradeUnsafe.address).encodeABI(),
+      ),
+      'UUPSUnauthorizedCallContext',
+      [],
+    );
+  });
+
+  it('calling upgradeToAndCall from a contract that is not an ERC1967 proxy (with the right implementation) reverts', async function () {
+    await expectRevertCustomError(
+      this.helper.$functionDelegateCall(
+        this.implUpgradeOk.address,
+        this.implUpgradeOk.contract.methods.upgradeToAndCall(this.implUpgradeUnsafe.address, '0x').encodeABI(),
+      ),
+      'UUPSUnauthorizedCallContext',
+      [],
+    );
+  });
+
+  it('rejects upgrading to an unsupported UUID', async function () {
+    await expectRevertCustomError(
+      this.instance.upgradeTo(this.implUnsupportedUUID.address),
+      'UUPSUnsupportedProxiableUUID',
+      [web3.utils.keccak256('invalid UUID')],
+    );
+  });
+
   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 });
+    expect(await getAddressInSlot(this.instance, ImplementationSlot)).to.be.equal(this.implUpgradeUnsafe.address);
   });
 
   // delegate to a non existing upgradeTo function causes a low level revert
@@ -63,24 +117,4 @@ contract('UUPSUpgradeable', function () {
       otherInstance.address,
     ]);
   });
-
-  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);
-  });
 });