Storage Collision in Proxies

1// Storage Layout Example
2contract Example {
3    uint256 a;      // Slot 0
4    uint256 b;      // Slot 1
5    address owner;  // Slot 2 (20 bytes, but occupies full slot)
6    bool active;    // Slot 3
7    
8    // Mappings and dynamic arrays use different slot calculation
9    mapping(address => uint256) balances;  // Slot 4 (base slot)
10    // Actual data at: keccak256(key . slot)
11    
12    uint256[] data; // Slot 5 (stores length)
13    // Actual data at: keccak256(slot) + index
14}
15
16// Storage visualization:
17// ┌────────┬─────────────────────────────────────┐
18// │ Slot 0 │ a (uint256)                         │
19// ├────────┼─────────────────────────────────────┤
20// │ Slot 1 │ b (uint256)                         │
21// ├────────┼─────────────────────────────────────┤
22// │ Slot 2 │ owner (address - 20 bytes)          │
23// ├────────┼─────────────────────────────────────┤
24// │ Slot 3 │ active (bool - 1 byte)              │
25// ├────────┼─────────────────────────────────────┤
26// │ Slot 4 │ balances mapping (base)             │
27// ├────────┼─────────────────────────────────────┤
28// │ Slot 5 │ data.length                         │
29// └────────┴─────────────────────────────────────┘

1// 🚨 VULNERABLE: Classic storage collision
2contract VulnerableProxy {
3    address public implementation;  // Slot 0
4    address public admin;           // Slot 1
5    
6    fallback() external payable {
7        (bool success,) = implementation.delegatecall(msg.data);
8        require(success);
9    }
10}
11
12contract Implementation {
13    address public owner;      // Slot 0 - COLLIDES with implementation!
14    uint256 public totalSupply; // Slot 1 - COLLIDES with admin!
15    
16    function initialize(address _owner) external {
17        owner = _owner;  // 💥 This OVERWRITES the proxy's implementation address!
18        // Proxy is now bricked - delegatecall will call random address
19    }
20}

1// Implementation V1
2contract TokenV1 {
3    address public owner;        // Slot 0
4    uint256 public totalSupply;  // Slot 1
5    mapping(address => uint256) public balances;  // Slot 2
6}
7
8// 🚨 VULNERABLE: Implementation V2 - WRONG WAY
9contract TokenV2 {
10    bool public paused;          // Slot 0 - NEW! Shifts everything!
11    address public owner;        // Slot 1 - Was at slot 0!
12    uint256 public totalSupply;  // Slot 2 - Was at slot 1!
13    mapping(address => uint256) public balances;  // Slot 3 - Was at slot 2!
14    
15    // After upgrade:
16    // - "paused" will be true/false based on low bytes of old owner address
17    // - "owner" will be garbage (the old totalSupply value)
18    // - "totalSupply" will be the base slot of old balances mapping
19    // - All balances are effectively lost (wrong mapping base)
20}
21
22// ✅ CORRECT: Implementation V2
23contract TokenV2Safe {
24    address public owner;        // Slot 0 - SAME
25    uint256 public totalSupply;  // Slot 1 - SAME
26    mapping(address => uint256) public balances;  // Slot 2 - SAME
27    bool public paused;          // Slot 3 - NEW! Added at the end
28}

1// Base contracts
2contract Ownable {
3    address public owner;  // Will be at some slot
4}
5
6contract Pausable {
7    bool public paused;    // Will be at some slot
8}
9
10// 🚨 V1: Ownable first
11contract TokenV1 is Ownable, Pausable {
12    uint256 public totalSupply;
13    // Storage layout:
14    // Slot 0: owner (from Ownable)
15    // Slot 1: paused (from Pausable)  
16    // Slot 2: totalSupply
17}
18
19// 🚨 V2: Pausable first - BREAKS EVERYTHING!
20contract TokenV2 is Pausable, Ownable {
21    uint256 public totalSupply;
22    // Storage layout:
23    // Slot 0: paused (from Pausable) - WAS owner!
24    // Slot 1: owner (from Ownable) - WAS paused!
25    // Slot 2: totalSupply
26    
27    // Result: owner and paused are SWAPPED!
28    // If old owner was 0x1234...5678 (nonzero), paused is now "true"
29    // owner is now address(0) or address(1) depending on old paused value
30}

1// 🚨 VULNERABLE: No storage gap
2contract BaseV1 {
3    address public owner;
4    uint256 public value;
5    // No gap! If we add variables in V2, child contracts break
6}
7
8contract ChildV1 is BaseV1 {
9    uint256 public childValue;  // Slot 2
10}
11
12// After upgrading BaseV1 to add a variable:
13contract BaseV2 {
14    address public owner;
15    uint256 public value;
16    bool public paused;  // NEW - slot 2
17    // 💥 childValue was at slot 2, now paused is there!
18}
19
20// ✅ CORRECT: Use storage gaps
21contract BaseV1Safe {
22    address public owner;
23    uint256 public value;
24    
25    // Reserve 50 slots for future upgrades
26    uint256[50] private __gap;
27}
28
29contract BaseV2Safe {
30    address public owner;
31    uint256 public value;
32    bool public paused;  // Uses one gap slot
33    
34    uint256[49] private __gap;  // Reduce gap by 1
35}

1// 🚨 VULNERABLE: Implementation stored at slot 0
2contract BadProxy {
3    address implementation;  // Slot 0 - will collide!
4    
5    fallback() external payable {
6        address impl = implementation;
7        assembly {
8            calldatacopy(0, 0, calldatasize())
9            let result := delegatecall(gas(), impl, 0, calldatasize(), 0, 0)
10            returndatacopy(0, 0, returndatasize())
11            switch result
12            case 0 { revert(0, returndatasize()) }
13            default { return(0, returndatasize()) }
14        }
15    }
16}
17
18// 💥 Any implementation with slot 0 variables will break:
19contract VictimImplementation {
20    address public owner;  // Slot 0 - COLLIDES with implementation!
21    
22    function setOwner(address _owner) external {
23        owner = _owner;  // Overwrites proxy's implementation pointer!
24    }
25}

1// V1: Uses packed storage
2contract TokenV1 {
3    address public owner;     // Slot 0 (20 bytes)
4    uint96 public totalSupply; // Slot 0 (12 bytes) - packed with owner
5    uint256 public maxSupply;  // Slot 1
6}
7
8// 🚨 V2: Changed type size - BREAKS LAYOUT!
9contract TokenV2 {
10    address public owner;      // Slot 0 (20 bytes)
11    uint256 public totalSupply; // Slot 1 - WAS packed in slot 0!
12    uint256 public maxSupply;   // Slot 2 - WAS at slot 1!
13    
14    // totalSupply now reads garbage (old slot 1 value)
15    // maxSupply is lost
16}

1// ✅ SECURE: EIP-1967 Compliant Proxy
2contract EIP1967Proxy {
3    // Standard slot for implementation address
4    bytes32 private constant IMPLEMENTATION_SLOT = 
5        bytes32(uint256(keccak256("eip1967.proxy.implementation")) - 1);
6    
7    // Standard slot for admin address  
8    bytes32 private constant ADMIN_SLOT = 
9        bytes32(uint256(keccak256("eip1967.proxy.admin")) - 1);
10    
11    constructor(address implementation, address admin) {
12        _setImplementation(implementation);
13        _setAdmin(admin);
14    }
15    
16    function _setImplementation(address newImplementation) private {
17        require(newImplementation.code.length > 0, "Not a contract");
18        
19        assembly {
20            sstore(IMPLEMENTATION_SLOT, newImplementation)
21        }
22    }
23    
24    function _getImplementation() internal view returns (address impl) {
25        assembly {
26            impl := sload(IMPLEMENTATION_SLOT)
27        }
28    }
29    
30    function _setAdmin(address newAdmin) private {
31        assembly {
32            sstore(ADMIN_SLOT, newAdmin)
33        }
34    }
35    
36    fallback() external payable {
37        address impl = _getImplementation();
38        assembly {
39            calldatacopy(0, 0, calldatasize())
40            let result := delegatecall(gas(), impl, 0, calldatasize(), 0, 0)
41            returndatacopy(0, 0, returndatasize())
42            switch result
43            case 0 { revert(0, returndatasize()) }
44            default { return(0, returndatasize()) }
45        }
46    }
47}

1// SPDX-License-Identifier: MIT
2pragma solidity ^0.8.19;
3
4import "@openzeppelin/contracts-upgradeable/proxy/utils/Initializable.sol";
5import "@openzeppelin/contracts-upgradeable/proxy/utils/UUPSUpgradeable.sol";
6import "@openzeppelin/contracts-upgradeable/access/OwnableUpgradeable.sol";
7
8contract TokenV1 is Initializable, UUPSUpgradeable, OwnableUpgradeable {
9    // ✅ State variables for V1
10    uint256 public totalSupply;
11    mapping(address => uint256) public balances;
12    
13    // ✅ Storage gap for future upgrades
14    uint256[48] private __gap;
15    
16    /// @custom:oz-upgrades-unsafe-allow constructor
17    constructor() {
18        _disableInitializers();  // ✅ Prevent implementation initialization
19    }
20    
21    function initialize(address owner_) external initializer {
22        __Ownable_init(owner_);
23        __UUPSUpgradeable_init();
24    }
25    
26    function _authorizeUpgrade(address) internal override onlyOwner {}
27    
28    function mint(address to, uint256 amount) external onlyOwner {
29        totalSupply += amount;
30        balances[to] += amount;
31    }
32}
33
34// ✅ V2: Correctly adds new variable at the end
35contract TokenV2 is Initializable, UUPSUpgradeable, OwnableUpgradeable {
36    // ✅ SAME as V1 - DO NOT CHANGE ORDER
37    uint256 public totalSupply;
38    mapping(address => uint256) public balances;
39    
40    // ✅ NEW variable added at the end (uses 1 gap slot)
41    bool public paused;
42    
43    // ✅ Reduced gap by 1
44    uint256[47] private __gap;
45    
46    /// @custom:oz-upgrades-unsafe-allow constructor
47    constructor() {
48        _disableInitializers();
49    }
50    
51    function initialize(address owner_) external initializer {
52        __Ownable_init(owner_);
53        __UUPSUpgradeable_init();
54    }
55    
56    function _authorizeUpgrade(address) internal override onlyOwner {}
57    
58    modifier whenNotPaused() {
59        require(!paused, "Paused");
60        _;
61    }
62    
63    function mint(address to, uint256 amount) external onlyOwner whenNotPaused {
64        totalSupply += amount;
65        balances[to] += amount;
66    }
67    
68    function setPaused(bool _paused) external onlyOwner {
69        paused = _paused;
70    }
71}

1// SPDX-License-Identifier: MIT
2pragma solidity ^0.8.19;
3
4import "forge-std/Test.sol";
5import "../src/TokenV1.sol";
6import "../src/TokenV2.sol";
7import "@openzeppelin/contracts/proxy/ERC1967/ERC1967Proxy.sol";
8
9contract StorageCollisionTest is Test {
10    TokenV1 implementationV1;
11    TokenV2 implementationV2;
12    ERC1967Proxy proxy;
13    TokenV1 tokenV1;
14    TokenV2 tokenV2;
15    
16    address owner = makeAddr("owner");
17    address user = makeAddr("user");
18    
19    function setUp() public {
20        implementationV1 = new TokenV1();
21        
22        bytes memory initData = abi.encodeCall(TokenV1.initialize, (owner));
23        proxy = new ERC1967Proxy(address(implementationV1), initData);
24        
25        tokenV1 = TokenV1(address(proxy));
26    }
27    
28    // ✅ Test: V1 state is preserved after V2 upgrade
29    function test_UpgradePreservesState() public {
30        // Setup V1 state
31        vm.prank(owner);
32        tokenV1.mint(user, 1000);
33        
34        assertEq(tokenV1.totalSupply(), 1000);
35        assertEq(tokenV1.balances(user), 1000);
36        
37        // Upgrade to V2
38        implementationV2 = new TokenV2();
39        vm.prank(owner);
40        tokenV1.upgradeToAndCall(address(implementationV2), "");
41        
42        tokenV2 = TokenV2(address(proxy));
43        
44        // ✅ Verify state is preserved
45        assertEq(tokenV2.totalSupply(), 1000, "totalSupply corrupted");
46        assertEq(tokenV2.balances(user), 1000, "balances corrupted");
47        
48        // ✅ New functionality works
49        assertEq(tokenV2.paused(), false, "paused should be false");
50    }
51    
52    // ✅ Test: Storage slots match between versions
53    function test_StorageLayoutCompatibility() public {
54        // Read V1 storage slots
55        bytes32 slot0V1 = vm.load(address(proxy), bytes32(uint256(0)));
56        bytes32 slot1V1 = vm.load(address(proxy), bytes32(uint256(1)));
57        
58        // Upgrade
59        implementationV2 = new TokenV2();
60        vm.prank(owner);
61        tokenV1.upgradeToAndCall(address(implementationV2), "");
62        
63        // Read same slots after upgrade
64        bytes32 slot0V2 = vm.load(address(proxy), bytes32(uint256(0)));
65        bytes32 slot1V2 = vm.load(address(proxy), bytes32(uint256(1)));
66        
67        // ✅ Slots should be unchanged
68        assertEq(slot0V1, slot0V2, "Slot 0 changed after upgrade");
69        assertEq(slot1V1, slot1V2, "Slot 1 changed after upgrade");
70    }
71    
72    // ✅ Test: EIP-1967 slots are used correctly
73    function test_EIP1967Compliance() public {
74        bytes32 implSlot = bytes32(uint256(
75            keccak256("eip1967.proxy.implementation")
76        ) - 1);
77        
78        bytes32 storedImpl = vm.load(address(proxy), implSlot);
79        
80        assertEq(
81            address(uint160(uint256(storedImpl))),
82            address(implementationV1),
83            "Implementation not at EIP-1967 slot"
84        );
85    }
86    
87    // ✅ Fuzz: Storage reads are consistent
88    function testFuzz_StorageConsistency(uint256 mintAmount) public {
89        vm.assume(mintAmount > 0 && mintAmount < type(uint128).max);
90        
91        vm.prank(owner);
92        tokenV1.mint(user, mintAmount);
93        
94        // Upgrade
95        implementationV2 = new TokenV2();
96        vm.prank(owner);
97        tokenV1.upgradeToAndCall(address(implementationV2), "");
98        tokenV2 = TokenV2(address(proxy));
99        
100        // Values should match
101        assertEq(tokenV2.totalSupply(), mintAmount);
102        assertEq(tokenV2.balances(user), mintAmount);
103    }
104}