test storage persistence after registry upgrade

README.md

@@ -12,9 +12,9 @@ A system for deploying and verifying proxy contracts with predictable storage la
 
 ### 2. TransparentVerifiableProxy
 - Transparent proxy pattern with verified storage layout
-- Fixed storage slots:
-  - Slot 0: `salt` (uint256)
-  - Slot 1: `owner` (address)
+- Fixed storage slots via Diamond Storage pattern:
+  - Slot 'proxy.verifiable.salt': `salt` (uint256)
+  - Slot 'proxy.verifiable.owner': `owner` (address)
 - Immutable `creator` field (set in bytecode)
 - Implements secure upgrade mechanism
 - Initializable to prevent implementation tampering

src/TransparentVerifiableProxy.sol

@@ -9,6 +9,21 @@ import {Proxy} from "@openzeppelin/contracts/proxy/Proxy.sol";
 import {ERC1967Utils} from "@openzeppelin/contracts/proxy/ERC1967/ERC1967Utils.sol";
 import {Initializable} from "@openzeppelin/contracts/proxy/utils/Initializable.sol";
 
+// EIP-2535 Diamond Storage pattern
+// ref: https://eips.ethereum.org/EIPS/eip-2535#storage
+library StorageSlot {
+    bytes32 constant SLOT_SALT = keccak256("proxy.verifiable.salt");
+    bytes32 constant SLOT_OWNER = keccak256("proxy.verifiable.owner");
+
+    function getSaltSlot() internal pure returns (bytes32) {
+        return SLOT_SALT;
+    }
+
+    function getOwnerSlot() internal pure returns (bytes32) {
+        return SLOT_OWNER;
+    }
+}
+
 interface ITransparentVerifiableProxy {
     /// @dev See {UUPSUpgradeable-upgradeToAndCall}
     function upgradeToAndCall(address newImplementation, bytes calldata data) external payable;
@@ -18,10 +33,6 @@ contract TransparentVerifiableProxy is Proxy, Initializable {
     // immutable variable (in bytecode)
     address public immutable creator;
 
-    // storage variables (in storage slots)
-    uint256 public salt; // Salt, being used creating the proxy (slot 0)
-    address public owner; // The owner of the proxy contract (slot 1)
-
     // ### EVENTS
     error ProxyDeniedOwnerAccess();
 
@@ -52,12 +63,34 @@ contract TransparentVerifiableProxy is Proxy, Initializable {
     {
         require(implementation != address(0), "New implementation cannot be the zero address");
 
-        salt = _salt;
-        owner = _owner;
+        bytes32 saltSlot = StorageSlot.getSaltSlot();
+        bytes32 ownerSlot = StorageSlot.getOwnerSlot();
 
+        assembly {
+            sstore(saltSlot, _salt)
+            sstore(ownerSlot, _owner)
+        }
         ERC1967Utils.upgradeToAndCall(implementation, data);
     }
 
+    function salt() public view returns (uint256) {
+        bytes32 slot = StorageSlot.getSaltSlot();
+        uint256 value;
+        assembly {
+            value := sload(slot)
+        }
+        return value;
+    }
+
+    function owner() public view returns (address) {
+        bytes32 slot = StorageSlot.getOwnerSlot();
+        address value;
+        assembly {
+            value := sload(slot)
+        }
+        return value;
+    }
+
     /**
      * @dev Returns the current implementation address.
      *

src/VerifiableFactory.sol

@@ -9,8 +9,6 @@ interface IProxy {
     function salt() external view returns (uint256);
 
     function owner() external view returns (address);
-
-    function creator() external view returns (address);
 }
 
 contract VerifiableFactory {
@@ -36,8 +34,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 +70,25 @@ 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) {
+                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/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;