From 10d9531cb20d4396149cc51255343b62a6c7697d Mon Sep 17 00:00:00 2001 From: codermaybe <1473434657@qq.com> Date: Sun, 23 Nov 2025 19:06:13 +0800 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20opt-in=20storage=20layout=20c?= =?UTF-8?q?ompatibility=20guard=20for=20UUPS?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/utils/UUPSUpgradeable.sol | 7 +++++ test/UUPSUpgradeable.t.sol | 15 +++++++++ test/utils/mocks/MockUUPSImplementation.sol | 31 +++++++++++++++++++ .../mocks/MockUUPSWithDifferentLayout.sol | 27 ++++++++++++++++ 4 files changed, 80 insertions(+) create mode 100644 test/utils/mocks/MockUUPSWithDifferentLayout.sol diff --git a/src/utils/UUPSUpgradeable.sol b/src/utils/UUPSUpgradeable.sol index 946e02a83..647cbcc3f 100644 --- a/src/utils/UUPSUpgradeable.sol +++ b/src/utils/UUPSUpgradeable.sol @@ -20,6 +20,8 @@ abstract contract UUPSUpgradeable is CallContextChecker { /// @dev The upgrade failed. error UpgradeFailed(); + /// @dev The storagelayout mismatch. + error StorageLayoutMismatch(); /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/ /* EVENTS */ @@ -52,6 +54,10 @@ abstract contract UUPSUpgradeable is CallContextChecker { /// ``` function _authorizeUpgrade(address newImplementation) internal virtual; + /// @dev Optional hook to guard against storage layout incompatibility. + /// Override in implementation if desired (e.g. compare a constant layout hash). + function _checkStorageLayout(address newImplementation) internal virtual {} + /// @dev Returns the storage slot used by the implementation, /// as specified in [ERC1822](https://eips.ethereum.org/EIPS/eip-1822). /// @@ -73,6 +79,7 @@ abstract contract UUPSUpgradeable is CallContextChecker { onlyProxy { _authorizeUpgrade(newImplementation); + _checkStorageLayout(newImplementation); /// @solidity memory-safe-assembly assembly { newImplementation := shr(96, shl(96, newImplementation)) // Clears upper 96 bits. diff --git a/test/UUPSUpgradeable.t.sol b/test/UUPSUpgradeable.t.sol index 9b8c44e4e..7f95cf4b6 100644 --- a/test/UUPSUpgradeable.t.sol +++ b/test/UUPSUpgradeable.t.sol @@ -5,6 +5,7 @@ import "./utils/SoladyTest.sol"; import {CallContextChecker, UUPSUpgradeable} from "../src/utils/UUPSUpgradeable.sol"; import {LibClone} from "../src/utils/LibClone.sol"; import {MockUUPSImplementation} from "../test/utils/mocks/MockUUPSImplementation.sol"; +import "./utils/mocks/MockUUPSWithDifferentLayout.sol"; contract UUPSUpgradeableTest is SoladyTest { MockUUPSImplementation impl1; @@ -70,4 +71,18 @@ contract UUPSUpgradeableTest is SoladyTest { vm.expectRevert(MockUUPSImplementation.Unauthorized.selector); MockUUPSImplementation(proxy).upgradeToAndCall(address(0xABCD), ""); } + + function testUpgradeRevertsOnStorageLayoutMismatch() public { + MockUUPSWithDifferentLayout bad = new MockUUPSWithDifferentLayout(); + + vm.prank(address(this)); + vm.expectRevert(UUPSUpgradeable.StorageLayoutMismatch.selector); + MockUUPSImplementation(proxy).upgradeToAndCall(address(bad), ""); + } + + function testUpgradeAllowsCompatibleImplementation() public { + MockUUPSImplementation compatible = new MockUUPSImplementation(); + MockUUPSImplementation(proxy).upgradeToAndCall(address(compatible), ""); + // passes silently + } } diff --git a/test/utils/mocks/MockUUPSImplementation.sol b/test/utils/mocks/MockUUPSImplementation.sol index 0d379f0c2..26792f049 100644 --- a/test/utils/mocks/MockUUPSImplementation.sol +++ b/test/utils/mocks/MockUUPSImplementation.sol @@ -11,6 +11,11 @@ contract MockUUPSImplementation is UUPSUpgradeable, Brutalizer { address public owner; + /// @dev Storage layout version hash for upgrade compatibility checks. + bytes32 private constant _STORAGE_LAYOUT_VERSION = + 0x7b86195f14c03aa21c7b6881091a51da39658e244432de468f62f9387ca79dba; + // keccak256("solady.mock.uups.v1") + error Unauthorized(); error CustomError(address owner_); @@ -26,6 +31,32 @@ contract MockUUPSImplementation is UUPSUpgradeable, Brutalizer { function _authorizeUpgrade(address) internal override onlyOwner {} + /// @dev Returns the storage layout version for compatibility verification. + function STORAGE_LAYOUT_VERSION() external pure returns (bytes32) { + return _STORAGE_LAYOUT_VERSION; + } + + /// @dev Checks storage layout compatibility by comparing version hashes. + function _checkStorageLayout(address newImpl) internal view override { + bytes32 expected = _STORAGE_LAYOUT_VERSION; + + /// @solidity memory-safe-assembly + assembly { + // Skip check if `newImpl` has no code to preserve standard UUPS errors. + if extcodesize(newImpl) { + mstore(0x00, 0xa19f05fc) // `STORAGE_LAYOUT_VERSION()` + if iszero(staticcall(gas(), newImpl, 0x1c, 0x04, 0x00, 0x20)) { + mstore(0x00, 0xc2fbd48f) // `StorageLayoutMismatch()` + revert(0x1c, 0x04) + } + if iszero(eq(mload(0x00), expected)) { + mstore(0x00, 0xc2fbd48f) // `StorageLayoutMismatch()` + revert(0x1c, 0x04) + } + } + } + } + function revertWithError() public view { revert CustomError(owner); } diff --git a/test/utils/mocks/MockUUPSWithDifferentLayout.sol b/test/utils/mocks/MockUUPSWithDifferentLayout.sol new file mode 100644 index 000000000..8e426c898 --- /dev/null +++ b/test/utils/mocks/MockUUPSWithDifferentLayout.sol @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: MIT +pragma solidity ^0.8.4; + +import {UUPSUpgradeable} from "../../../src/utils/UUPSUpgradeable.sol"; + +contract MockUUPSWithDifferentLayout is UUPSUpgradeable { + bool public extraVar; // Breaks layout compatibility + uint256 public value; + address public owner; + + /// @dev Intentionally different version to trigger mismatch. + bytes32 private constant _STORAGE_LAYOUT_VERSION = + 0xbe1c9475ae85c040745b734fec444fcbc3d6e069d8e0b5027394bd8b0a614f94; + // keccak256("solady.mock.uups.incompatible.v1") + + function _authorizeUpgrade(address) internal view override { + require(msg.sender == address(this), "unauthorized"); + } + + /// @dev Returns incompatible version hash to test mismatch detection. + function STORAGE_LAYOUT_VERSION() external pure returns (bytes32) { + return _STORAGE_LAYOUT_VERSION; + } + + /// @dev No-op override to skip layout check (for negative testing). + function _checkStorageLayout(address) internal pure override {} +}