test storage persistence after registry upgrade

ensdomains · Nov 19, 2024 · 195f151 · 195f151
1 parent 4678b4d
commit 195f151
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 17 deletions.
diff --git a/src/VerifiableFactory.sol b/src/VerifiableFactory.sol
@@ -36,8 +36,6 @@ contract VerifiableFactory {
      * @return proxy The address of the deployed `TransparentVerifiableProxy`.
      */
     function deployProxy(address implementation, uint256 salt) external returns (address) {
-        console.log("deploys");
-        console.logAddress(msg.sender);
         bytes32 outerSalt = keccak256(abi.encode(msg.sender, salt));
 
         TransparentVerifiableProxy proxy = new TransparentVerifiableProxy{salt: outerSalt}(address(this));
@@ -74,26 +72,32 @@ contract VerifiableFactory {
             return false;
         }
         try IProxy(proxy).salt() returns (uint256 salt) {
-            try IProxy(proxy).creator() returns (address creator) {
-                // verify the creator matches this factory
-                if (address(this) != creator) {
-                    return false;
-                }
+            try IProxy(proxy).owner() returns (address owner) {
+                try IProxy(proxy).creator() returns (address creator) {
+                    return _verifyContract(proxy, creator, owner, salt);
+                } catch {}
+            } catch {}
+        } catch {}
+
+        return false;
+    }
 
-                // reconstruct the address using CREATE2 and verify it matches
-                bytes32 outerSalt = keccak256(abi.encode(msg.sender, salt));
+    function _verifyContract(address proxy, address creator, address owner, uint256 salt) private view returns (bool) {
+        // verify the creator matches this factory
+        if (address(this) != creator) {
+            return false;
+        }
 
-                // get creation bytecode with constructor arguments
-                bytes memory bytecode =
-                    abi.encodePacked(type(TransparentVerifiableProxy).creationCode, abi.encode(address(this)));
+        // reconstruct the address using CREATE2 and verify it matches
+        bytes32 outerSalt = keccak256(abi.encode(owner, salt));
 
-                address expectedProxyAddress = Create2.computeAddress(outerSalt, keccak256(bytecode), address(this));
+        // get creation bytecode with constructor arguments
+        bytes memory bytecode =
+            abi.encodePacked(type(TransparentVerifiableProxy).creationCode, abi.encode(address(this)));
 
-                return expectedProxyAddress == proxy;
-            } catch {}
-        } catch {}
+        address expectedProxyAddress = Create2.computeAddress(outerSalt, keccak256(bytecode), address(this));
 
-        return false;
+        return expectedProxyAddress == proxy;
     }
 
     function isContract(address account) internal view returns (bool) {

diff --git a/test/VerifiableFactory.t.sol b/test/VerifiableFactory.t.sol
@@ -142,6 +142,47 @@ contract VerifiableFactoryTest is Test {
         assertEq(proxy.creator(), address(factory), "Wrong creator");
     }
 
+    function test_StoragePersistenceAfterUpgrade() public {
+        uint256 salt = 1;
+        address testAccount = makeAddr("testAccount");
+
+        // deploy proxy
+        vm.prank(owner);
+        address proxyAddress = factory.deployProxy(address(implementation), salt);
+
+        // initialize v1 implementation
+        MockRegistry proxyV1 = MockRegistry(proxyAddress);
+
+        // initialize registry
+        vm.prank(owner);
+        proxyV1.initialize(owner);
+        assertEq(proxyV1.admin(), owner, "Admin should be set");
+
+        // register an address
+        vm.prank(owner);
+        proxyV1.register(testAccount);
+        assertTrue(proxyV1.registeredAddresses(testAccount), "Address should be registered in V1");
+        assertEq(proxyV1.getRegistryVersion(), 1, "Should be V1 implementation");
+
+        // upgrade to v2
+        vm.prank(owner);
+        factory.upgradeImplementation(proxyAddress, address(implementationV2), "");
+
+        // verify state persists after upgrade
+        MockRegistryV2 proxyV2 = MockRegistryV2(proxyAddress);
+
+        // check storage persistence
+        assertTrue(proxyV2.registeredAddresses(testAccount), "Address registration should persist after upgrade");
+        assertEq(proxyV2.admin(), owner, "Admin should persist after upgrade");
+        assertEq(proxyV2.getRegistryVersion(), 2, "Should be V2 implementation");
+
+        // verify v2 functionality still works as it should be
+        address newTestAccount = makeAddr("newTestAccount");
+        vm.prank(owner);
+        proxyV2.register(newTestAccount);
+        assertTrue(proxyV2.registeredAddresses(newTestAccount), "Should be able to register new address in V2");
+    }
+
     // ### Helpers
     function isContract(address account) internal view returns (bool) {
         uint256 size;