✨ Add _hasOwnershipCycle to ERC6551 (#944)

.gas-snapshot

@@ -176,8 +176,8 @@ ERC1155Test:testSafeTransferFromToZeroReverts(uint256) (runs: 278, μ: 71146, ~:
 ERC1155Test:test__codesize() (gas: 42341)
 ERC1271Test:testBasefeeBytecodeContract() (gas: 45430)
 ERC1271Test:testIsValidSignature() (gas: 2464705)
-ERC1271Test:testIsValidSignature(uint256) (runs: 278, μ: 245649, ~: 213657)
-ERC1271Test:test__codesize() (gas: 30392)
+ERC1271Test:testIsValidSignature(uint256) (runs: 278, μ: 251606, ~: 214580)
+ERC1271Test:test__codesize() (gas: 30438)
 ERC1967FactoryTest:testChangeAdmin() (gas: 266356)
 ERC1967FactoryTest:testChangeAdminUnauthorized() (gas: 257653)
 ERC1967FactoryTest:testDeploy() (gas: 257202)
@@ -194,7 +194,7 @@ ERC1967FactoryTest:testUpgradeAndCallWithRevert() (gas: 265824)
 ERC1967FactoryTest:testUpgradeUnauthorized() (gas: 270340)
 ERC1967FactoryTest:testUpgradeWithCorruptedProxy() (gas: 263294)
 ERC1967FactoryTest:test__codesize() (gas: 32011)
-ERC20Invariants:invariantBalanceSum() (runs: 256, calls: 128000, reverts: 78216)
+ERC20Invariants:invariantBalanceSum() (runs: 256, calls: 128000, reverts: 78398)
 ERC20Invariants:test__codesize() (gas: 7532)
 ERC20Test:testApprove() (gas: 35730)
 ERC20Test:testApprove(address,uint256) (runs: 278, μ: 30751, ~: 31181)
@@ -275,24 +275,24 @@ ERC4626Test:testWithdrawWithNotEnoughUnderlyingAmountReverts() (gas: 144343)
 ERC4626Test:testWithdrawZero() (gas: 51874)
 ERC4626Test:test__codesize() (gas: 37175)
 ERC6551Test:testBaseFeeMini() (gas: 39514)
-ERC6551Test:testCdFallback() (gas: 895599)
-ERC6551Test:testDeployERC6551(uint256) (runs: 278, μ: 170165, ~: 168891)
+ERC6551Test:testCdFallback() (gas: 895594)
+ERC6551Test:testDeployERC6551(uint256) (runs: 278, μ: 170456, ~: 168890)
 ERC6551Test:testDeployERC6551Proxy() (gas: 80395)
 ERC6551Test:testExecute() (gas: 507518)
 ERC6551Test:testExecuteBatch() (gas: 817605)
-ERC6551Test:testExecuteBatch(uint256) (runs: 278, μ: 639638, ~: 763255)
+ERC6551Test:testExecuteBatch(uint256) (runs: 278, μ: 612344, ~: 762410)
 ERC6551Test:testInitializeERC6551ProxyImplementation() (gas: 189914)
-ERC6551Test:testIsValidSigner(address) (runs: 278, μ: 167471, ~: 167467)
-ERC6551Test:testOnERC1155BatchReceived() (gas: 1702522)
-ERC6551Test:testOnERC1155Received() (gas: 1699885)
-ERC6551Test:testOnERC721Received() (gas: 1738813)
-ERC6551Test:testOnERC721ReceivedCycles() (gas: 1728089)
-ERC6551Test:testOnERC721ReceivedCyclesWithDifferentChainIds(uint256) (runs: 278, μ: 450735, ~: 455448)
-ERC6551Test:testOwnerWorksWithChainIdChange(uint256,uint256) (runs: 278, μ: 1364052, ~: 1364067)
+ERC6551Test:testIsValidSigner(address) (runs: 278, μ: 167472, ~: 167467)
+ERC6551Test:testOnERC1155BatchReceived() (gas: 1702519)
+ERC6551Test:testOnERC1155Received() (gas: 1699882)
+ERC6551Test:testOnERC721Received() (gas: 1738866)
+ERC6551Test:testOnERC721ReceivedCycles() (gas: 2914176)
+ERC6551Test:testOnERC721ReceivedCyclesWithDifferentChainIds(uint256) (runs: 278, μ: 450409, ~: 455548)
+ERC6551Test:testOwnerWorksWithChainIdChange(uint256,uint256) (runs: 278, μ: 1368677, ~: 1368689)
 ERC6551Test:testSupportsInterface() (gas: 169450)
-ERC6551Test:testUpdateState(uint256) (runs: 278, μ: 235212, ~: 235137)
+ERC6551Test:testUpdateState(uint256) (runs: 278, μ: 235193, ~: 235130)
 ERC6551Test:testUpgrade() (gas: 248450)
-ERC6551Test:test__codesize() (gas: 51546)
+ERC6551Test:test__codesize() (gas: 52354)
 ERC6909Test:testApprove() (gas: 36872)
 ERC6909Test:testApprove(address,uint256,uint256) (runs: 278, μ: 37442, ~: 37514)
 ERC6909Test:testBurn() (gas: 40870)
@@ -1191,7 +1191,7 @@ UpgradeableBeaconTest:testUpgradeableSolidityBeaconOnlyOwnerFunctions() (gas: 26
 UpgradeableBeaconTest:testUpgradeableYulBeaconOnlyFnSelectorNotRecognised() (gas: 172796)
 UpgradeableBeaconTest:testUpgradeableYulBeaconOnlyOwnerFunctions() (gas: 198398)
 UpgradeableBeaconTest:test__codesize() (gas: 8808)
-WETHInvariants:invariantTotalSupplyEqualsBalance() (runs: 256, calls: 128000, reverts: 59410)
+WETHInvariants:invariantTotalSupplyEqualsBalance() (runs: 256, calls: 128000, reverts: 59795)
 WETHInvariants:test__codesize() (gas: 5264)
 WETHTest:testDeposit() (gas: 68090)
 WETHTest:testDeposit(uint256) (runs: 278, μ: 67950, ~: 68384)

src/accounts/ERC6551.sol

@@ -299,6 +299,46 @@ abstract contract ERC6551 is UUPSUpgradeable, Receiver, ERC1271 {
  }
  }
 
+ /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
+ /* INTERNAL HELPERS */
+ /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
+
+ /// @dev Returns whether there is an ownership cycle.
+ function _hasOwnershipCycle() internal view returns (bool result) {
+ uint256 cachedChainId = _cachedChainId;
+ /// @solidity memory-safe-assembly
+ assembly {
+ let m := mload(0x40) // Cache the free memory pointer.
+ extcodecopy(address(), 0x00, 0x4d, 0x60) // `(chainId, tokenContract, tokenId)`.
+ mstore(0x60, 0xfc0c546a) // `token()`.
+ for {} 1 {} {
+ let tokenContract := mload(0x20)
+ mstore(0x20, 0x6352211e) // `ownerOf(uint256)`.
+ let currentOwner :=
+ mul( // `chainId == cachedChainId ? tokenContract.ownerOf(tokenId) : address(0)`.
+ mload(0x20),
+ and(
+ and(gt(returndatasize(), 0x1f), eq(mload(0x00), cachedChainId)),
+ staticcall(gas(), tokenContract, 0x3c, 0x24, 0x20, 0x20)
+ )
+ )
+ if iszero(eq(currentOwner, address())) {
+ if iszero(
+ and( // `(chainId, tokenContract, tokenId) = currentOwner.token()`.
+ gt(returndatasize(), 0x5f),
+ staticcall(gas(), currentOwner, 0x7c, 0x04, 0x00, 0x60)
+ )
+ ) { break }
+ continue
+ }
+ result := 1
+ break
+ }
+ mstore(0x40, m) // Restore the free memory pointer.
+ mstore(0x60, 0) // Restore the zero pointer.
+ }
+ }
+
  /*´:°•.°+.*•´.*:˚.°*.˚•´.°:°•.°•.*•´.*:˚.°*.˚•´.°:°•.°+.*•´.*:*/
  /* OVERRIDES */
  /*.•°:°.´+˚.*°.˚:*.´•*.+°.•°:´*.´•*.•°.•°:°.´:•˚°.*°.˚:*.´+°.•*/
@@ -322,51 +362,24 @@ abstract contract ERC6551 is UUPSUpgradeable, Receiver, ERC1271 {
  /// @dev For handling token callbacks.
  /// Safe-transferred ERC721 tokens will trigger a ownership cycle check.
  modifier receiverFallback() override(Receiver) {
- uint256 cachedChainId = _cachedChainId;
- /// @solidity memory-safe-assembly
- assembly {
- let s := shr(224, calldataload(0x00))
- // 0x150b7a02: `onERC721Received(address,address,uint256,bytes)`.
- if eq(s, 0x150b7a02) {
- extcodecopy(address(), 0x00, 0x4d, 0x60) // `chainId`, `tokenContract`, `tokenId`.
- mstore(0x60, 0xfc0c546a) // `token()`.
- for {} 1 {} {
- let tokenContract := mload(0x20)
- // `tokenId` is already at 0x40.
- mstore(0x20, 0x6352211e) // `ownerOf(uint256)`.
- let chainsEq := eq(mload(0x00), cachedChainId)
- let currentOwner :=
- mul(
- mload(0x20),
- and(
- and(gt(returndatasize(), 0x1f), chainsEq),
- staticcall(gas(), tokenContract, 0x3c, 0x24, 0x20, 0x20)
- )
- )
- if iszero(
- or(
- eq(currentOwner, address()),
- and(
- and(chainsEq, eq(tokenContract, caller())),
- eq(mload(0x40), calldataload(0x44))
- )
- )
- ) {
- if iszero(
- and(
- gt(returndatasize(), 0x5f),
- staticcall(gas(), currentOwner, 0x7c, 0x04, 0x00, 0x60)
- )
- ) {
- mstore(0x40, s) // Store `msg.sig`.
- return(0x5c, 0x20) // Return `msg.sig`.
- }
- continue
- }
- mstore(0x00, 0xaed146d3) // `SelfOwnDetected()`.
- revert(0x1c, 0x04)
+ uint256 s = uint256(bytes32(msg.sig)) >> 224;
+ // 0x150b7a02: `onERC721Received(address,address,uint256,bytes)`.
+ if (s == 0x150b7a02) {
+ if (!_hasOwnershipCycle()) {
+ /// @solidity memory-safe-assembly
+ assembly {
+ mstore(0x20, s) // Store `msg.sig`.
+ return(0x3c, 0x20) // Return `msg.sig`.
  }
  }
+ /// @solidity memory-safe-assembly
+ assembly {
+ mstore(0x00, 0xaed146d3) // `SelfOwnDetected()`.
+ revert(0x1c, 0x04)
+ }
+ }
+ /// @solidity memory-safe-assembly
+ assembly {
  // 0xf23a6e61: `onERC1155Received(address,address,uint256,uint256,bytes)`.
  // 0xbc197c81: `onERC1155BatchReceived(address,address,uint256[],uint256[],bytes)`.
  if or(eq(s, 0xf23a6e61), eq(s, 0xbc197c81)) {

test/ERC6551.t.sol

@@ -172,6 +172,16 @@ contract ERC6551Test is SoladyTest {
  t[i].owner, address(t[j].account), t[i].tokenId
  );
  }
+ for (uint256 j; j != i; ++j) {
+ vm.prank(t[i].owner);
+ MockERC721(_erc721).safeTransferFrom(
+ t[i].owner, address(t[j].account), t[i].tokenId
+ );
+ vm.prank(address(t[j].account));
+ MockERC721(_erc721).safeTransferFrom(
+ address(t[j].account), t[i].owner, t[i].tokenId
+ );
+ }
  }
 
  _TestTemps memory u = _testTemps();