Simplify UUPSUpgradeable along the lines of ERC1822

contracts/proxy/ERC1822/IProxiable.sol

@@ -0,0 +1,8 @@
+// SPDX-License-Identifier: MIT
+// OpenZeppelin Contracts v4.x.0 (proxy/ERC1822/IProxiable.sol)
+
+pragma solidity ^0.8.0;
+
+interface IProxiable {
+    function proxiableUUID() external view returns (bytes32);
+}

contracts/proxy/ERC1967/ERC1967Upgrade.sol

@@ -4,6 +4,7 @@
 pragma solidity ^0.8.2;
 
 import "../beacon/IBeacon.sol";
+import "../ERC1822/IProxiable.sol";
 import "../../utils/Address.sol";
 import "../../utils/StorageSlot.sol";
 
@@ -82,28 +83,14 @@ abstract contract ERC1967Upgrade {
         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)
+        if (StorageSlot.getBooleanSlot(_ROLLBACK_SLOT).value) {
+            _setImplementation(newImplementation);
+        } else {
+            require(
+                IProxiable(newImplementation).proxiableUUID() == _IMPLEMENTATION_SLOT,
+                "Invalid proxiableUUID value"
             );
-            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);
+            _upgradeToAndCall(newImplementation, data, forceCall);
         }
     }
 

contracts/proxy/utils/UUPSUpgradeable.sol

@@ -3,6 +3,7 @@
 
 pragma solidity ^0.8.0;
 
+import "../ERC1822/IProxiable.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 IProxiable, ERC1967Upgrade {
     /// @custom:oz-upgrades-unsafe-allow state-variable-immutable state-variable-assignment
     address private immutable __self = address(this);
 
@@ -34,6 +35,23 @@ 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, "Function 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.
+     */
+    function proxiableUUID() external view virtual override notDelegated returns (bytes32) {
+        return _IMPLEMENTATION_SLOT;
+    }
+
     /**
      * @dev Upgrade the implementation of the proxy to `newImplementation`.
      *

test/proxy/utils/UUPSUpgradeable.test.js

@@ -6,6 +6,8 @@ const UUPSUpgradeableUnsafeMock = artifacts.require('UUPSUpgradeableUnsafeMock')
 const UUPSUpgradeableBrokenMock = artifacts.require('UUPSUpgradeableBrokenMock');
 const CountersImpl = artifacts.require('CountersImpl');
 
+const Legacy = [ './legacy/UUPSUpgradeableMockV1' ];
+
 contract('UUPSUpgradeable', function (accounts) {
   before(async function () {
     this.implInitial = await UUPSUpgradeableMock.new();
@@ -44,7 +46,7 @@ contract('UUPSUpgradeable', function (accounts) {
     expectEvent(receipt, 'Upgraded', { implementation: this.implUpgradeUnsafe.address });
   });
 
-  it('reject upgrade to broken upgradeable implementation', async function () {
+  it.skip('reject upgrade to broken upgradeable implementation (no longer supported)', async function () {
     await expectRevert(
       this.instance.upgradeTo(this.implUpgradeBroken.address),
       'ERC1967Upgrade: upgrade breaks further upgrades',
@@ -55,7 +57,7 @@ contract('UUPSUpgradeable', function (accounts) {
   it('reject upgrade to non uups implementation', async function () {
     await expectRevert(
       this.instance.upgradeTo(this.implUpgradeNonUUPS.address),
-      'Address: low-level delegate call failed',
+      'Transaction reverted: function selector was not recognized and there\'s no fallback function',
     );
   });
 
@@ -66,7 +68,27 @@ contract('UUPSUpgradeable', function (accounts) {
     // infinite loop reverts when a nested call is out-of-gas
     await expectRevert(
       this.instance.upgradeTo(otherInstance.address),
-      'Address: low-level delegate call failed',
+      'Function must not be called through delegatecall',
     );
   });
+
+  describe('compatibility with legacy version of UUPSUpgradeable', function () {
+    // This could possibly be cleaner/simpler with ethers.js
+    for (const path of Legacy) {
+      it(`can upgrade from ${path}`, async function () {
+        const artefact = require(path);
+        // deploy legacy implementation
+        const legacy = new web3.eth.Contract(artefact.abi);
+        const { _address: implAddr } = await legacy.deploy({ data: artefact.bytecode }).send({ from: accounts[0] });
+
+        // deploy proxy
+        const { address: proxyAddr } = await ERC1967Proxy.new(implAddr, '0x');
+        const proxy = new web3.eth.Contract(artefact.abi, proxyAddr);
+
+        // perform upgrade
+        const receipt = await proxy.methods.upgradeTo(this.implInitial.address).send({ from: accounts[0] });
+        expectEvent(receipt, 'Upgraded', { implementation: this.implInitial.address });
+      });
+    }
+  });
 });