mg/uupl-smart-contract
2
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: demonstrate state-preserving V2 upgrade

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files
test/BankInvariant.t.sol
+23 -4
View File
@@ -2,26 +2,45 @@
pragma solidity 0.8.35;
import {BankTestBase} from "./helpers/BankTestBase.sol";
import {BankHandler} from "./helpers/BankHandler.sol";
import {BankV2} from "../src/BankV2.sol";
import {Options, Upgrades} from "openzeppelin-foundry-upgrades/Upgrades.sol";
import {BankV2Handler} from "./helpers/BankV2Handler.sol";
contract BankInvariantTest is BankTestBase {
BankHandler internal handler;
BankV2Handler internal handler;
function setUp() public override {
super.setUp();
Options memory opts;
opts.referenceContract = "BankV1.sol:BankV1";
Upgrades.upgradeProxy(proxy, "BankV2.sol:BankV2", "", opts, owner);
BankV2 bankV2 = BankV2(proxy);
handler = new BankHandler(token, bank);
handler = new BankV2Handler(token, bankV2);
vm.prank(owner);
token.transferOwnership(address(handler));
targetContract(address(handler));
bytes4[] memory selectors = new bytes4[](3);
bytes4[] memory selectors = new bytes4[](4);
selectors[0] = handler.deposit.selector;
selectors[1] = handler.withdraw.selector;
selectors[2] = handler.donate.selector;
selectors[3] = handler.transfer.selector;
targetSelector(FuzzSelector({addr: address(handler), selectors: selectors}));
}
function testHandlerExecutesDeterministicTrackedTransfer() public {
handler.deposit(0, 40e6);
handler.transfer(0, 0, 15e6);
assertEq(bank.balanceOf(handler.actorAt(0)), 25e6);
assertEq(bank.balanceOf(handler.actorAt(1)), 15e6);
assertEq(handler.ghostTransferred(), 15e6);
assertEq(bank.totalLiabilities(), 40e6);
assertEq(token.balanceOf(address(bank)), 40e6);
}
function invariant_liabilitiesEqualTrackedBalances() public view {
uint256 sum;
for (uint256 i; i < handler.actorCount(); ++i) {
test/ScriptPreflight.t.sol
+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(
test/helpers/BankV2Handler.sol
+32
View File
@@ -0,0 +1,32 @@
// SPDX-License-Identifier: MIT
pragma solidity 0.8.35;
import {BankV2} from "../../src/BankV2.sol";
import {MockUSDC} from "../../src/MockUSDC.sol";
import {BankHandler} from "./BankHandler.sol";
contract BankV2Handler is BankHandler {
BankV2 internal immutable bankV2;
uint256 public ghostTransferred;
constructor(MockUSDC token_, BankV2 bank_) BankHandler(token_, bank_) {
bankV2 = bank_;
}
function transfer(uint256 fromSeed, uint256 toSeed, uint256 amount) external {
uint256 count = actorCount();
uint256 fromIndex = fromSeed % count;
address from = actorAt(fromIndex);
uint256 balance = bankV2.balanceOf(from);
if (balance == 0) return;
uint256 toIndex = (fromIndex + 1 + (toSeed % (count - 1))) % count;
address to = actorAt(toIndex);
amount = bound(amount, 1, balance);
vm.prank(from);
bankV2.transferBalance(to, amount);
ghostTransferred += amount;
}
}