lidofinance · TheDZhon · Jun 17, 2025 · Jun 16, 2025 · Jun 16, 2025 · Jun 17, 2025
@@ -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;
 
     /**

@@ -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 public constant UPGRADE_STARTED_SLOT =
+        0x058d69f67a3d86c424c516d23a070ff8bed34431617274caa2049bd702675e3f;
+
 
     /// @param _params Params required to initialize the addresses contract
     constructor(V3AddressesParams memory _params) V3Addresses(_params) {
@@ -124,8 +133,11 @@ contract V3Template is V3Addresses {
     function startUpgrade() external {
         if (msg.sender != AGENT) revert OnlyAgentCanUpgrade();
         if (block.timestamp >= EXPIRE_SINCE_INCLUSIVE) revert Expired();
+        if (isUpgradeFinished) revert UpgradeAlreadyFinished();
+        if (_isStartCalledInThisTx()) revert StartAlreadyCalledInThisTx();
         if (upgradeBlockNumber != UPGRADE_NOT_STARTED) revert UpgradeAlreadyStarted();
 
+        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);

@@ -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<void> {
+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<void> {
   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 };
 }
@@ -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,49 +58,109 @@ describe("Integration: Upgrade Template V3 tests", () => {
   afterEach(async () => await Snapshot.restore(snapshot));
 
   function it_(title: string, fn: () => Promise<void>) {
-    it(title, async function () {
+    return it(title, async function () {
       if (needToSkipTemplateTests()) {
         this.skip();
       }
       await fn();
     });
   }
 
-  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>("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>("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",
+      );
+    });
   });
 });
@@ -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();
+    }
+}