solady icon indicating copy to clipboard operation
solady copied to clipboard

✨ Add opt-in storage layout compatibility guard for UUPS

Open codermaybe opened this issue 3 months ago • 3 comments

Description

Adds an optional storage layout compatibility check mechanism for UUPS upgrades to prevent storage collision vulnerabilities.

Changes

  • Added StorageLayoutMismatch() error to UUPSUpgradeable.sol
  • Added _checkStorageLayout(address newImplementation) virtual hook (empty by default for backwards compatibility)
  • Provided reference implementation in MockUUPSImplementation.sol using version hash comparison
  • Added test cases for both compatible and incompatible layout scenarios

Usage Pattern

Implementations can override _checkStorageLayout to enforce storage compatibility:

bytes32 private constant _STORAGE_LAYOUT_VERSION = keccak256("MyContract.v1");

function STORAGE_LAYOUT_VERSION() external pure returns (bytes32) {
    return _STORAGE_LAYOUT_VERSION;
}

function _checkStorageLayout(address newImpl) internal view override {
    // Check version hash matches - see MockUUPSImplementation.sol for full example
}

Design Decisions

  • Opt-in mechanism: Default empty implementation preserves existing behavior
  • Non-breaking: Zero impact on contracts that don't override the hook
  • Gas-optimized: Uses assembly for efficient staticcall and error handling
  • Flexible: Implementations can choose their own validation strategy

Addresses the storage layout safety concern raised in the original UUPS test demo.

Closes #1489

Checklist

Ensure you completed all of the steps below before submitting your pull request:

  • [x] Ran forge fmt?
  • [x] Ran forge test?

codermaybe avatar Nov 23 '25 12:11 codermaybe