Michael.W基于Foundry精读Openzeppelin第64期——UUPSUpgradeable.sol

0. 版本

[openzeppelin]：v4.8.3，[forge-std]：v1.5.6

0.1 UUPSUpgradeable.sol

Github: https://github.com/OpenZeppelin/openzeppelin-contracts/blob/v4.8.3/contracts/proxy/utils/UUPSUpgradeable.sol

UUPSUpgradeable库是专为UUPS代理设计的一种合约升级机制的实现。当本合约被设置为ERC1967Proxy代理合约背后的逻辑合约后，可以对其进行合约升级操作。作为逻辑合约的父合约，本库的安全机制可保证不会因某次错误的升级而打破合约的可升级性。

ERC1967Proxy库详解参见：https://learnblockchain.cn/article/8594

1. 目标合约

继承UUPSUpgradeable合约：

Github: https://github.com/RevelationOfTuring/foundry-openzeppelin-contracts/blob/master/src/proxy/utils/MockUUPSUpgradeable.sol

// SPDX-License-Identifier: UNLICENSED
pragma solidity ^0.8.0;

import "openzeppelin-contracts/contracts/proxy/utils/UUPSUpgradeable.sol";
import "openzeppelin-contracts/contracts/access/Ownable.sol";

interface IImplementation {
    event StorageChanged(uint preValue, uint newValue);
}

// specific logic of implementation
contract Implementation is Ownable, IImplementation {
    // storage
    uint public i;

    // initializer
    function __Implementation_init(address owner) external {
        _transferOwnership(owner);
    }

    // logic function
    function setI(uint newI) external virtual {
        emit StorageChanged(i, newI);
        i = newI;
    }
}

contract MockUUPSUpgradeable is UUPSUpgradeable, Implementation {
    // modified by `onlyOwner`
    function _authorizeUpgrade(address newImplementation) internal override onlyOwner {}
}

contract MockUUPSUpgradeableNew is MockUUPSUpgradeable {
    // change the logic
    function setI(uint newI) external override {
        uint currentI = i;
        i = currentI + newI;
        emit StorageChanged(currentI, newI);
    }
}

contract MockUUPSUpgradeableWithWrongProxiableUUID is MockUUPSUpgradeable {
    // inconsistent proxiable uuid
    function proxiableUUID() external view virtual override notDelegated returns (bytes32) {
        return bytes32(uint(_IMPLEMENTATION_SLOT) - 1);
    }
}

contract MockUUPSUpgradeableWithRollbackTest is MockUUPSUpgradeable {
    bytes32 private constant _ROLLBACK_SLOT = bytes32(uint(keccak256("eip1967.proxy.rollback")) - 1);

    function upgradeTo(address newImplementation) external override onlyProxy {
        _authorizeUpgrade(newImplementation);
        _upgradeToAndCallUUPSWithRollbackTest(
            newImplementation,
            "",
            false
        );
    }

    function upgradeToAndCall(address newImplementation, bytes memory data) external payable override onlyProxy {
        _authorizeUpgrade(newImplementation);
        _upgradeToAndCallUUPSWithRollbackTest(
            newImplementation,
            data,
            true
        );
    }

    // do rollback test both in {upgradeTo} and {upgradeToAndCall}
    function _upgradeToAndCallUUPSWithRollbackTest(
        address newImplementation,
        bytes memory data,
        bool forceCall
    ) private {
        address preImplementation = _getImplementation();
        // upgrade to and check UUPS on the new implementation with setup call first
        _upgradeToAndCallUUPS(newImplementation, data, forceCall);

        StorageSlot.BooleanSlot storage isInRollbackTest = StorageSlot.getBooleanSlot(_ROLLBACK_SLOT);
        // go into rollback test
        isInRollbackTest.value = true;
        // upgrade to the pre-implementation from the new implementation without setup call
        _upgradeToAndCallUUPS(preImplementation, "", false);
        require(preImplementation == _getImplementation(), "fail to recover pre-implementation in rollback test");
        // upgrade to the new implementation again without setup call
        _upgradeToAndCallUUPS(newImplementation, "", false);
        // get out of rollback test
        isInRollbackTest.value = false;
    }
}

contract MockUUPSUpgradeableWithRollbackTestNew is MockUUPSUpgradeableWithRollbackTest {
    // change the logic
    function setI(uint newI) external override {
        uint currentI = i;
        i = currentI + newI;
        emit StorageChanged(currentI, newI);
    }
}

全部foundry测试合约：

Github: https://github.com/RevelationOfTuring/foundry-openzeppelin-contracts/blob/master/test/proxy/utils/UUPSUpgradeable.t.sol

2. 代码精读

2.1 modifier onlyProxy()

该修饰器用于检查调用其所修饰函数的方式是delegatecall，并校验调用的上下文是一个类ERC1967的代理合约且本逻辑合约是其背后设置的逻辑合约。

注：这个类ERC1967的代理合约只包括将本逻辑合约设置为其背后逻辑合约的UUPS代理或透明代理，而ERC1167最小代理合约（即Clones）通常是无法通过该修饰器的检验的。

ERC1167最小代理合约详解参见：https://learnblockchain.cn/article/8507

    // 本逻辑合约地址
    // 注：由于是immutable类型常量，在逻辑合约部署时就会将逻辑合约地址以字节码的形式存在链上
    address private immutable __self = address(this);

    modifier onlyProxy() {
        // 校验调用上下文中的本合约地址不是本逻辑合约地址，
        // 即校验该调用为delegatecall
        require(address(this) != __self, "Function must be called through delegatecall");
        // 要求调用上下文中调用_getImplementation()返回地址为本逻辑合约地址，
        // 即校验调用上下文是一个类ERC1967的代理合约且本逻辑合约就是其背后的逻辑合约
        require(_getImplementation() == __self, "Function must be called through active proxy");
        // 执行被修饰函数的逻辑
        _;
    }

2.2 modifier notDelegated()

该修饰器用于检查调用其所修饰函数的方式非delegatecall，即修饰一个只能在本逻辑合约直接调用而无法从代理合约委托调用的函数。

    modifier notDelegated() {
        // 校验调用上下文中的本合约地址是本逻辑合约地址，
        // 即校验该调用为非delegatecall
        require(address(this) == __self, "UUPSUpgradeable: must not be called through delegatecall");
        // 执行被修饰函数的逻辑
        _;
    }

2.3 proxiableUUID()

返回用于存储逻辑合约地址的slot号。本函数是ERC1822中proxiableUUID()的实现，用于检查合约升级时新旧逻辑合约的兼容性。

注：指向逻辑合约的代理合约不能做逻辑合约，因为这可能在升级过程中引发自己delegatecall自己的情况。该情况会导致gas耗尽从而阻塞合约升级。

例子：

1. 假设代理合约A -> 逻辑合约B，现在A要将逻辑合约升级更换为代理合约A自身；

2. 假设代理合约A -> 逻辑合约B（B也是一个代理合约），代理合约B -> 逻辑合约C，现在B要将其逻辑合约升级更换为A。

为了防止在合约升级过程中出现自己delegatecall自己的情况，将proxiableUUID()通过修饰器notDelegated限定为不可delegatecall。

    function proxiableUUID() external view virtual override notDelegated returns (bytes32) {
        // 返回ERC1967中规定的用于存储逻辑合约地址的slot号
        // 注：ERC1967Upgrade._IMPLEMENTATION_SLOT详解参见：https://learnblockchain.cn/article/8581
        return _IMPLEMENTATION_SLOT;
    }

foundry代码验证：

contract UUPSUpgradeableTest is Test {
    bytes32 private constant _IMPLEMENTATION_SLOT = bytes32(uint(keccak256("eip1967.proxy.implementation")) - 1);
    MockUUPSUpgradeable private _testing = new MockUUPSUpgradeable();

    function test_ProxiableUUID() external {
        assertEq(_testing.proxiableUUID(), _IMPLEMENTATION_SLOT);

        // revert with a delegatecall
        MockUUPSUpgradeable proxy = MockUUPSUpgradeable(address(new ERC1967Proxy(address(_testing), "")));
        vm.expectRevert("UUPSUpgradeable: must not be called through delegatecall");
        proxy.proxiableUUID();
    }
}

2.4 upgradeTo(address newImplementation) && _authorizeUpgrade(address newImplementation) internal

• upgradeTo(address newImplementation)：调用类ERC1967的代理合约，将代理合约背后的原逻辑合约升级为newImplementation；
• _authorizeUpgrade(address newImplementation) internal：用于校验msg.sender是否具有合约升级权限的函数。该函数为virtual，需要在子类合约中进行函数实现。通常可使用Openzeppelin的AccessControl库或Ownable库对升级权限进行管理和约束。

注：

• AccessControl库详解参见：https://learnblockchain.cn/article/6632
• Ownable库详解参见：https://learnblockchain.cn/article/6576

    function upgradeTo(address newImplementation) external virtual onlyProxy {
        // 校验msg.sender是否具有合约升级权限
        _authorizeUpgrade(newImplementation);
        // 调用ERC1967Upgrade._upgradeToAndCallUUPS()，将类ERC1967代理合约背后的逻辑合约地址升级更换为newImplementation并对其进行UUPS安全检查。
        // 注：ERC1967Upgrade._upgradeToAndCallUUPS()详解参见：https://learnblockchain.cn/article/8581
        _upgradeToAndCallUUPS(newImplementation, new bytes(0), false);
    }

    function _authorizeUpgrade(address newImplementation) internal virtual;

foundry代码验证：

contract UUPSUpgradeableTest is Test, IImplementation, IERC1967 {
    MockUUPSUpgradeable private _testing = new MockUUPSUpgradeable();
    MockUUPSUpgradeableNew private _implementationNew = new MockUUPSUpgradeableNew();
    MockUUPSUpgradeableWithWrongProxiableUUID private _implementationNewWithWrongProxiableUUID = new MockUUPSUpgradeableWithWrongProxiableUUID();

    function test_UpgradeTo() external {
        // case 1: pass
        // deploy proxy with a setup call
        address proxyAddress = address(
            new ERC1967Proxy(
                address(_testing),
                abi.encodeCall(_testing.__Implementation_init, (address(this)))
            )
        );

        MockUUPSUpgradeable proxyAsMockUUPSUpgradeable = MockUUPSUpgradeable(proxyAddress);
        // check setup call
        assertEq(proxyAsMockUUPSUpgradeable.owner(), address(this));
        // call on proxy
        uint newI = 1024;
        vm.expectEmit(proxyAddress);
        emit IImplementation.StorageChanged(0, newI);

        proxyAsMockUUPSUpgradeable.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeable.i(), newI);

        // upgrade to new implementation
        vm.expectEmit(proxyAddress);
        emit IERC1967.Upgraded(address(_implementationNew));

        proxyAsMockUUPSUpgradeable.upgradeTo(address(_implementationNew));
        MockUUPSUpgradeableNew proxyAsMockUUPSUpgradeableNew = MockUUPSUpgradeableNew(proxyAddress);
        // check storage
        assertEq(proxyAsMockUUPSUpgradeableNew.i(), newI);
        // call new logic
        vm.expectEmit(proxyAddress);
        emit IImplementation.StorageChanged(newI, newI);

        proxyAsMockUUPSUpgradeableNew.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeableNew.i(), newI + newI);

        // case 2: revert if not pass {_authorizeUpgrade}
        address auth = address(1024);
        proxyAsMockUUPSUpgradeable = MockUUPSUpgradeable(address(
            new ERC1967Proxy(
                address(_testing),
                abi.encodeCall(_testing.__Implementation_init, (auth))
            )
        ));

        vm.expectRevert("Ownable: caller is not the owner");
        proxyAsMockUUPSUpgradeable.upgradeTo(address(_implementationNew));

        // case 3: revert if the new implementation has an inconsistent proxiable uuid
        proxyAsMockUUPSUpgradeable = MockUUPSUpgradeable(address(
            new ERC1967Proxy(
                address(_testing),
                abi.encodeCall(_testing.__Implementation_init, (address(this)))
            )
        ));

        vm.expectRevert("ERC1967Upgrade: unsupported proxiableUUID");
        proxyAsMockUUPSUpgradeable.upgradeTo(address(_implementationNewWithWrongProxiableUUID));

        // case 4: revert if the new implementation has no {proxiableUUID}
        vm.expectRevert("ERC1967Upgrade: new implementation is not UUPS");
        proxyAsMockUUPSUpgradeable.upgradeTo(address(this));

        // case 5: revert with no msg if the new implementation address is an EOA
        vm.expectRevert();
        proxyAsMockUUPSUpgradeable.upgradeTo(address(1024));

        // case 6: revert if call {upgradeTo} directly
        vm.expectRevert("Function must be called through delegatecall");
        _testing.upgradeTo(address(this));

        // case 7: revert if the context of delegatecall is not an active proxy
        vm.expectRevert("Function must be called through active proxy");
        Address.functionDelegateCall(
            address(_testing),
            abi.encodeCall(_testing.upgradeTo, (address(_implementationNew)))
        );
    }
}

2.5 upgradeToAndCall(address newImplementation, bytes memory data)

调用类ERC1967的代理合约，将代理合约背后的原逻辑合约升级为newImplementation并随后再执行一个额外的到新逻辑合约的delegatecall。

    function upgradeToAndCall(address newImplementation, bytes memory data) external payable virtual onlyProxy {
        // 校验msg.sender是否具有合约升级权限
        _authorizeUpgrade(newImplementation);
        // 调用ERC1967Upgrade._upgradeToAndCallUUPS()，将类ERC1967代理合约背后的逻辑合约地址升级更换为newImplementation并对其进行UUPS安全检查。
        // 随后再以data为calldata执行一个额外的到新逻辑合约的delegatecall。
        // 注：ERC1967Upgrade._upgradeToAndCallUUPS()详解参见：https://learnblockchain.cn/article/8581
        _upgradeToAndCallUUPS(newImplementation, data, true);
    }

foundry代码验证：

contract UUPSUpgradeableTest is Test, IImplementation, IERC1967 {
    MockUUPSUpgradeable private _testing = new MockUUPSUpgradeable();
    MockUUPSUpgradeableNew private _implementationNew = new MockUUPSUpgradeableNew();
    MockUUPSUpgradeableWithWrongProxiableUUID private _implementationNewWithWrongProxiableUUID = new MockUUPSUpgradeableWithWrongProxiableUUID();

    function test_UpgradeToAndCall() external {
        // case 1: pass
        // deploy proxy with a setup call
        address proxyAddress = address(
            new ERC1967Proxy(
                address(_testing),
                abi.encodeCall(_testing.__Implementation_init, (address(this)))
            )
        );

        MockUUPSUpgradeable proxyAsMockUUPSUpgradeable = MockUUPSUpgradeable(proxyAddress);
        assertEq(proxyAsMockUUPSUpgradeable.owner(), address(this));
        uint newI = 1024;
        proxyAsMockUUPSUpgradeable.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeable.i(), newI);

        // upgrade to new implementation with a setup call
        vm.expectEmit(proxyAddress);
        emit IERC1967.Upgraded(address(_implementationNew));
        emit IImplementation.StorageChanged(newI, newI);

        bytes memory data = abi.encodeCall(
            _implementationNew.setI,
            (newI)
        );
        proxyAsMockUUPSUpgradeable.upgradeToAndCall(
            address(_implementationNew),
            data
        );
        MockUUPSUpgradeableNew proxyAsMockUUPSUpgradeableNew = MockUUPSUpgradeableNew(proxyAddress);
        // check storage
        assertEq(proxyAsMockUUPSUpgradeableNew.i(), newI + newI);
        // call new logic
        vm.expectEmit(proxyAddress);
        emit IImplementation.StorageChanged(newI + newI, newI);

        proxyAsMockUUPSUpgradeableNew.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeableNew.i(), newI + newI + newI);

        // case 2: revert if not pass {_authorizeUpgrade}
        address auth = address(1024);
        proxyAsMockUUPSUpgradeable = MockUUPSUpgradeable(address(
            new ERC1967Proxy(
                address(_testing),
                abi.encodeCall(_testing.__Implementation_init, (auth))
            )
        ));

        vm.expectRevert("Ownable: caller is not the owner");
        proxyAsMockUUPSUpgradeable.upgradeToAndCall(
            address(_implementationNew),
            data
        );

        // case 3: revert if the new implementation has an inconsistent proxiable uuid
        proxyAsMockUUPSUpgradeable = MockUUPSUpgradeable(address(
            new ERC1967Proxy(
                address(_testing),
                abi.encodeCall(_testing.__Implementation_init, (address(this)))
            )
        ));

        vm.expectRevert("ERC1967Upgrade: unsupported proxiableUUID");
        proxyAsMockUUPSUpgradeable.upgradeToAndCall(
            address(_implementationNewWithWrongProxiableUUID),
            data
        );

        // case 4: revert if the new implementation has no {proxiableUUID}
        vm.expectRevert("ERC1967Upgrade: new implementation is not UUPS");
        proxyAsMockUUPSUpgradeable.upgradeToAndCall(
            address(this),
            data
        );

        // case 5: revert with no msg if the new implementation address is an EOA
        vm.expectRevert();
        proxyAsMockUUPSUpgradeable.upgradeToAndCall(
            address(1024),
            data
        );

        // case 6: revert if call {upgradeToAndCall} directly
        vm.expectRevert("Function must be called through delegatecall");
        _testing.upgradeToAndCall(
            address(this),
            data
        );

        // case 7: revert if the context of delegatecall is not an active proxy
        vm.expectRevert("Function must be called through active proxy");
        Address.functionDelegateCall(
            address(_testing),
            abi.encodeCall(
                _testing.upgradeToAndCall,
                (address(_implementationNew), data)
            )
        );
    }
}

2.6 测试本库带有rollback test的UUPS升级机制

foundry代码验证：

contract UUPSUpgradeableTest is Test, IImplementation, IERC1967 {
    MockUUPSUpgradeableWithRollbackTest private _implementationWithRollbackTest = new MockUUPSUpgradeableWithRollbackTest();
    MockUUPSUpgradeableWithRollbackTestNew private _implementationWithRollbackTestNew = new MockUUPSUpgradeableWithRollbackTestNew();

    function test_UpgradeWithUUPSRollbackTest() external {
        // case 1: test {upgradeTo} with UUPS rollback test
        // deploy proxy with a setup call
        address proxyAddress = address(
            new ERC1967Proxy(
                address(_implementationWithRollbackTest),
                abi.encodeCall(_implementationWithRollbackTest.__Implementation_init, (address(this)))
            )
        );

        MockUUPSUpgradeableWithRollbackTest proxyAsMockUUPSUpgradeableWithRollbackTest = MockUUPSUpgradeableWithRollbackTest(proxyAddress);
        // check setup call
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTest.owner(), address(this));
        // call on proxy
        uint newI = 1024;
        vm.expectEmit(proxyAddress);
        emit IImplementation.StorageChanged(0, newI);

        proxyAsMockUUPSUpgradeableWithRollbackTest.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTest.i(), newI);

        // upgrade to new implementation
        vm.expectEmit(proxyAddress);
        emit IERC1967.Upgraded(address(_implementationWithRollbackTestNew));

        proxyAsMockUUPSUpgradeableWithRollbackTest.upgradeTo(address(_implementationWithRollbackTestNew));
        MockUUPSUpgradeableWithRollbackTestNew proxyAsMockUUPSUpgradeableWithRollbackTestNew = MockUUPSUpgradeableWithRollbackTestNew(proxyAddress);
        // check storage
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTestNew.i(), newI);
        // call new logic
        vm.expectEmit(proxyAddress);
        emit IImplementation.StorageChanged(newI, newI);

        proxyAsMockUUPSUpgradeableWithRollbackTestNew.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTestNew.i(), newI + newI);

        // case 2: test {upgradeToAndCall} with UUPS rollback test
        // deploy proxy with a setup call
        proxyAddress = address(
            new ERC1967Proxy(
                address(_implementationWithRollbackTest),
                abi.encodeCall(_implementationWithRollbackTest.__Implementation_init, (address(this)))
            )
        );

        proxyAsMockUUPSUpgradeableWithRollbackTest = MockUUPSUpgradeableWithRollbackTest(proxyAddress);
        // check setup call
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTest.owner(), address(this));
        // call on proxy
        newI = 1024;
        vm.expectEmit(proxyAddress);
        emit IImplementation.StorageChanged(0, newI);

        proxyAsMockUUPSUpgradeableWithRollbackTest.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTest.i(), newI);

        // upgrade to new implementation with a setup call
        vm.expectEmit(proxyAddress);
        emit IERC1967.Upgraded(address(_implementationWithRollbackTestNew));
        emit IImplementation.StorageChanged(newI, newI);

        proxyAsMockUUPSUpgradeableWithRollbackTest.upgradeToAndCall(
            address(_implementationWithRollbackTestNew),
            abi.encodeCall(
                _implementationWithRollbackTestNew.setI,
                (newI)
            )
        );
        proxyAsMockUUPSUpgradeableWithRollbackTestNew = MockUUPSUpgradeableWithRollbackTestNew(proxyAddress);
        // check storage
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTestNew.i(), newI + newI);
        // call new logic
        vm.expectEmit(proxyAddress);
        emit IImplementation.StorageChanged(newI + newI, newI);

        proxyAsMockUUPSUpgradeableWithRollbackTestNew.setI(newI);
        assertEq(proxyAsMockUUPSUpgradeableWithRollbackTestNew.i(), newI + newI + newI);
    }
}

ps: 本人热爱图灵，热爱中本聪，热爱V神。 以下是我个人的公众号，如果有技术问题可以关注我的公众号来跟我交流。 同时我也会在这个公众号上每周更新我的原创文章，喜欢的小伙伴或者老伙计可以支持一下！ 如果需要转发，麻烦注明作者。十分感谢！

公众号名称：后现代泼痞浪漫主义奠基人