move upgradeToAndCallUUPS to UUPSUpgradeable

OpenZeppelin · Jun 15, 2023 · c3aec1c · c3aec1c
1 parent ff85c7b
commit c3aec1c
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 38 deletions.
diff --git a/contracts/proxy/ERC1967/ERC1967Utils.sol b/contracts/proxy/ERC1967/ERC1967Utils.sol
@@ -5,7 +5,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";
 
@@ -59,11 +58,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 +97,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

diff --git a/contracts/proxy/utils/UUPSUpgradeable.sol b/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
@@ -81,7 +86,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 +101,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 +115,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);
+            }
+        } catch {
+            // The implementation is not UUPS
+            revert ERC1967Utils.ERC1967InvalidImplementation(newImplementation);
+        }
+        ERC1967Utils.upgradeToAndCall(newImplementation, data, forceCall);
+    }
 }
diff --git a/test/proxy/utils/UUPSUpgradeable.test.js b/test/proxy/utils/UUPSUpgradeable.test.js
@@ -71,12 +71,6 @@ contract('UUPSUpgradeable', function () {
     );
 
     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);