Fix salt sniping. Lean ThirdwebContract. Test.

Author: jakeloo

contracts/ByocFactory.sol

@@ -5,6 +5,7 @@ pragma solidity ^0.8.11;
 import "@openzeppelin/contracts/utils/Create2.sol";
 import "@openzeppelin/contracts/proxy/Clones.sol";
 import "@openzeppelin/contracts/utils/Address.sol";
+import "@openzeppelin/contracts/utils/Multicall.sol";
 import "@openzeppelin/contracts/access/AccessControlEnumerable.sol";
 import "@openzeppelin/contracts/metatx/ERC2771Context.sol";
 
@@ -13,7 +14,7 @@ import { IByocFactory } from "./interfaces/IByocFactory.sol";
 import { TWRegistry } from "./TWRegistry.sol";
 import "./ThirdwebContract.sol";
 
-contract ByocFactory is IByocFactory, ERC2771Context, AccessControlEnumerable {
+contract ByocFactory is IByocFactory, ERC2771Context, Multicall, AccessControlEnumerable {
     /*///////////////////////////////////////////////////////////////
                             State variables
     //////////////////////////////////////////////////////////////*/
@@ -24,8 +25,8 @@ contract ByocFactory is IByocFactory, ERC2771Context, AccessControlEnumerable {
     /// @dev Whether the registry is paused.
     bool public isPaused;
 
-    /// @dev Empty var used in deployment.
-    address public deployer;
+    /// @dev contract address deployed through the factory => deployer
+    mapping(address => address) public getContractDeployer;
 
     /*///////////////////////////////////////////////////////////////
                     Constructor + modifiers
@@ -56,12 +57,15 @@ contract ByocFactory is IByocFactory, ERC2771Context, AccessControlEnumerable {
         uint256 _value,
         string memory publishMetadataUri
     ) external onlyUnpausedOrAdmin returns (address deployedAddress) {
-        deployer = _msgSender();
-
         require(bytes(publishMetadataUri).length > 0, "No publish metadata");
 
         bytes memory contractBytecode = abi.encodePacked(_contractBytecode, _constructorArgs);
-        bytes32 salt = _salt == "" ? keccak256(abi.encodePacked(_msgSender(), block.number)) : _salt;
+        bytes32 salt = _salt == ""
+            ? keccak256(abi.encodePacked(_msgSender(), block.number, keccak256(contractBytecode)))
+            : keccak256(abi.encodePacked(_msgSender(), _salt));
+
+        address computedContractAddress = Create2.computeAddress(salt, keccak256(contractBytecode), address(this));
+        getContractDeployer[computedContractAddress] = _msgSender();
 
         deployedAddress = Create2.deploy(_value, salt, contractBytecode);
 
@@ -74,8 +78,6 @@ contract ByocFactory is IByocFactory, ERC2771Context, AccessControlEnumerable {
 
         registry.add(_publisher, deployedAddress);
 
-        delete deployer;
-
         emit ContractDeployed(_msgSender(), _publisher, deployedAddress);
     }
 
@@ -88,7 +90,13 @@ contract ByocFactory is IByocFactory, ERC2771Context, AccessControlEnumerable {
         uint256 _value,
         string memory publishMetadataUri
     ) external onlyUnpausedOrAdmin returns (address deployedAddress) {
-        bytes32 salt = _salt == "" ? keccak256(abi.encodePacked(_msgSender(), block.number)) : _salt;
+        bytes32 salt = _salt == ""
+            ? keccak256(abi.encodePacked(_msgSender(), block.number, _implementation, _initializeData))
+            : keccak256(abi.encodePacked(_msgSender(), _salt));
+
+        address computedContractAddress = Clones.predictDeterministicAddress(_implementation, salt, address(this));
+        getContractDeployer[computedContractAddress] = _msgSender();
+
         deployedAddress = Clones.cloneDeterministic(_implementation, salt);
 
         ThirdwebContract(deployedAddress).setPublishMetadataUri(publishMetadataUri);

contracts/ThirdwebContract.sol

@@ -1,19 +1,15 @@
 // SPDX-License-Identifier: Apache-2.0
 pragma solidity ^0.8.0;
 
-interface IDeployer {
-    function deployer() external view returns (address);
+interface IContractDeployer {
+    function getContractDeployer(address _contract) external view returns (address);
 }
 
+error ThirdwebContract_MetadataAlreadyInitialized();
+
 contract ThirdwebContract {
     /// @dev The publish metadata of the contract of which this contract is an instance.
     string private publishMetadataUri;
-    /// @dev Address of the contract deployer.
-    address private deployer;
-
-    constructor() {
-        deployer = IDeployer(msg.sender).deployer();
-    }
 
     /// @dev Returns the publish metadata for this contract.
     function getPublishMetadataUri() external view returns (string memory) {
@@ -22,12 +18,22 @@ contract ThirdwebContract {
 
     /// @dev Initializes the publish metadata and at deploy time.
     function setPublishMetadataUri(string memory uri) external {
-        require(bytes(publishMetadataUri).length == 0, "Published metadata already initialized");
+        if (bytes(publishMetadataUri).length != 0) {
+            revert ThirdwebContract_MetadataAlreadyInitialized();
+        }
         publishMetadataUri = uri;
     }
 
-    /// @dev Returns msg.sender, if caller is not thirdweb factory. Returns the intended msg.sender if caller is factory.
+    /// @dev Enable access to the original contract deployer in the constructor. If this function is called outside of a constructor, it will return address(0) instead.
+    ///      Save 1 storage slot from not storing the factory address and not having to hardcode the factory address.
     function _contractDeployer() internal view returns (address) {
-        return deployer;
+        if (address(this).code.length == 0) {
+            try IContractDeployer(msg.sender).getContractDeployer(address(this)) returns (address deployer) {
+                return deployer;
+            } catch {
+                return address(0);
+            }
+        }
+        return address(0);
     }
 }

src/test/ThirdwebContract.t.sol

@@ -0,0 +1,48 @@
+// SPDX-License-Identifier: Apache-2.0
+pragma solidity ^0.8.11;
+
+// Test imports
+import "./utils/BaseTest.sol";
+import "contracts/TWFactory.sol";
+import "contracts/TWRegistry.sol";
+
+// Helpers
+import "@openzeppelin/contracts/utils/Create2.sol";
+import "@openzeppelin/contracts/proxy/Clones.sol";
+import "contracts/TWProxy.sol";
+import "contracts/ThirdwebContract.sol";
+
+contract MyThirdwebContract is ThirdwebContract {
+    address public contractDeployerFromConstructor;
+
+    constructor() {
+        contractDeployerFromConstructor = _contractDeployer();
+    }
+
+    function getContractDeployerOutsideFromConstructor() public view returns (address) {
+        return _contractDeployer();
+    }
+}
+
+contract ThirdwebContractTest is BaseTest {
+    function setUp() public override {
+        super.setUp();
+    }
+
+    function getContractDeployer(address) public pure returns (address) {
+        return address(42);
+    }
+
+    function deployContract() public returns (address) {
+        return Create2.deploy(0, keccak256(abi.encode(0)), type(MyThirdwebContract).creationCode);
+    }
+
+    function test_ThirdwebContract_ContractDeployerConstructor() external {
+        address addy = deployContract();
+        address contractDeployer = MyThirdwebContract(addy).contractDeployerFromConstructor();
+        address contractDeployerNotConstructor = MyThirdwebContract(addy).getContractDeployerOutsideFromConstructor();
+
+        assertEq(contractDeployer, getContractDeployer(addy));
+        assertEq(contractDeployerNotConstructor, address(0));
+    }
+}