coderabbit

zeta-chain · Nov 8, 2024 · 8eaa303 · 8eaa303
1 parent a94c307
commit 8eaa303
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 3 deletions.
diff --git a/v2/scripts/upgrade/SimulateERC20CustodyUpgrade.s.sol b/v2/scripts/upgrade/SimulateERC20CustodyUpgrade.s.sol
@@ -15,12 +15,22 @@ contract UpgradeSimulation is Script {
         bytes32 implSalt = keccak256("ERC20Custody");
         address erc20CustodyImpl = address(new ERC20Custody{salt: implSalt}());
 
-        // Simulate the upgrade
         ERC20Custody proxy = ERC20Custody(proxyAddress);
+
+        // Get the current state
+        address gateway = proxy.gateway();
+        address tssAddress = proxy.tssAddress();
+        bool supportsLegacy = proxy.supportsLegacy();
+
+        // Simulate the upgrade
         vm.prank(adminAddress);
         proxy.upgradeToAndCall(erc20CustodyImpl, "");
 
         // After upgrade, verify that the state is intact
+        require(gateway == proxy.gateway(), "gateway address mismatch");
+        require(tssAddress == proxy.tssAddress(), "tss address mismatch");
+        require(supportsLegacy == proxy.supportsLegacy(), "supportsLegacy mismatch");
+
         console.log("Upgraded contract state:");
         console.log("gateway address:", address(proxy.gateway()));
         console.log("tss address:", proxy.tssAddress());

diff --git a/v2/scripts/upgrade/SimulateGatewayEVMUpgrade.s.sol b/v2/scripts/upgrade/SimulateGatewayEVMUpgrade.s.sol
@@ -15,12 +15,24 @@ contract UpgradeSimulation is Script {
         bytes32 implSalt = keccak256("GatewayEVM");
         address gatewayImpl = address(new GatewayEVM{salt: implSalt}());
 
-        // Simulate the upgrade
         GatewayEVM proxy = GatewayEVM(proxyAddress);
+
+        // Get the current state
+        address custody = proxy.custody();
+        address tssAddress = proxy.tssAddress();
+        address zetaConnector = proxy.zetaConnector();
+        address zetaToken = proxy.zetaToken();
+
+        // Simulate the upgrade
         vm.prank(adminAddress);
         proxy.upgradeToAndCall(gatewayImpl, "");
 
         // After upgrade, verify that the state is intact
+        require(custody == proxy.custody(), "custody address mismatch");
+        require(tssAddress == proxy.tssAddress(), "tss address mismatch");
+        require(zetaConnector == proxy.zetaConnector(), "zetaConnector address mismatch");
+        require(zetaToken == proxy.zetaToken(), "zetaToken address mismatch");
+
         console.log("Upgraded contract state:");
         console.log("custody address:", proxy.custody());
         console.log("tss address:", proxy.tssAddress());

diff --git a/v2/scripts/upgrade/SimulateGatewayZEVMUpgrade.s.sol b/v2/scripts/upgrade/SimulateGatewayZEVMUpgrade.s.sol
@@ -15,12 +15,18 @@ contract UpgradeSimulation is Script {
         bytes32 implSalt = keccak256("GatewayZEVM");
         address gatewayImpl = address(new GatewayZEVM{salt: implSalt}());
 
-        // Simulate the upgrade
         GatewayZEVM proxy = GatewayZEVM(proxyAddress);
+
+        // Get the current state
+        address zetaToken = proxy.zetaToken();
+
+        // Simulate the upgrade
         vm.prank(adminAddress);
         proxy.upgradeToAndCall(gatewayImpl, "");
 
         // After upgrade, verify that the state is intact
+        require(zetaToken == proxy.zetaToken(), "zetaToken address mismatch");
+
         console.log("Upgraded contract state:");
         console.log("zetaToken address:", proxy.zetaToken());
     }