Fix dirty bits in upper bits in implementation address in Clones.sol

contracts/proxy/Clones.sol

@@ -39,12 +39,11 @@ library Clones {
         }
         /// @solidity memory-safe-assembly
         assembly {
-            // Stores the bytecode after address
-            mstore(0x20, 0x5af43d82803e903d91602b57fd5bf3)
-            // implementation address
-            mstore(0x11, implementation)
-            // Packs the first 3 bytes of the `implementation` address with the bytecode before the address.
-            mstore(0x00, or(shr(0x88, implementation), 0x3d602d80600a3d3981f3363d3d373d3d3d363d73000000))
+            // Cleans the upper 96 bits of the `implementation` word, then packs the first 3 bytes
+            // of the `implementation` address with the bytecode before the address.
+            mstore(0x00, or(shr(0xe8, shl(0x60, implementation)), 0x3d602d80600a3d3981f3363d3d373d3d3d363d73000000))
+            // Packs the remaining 17 bytes of `implementation` with the bytecode after the address.
+            mstore(0x20, or(shl(0x78, implementation), 0x5af43d82803e903d91602b57fd5bf3))
             instance := create(value, 0x09, 0x37)
         }
         if (instance == address(0)) {
@@ -80,12 +79,11 @@ library Clones {
         }
         /// @solidity memory-safe-assembly
         assembly {
-            // Stores the bytecode after address
-            mstore(0x20, 0x5af43d82803e903d91602b57fd5bf3)
-            // implementation address
-            mstore(0x11, implementation)
-            // Packs the first 3 bytes of the `implementation` address with the bytecode before the address.
-            mstore(0x00, or(shr(0x88, implementation), 0x3d602d80600a3d3981f3363d3d373d3d3d363d73000000))
+            // Cleans the upper 96 bits of the `implementation` word, then packs the first 3 bytes
+            // of the `implementation` address with the bytecode before the address.
+            mstore(0x00, or(shr(0xe8, shl(0x60, implementation)), 0x3d602d80600a3d3981f3363d3d373d3d3d363d73000000))
+            // Packs the remaining 17 bytes of `implementation` with the bytecode after the address.
+            mstore(0x20, or(shl(0x78, implementation), 0x5af43d82803e903d91602b57fd5bf3))
             instance := create2(value, 0x09, 0x37, salt)
         }
         if (instance == address(0)) {

test/proxy/Clones.t.sol

@@ -2,10 +2,54 @@
 
 pragma solidity ^0.8.20;
 
+// solhint-disable func-name-mixedcase
+
 import {Test} from "forge-std/Test.sol";
 import {Clones} from "@openzeppelin/contracts/proxy/Clones.sol";
 
+contract DummyContract {
+    function getNumber() external pure returns (uint256) {
+        return 42;
+    }
+}
+
+contract ClonesMock {
+    function cloneContract(address target) external returns (address) {
+        return Clones.clone(target);
+    }
+
+    function cloneDirtyTarget(address target) external returns (address) {
+        // Add dirty bits; will start with (0xff..)
+        assembly {
+            // Get type(uint256).max and shift to the left by 20 bytes => `or` with target
+            target := or(shl(160, not(0)), target)
+        }
+        return Clones.clone(target);
+    }
+
+    function cloneContractDeterministic(address target) external returns (address) {
+        return Clones.clone(target);
+    }
+
+    function cloneDirtyTargetDeterministic(address target) external returns (address) {
+        // Add dirty bits; will start with (0xff..)
+        assembly {
+            // Get type(uint256).max and shift to the left by 20 bytes => `or` with target
+            target := or(shl(160, not(0)), target)
+        }
+        return Clones.clone(target);
+    }
+}
+
 contract ClonesTest is Test {
+    DummyContract private dummy;
+    ClonesMock private clonesWrapper;
+
+    function setUp() public {
+        dummy = new DummyContract();
+        clonesWrapper = new ClonesMock();
+    }
+
     function testSymbolicPredictDeterministicAddressSpillage(address implementation, bytes32 salt) public {
         address predicted = Clones.predictDeterministicAddress(implementation, salt);
         bytes32 spillage;
@@ -15,4 +59,48 @@ contract ClonesTest is Test {
         }
         assertEq(spillage, bytes32(0));
     }
+
+    function test_DummyContractReturnsCorrectNumber() external {
+        assertEq(dummy.getNumber(), 42);
+    }
+
+    /// Clone
+
+    function test_ValidCloneOfDummyContractReturnsCorrectNumber() external {
+        DummyContract dummyClone = DummyContract(clonesWrapper.cloneContract(address(dummy)));
+        assertEq(dummyClone.getNumber(), dummy.getNumber());
+    }
+
+    function test_ClonesOfSameAddressHaveSameCodeDirty() external {
+        address clone = clonesWrapper.cloneContract(address(dummy));
+        address invalidClone = clonesWrapper.cloneDirtyTarget(address(dummy));
+
+        assertEq(keccak256(clone.code), keccak256(invalidClone.code));
+    }
+
+    function testFail_CallToInvalidCloneWillRevertBecauseItHasNoCode() external {
+        DummyContract invalidDummyClone = DummyContract(clonesWrapper.cloneDirtyTarget(address(dummy)));
+        vm.expectRevert();
+        invalidDummyClone.getNumber();
+    }
+
+    /// Clone Deterministic
+
+    function test_ValidCloneOfDummyContractReturnsCorrectNumber_Deterministic() external {
+        DummyContract dummyClone = DummyContract(clonesWrapper.cloneContractDeterministic(address(dummy)));
+        assertEq(dummyClone.getNumber(), dummy.getNumber());
+    }
+
+    function test_ClonesOfSameAddressHaveSameCodeDirty_Deterministic() external {
+        address clone = clonesWrapper.cloneContractDeterministic(address(dummy));
+        address invalidClone = clonesWrapper.cloneDirtyTargetDeterministic(address(dummy));
+
+        assertEq(keccak256(clone.code), keccak256(invalidClone.code));
+    }
+
+    function testFail_CallToInvalidCloneWillRevertBecauseItHasNoCode_Deterministic() external {
+        DummyContract invalidDummyClone = DummyContract(clonesWrapper.cloneDirtyTargetDeterministic(address(dummy)));
+        vm.expectRevert();
+        invalidDummyClone.getNumber();
+    }
 }