feat: demonstrate state-preserving V2 upgrade

This commit is contained in:
golem
2026-08-21 15:31:40 -06:00
parent ca6a99b913
commit 94f2ad6e09
16 changed files with 1082 additions and 36 deletions
+244 -7
View File
@@ -2,11 +2,15 @@
pragma solidity 0.8.35;
import {Test} from "forge-std/Test.sol";
import {Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";
import {DemoScript} from "../script/lib/DemoScript.sol";
import {DeployV1} from "../script/DeployV1.s.sol";
import {SeedV1Demo} from "../script/SeedV1Demo.s.sol";
import {CheckState} from "../script/CheckState.s.sol";
import {UpgradeV2} from "../script/UpgradeV2.s.sol";
import {TransferV2Demo} from "../script/TransferV2Demo.s.sol";
import {BankV1} from "../src/BankV1.sol";
import {BankV2} from "../src/BankV2.sol";
import {MockUSDC} from "../src/MockUSDC.sol";
contract ScriptPreflightHarness is DemoScript {
@@ -35,12 +39,39 @@ contract ScriptPreflightHarness is DemoScript {
}
}
contract UpgradeV2Harness is UpgradeV2 {
function assertSnapshotUnchanged(Snapshot calldata before_, Snapshot calldata after_) external pure {
_assertSnapshotUnchanged(before_, after_);
}
function runWithPaths(string calldata activePath, string calldata pendingPath, address sender)
external
returns (bool upgraded, address implementation)
{
_requireSupportedChain(block.chainid);
return _run(activePath, pendingPath, sender);
}
}
contract TransferV2DemoHarness is TransferV2Demo {
function runWithPath(string calldata manifestPath) external {
_run(manifestPath);
}
}
contract CheckStateHarness is CheckState {
function runWithPath(string calldata manifestPath, string calldata stage) external view {
_run(manifestPath, stage);
}
}
contract ScriptPreflightTest is Test {
ScriptPreflightHarness internal harness;
string internal fixtureDir;
address internal constant OWNER = 0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266;
address internal constant ALICE = 0x70997970C51812dc3A010C7d01b50e0d17dc79C8;
address internal constant BOB = 0x3C44CdDdB6a900fa2b585dd299e03d12FA4293BC;
address internal constant TOKEN = 0x1000000000000000000000000000000000000001;
address internal constant PROXY = 0x2000000000000000000000000000000000000002;
address internal constant IMPLEMENTATION = 0x3000000000000000000000000000000000000003;
@@ -302,7 +333,7 @@ contract ScriptPreflightTest is Test {
}
function testSeedV1DemoExecutesExactActOneStateAndCheckStateAcceptsIt() public {
(string memory path, DemoScript.Manifest memory manifest) = _deployFixture();
(string memory path, DemoScript.Manifest memory manifest) = _deployFixtureNamed("seed-v1");
vm.setEnv("DEPLOYMENT_MANIFEST_PATH", path);
new SeedV1Demo().run();
@@ -320,18 +351,176 @@ contract ScriptPreflightTest is Test {
}
function testCheckStateRejectsManifestImplementationMismatch() public {
(string memory path, DemoScript.Manifest memory manifest) = _deployFixture();
(string memory path, DemoScript.Manifest memory manifest) = _deployUpgradeFixtureNamed("check-mismatch");
vm.writeJson(string.concat('"', vm.toString(manifest.token), '"'), path, ".implementation");
vm.setEnv("DEPLOYMENT_MANIFEST_PATH", path);
vm.setEnv("DEMO_EXPECTED_STAGE", "deployed");
CheckState checker = new CheckState();
CheckStateHarness checker = new CheckStateHarness();
vm.expectRevert(
abi.encodeWithSelector(
DemoScript.UnexpectedAddress.selector, "implementation", manifest.token, manifest.implementation
)
);
checker.run();
checker.runWithPath(path, "deployed");
}
function testUpgradeV2RequiresManifestIdentityCodeOwnerSenderAndVersionOne() public {
(string memory path, DemoScript.Manifest memory manifest) = _deployUpgradeFixtureNamed("upgrade-preflight");
string memory pending = string.concat(fixtureDir, "/upgrade-pending.json");
UpgradeV2Harness upgrader = new UpgradeV2Harness();
vm.expectRevert(
abi.encodeWithSelector(DemoScript.UnexpectedAddress.selector, "SCRIPT_SENDER", manifest.owner, ALICE)
);
upgrader.runWithPaths(path, pending, ALICE);
vm.writeJson(string.concat('"', vm.toString(manifest.token), '"'), path, ".implementation");
vm.expectRevert(
abi.encodeWithSelector(
DemoScript.UnexpectedAddress.selector, "implementation", manifest.token, manifest.implementation
)
);
upgrader.runWithPaths(path, pending, OWNER);
}
function testUpgradeV2RejectsUnsupportedChainMissingProxyCodeWrongOwnerAndUnexpectedVersion() public {
UpgradeV2Harness upgrader = new UpgradeV2Harness();
vm.chainId(1);
vm.expectRevert(abi.encodeWithSelector(DemoScript.UnsupportedChain.selector, uint256(1)));
upgrader.runWithPaths("unused", "unused", OWNER);
vm.chainId(31337);
(string memory missingPath,) = _deployUpgradeFixtureNamed("upgrade-missing-proxy");
vm.writeJson('"0x4000000000000000000000000000000000000004"', missingPath, ".proxy");
vm.expectRevert(
abi.encodeWithSelector(DemoScript.MissingCode.selector, "proxy", 0x4000000000000000000000000000000000000004)
);
upgrader.runWithPaths(missingPath, string.concat(fixtureDir, "/missing-pending.json"), OWNER);
(string memory ownerPath, DemoScript.Manifest memory ownerManifest) =
_deployUpgradeFixtureNamed("upgrade-owner");
vm.prank(OWNER);
BankV1(ownerManifest.proxy).transferOwnership(ALICE);
vm.expectRevert(abi.encodeWithSelector(DemoScript.UnexpectedAddress.selector, "owner", OWNER, ALICE));
upgrader.runWithPaths(ownerPath, string.concat(fixtureDir, "/owner-pending.json"), ALICE);
(string memory versionPath, DemoScript.Manifest memory versionManifest) =
_deployUpgradeFixtureNamed("upgrade-version");
vm.mockCall(
versionManifest.proxy, abi.encodeWithSelector(BankV1.contractVersion.selector), abi.encode(uint256(3))
);
vm.expectRevert(abi.encodeWithSelector(UpgradeV2.UnexpectedVersion.selector, uint256(3)));
upgrader.runWithPaths(versionPath, string.concat(fixtureDir, "/version-pending.json"), OWNER);
}
function testUpgradeV2PreservesCompleteSnapshotAndWritesOnlyStagingRecord() public {
(string memory path, DemoScript.Manifest memory manifest) = _deployUpgradeFixtureNamed("upgrade-preserve");
_seedState(manifest);
string memory beforeManifest = vm.readFile(path);
string memory pending = string.concat(fixtureDir, "/upgrade-pending.json");
(bool upgraded, address implementation) = new UpgradeV2Harness().runWithPaths(path, pending, OWNER);
assertTrue(upgraded);
assertNotEq(implementation, manifest.implementation);
assertEq(vm.readFile(path), beforeManifest);
BankV2 bankV2 = BankV2(manifest.proxy);
assertEq(bankV2.contractVersion(), 2);
assertEq(bankV2.owner(), manifest.owner);
assertEq(address(bankV2.asset()), manifest.token);
assertFalse(bankV2.paused());
assertEq(bankV2.balanceOf(ALICE), 900e6);
assertEq(bankV2.balanceOf(BOB), 500e6);
assertEq(bankV2.totalLiabilities(), 1_400e6);
assertEq(MockUSDC(manifest.token).balanceOf(manifest.proxy), 1_400e6);
string memory marker = vm.readFile(pending);
assertEq(vm.parseJsonString(marker, ".mode"), "upgrade");
assertEq(vm.parseJsonAddress(marker, ".implementation"), implementation);
assertEq(vm.parseJsonAddress(marker, ".snapshot.proxy"), manifest.proxy);
assertEq(vm.parseJsonUint(marker, ".snapshot.balances[1].balance"), 900e6);
assertEq(vm.parseJsonUint(marker, ".snapshot.balances[2].balance"), 500e6);
}
function testUpgradeV2AlreadyActiveWritesVerifiableNoopWithoutChangingImplementation() public {
(string memory path, DemoScript.Manifest memory manifest) = _deployUpgradeFixtureNamed("upgrade-noop");
string memory pending = string.concat(fixtureDir, "/upgrade-pending.json");
UpgradeV2Harness upgrader = new UpgradeV2Harness();
(, address implementation) = upgrader.runWithPaths(path, pending, OWNER);
vm.writeJson(string.concat('"', vm.toString(implementation), '"'), path, ".implementation");
uint256 nonceBefore = vm.getNonce(OWNER);
(bool upgraded, address observedImplementation) = upgrader.runWithPaths(path, pending, OWNER);
assertFalse(upgraded);
assertEq(observedImplementation, implementation);
assertEq(vm.getNonce(OWNER), nonceBefore);
string memory marker = vm.readFile(pending);
assertEq(vm.parseJsonString(marker, ".mode"), "noop");
assertEq(vm.parseJsonUint(marker, ".chainId"), 31337);
assertEq(vm.parseJsonUint(marker, ".observedBlock"), block.number);
assertEq(vm.parseJsonUint(marker, ".ownerNonce"), nonceBefore);
assertEq(vm.parseJsonAddress(marker, ".proxy"), manifest.proxy);
assertEq(vm.parseJsonAddress(marker, ".implementation"), implementation);
}
function testSnapshotComparisonRejectsApplicationMutationButAllowsImplementationAndVersionChange() public {
UpgradeV2Harness upgradeHarness = new UpgradeV2Harness();
UpgradeV2.Snapshot memory before_ = _snapshot();
UpgradeV2.Snapshot memory after_ = _snapshot();
after_.implementation = address(0x9999);
after_.version = 2;
upgradeHarness.assertSnapshotUnchanged(before_, after_);
for (uint256 mutation; mutation < 10; ++mutation) {
before_ = _snapshot();
after_ = _snapshot();
if (mutation == 0) after_.proxy = address(0x9999);
else if (mutation == 1) after_.owner = address(0x9999);
else if (mutation == 2) after_.asset = address(0x9999);
else if (mutation == 3) after_.paused = true;
else if (mutation == 4) after_.balances = new uint256[](2);
else if (mutation == 5) after_.balances[1] += 1;
else if (mutation == 6) after_.liabilities += 1;
else if (mutation == 7) after_.reserves += 1;
else if (mutation == 8) after_.surplus += 1;
else after_.deploymentBlock += 1;
vm.expectRevert();
upgradeHarness.assertSnapshotUnchanged(before_, after_);
}
}
function testTransferV2DemoExecutesExactActThreeAndCheckStateAcceptsBothV2Stages() public {
(string memory path, DemoScript.Manifest memory manifest) = _deployUpgradeFixtureNamed("transfer-v2");
_seedState(manifest);
string memory pending = string.concat(fixtureDir, "/upgrade-pending.json");
(, address implementation) = new UpgradeV2Harness().runWithPaths(path, pending, OWNER);
vm.writeJson(string.concat('"', vm.toString(implementation), '"'), path, ".implementation");
new CheckStateHarness().runWithPath(path, "upgraded");
new TransferV2DemoHarness().runWithPath(path);
new CheckStateHarness().runWithPath(path, "v2");
BankV2 bankV2 = BankV2(manifest.proxy);
assertEq(bankV2.balanceOf(ALICE), 650e6);
assertEq(bankV2.balanceOf(BOB), 750e6);
assertEq(bankV2.totalLiabilities(), 1_400e6);
assertEq(MockUSDC(manifest.token).balanceOf(manifest.proxy), 1_400e6);
}
function _snapshot() internal pure returns (UpgradeV2.Snapshot memory snapshot) {
snapshot.proxy = PROXY;
snapshot.implementation = IMPLEMENTATION;
snapshot.owner = OWNER;
snapshot.asset = TOKEN;
snapshot.paused = false;
snapshot.balances = new uint256[](3);
snapshot.balances[0] = 10;
snapshot.balances[1] = 20;
snapshot.balances[2] = 30;
snapshot.liabilities = 60;
snapshot.reserves = 70;
snapshot.surplus = 10;
snapshot.deploymentBlock = 1;
snapshot.version = 1;
}
function _writeManifest(
@@ -365,7 +554,14 @@ contract ScriptPreflightTest is Test {
}
function _deployFixture() internal returns (string memory path, DemoScript.Manifest memory manifest) {
path = string.concat(fixtureDir, "/deployed.json");
return _deployFixtureNamed("deployed");
}
function _deployFixtureNamed(string memory name)
internal
returns (string memory path, DemoScript.Manifest memory manifest)
{
path = string.concat(fixtureDir, "/", name, ".json");
vm.chainId(31337);
vm.setEnv("SCRIPT_SENDER", vm.toString(OWNER));
vm.setEnv("DEPLOYMENT_MANIFEST_PATH", path);
@@ -374,6 +570,47 @@ contract ScriptPreflightTest is Test {
manifest = harness.readManifest(path, true);
}
function _deployUpgradeFixtureNamed(string memory name)
internal
returns (string memory path, DemoScript.Manifest memory manifest)
{
vm.chainId(31337);
MockUSDC deployedToken = new MockUSDC(OWNER);
address deployedProxy = Upgrades.deployUUPSProxy(
"BankV1.sol:BankV1", abi.encodeCall(BankV1.initialize, (address(deployedToken), OWNER))
);
manifest.schemaVersion = 1;
manifest.network = "anvil";
manifest.chainId = 31337;
manifest.deploymentBlock = 1;
manifest.rpcUrl = "http://127.0.0.1:8545";
manifest.token = address(deployedToken);
manifest.proxy = deployedProxy;
manifest.implementation = Upgrades.getImplementationAddress(deployedProxy);
manifest.owner = OWNER;
manifest.actors = _actors();
path = string.concat(fixtureDir, "/", name, ".json");
vm.writeFile(path, harness.serializeManifest(manifest));
}
function _seedState(DemoScript.Manifest memory manifest) internal {
MockUSDC deployedToken = MockUSDC(manifest.token);
BankV1 deployedBank = BankV1(manifest.proxy);
vm.startPrank(OWNER);
deployedToken.mint(ALICE, 2_000e6);
deployedToken.mint(BOB, 1_000e6);
vm.stopPrank();
vm.startPrank(ALICE);
deployedToken.approve(manifest.proxy, 1_000e6);
deployedBank.deposit(1_000e6);
deployedBank.withdraw(100e6);
vm.stopPrank();
vm.startPrank(BOB);
deployedToken.approve(manifest.proxy, 500e6);
deployedBank.deposit(500e6);
vm.stopPrank();
}
function _writePublicAnvilManifest() internal returns (string memory path) {
path = string.concat(fixtureDir, "/public-anvil-manifest.json");
vm.writeFile(