test storage persistence after registry upgrade

src/ITransparentVerifiableProxy.sol

@@ -0,0 +1,13 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+interface ITransparentVerifiableProxy {
+    function salt() external view returns (uint256);
+
+    function owner() external view returns (address);
+
+    function creator() external view returns (address);
+
+    /// @dev See {UUPSUpgradeable-upgradeToAndCall}
+    function upgradeToAndCall(address newImplementation, bytes calldata data) external payable;
+}

src/TransparentVerifiableProxy.sol

@@ -10,11 +10,7 @@ import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.s
 import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
 import {StorageSlot} from "@openzeppelin/contracts/utils/StorageSlot.sol";
 import {SlotDerivation} from "@openzeppelin/contracts/utils/SlotDerivation.sol";
-
-interface ITransparentVerifiableProxy {
-    /// @dev See {UUPSUpgradeable-upgradeToAndCall}
-    function upgradeToAndCall(address newImplementation, bytes calldata data) external payable;
-}
+import {ITransparentVerifiableProxy} from "./ITransparentVerifiableProxy.sol";
 
 contract TransparentVerifiableProxy is Proxy, Initializable {
     using StorageSlot for bytes32;
@@ -31,9 +27,9 @@ contract TransparentVerifiableProxy is Proxy, Initializable {
     // ### EVENTS
     error ProxyDeniedOwnerAccess();
 
-    // // Modifier that allows only the owner to call certain functions
-    // modifier onlyOwner() {
-    //     require(msg.sender == owner, "Caller is not the owner");
+    // modifier that allows only the creator address to call certain functions
+    // modifier onlyCreator() {
+    //     require(msg.sender == creator, "Caller is not the creator address");
     //     _;
     // }
 

src/VerifiableFactory.sol

@@ -3,15 +3,8 @@ pragma solidity ^0.8.20;
 
 import {console} from "forge-std/console.sol";
 import {Create2} from "@openzeppelin/contracts/utils/Create2.sol";
-import {ITransparentVerifiableProxy, TransparentVerifiableProxy} from "./TransparentVerifiableProxy.sol";
-
-interface IProxy {
-    function salt() external view returns (uint256);
-
-    function owner() external view returns (address);
-
-    function creator() external view returns (address);
-}
+import {TransparentVerifiableProxy} from "./TransparentVerifiableProxy.sol";
+import {ITransparentVerifiableProxy} from "./ITransparentVerifiableProxy.sol";
 
 contract VerifiableFactory {
     event ProxyDeployed(address indexed sender, address indexed proxyAddress, uint256 salt, address implementation);
@@ -36,8 +29,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));
@@ -52,7 +43,7 @@ contract VerifiableFactory {
 
     // Function to upgrade the proxy's implementation (only owner of proxy can call this)
     function upgradeImplementation(address proxyAddress, address newImplementation, bytes memory data) external {
-        address owner = IProxy(proxyAddress).owner();
+        address owner = ITransparentVerifiableProxy(proxyAddress).owner();
         require(owner == msg.sender, "Only the owner can upgrade");
 
         // Upgrade the proxy to point to the new implementation
@@ -73,27 +64,26 @@ contract VerifiableFactory {
         if (!isContract(proxy)) {
             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 ITransparentVerifiableProxy(proxy).salt() returns (uint256 salt) {
+            try ITransparentVerifiableProxy(proxy).owner() returns (address owner) {
+                return _verifyContract(proxy, owner, salt);
+            } catch {}
+        } catch {}
 
-                // reconstruct the address using CREATE2 and verify it matches
-                bytes32 outerSalt = keccak256(abi.encode(msg.sender, salt));
+        return false;
+    }
 
-                // get creation bytecode with constructor arguments
-                bytes memory bytecode =
-                    abi.encodePacked(type(TransparentVerifiableProxy).creationCode, abi.encode(address(this)));
+    function _verifyContract(address proxy, address owner, uint256 salt) private view returns (bool) {
+        // 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) {

test/TransparentVerifiableProxy.t.sol

@@ -0,0 +1,72 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.20;
+
+import "forge-std/Test.sol";
+import "../src/TransparentVerifiableProxy.sol";
+import {SlotDerivation} from "@openzeppelin/contracts/utils/SlotDerivation.sol";
+import {MockRegistry} from "../src/mock/MockRegistry.sol";
+
+contract TransparentVerifiableProxyTest is Test {
+    using SlotDerivation for bytes32;
+
+    TransparentVerifiableProxy proxy;
+
+    address creator = address(0x1);
+    address owner = address(0x2);
+    address implementation = address(new MockRegistry());
+    uint256 salt = 12345;
+    bytes emptyData;
+
+    string internal constant _VERIFICATION_SLOT = "proxy.verifiable";
+    string internal constant _SALT = "salt";
+    string internal constant _OWNER = "owner";
+
+    function setUp() public {
+        proxy = new TransparentVerifiableProxy(creator);
+    }
+
+    function testInitialize() public {
+        // initialize the proxy
+        proxy.initialize(salt, owner, implementation, emptyData);
+
+        // check salt and owner values
+        assertEq(proxy.salt(), salt, "Salt mismatch");
+        assertEq(proxy.owner(), owner, "Owner mismatch");
+    }
+
+    function testSaltStorage() public {
+        // initialize the proxy
+        proxy.initialize(salt, owner, implementation, emptyData);
+
+        // compute the base slot
+        bytes32 baseSlot = SlotDerivation.erc7201Slot(_VERIFICATION_SLOT);
+
+        // use SlotDerivation to compute the salt slot
+        bytes32 saltSlot = baseSlot.deriveMapping(_SALT);
+
+        // directly manipulate the storage for the salt
+        uint256 newSalt = 54321;
+        vm.store(address(proxy), saltSlot, bytes32(newSalt));
+
+        // verify the updated salt
+        assertEq(proxy.salt(), newSalt, "Salt update failed");
+    }
+
+    function testOwnerStorage() public {
+        // initialize the proxy
+        proxy.initialize(salt, owner, implementation, emptyData);
+
+        // compute the base slot
+        bytes32 baseSlot = SlotDerivation.erc7201Slot(_VERIFICATION_SLOT);
+
+        // use SlotDerivation to compute the owner slot
+        bytes32 ownerSlot = baseSlot.deriveMapping(_OWNER);
+
+        // directly manipulate the storage for the owner
+        address newOwner = address(0x4);
+        vm.store(address(proxy), ownerSlot, bytes32(uint256(uint160(newOwner))));
+
+        // verify the updated owner
+        assertEq(proxy.owner(), newOwner, "Owner update failed");
+    }
+}

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;