From db04ec2708f866d52de1d0323ef4553c73cac90e Mon Sep 17 00:00:00 2001 From: Artyom Veremeenko Date: Mon, 16 Jun 2025 10:49:28 +0300 Subject: [PATCH 1/3] feat(V3Template): check start and finish in same tx (transient storage) --- contracts/upgrade/V3Template.sol | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/contracts/upgrade/V3Template.sol b/contracts/upgrade/V3Template.sol index e4ebbd2fc..b61099d75 100644 --- a/contracts/upgrade/V3Template.sol +++ b/contracts/upgrade/V3Template.sol @@ -110,6 +110,15 @@ contract V3Template is V3Addresses { uint256 public initialTotalPooledEther; address[] public contractsWithBurnerAllowances; + // + // Slots for transient storage + // + + // Slot for the upgrade started flag + // keccak256("V3Template.upgradeStartedFlag") + bytes32 private constant UPGRADE_STARTED_SLOT = + 0x058d69f67a3d86c424c516d23a070ff8bed34431617274caa2049bd702675e3f; + /// @param _params Params required to initialize the addresses contract constructor(V3AddressesParams memory _params) V3Addresses(_params) { @@ -125,7 +134,10 @@ contract V3Template is V3Addresses { if (msg.sender != AGENT) revert OnlyAgentCanUpgrade(); if (block.timestamp >= EXPIRE_SINCE_INCLUSIVE) revert Expired(); if (upgradeBlockNumber != UPGRADE_NOT_STARTED) revert UpgradeAlreadyStarted(); + if (isUpgradeFinished) revert UpgradeAlreadyFinished(); + if (_isStartCalledInThisTx()) revert StartAlreadyCalledInThisTx(); + assembly { tstore(UPGRADE_STARTED_SLOT, 1) } upgradeBlockNumber = block.number; initialTotalShares = ILidoWithFinalizeUpgrade(LIDO).getTotalShares(); @@ -141,8 +153,8 @@ contract V3Template is V3Addresses { function finishUpgrade() external { if (msg.sender != AGENT) revert OnlyAgentCanUpgrade(); - if (upgradeBlockNumber != block.number) revert StartAndFinishMustBeInSameBlock(); if (isUpgradeFinished) revert UpgradeAlreadyFinished(); + if (!_isStartCalledInThisTx()) revert StartAndFinishMustBeInSameTx(); isUpgradeFinished = true; @@ -356,6 +368,12 @@ contract V3Template is V3Addresses { } } + function _isStartCalledInThisTx() internal view returns (bool isStartCalledInThisTx) { + assembly { + isStartCalledInThisTx := tload(UPGRADE_STARTED_SLOT) + } + } + error OnlyAgentCanUpgrade(); error UpgradeAlreadyStarted(); error UpgradeAlreadyFinished(); @@ -366,6 +384,8 @@ contract V3Template is V3Addresses { error NonZeroRoleHolders(address contractAddress, bytes32 role); error IncorrectAragonAppImplementation(address repo, address implementation); error StartAndFinishMustBeInSameBlock(); + error StartAndFinishMustBeInSameTx(); + error StartAlreadyCalledInThisTx(); error Expired(); error IncorrectBurnerSharesMigration(); error IncorrectBurnerAllowance(address contractAddress, address burner); From 1c05e3de9e69873725b2aad0fa731bd1f4b51438 Mon Sep 17 00:00:00 2001 From: Artyom Veremeenko Date: Mon, 16 Jun 2025 13:42:57 +0300 Subject: [PATCH 2/3] fix: code review --- contracts/0.8.25/vaults/dashboard/Dashboard.sol | 2 +- contracts/upgrade/V3Template.sol | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/contracts/0.8.25/vaults/dashboard/Dashboard.sol b/contracts/0.8.25/vaults/dashboard/Dashboard.sol index 9b16a9485..a1336c021 100644 --- a/contracts/0.8.25/vaults/dashboard/Dashboard.sol +++ b/contracts/0.8.25/vaults/dashboard/Dashboard.sol @@ -53,7 +53,7 @@ contract Dashboard is NodeOperatorFee { * @notice Slot for the fund-on-receive flag * keccak256("vaults.Dashboard.fundOnReceive") */ - bytes32 private constant FUND_ON_RECEIVE_FLAG_SLOT = + bytes32 public constant FUND_ON_RECEIVE_FLAG_SLOT = 0x7408b7b034fda7051615c19182918ecb91d753231cffd86f81a45d996d63e038; /** diff --git a/contracts/upgrade/V3Template.sol b/contracts/upgrade/V3Template.sol index b61099d75..90db7672a 100644 --- a/contracts/upgrade/V3Template.sol +++ b/contracts/upgrade/V3Template.sol @@ -116,7 +116,7 @@ contract V3Template is V3Addresses { // Slot for the upgrade started flag // keccak256("V3Template.upgradeStartedFlag") - bytes32 private constant UPGRADE_STARTED_SLOT = + bytes32 public constant UPGRADE_STARTED_SLOT = 0x058d69f67a3d86c424c516d23a070ff8bed34431617274caa2049bd702675e3f; From 93608ee019ea2db0ae8e2b5664379bc92c669d7e Mon Sep 17 00:00:00 2001 From: Artyom Veremeenko Date: Tue, 17 Jun 2025 09:59:04 +0300 Subject: [PATCH 3/3] feat: more V3Template tests --- contracts/upgrade/V3Template.sol | 2 +- .../upgrade/steps/0500-mock-aragon-voting.ts | 10 +- .../contracts/VaultHub__MockForDashboard.sol | 32 ++++- .../upgrade-template-v3.integration.ts | 124 ++++++++++++++---- test/upgrade/V3Template_Harness.sol | 17 +++ 5 files changed, 149 insertions(+), 36 deletions(-) create mode 100644 test/upgrade/V3Template_Harness.sol diff --git a/contracts/upgrade/V3Template.sol b/contracts/upgrade/V3Template.sol index 90db7672a..bd4057190 100644 --- a/contracts/upgrade/V3Template.sol +++ b/contracts/upgrade/V3Template.sol @@ -133,9 +133,9 @@ contract V3Template is V3Addresses { function startUpgrade() external { if (msg.sender != AGENT) revert OnlyAgentCanUpgrade(); if (block.timestamp >= EXPIRE_SINCE_INCLUSIVE) revert Expired(); - if (upgradeBlockNumber != UPGRADE_NOT_STARTED) revert UpgradeAlreadyStarted(); if (isUpgradeFinished) revert UpgradeAlreadyFinished(); if (_isStartCalledInThisTx()) revert StartAlreadyCalledInThisTx(); + if (upgradeBlockNumber != UPGRADE_NOT_STARTED) revert UpgradeAlreadyStarted(); assembly { tstore(UPGRADE_STARTED_SLOT, 1) } upgradeBlockNumber = block.number; diff --git a/scripts/upgrade/steps/0500-mock-aragon-voting.ts b/scripts/upgrade/steps/0500-mock-aragon-voting.ts index 894e6f03f..71c9ebf8c 100644 --- a/scripts/upgrade/steps/0500-mock-aragon-voting.ts +++ b/scripts/upgrade/steps/0500-mock-aragon-voting.ts @@ -5,7 +5,7 @@ import { impersonate } from "lib/account"; import { loadContract } from "lib/contract"; import { readNetworkState, Sk } from "lib/state-file"; -export async function main(): Promise { +export async function createVoteAndGetExecuteTxPromise() { const state = readNetworkState(); const agentAddress = state[Sk.appAgent].proxy.address; const votingAddress = state[Sk.appVoting].proxy.address; @@ -24,8 +24,14 @@ export async function main(): Promise { if (!(await voteScript.isValidVoteScript(voteId))) throw new Error("Vote script is not valid"); await voting.connect(deployer).vote(voteId, true, false); await advanceChainTime(await voting.voteTime()); - const executeTx = await voting.executeVote(voteId); + return voting.executeVote(voteId); +} + +export async function main() { + const executeTx = await createVoteAndGetExecuteTxPromise(); const executeReceipt = await executeTx.wait(); log.success("Voting executed: gas used", executeReceipt!.gasUsed); + + return { executeReceipt }; } diff --git a/test/0.8.25/vaults/dashboard/contracts/VaultHub__MockForDashboard.sol b/test/0.8.25/vaults/dashboard/contracts/VaultHub__MockForDashboard.sol index 2c8cf1ff4..f93cda3eb 100644 --- a/test/0.8.25/vaults/dashboard/contracts/VaultHub__MockForDashboard.sol +++ b/test/0.8.25/vaults/dashboard/contracts/VaultHub__MockForDashboard.sol @@ -85,7 +85,7 @@ contract VaultHub__MockForDashboard { reservationFeeBP: 100, isBeaconDepositsManuallyPaused: false }); - + emit Mock__VaultConnected(vault); } @@ -167,9 +167,25 @@ contract VaultHub__MockForDashboard { return vaultConnections[_vault].vaultIndex != 0; } - function updateConnection(address _vault, uint256 _shareLimit, uint256 _reserveRatioBP, uint256 _forcedRebalanceThresholdBP, uint256 _infraFeeBP, uint256 _liquidityFeeBP, uint256 _reservationFeeBP) external { + function updateConnection( + address _vault, + uint256 _shareLimit, + uint256 _reserveRatioBP, + uint256 _forcedRebalanceThresholdBP, + uint256 _infraFeeBP, + uint256 _liquidityFeeBP, + uint256 _reservationFeeBP + ) external { if (!isVaultConnected(_vault)) revert NotConnectedToHub(_vault); - emit Mock__VaultConnectionUpdated(_vault, _shareLimit, _reserveRatioBP, _forcedRebalanceThresholdBP, _infraFeeBP, _liquidityFeeBP, _reservationFeeBP); + emit Mock__VaultConnectionUpdated( + _vault, + _shareLimit, + _reserveRatioBP, + _forcedRebalanceThresholdBP, + _infraFeeBP, + _liquidityFeeBP, + _reservationFeeBP + ); } event Mock__ValidatorExitRequested(address vault, bytes pubkeys); @@ -186,7 +202,15 @@ contract VaultHub__MockForDashboard { event Mock__VaultDisconnectInitiated(address vault); event Mock__Rebalanced(address vault, uint256 amount); event Mock__VaultConnected(address vault); - event Mock__VaultConnectionUpdated(address vault, uint256 shareLimit, uint256 reserveRatioBP, uint256 forcedRebalanceThresholdBP, uint256 infraFeeBP, uint256 liquidityFeeBP, uint256 reservationFeeBP); + event Mock__VaultConnectionUpdated( + address vault, + uint256 shareLimit, + uint256 reserveRatioBP, + uint256 forcedRebalanceThresholdBP, + uint256 infraFeeBP, + uint256 liquidityFeeBP, + uint256 reservationFeeBP + ); error ZeroArgument(string argument); error NotConnectedToHub(address vault); diff --git a/test/integration/upgrade/upgrade-template-v3.integration.ts b/test/integration/upgrade/upgrade-template-v3.integration.ts index 800af6872..0cdbd17c5 100644 --- a/test/integration/upgrade/upgrade-template-v3.integration.ts +++ b/test/integration/upgrade/upgrade-template-v3.integration.ts @@ -1,11 +1,12 @@ import { expect } from "chai"; import hre from "hardhat"; import { beforeEach } from "mocha"; +import { createVoteAndGetExecuteTxPromise } from "scripts/upgrade/steps/0500-mock-aragon-voting"; import { HardhatEthersSigner } from "@nomicfoundation/hardhat-ethers/signers"; import { time } from "@nomicfoundation/hardhat-network-helpers"; -import { OssifiableProxy, V3Template } from "typechain-types"; +import { OssifiableProxy, V3Template, V3Template__Harness, V3Template__Harness__factory } from "typechain-types"; import { deployUpgrade, loadContract, readNetworkState, Sk } from "lib"; import { getProtocolContext, ProtocolContext } from "lib/protocol"; @@ -23,6 +24,7 @@ describe("Integration: Upgrade Template V3 tests", () => { let template: V3Template; let deployer: HardhatEthersSigner; let agentSigner: HardhatEthersSigner; + let agentMock: V3Template__Harness; before(async () => { originalSnapshot = await Snapshot.take(); @@ -41,6 +43,10 @@ describe("Integration: Upgrade Template V3 tests", () => { ctx = await getProtocolContext(true); agentSigner = await ctx.getSigner("agent"); + + // const agentMockFactory = new V3Template__Harness__factory(deployer); + agentMock = await new V3Template__Harness__factory(deployer).deploy(await template.getAddress()); + await agentMock.waitForDeployment(); }); after(async () => await Snapshot.restore(originalSnapshot)); @@ -52,7 +58,7 @@ describe("Integration: Upgrade Template V3 tests", () => { afterEach(async () => await Snapshot.restore(snapshot)); function it_(title: string, fn: () => Promise) { - it(title, async function () { + return it(title, async function () { if (needToSkipTemplateTests()) { this.skip(); } @@ -60,41 +66,101 @@ describe("Integration: Upgrade Template V3 tests", () => { }); } - it_("should revert when startUpgrade is called by non-voting address", async function () { - await expect(template.connect(deployer).startUpgrade()).to.be.revertedWithCustomError( - template, - "OnlyAgentCanUpgrade", - ); + it_("happy path", async function () { + await expect(createVoteAndGetExecuteTxPromise()) + .to.emit(template, "UpgradeStarted") + .and.to.emit(template, "UpgradeFinished"); + expect(await template.upgradeBlockNumber()).to.not.equal(0); + expect(await template.isUpgradeFinished()).to.equal(true); }); - it_("should revert when startUpgrade is called after expiration", async function () { - await time.setNextBlockTimestamp(await template.EXPIRE_SINCE_INCLUSIVE()); - await expect(template.connect(agentSigner).startUpgrade()).to.be.revertedWithCustomError(template, "Expired"); - }); + describe("startUpgrade", () => { + it_("should revert when startUpgrade is called by non-agent address", async function () { + await expect(template.connect(deployer).startUpgrade()).to.be.revertedWithCustomError( + template, + "OnlyAgentCanUpgrade", + ); + }); - it_( - "should revert with IncorrectProxyImplementation when startUpgrade is called with incorrect proxy implementation", - async function () { - const locatorProxy = await loadContract("OssifiableProxy", ctx.contracts.locator.address); - const unexpectedLocatorImplementation = ctx.contracts.burner.address; - await locatorProxy.connect(agentSigner).proxy__upgradeTo(unexpectedLocatorImplementation); + it_("should revert when startUpgrade is called after expiration", async function () { + await time.setNextBlockTimestamp(await template.EXPIRE_SINCE_INCLUSIVE()); + await expect(template.connect(agentSigner).startUpgrade()).to.be.revertedWithCustomError(template, "Expired"); + }); - // Attempt to start the upgrade, which should revert with IncorrectProxyImplementation + it_( + "should revert with IncorrectProxyImplementation when startUpgrade is called with incorrect proxy implementation for locator and accountingOracle", + async function () { + const unexpectedImpl = ctx.contracts.kernel.address; + const testCases = [ + { + address: ctx.contracts.locator.address, + }, + { + address: ctx.contracts.accountingOracle.address, + }, + ]; + + for (const { address } of testCases) { + const proxy = await loadContract("OssifiableProxy", address); + await proxy.connect(agentSigner).proxy__upgradeTo(unexpectedImpl); + + // Attempt to start the upgrade, which should revert with IncorrectProxyImplementation + await expect(template.connect(agentSigner).startUpgrade()).to.be.revertedWithCustomError( + template, + "IncorrectProxyImplementation", + ); + } + }, + ); + + it_("should revert when startUpgrade is called after it has already been started", async function () { + await template.connect(agentSigner).startUpgrade(); await expect(template.connect(agentSigner).startUpgrade()).to.be.revertedWithCustomError( template, - "IncorrectProxyImplementation", + "UpgradeAlreadyStarted", ); - }, - ); + }); - it_("should revert when startUpgrade is called after it has already been started", async function () { - // First call should succeed - await template.connect(agentSigner).startUpgrade(); + it_("should revert when startUpgrade is called after upgrade is already finished", async function () { + await createVoteAndGetExecuteTxPromise(); + await expect(template.connect(agentSigner).startUpgrade()).to.be.revertedWithCustomError( + template, + "UpgradeAlreadyFinished", + ); + }); - // Second call should revert with UpgradeAlreadyStarted - await expect(template.connect(agentSigner).startUpgrade()).to.be.revertedWithCustomError( - template, - "UpgradeAlreadyStarted", - ); + it_("should revert when startUpgrade is called twice in the same transaction", async function () { + await hre.ethers.provider.send("hardhat_setCode", [agentSigner.address, await agentMock.getDeployedCode()]); + const harness = (await new V3Template__Harness__factory(deployer).attach( + agentSigner.address, + )) as V3Template__Harness; + + await expect(harness.startUpgradeTwice()).to.be.revertedWithCustomError(template, "StartAlreadyCalledInThisTx"); + }); + }); + + describe("finishUpgrade", () => { + it_("should revert when finishUpgrade is called by non-agent address", async function () { + await template.connect(agentSigner).startUpgrade(); + await expect(template.connect(deployer).finishUpgrade()).to.be.revertedWithCustomError( + template, + "OnlyAgentCanUpgrade", + ); + }); + + it_("should revert when finishUpgrade is called before startUpgrade", async function () { + await expect(template.connect(agentSigner).finishUpgrade()).to.be.revertedWithCustomError( + template, + "StartAndFinishMustBeInSameTx", + ); + }); + + it_("should revert when finishUpgrade is called after upgrade is already finished", async function () { + await createVoteAndGetExecuteTxPromise(); + await expect(template.connect(agentSigner).finishUpgrade()).to.be.revertedWithCustomError( + template, + "UpgradeAlreadyFinished", + ); + }); }); }); diff --git a/test/upgrade/V3Template_Harness.sol b/test/upgrade/V3Template_Harness.sol new file mode 100644 index 000000000..8fcc2a453 --- /dev/null +++ b/test/upgrade/V3Template_Harness.sol @@ -0,0 +1,17 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.8.25; + +import {V3Template} from "contracts/upgrade/V3Template.sol"; + +contract V3Template__Harness { + V3Template public immutable TEMPLATE; + + constructor(address _template) { + TEMPLATE = V3Template(_template); + } + + function startUpgradeTwice() external { + TEMPLATE.startUpgrade(); + TEMPLATE.startUpgrade(); + } +}