Skip to content

Commit fd65c58

Browse files
coffeexcoincoffeexcoin
authored
feat: add breakable beacon proxy, move proxies to dedicated folder (#31)
1 parent 13b799b commit fd65c58

9 files changed

Lines changed: 354 additions & 11 deletions

contracts/AccountProxy.sol renamed to contracts/proxy/AccountProxy.sol

File renamed without changes.

‎contracts/proxy/BreakableBeaconProxy.sol‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
// SPDX-License-Identifier: GPL-3.0
2+
pragma solidity ^0.8.24;
3+
4+
import {Proxy} from "@openzeppelin/contracts/proxy/Proxy.sol";
5+
import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
6+
import {IBeacon} from "@openzeppelin/contracts/proxy/beacon/IBeacon.sol";
7+
8+
contract BreakableBeaconProxy is Proxy {
9+
constructor(address beacon) {
10+
ERC1967Utils.upgradeBeaconToAndCall(beacon, new bytes(0));
11+
}
12+
13+
function _implementation() internal view override returns (address) {
14+
address beacon = ERC1967Utils.getBeacon();
15+
if (beacon != address(0)) {
16+
return IBeacon(beacon).implementation();
17+
}
18+
return ERC1967Utils.getImplementation();
19+
}
20+
}

‎contracts/proxy/BreakableUpgradeable.sol‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
// SPDX-License-Identifier: GPL-3.0
2+
pragma solidity ^0.8.24;
3+
4+
import {IERC1822Proxiable} from "@openzeppelin/contracts/interfaces/draft-IERC1822.sol";
5+
import {IBeacon} from "@openzeppelin/contracts/proxy/beacon/IBeacon.sol";
6+
import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
7+
import {UUPSUpgradeable} from "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
8+
import {StorageSlot} from "@openzeppelin/contracts/utils/StorageSlot.sol";
9+
10+
abstract contract BreakableUpgradeable is UUPSUpgradeable {
11+
/// @custom:oz-upgrades-unsafe-allow state-variable-immutable
12+
address private immutable __breakableSelf = address(this);
13+
14+
error InvalidBeaconImplementation(address implementation);
15+
event BeaconBroken(address indexed beacon, address indexed implementation);
16+
17+
/// @dev Authorizes the upgrade, detaches from a beacon if needed, and then performs a standard UUPS upgrade.
18+
function upgradeToAndCall(
19+
address newImplementation,
20+
bytes memory data
21+
) public payable virtual override onlyProxy {
22+
_authorizeUpgrade(newImplementation);
23+
_breakBeacon();
24+
_upgradeToAndCallUUPSBreakable(newImplementation, data);
25+
}
26+
27+
/// @dev Detaches a BreakableBeaconProxy by freezing this implementation in the ERC-1967 implementation slot.
28+
function _breakBeacon() internal virtual {
29+
_checkProxy();
30+
31+
address beacon = ERC1967Utils.getBeacon();
32+
if (beacon == address(0)) {
33+
return;
34+
}
35+
36+
address implementation = IBeacon(beacon).implementation();
37+
if (implementation != __breakableSelf) {
38+
revert InvalidBeaconImplementation(implementation);
39+
}
40+
41+
StorageSlot
42+
.getAddressSlot(ERC1967Utils.IMPLEMENTATION_SLOT)
43+
.value = implementation;
44+
StorageSlot.getAddressSlot(ERC1967Utils.BEACON_SLOT).value = address(0);
45+
46+
emit BeaconBroken(beacon, implementation);
47+
}
48+
49+
function _checkProxy() internal view virtual override {
50+
address beacon = ERC1967Utils.getBeacon();
51+
if (address(this) == __breakableSelf) {
52+
revert UUPSUnauthorizedCallContext();
53+
} else if (beacon == address(0)) {
54+
super._checkProxy();
55+
} else if (IBeacon(beacon).implementation() != __breakableSelf) {
56+
revert UUPSUnauthorizedCallContext();
57+
}
58+
}
59+
60+
function _upgradeToAndCallUUPSBreakable(
61+
address newImplementation,
62+
bytes memory data
63+
) private {
64+
try IERC1822Proxiable(newImplementation).proxiableUUID() returns (
65+
bytes32 slot
66+
) {
67+
if (slot != ERC1967Utils.IMPLEMENTATION_SLOT) {
68+
revert UUPSUnsupportedProxiableUUID(slot);
69+
}
70+
ERC1967Utils.upgradeToAndCall(newImplementation, data);
71+
} catch {
72+
revert ERC1967Utils.ERC1967InvalidImplementation(newImplementation);
73+
}
74+
}
75+
}

contracts/ERC1967Proxy.sol renamed to contracts/proxy/ERC1967Proxy.sol

File renamed without changes.

‎contracts/test/MockBeacon.sol‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity ^0.8.24;
3+
4+
import {IBeacon} from "@openzeppelin/contracts/proxy/beacon/IBeacon.sol";
5+
6+
contract MockBeacon is IBeacon {
7+
error InvalidImplementation(address implementation);
8+
9+
event Upgraded(address indexed implementation);
10+
11+
address private _implementation;
12+
13+
constructor(address implementation_) {
14+
_setImplementation(implementation_);
15+
}
16+
17+
function implementation() external view override returns (address) {
18+
return _implementation;
19+
}
20+
21+
function upgradeTo(address newImplementation) external {
22+
_setImplementation(newImplementation);
23+
}
24+
25+
function _setImplementation(address newImplementation) private {
26+
if (newImplementation.code.length == 0) {
27+
revert InvalidImplementation(newImplementation);
28+
}
29+
30+
_implementation = newImplementation;
31+
emit Upgraded(newImplementation);
32+
}
33+
}

‎contracts/test/MockBreakableUpgradeable.sol‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
// SPDX-License-Identifier: MIT
2+
pragma solidity ^0.8.24;
3+
4+
import {BreakableUpgradeable} from "../proxy/BreakableUpgradeable.sol";
5+
6+
contract MockBreakableUpgradeable is BreakableUpgradeable {
7+
error AlreadyInitialized();
8+
error NotOwner();
9+
10+
uint256 public value;
11+
address public owner;
12+
13+
uint256 private immutable _version;
14+
15+
constructor(uint256 version_) {
16+
_version = version_;
17+
}
18+
19+
function initialize(address owner_, uint256 value_) external {
20+
if (owner != address(0)) {
21+
revert AlreadyInitialized();
22+
}
23+
24+
owner = owner_;
25+
value = value_;
26+
}
27+
28+
function version() external view returns (uint256) {
29+
return _version;
30+
}
31+
32+
function setValue(uint256 value_) external {
33+
value = value_;
34+
}
35+
36+
function _authorizeUpgrade(address) internal view override {
37+
if (msg.sender != owner) {
38+
revert NotOwner();
39+
}
40+
}
41+
}

‎package-lock.json‎

Lines changed: 9 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@
55
"dependencies": {
66
"@matterlabs/zksync-contracts": "28.0.1",
77
"@nomad-xyz/excessively-safe-call": "github:nomad-xyz/ExcessivelySafeCall",
8-
"@openzeppelin/contracts": "5.1.0",
9-
"@openzeppelin/contracts-upgradeable": "5.1.0",
8+
"@openzeppelin/contracts": "5.6.1",
9+
"@openzeppelin/contracts-upgradeable": "5.6.1",
1010
"solady": "^0.1.21",
1111
"ts-morph": "^19.0.0"
1212
},

‎test/proxy/breakablebeacon.test.ts‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,174 @@
1+
import { expect } from "chai";
2+
import { getAddress, ZeroAddress } from "ethers";
3+
import * as hre from "hardhat";
4+
import type { Wallet } from "zksync-ethers";
5+
import { Contract, Provider } from "zksync-ethers";
6+
7+
import { LOCAL_RICH_WALLETS, getWallet } from "../../deploy/utils";
8+
import { ClaveDeployer } from "../utils/deployer";
9+
10+
const IMPLEMENTATION_SLOT =
11+
"0x360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc";
12+
const BEACON_SLOT =
13+
"0xa3f0ad74e5423aebfd80d3ef4346578335a9a72aeaee59ff6cb3582b35133d50";
14+
15+
const storageAddress = (value: string): string => {
16+
return getAddress(`0x${value.slice(-40)}`);
17+
};
18+
19+
describe("BreakableBeaconProxy", () => {
20+
let deployer: ClaveDeployer;
21+
let provider: Provider;
22+
let ownerWallet: Wallet;
23+
let otherWallet: Wallet;
24+
let mockAbi: unknown[];
25+
26+
before(async () => {
27+
ownerWallet = getWallet(hre, LOCAL_RICH_WALLETS[0].privateKey);
28+
otherWallet = getWallet(hre, LOCAL_RICH_WALLETS[1].privateKey);
29+
deployer = new ClaveDeployer(hre, ownerWallet);
30+
provider = new Provider(hre.network.config.url, undefined, {
31+
cacheTimeout: -1,
32+
});
33+
34+
mockAbi = (await hre.zksyncEthers.loadArtifact("MockBreakableUpgradeable"))
35+
.abi;
36+
});
37+
38+
const deployFixture = async (): Promise<{
39+
owner: string;
40+
implementationV2: Contract;
41+
implementationV3: Contract;
42+
beacon: Contract;
43+
proxy: Contract;
44+
proxied: Contract;
45+
}> => {
46+
const owner = await ownerWallet.getAddress();
47+
const implementationV1 = await deployer.deployCustomContract(
48+
"MockBreakableUpgradeable",
49+
[1],
50+
);
51+
const implementationV2 = await deployer.deployCustomContract(
52+
"MockBreakableUpgradeable",
53+
[2],
54+
);
55+
const implementationV3 = await deployer.deployCustomContract(
56+
"MockBreakableUpgradeable",
57+
[3],
58+
);
59+
const beacon = await deployer.deployCustomContract("MockBeacon", [
60+
await implementationV1.getAddress(),
61+
]);
62+
const proxy = await deployer.deployCustomContract("BreakableBeaconProxy", [
63+
await beacon.getAddress(),
64+
]);
65+
const proxied = new Contract(
66+
await proxy.getAddress(),
67+
mockAbi,
68+
ownerWallet,
69+
);
70+
71+
await (await proxied.initialize(owner, 123)).wait();
72+
73+
return {
74+
owner,
75+
implementationV2,
76+
implementationV3,
77+
beacon,
78+
proxy,
79+
proxied,
80+
};
81+
};
82+
83+
const expectProxySlots = async (
84+
proxy: Contract,
85+
expectedBeacon: string,
86+
expectedImplementation: string,
87+
): Promise<void> => {
88+
const proxyAddress = await proxy.getAddress();
89+
const [beaconSlot, implementationSlot] = await Promise.all([
90+
provider.getStorage(proxyAddress, BEACON_SLOT),
91+
provider.getStorage(proxyAddress, IMPLEMENTATION_SLOT),
92+
]);
93+
94+
expect(storageAddress(beaconSlot)).to.eq(getAddress(expectedBeacon));
95+
expect(storageAddress(implementationSlot)).to.eq(
96+
getAddress(expectedImplementation),
97+
);
98+
};
99+
100+
it("delegates through the beacon before it is broken", async () => {
101+
const { owner, implementationV2, beacon, proxy, proxied } =
102+
await deployFixture();
103+
104+
expect(await proxied.version()).to.eq(1n);
105+
expect(await proxied.value()).to.eq(123n);
106+
expect(await proxied.owner()).to.eq(owner);
107+
108+
await (await beacon.upgradeTo(await implementationV2.getAddress())).wait();
109+
110+
expect(await proxied.version()).to.eq(2n);
111+
expect(await proxied.value()).to.eq(123n);
112+
await expectProxySlots(proxy, await beacon.getAddress(), ZeroAddress);
113+
});
114+
115+
it("breaks away from the beacon during a UUPS upgrade", async () => {
116+
const {
117+
owner,
118+
implementationV2,
119+
implementationV3,
120+
beacon,
121+
proxy,
122+
proxied,
123+
} = await deployFixture();
124+
125+
await (
126+
await proxied.upgradeToAndCall(await implementationV2.getAddress(), "0x")
127+
).wait();
128+
129+
expect(await proxied.version()).to.eq(2n);
130+
expect(await proxied.value()).to.eq(123n);
131+
expect(await proxied.owner()).to.eq(owner);
132+
await expectProxySlots(
133+
proxy,
134+
ZeroAddress,
135+
await implementationV2.getAddress(),
136+
);
137+
138+
await (await beacon.upgradeTo(await implementationV3.getAddress())).wait();
139+
140+
expect(await proxied.version()).to.eq(2n);
141+
});
142+
143+
it("continues to support UUPS upgrades after it is broken", async () => {
144+
const { implementationV2, implementationV3, proxy, proxied } =
145+
await deployFixture();
146+
147+
await (
148+
await proxied.upgradeToAndCall(await implementationV2.getAddress(), "0x")
149+
).wait();
150+
await (
151+
await proxied.upgradeToAndCall(await implementationV3.getAddress(), "0x")
152+
).wait();
153+
154+
expect(await proxied.version()).to.eq(3n);
155+
await expectProxySlots(
156+
proxy,
157+
ZeroAddress,
158+
await implementationV3.getAddress(),
159+
);
160+
});
161+
162+
it("does not break the beacon when authorization fails", async () => {
163+
const { implementationV2, beacon, proxy, proxied } = await deployFixture();
164+
165+
await expect(
166+
proxied
167+
.connect(otherWallet)
168+
.upgradeToAndCall(await implementationV2.getAddress(), "0x"),
169+
).to.be.revertedWithCustomError(proxied, "NotOwner");
170+
171+
expect(await proxied.version()).to.eq(1n);
172+
await expectProxySlots(proxy, await beacon.getAddress(), ZeroAddress);
173+
});
174+
});

0 commit comments

Comments
 (0)