✨ Allow ERC1967 bootstrap to do upgrade and call in a single call (#1118

)
Vectorized · Oct 11, 2024 · ee879db · ee879db
1 parent e1c147d
commit ee879db
Show file tree

Hide file tree

Showing 2 changed files with 59 additions and 10 deletions.
diff --git a/src/utils/LibClone.sol b/src/utils/LibClone.sol
@@ -1714,14 +1714,34 @@ library LibClone {
     function bootstrapERC1967(address instance, address implementation) internal {
         /// @solidity memory-safe-assembly
         assembly {
-            mstore(0x00, shr(96, shl(96, implementation)))
-            if iszero(call(gas(), instance, 0, 0x00, 0x20, codesize(), 0x00)) {
+            mstore(0x00, implementation)
+            if iszero(call(gas(), instance, 0, 0x0c, 0x14, codesize(), 0x00)) {
                 mstore(0x00, 0x30116425) // `DeploymentFailed()`.
                 revert(0x1c, 0x04)
             }
         }
     }
 
+    /// @dev Replaces the implementation at `instance`, and then call it with `data`.
+    function bootstrapERC1967AndCall(address instance, address implementation, bytes memory data)
+        internal
+    {
+        /// @solidity memory-safe-assembly
+        assembly {
+            let n := mload(data)
+            mstore(data, implementation)
+            if iszero(call(gas(), instance, 0, add(data, 0x0c), add(n, 0x14), codesize(), 0x00)) {
+                if iszero(returndatasize()) {
+                    mstore(0x00, 0x30116425) // `DeploymentFailed()`.
+                    revert(0x1c, 0x04)
+                }
+                returndatacopy(mload(0x40), 0x00, returndatasize())
+                revert(mload(0x40), returndatasize())
+            }
+            mstore(data, n) // Restore the length of `data`.
+        }
+    }
+
     /// @dev Returns the implementation address of the ERC1967 bootstrap for this contract.
     function predictDeterministicAddressERC1967Bootstrap() internal view returns (address) {
         return predictDeterministicAddressERC1967Bootstrap(address(this), address(this));
@@ -1745,12 +1765,13 @@ library LibClone {
         /// @solidity memory-safe-assembly
         assembly {
             c := mload(0x40)
-            mstore(add(c, 0x60), 0xca505d382bbc5500000000000000000000000000000000000000000000000000)
-            mstore(add(c, 0x40), 0x0338573d357f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3)
+            mstore(add(c, 0x80), 0x3d3560601c5af46047573d6000383e3d38fd0000000000000000000000000000)
+            mstore(add(c, 0x60), 0xa920a3ca505d382bbc55601436116049575b005b363d3d373d3d601436036014)
+            mstore(add(c, 0x40), 0x0338573d3560601c7f360894a13ba1a3210667c828492db98dca3e2076cc3735)
             mstore(add(c, 0x20), authorizedUpgrader)
-            mstore(add(c, 0x0c), 0x603d80600a3d393df3fe3373)
-            mstore(c, 0x47)
-            mstore(0x40, add(c, 0x80))
+            mstore(add(c, 0x0c), 0x606880600a3d393df3fe3373)
+            mstore(c, 0x72)
+            mstore(0x40, add(c, 0xa0))
         }
     }
 

diff --git a/test/LibClone.t.sol b/test/LibClone.t.sol
@@ -1010,9 +1010,9 @@ contract LibCloneTest is SoladyTest {
     function testERC1967BootstrapInitCode(address authorizedUpgrader) public {
         bytes memory c = LibClone.initCodeERC1967Bootstrap(authorizedUpgrader);
         bytes memory expected = abi.encodePacked(
-            hex"603d80600a3d393df3fe3373",
+            hex"606880600a3d393df3fe3373",
             authorizedUpgrader,
-            hex"0338573d357f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc55"
+            hex"0338573d3560601c7f360894a13ba1a3210667c828492db98dca3e2076cc3735a920a3ca505d382bbc55601436116049575b005b363d3d373d3d6014360360143d3560601c5af46047573d6000383e3d38fd"
         );
         assertEq(c, expected);
     }
@@ -1022,6 +1022,8 @@ contract LibCloneTest is SoladyTest {
         address bootstrap = LibClone.erc1967Bootstrap(authorizedUpgrader);
         address instance = this.deployDeterministicERC1967I(bootstrap, salt);
         assertEq(LibClone.implementationOf(instance), bootstrap);
+        bytes memory bootstrapCode = address(bootstrap).code;
+        assertEq(uint8(bootstrapCode[bootstrapCode.length - 1]), uint8(0xfd));
         if (_randomChance(2)) {
             vm.prank(authorizedUpgrader);
             LibClone.bootstrapERC1967(instance, implementation);
@@ -1054,16 +1056,42 @@ contract LibCloneTest is SoladyTest {
                 vm.load(instance, _ERC1967_IMPLEMENTATION_SLOT),
                 bytes32(uint256(uint160(implementation)))
             );
-        } else {
+        } else if (_randomChance(2)) {
             LibClone.bootstrapERC1967(instance, address(this));
             assertEq(
                 vm.load(instance, _ERC1967_IMPLEMENTATION_SLOT),
                 bytes32(uint256(uint160(address(this))))
             );
             _checkBehavesLikeProxy(instance);
+        } else if (_randomChance(2)) {
+            uint256 x = _random();
+            this.bootstrapERC1967AndCall(
+                instance, address(this), abi.encodeWithSignature("setValue(uint256)", x)
+            );
+            assertEq(LibCloneTest(instance).value(), x);
+            assertEq(
+                vm.load(instance, _ERC1967_IMPLEMENTATION_SLOT),
+                bytes32(uint256(uint160(address(this))))
+            );
+            _checkBehavesLikeProxy(instance);
+        } else {
+            vm.expectRevert(abi.encodeWithSelector(CustomError.selector, uint256(0)));
+            this.bootstrapERC1967AndCall(
+                instance, address(this), abi.encodeWithSignature("revertWithError()")
+            );
+            assertEq(
+                vm.load(instance, _ERC1967_IMPLEMENTATION_SLOT),
+                bytes32(uint256(uint160(bootstrap)))
+            );
         }
     }
 
+    function bootstrapERC1967AndCall(address instance, address implementation, bytes memory data)
+        public
+    {
+        LibClone.bootstrapERC1967AndCall(instance, implementation, data);
+    }
+
     function _beacon() internal returns (address result) {
         if (_deployedBeacon != address(0)) return _deployedBeacon;
         if (_randomChance(2)) {