feat: add breakable beacon proxy, move proxies to dedicated folder

contracts/proxy/BreakableBeaconProxy.sol

@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.24;
+
+import {Proxy} from "@openzeppelin/contracts/proxy/Proxy.sol";
+import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
+import {IBeacon} from "@openzeppelin/contracts/proxy/beacon/IBeacon.sol";
+
+contract BreakableBeaconProxy is Proxy {
+    constructor(address beacon) {
+        ERC1967Utils.upgradeBeaconToAndCall(beacon, new bytes(0));
+    }
+
+    function _implementation() internal view override returns (address) {
+        address beacon = ERC1967Utils.getBeacon();
+        if (beacon != address(0)) {
+            return IBeacon(beacon).implementation();
+        }
+        return ERC1967Utils.getImplementation();
+    }
+}

contracts/proxy/BreakableUpgradeable.sol

@@ -0,0 +1,75 @@
+// SPDX-License-Identifier: GPL-3.0
+pragma solidity ^0.8.24;
+
+import {IERC1822Proxiable} from "@openzeppelin/contracts/interfaces/draft-IERC1822.sol";
+import {IBeacon} from "@openzeppelin/contracts/proxy/beacon/IBeacon.sol";
+import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
+import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
+import {StorageSlot} from "@openzeppelin/contracts/utils/StorageSlot.sol";
+
+abstract contract BreakableUpgradeable is UUPSUpgradeable {
+    /// @custom:oz-upgrades-unsafe-allow state-variable-immutable
+    address private immutable __breakableSelf = address(this);
+
+    error InvalidBeaconImplementation(address implementation);
+    event BeaconBroken(address indexed beacon, address indexed implementation);
+
+    /// @dev Authorizes the upgrade, detaches from a beacon if needed, and then performs a standard UUPS upgrade.
+    function upgradeToAndCall(
+        address newImplementation,
+        bytes memory data
+    ) public payable virtual override onlyProxy {
+        _authorizeUpgrade(newImplementation);
+        _breakBeacon();
+        _upgradeToAndCallUUPSBreakable(newImplementation, data);
+    }
+
+    /// @dev Detaches a BreakableBeaconProxy by freezing this implementation in the ERC-1967 implementation slot.
+    function _breakBeacon() internal virtual {
+        _checkProxy();
+
+        address beacon = ERC1967Utils.getBeacon();
+        if (beacon == address(0)) {
+            return;
+        }
+
+        address implementation = IBeacon(beacon).implementation();
+        if (implementation != __breakableSelf) {
+            revert InvalidBeaconImplementation(implementation);
+        }
+
+        StorageSlot
+            .getAddressSlot(ERC1967Utils.IMPLEMENTATION_SLOT)
+            .value = implementation;
+        StorageSlot.getAddressSlot(ERC1967Utils.BEACON_SLOT).value = address(0);
+
+        emit BeaconBroken(beacon, implementation);
+    }
+
+    function _checkProxy() internal view virtual override {
+        address beacon = ERC1967Utils.getBeacon();
+        if (address(this) == __breakableSelf) {
+            revert UUPSUnauthorizedCallContext();
+        } else if (beacon == address(0)) {
+            super._checkProxy();
+        } else if (IBeacon(beacon).implementation() != __breakableSelf) {
+            revert UUPSUnauthorizedCallContext();
+        }
+    }
+
+    function _upgradeToAndCallUUPSBreakable(
+        address newImplementation,
+        bytes memory data
+    ) private {
+        try IERC1822Proxiable(newImplementation).proxiableUUID() returns (
+            bytes32 slot
+        ) {
+            if (slot != ERC1967Utils.IMPLEMENTATION_SLOT) {
+                revert UUPSUnsupportedProxiableUUID(slot);
+            }
+            ERC1967Utils.upgradeToAndCall(newImplementation, data);
+        } catch {
+            revert ERC1967Utils.ERC1967InvalidImplementation(newImplementation);
+        }
+    }
+}

contracts/test/MockBeacon.sol

@@ -0,0 +1,33 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {IBeacon} from "@openzeppelin/contracts/proxy/beacon/IBeacon.sol";
+
+contract MockBeacon is IBeacon {
+    error InvalidImplementation(address implementation);
+
+    event Upgraded(address indexed implementation);
+
+    address private _implementation;
+
+    constructor(address implementation_) {
+        _setImplementation(implementation_);
+    }
+
+    function implementation() external view override returns (address) {
+        return _implementation;
+    }
+
+    function upgradeTo(address newImplementation) external {
+        _setImplementation(newImplementation);
+    }
+
+    function _setImplementation(address newImplementation) private {
+        if (newImplementation.code.length == 0) {
+            revert InvalidImplementation(newImplementation);
+        }
+
+        _implementation = newImplementation;
+        emit Upgraded(newImplementation);
+    }
+}

contracts/test/MockBreakableUpgradeable.sol

@@ -0,0 +1,41 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.24;
+
+import {BreakableUpgradeable} from "../proxy/BreakableUpgradeable.sol";
+
+contract MockBreakableUpgradeable is BreakableUpgradeable {
+    error AlreadyInitialized();
+    error NotOwner();
+
+    uint256 public value;
+    address public owner;
+
+    uint256 private immutable _version;
+
+    constructor(uint256 version_) {
+        _version = version_;
+    }
+
+    function initialize(address owner_, uint256 value_) external {
+        if (owner != address(0)) {
+            revert AlreadyInitialized();
+        }
+
+        owner = owner_;
+        value = value_;
+    }
+
+    function version() external view returns (uint256) {
+        return _version;
+    }
+
+    function setValue(uint256 value_) external {
+        value = value_;
+    }
+
+    function _authorizeUpgrade(address) internal view override {
+        if (msg.sender != owner) {
+            revert NotOwner();
+        }
+    }
+}

package.json

@@ -5,8 +5,8 @@
   "dependencies": {
     "@matterlabs/zksync-contracts": "28.0.1",
     "@nomad-xyz/excessively-safe-call": "github:nomad-xyz/ExcessivelySafeCall",
-    "@openzeppelin/contracts": "5.1.0",
-    "@openzeppelin/contracts-upgradeable": "5.1.0",
+    "@openzeppelin/contracts": "5.6.1",
+    "@openzeppelin/contracts-upgradeable": "5.6.1",
     "solady": "^0.1.21",
     "ts-morph": "^19.0.0"
   },

test/proxy/breakablebeacon.test.ts

@@ -0,0 +1,174 @@
+import { expect } from "chai";
+import { getAddress, ZeroAddress } from "ethers";
+import * as hre from "hardhat";
+import type { Wallet } from "zksync-ethers";
+import { Contract, Provider } from "zksync-ethers";
+
+import { LOCAL_RICH_WALLETS, getWallet } from "../../deploy/utils";
+import { ClaveDeployer } from "../utils/deployer";
+
+const IMPLEMENTATION_SLOT =
+  "0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc";
+const BEACON_SLOT =
+  "0xa3f0ad74e5423aebfd80d3ef4346578335a9a72aeaee59ff6cb3582b35133d50";
+
+const storageAddress = (value: string): string => {
+  return getAddress(`0x${value.slice(-40)}`);
+};
+
+describe("BreakableBeaconProxy", () => {
+  let deployer: ClaveDeployer;
+  let provider: Provider;
+  let ownerWallet: Wallet;
+  let otherWallet: Wallet;
+  let mockAbi: unknown[];
+
+  before(async () => {
+    ownerWallet = getWallet(hre, LOCAL_RICH_WALLETS[0].privateKey);
+    otherWallet = getWallet(hre, LOCAL_RICH_WALLETS[1].privateKey);
+    deployer = new ClaveDeployer(hre, ownerWallet);
+    provider = new Provider(hre.network.config.url, undefined, {
+      cacheTimeout: -1,
+    });
+
+    mockAbi = (await hre.zksyncEthers.loadArtifact("MockBreakableUpgradeable"))
+      .abi;
+  });
+
+  const deployFixture = async (): Promise<{
+    owner: string;
+    implementationV2: Contract;
+    implementationV3: Contract;
+    beacon: Contract;
+    proxy: Contract;
+    proxied: Contract;
+  }> => {
+    const owner = await ownerWallet.getAddress();
+    const implementationV1 = await deployer.deployCustomContract(
+      "MockBreakableUpgradeable",
+      [1],
+    );
+    const implementationV2 = await deployer.deployCustomContract(
+      "MockBreakableUpgradeable",
+      [2],
+    );
+    const implementationV3 = await deployer.deployCustomContract(
+      "MockBreakableUpgradeable",
+      [3],
+    );
+    const beacon = await deployer.deployCustomContract("MockBeacon", [
+      await implementationV1.getAddress(),
+    ]);
+    const proxy = await deployer.deployCustomContract("BreakableBeaconProxy", [
+      await beacon.getAddress(),
+    ]);
+    const proxied = new Contract(
+      await proxy.getAddress(),
+      mockAbi,
+      ownerWallet,
+    );
+
+    await (await proxied.initialize(owner, 123)).wait();
+
+    return {
+      owner,
+      implementationV2,
+      implementationV3,
+      beacon,
+      proxy,
+      proxied,
+    };
+  };
+
+  const expectProxySlots = async (
+    proxy: Contract,
+    expectedBeacon: string,
+    expectedImplementation: string,
+  ): Promise<void> => {
+    const proxyAddress = await proxy.getAddress();
+    const [beaconSlot, implementationSlot] = await Promise.all([
+      provider.getStorage(proxyAddress, BEACON_SLOT),
+      provider.getStorage(proxyAddress, IMPLEMENTATION_SLOT),
+    ]);
+
+    expect(storageAddress(beaconSlot)).to.eq(getAddress(expectedBeacon));
+    expect(storageAddress(implementationSlot)).to.eq(
+      getAddress(expectedImplementation),
+    );
+  };
+
+  it("delegates through the beacon before it is broken", async () => {
+    const { owner, implementationV2, beacon, proxy, proxied } =
+      await deployFixture();
+
+    expect(await proxied.version()).to.eq(1n);
+    expect(await proxied.value()).to.eq(123n);
+    expect(await proxied.owner()).to.eq(owner);
+
+    await (await beacon.upgradeTo(await implementationV2.getAddress())).wait();
+
+    expect(await proxied.version()).to.eq(2n);
+    expect(await proxied.value()).to.eq(123n);
+    await expectProxySlots(proxy, await beacon.getAddress(), ZeroAddress);
+  });
+
+  it("breaks away from the beacon during a UUPS upgrade", async () => {
+    const {
+      owner,
+      implementationV2,
+      implementationV3,
+      beacon,
+      proxy,
+      proxied,
+    } = await deployFixture();
+
+    await (
+      await proxied.upgradeToAndCall(await implementationV2.getAddress(), "0x")
+    ).wait();
+
+    expect(await proxied.version()).to.eq(2n);
+    expect(await proxied.value()).to.eq(123n);
+    expect(await proxied.owner()).to.eq(owner);
+    await expectProxySlots(
+      proxy,
+      ZeroAddress,
+      await implementationV2.getAddress(),
+    );
+
+    await (await beacon.upgradeTo(await implementationV3.getAddress())).wait();
+
+    expect(await proxied.version()).to.eq(2n);
+  });
+
+  it("continues to support UUPS upgrades after it is broken", async () => {
+    const { implementationV2, implementationV3, proxy, proxied } =
+      await deployFixture();
+
+    await (
+      await proxied.upgradeToAndCall(await implementationV2.getAddress(), "0x")
+    ).wait();
+    await (
+      await proxied.upgradeToAndCall(await implementationV3.getAddress(), "0x")
+    ).wait();
+
+    expect(await proxied.version()).to.eq(3n);
+    await expectProxySlots(
+      proxy,
+      ZeroAddress,
+      await implementationV3.getAddress(),
+    );
+  });
+
+  it("does not break the beacon when authorization fails", async () => {
+    const { implementationV2, beacon, proxy, proxied } = await deployFixture();
+
+    await expect(
+      proxied
+        .connect(otherWallet)
+        .upgradeToAndCall(await implementationV2.getAddress(), "0x"),
+    ).to.be.revertedWithCustomError(proxied, "NotOwner");
+
+    expect(await proxied.version()).to.eq(1n);
+    await expectProxySlots(proxy, await beacon.getAddress(), ZeroAddress);
+  });
+});