refactor: update UserNonces

test: add DeployTwice, DeployAndUpdateNonce, getBytecodeHash tests
chainrule-labs · Oct 25, 2023 · 3d9c062 · 3d9c062
1 parent 6526a75
commit 3d9c062
Show file tree

Hide file tree

Showing 5 changed files with 120 additions and 23 deletions.
diff --git a/README.md b/README.md
@@ -22,6 +22,12 @@ WE MANDATE THAT ONLY SIGNER AND XSAFE CAN SUBMIT SIGNATURES FOR DEPLOYMENT
 
 DEPLOYMENT REQUIRES SIGNATURE, NOT ADDRESS DIRECTLY
 
--   SALT IS A FUNCTION OF ADDRESS - ANY ONE CAN SUBMIT AN ADDRESS. ONLY PRIVATE KEY HOLDER CAN SUBMIT SIGNATURE
+-   SALT IS A FUNCTION OF ADDRESS - ANYONE CAN SUBMIT AN ADDRESS. ONLY PRIVATE KEYHOLDER CAN SUBMIT SIGNATURE
 -   WITH SIGNATURE METHOD, THE ONLY WAY TO OBTAIN A SALT, TIED TO AN ADDRESS, IS TO OWN THAT ADDRESS' PRIVATE KEY
 -   ONLY RIGHTFUL OWNERS CAN DEPLOY THEIR CONTRACTS AT THEIR DETERMINISTIC ADDRESSES
+
+USER NONCE IS DEPENDENT ON THE BYTECODE OF THE FUNCTION
+-   RATIONALE: USER DOES NOT NEED TO TRACK THE ORDER IN WHICH THEY DEPLOY CONTRACTS A, B, C... ON CHAIN 1 | B, A, C... ON CHAIN 2
+-   IN AN EFFORT TO KEEP THE SIGNABLE TRANSACTION HASH UNIQUE FOR EACH TRANSACTION, BYTECODE IS INCLUDED IN THE HAS
+-   WE CHOSE NOT TO INCLUDE BYTECODE IN THE SALT BECAUSE CREATE2 IS A HASH OF THE BYTECODE AND THE SALT. 
+    -   EVEN IF THE SALT IS THE SAME FOR CONTRACT A AND B, THEY WILL HAVE DIFFERENT BYTECODES AND THUS DIFFERENT ADDRESSES.
diff --git a/script/PredictiveDeployer.sol → script/PredictiveDeployer.s.sol b/script/PredictiveDeployer.sol → script/PredictiveDeployer.s.sol
diff --git a/src/PredictiveDeployer.sol b/src/PredictiveDeployer.sol
@@ -16,7 +16,9 @@ contract PredictiveDeployer is Initializable, UUPSUpgradeable, Ownable {
     bytes32 internal domainSeparator;
 
     // Factory Storage
-    mapping(address => uint256) public userNonces;
+    //mapping(address => uint256) public userNonces;
+    mapping(address => mapping(bytes32 => uint256)) public userNonces;
+    mapping(address => address[]) public deploymentHistory;
 
     // Events
     event Deploy(address indexed principal, address indexed child, bytes32 hashedBytecode, uint256 nonce);
@@ -52,8 +54,12 @@ contract PredictiveDeployer is Initializable, UUPSUpgradeable, Ownable {
      * @param _nonce The principal account's current nonce on this factory contract.
      * @return messageHash The expected signed message hash.
      */
-    function _computeMessageHash(address _principal, uint256 _nonce) internal view returns (bytes32) {
-        bytes32 txHash = getTransactionHash(_principal, _nonce);
+    function _computeMessageHash(address _principal, bytes memory _bytecode, uint256 _nonce)
+        internal
+        view
+        returns (bytes32)
+    {
+        bytes32 txHash = getTransactionHash(_principal, _bytecode, _nonce);
         return keccak256(abi.encodePacked("\x19Ethereum Signed Message:\n32", txHash));
     }
 
@@ -63,8 +69,12 @@ contract PredictiveDeployer is Initializable, UUPSUpgradeable, Ownable {
      * @param _nonce The principal account's current nonce on this factory contract.
      * @return txHash The unique deployment transaction hash to be signed.
      */
-    function getTransactionHash(address _principal, uint256 _nonce) public view returns (bytes32) {
-        return keccak256(abi.encodePacked(domainSeparator, _principal, _nonce));
+    function getTransactionHash(address _principal, bytes memory _bytecode, uint256 _nonce)
+        public
+        view
+        returns (bytes32)
+    {
+        return keccak256(abi.encodePacked(domainSeparator, _principal, _bytecode, _nonce));
     }
 
     /**
@@ -74,11 +84,20 @@ contract PredictiveDeployer is Initializable, UUPSUpgradeable, Ownable {
      * @return child The predicted address of the contract to be deployed.
      */
     function getAddress(address _principal, bytes memory _bytecode) public view returns (address) {
-        uint256 salt = uint256(uint160(_principal)) + userNonces[_principal];
+        uint256 salt = uint256(uint160(_principal)) + userNonces[_principal][keccak256(_bytecode)];
         bytes32 hash = keccak256(abi.encodePacked(bytes1(0xff), address(this), salt, keccak256(_bytecode)));
         return address(uint160(uint256(hash)));
     }
 
+    /**
+     * @dev Returns the keccak256 hash of the provided bytecode.
+     * @param _bytecode The bytecode of the contract to be deployed.
+     * @return hash The keccak256 hash of the provided bytecode.
+     */
+    function getBytecodeHash(bytes memory _bytecode) public pure returns (bytes32) {
+        return keccak256(_bytecode);
+    }
+
     /**
      * @dev Deploys arbitrary child contracts at predictable addresses, derived from account signatures.
      * @param _principal The address of the account that signed the message hash.
@@ -87,8 +106,8 @@ contract PredictiveDeployer is Initializable, UUPSUpgradeable, Ownable {
      * @return child The address of the deployed contract.
      */
     function deploy(address _principal, bytes memory _signature, bytes memory _bytecode) public returns (address) {
-        uint256 currentNonce = userNonces[_principal];
-        bytes32 expectedMessageHash = _computeMessageHash(_principal, currentNonce);
+        uint256 currentNonce = userNonces[_principal][keccak256(_bytecode)];
+        bytes32 expectedMessageHash = _computeMessageHash(_principal, _bytecode, currentNonce);
 
         // Recover principal address
         address recoveredSigner = ECDSA.recover(expectedMessageHash, _signature);
@@ -100,7 +119,7 @@ contract PredictiveDeployer is Initializable, UUPSUpgradeable, Ownable {
         uint256 salt = uint256(uint160(_principal)) + currentNonce;
 
         // Update nonce state
-        userNonces[_principal]++;
+        userNonces[_principal][keccak256(_bytecode)]++;
 
         // Deploy contract
         address child;
@@ -109,6 +128,9 @@ contract PredictiveDeployer is Initializable, UUPSUpgradeable, Ownable {
             if iszero(extcodesize(child)) { revert(0, 0) }
         }
 
+        // Update deployment history
+        deploymentHistory[_principal].push(child);
+
         emit Deploy(_principal, child, keccak256(_bytecode), currentNonce);
         return child;
     }

diff --git a/src/interfaces/IPredictiveDeployer.sol b/src/interfaces/IPredictiveDeployer.sol
@@ -2,8 +2,11 @@
 pragma solidity ^0.8.21;
 
 interface IPredictiveDeployer {
-    function userNonces(address _principal) external returns (uint256);
-    function getTransactionHash(address _principal, uint256 _nonce) external returns (bytes32);
+    function userNonces(address _principal, bytes32 _hashedBytecode) external returns (uint256);
+    function getTransactionHash(address _principal, bytes memory _bytecode, uint256 _nonce)
+        external
+        returns (bytes32);
+    function getBytecodeHash(bytes memory _bytecode) external pure returns (bytes32);
     function getAddress(address _principal, bytes memory _bytecode) external returns (address);
     function deploy(address _principal, bytes memory _signature, bytes memory _bytecode) external returns (address);
 }
diff --git a/test/PredictiveDeployer.t.sol b/test/PredictiveDeployer.t.sol
@@ -16,6 +16,7 @@ contract PredictiveDeployerTest is Test {
     ERC1967Proxy public proxy;
     Child public child;
     bytes public childBytecode;
+    bytes32 public hashedChildBytecode;
 
     // Events
     event Deploy(address indexed sender, address indexed child, bytes32 hashedBytecode, uint256 nonce);
@@ -29,16 +30,18 @@ contract PredictiveDeployerTest is Test {
         // Get Bytecode
         bytes memory bytecode = type(Child).creationCode;
         childBytecode = abi.encodePacked(bytecode, abi.encode(address(this)));
+        hashedChildBytecode = keccak256(childBytecode);
     }
 
     function testFuzz_GetAddress(uint256 pkNum, address sender) public {
         // Setup
         VmSafe.Wallet memory wallet = vm.createWallet(uint256(keccak256(abi.encodePacked(uint256(pkNum)))));
 
-        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr);
+        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr, hashedChildBytecode);
 
         // Get signature information
-        bytes32 txHash = IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, currentNonce);
+        bytes32 txHash =
+            IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, childBytecode, currentNonce);
 
         bytes32 messageHash = keccak256(abi.encodePacked("\x19Ethereum Signed Message:\n32", txHash));
 
@@ -66,10 +69,11 @@ contract PredictiveDeployerTest is Test {
         // Setup
         VmSafe.Wallet memory wallet = vm.createWallet(uint256(keccak256(abi.encodePacked(uint256(pkNum)))));
 
-        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr);
+        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr, hashedChildBytecode);
 
         // Get signature information
-        bytes32 txHash = IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, currentNonce);
+        bytes32 txHash =
+            IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, childBytecode, currentNonce);
 
         bytes32 messageHash = keccak256(abi.encodePacked("\x19Ethereum Signed Message:\n32", txHash));
 
@@ -90,14 +94,42 @@ contract PredictiveDeployerTest is Test {
         assertEq(actualChild, expectedChild);
     }
 
-    function test_CannotDeployReplay(uint256 pkNum) public {
+    // helper function
+    function deployChild(VmSafe.Wallet memory wallet, bytes memory bytecode, bytes32 hashedBytecode)
+        internal
+        returns (address)
+    {
+        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr, hashedBytecode);
+        bytes32 txHash = IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, bytecode, currentNonce);
+        bytes32 messageHash = keccak256(abi.encodePacked("\x19Ethereum Signed Message:\n32", txHash));
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(wallet.privateKey, messageHash);
+        bytes memory signature = abi.encodePacked(r, s, v);
+        return IPredictiveDeployer(address(proxy)).deploy(wallet.addr, signature, bytecode);
+    }
+
+    function testFuzz_DeployTwice(uint256 pkNum) public {
         // Setup
         VmSafe.Wallet memory wallet = vm.createWallet(uint256(keccak256(abi.encodePacked(uint256(pkNum)))));
 
-        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr);
+        // First Deployment
+        address actualChild1 = deployChild(wallet, childBytecode, hashedChildBytecode);
+
+        // Second Deployment
+        address actualChild2 = deployChild(wallet, childBytecode, hashedChildBytecode);
+
+        // Assertions
+        require(actualChild1 != actualChild2, "Addresses should not be equal");
+    }
+
+    function testFuzz_CannotDeployReplay(uint256 pkNum) public {
+        // Setup
+        VmSafe.Wallet memory wallet = vm.createWallet(uint256(keccak256(abi.encodePacked(uint256(pkNum)))));
+
+        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr, hashedChildBytecode);
 
         // Get signature information
-        bytes32 txHash = IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, currentNonce);
+        bytes32 txHash =
+            IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, childBytecode, currentNonce);
 
         bytes32 messageHash = keccak256(abi.encodePacked("\x19Ethereum Signed Message:\n32", txHash));
 
@@ -112,17 +144,18 @@ contract PredictiveDeployerTest is Test {
         IPredictiveDeployer(address(proxy)).deploy(wallet.addr, signature, childBytecode);
     }
 
-    function test_CannotDeployWithoutApproval(uint256 pkNum, address invalidPrincipal) public {
+    function testFuzz_CannotDeployWithoutApproval(uint256 pkNum, address invalidPrincipal) public {
         // Setup
         VmSafe.Wallet memory wallet = vm.createWallet(uint256(keccak256(abi.encodePacked(uint256(pkNum)))));
 
         // Failing condition: principal is not the signer
         vm.assume(invalidPrincipal != wallet.addr);
 
-        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr);
+        uint256 currentNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr, hashedChildBytecode);
 
         // Get signature information
-        bytes32 txHash = IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, currentNonce);
+        bytes32 txHash =
+            IPredictiveDeployer(address(proxy)).getTransactionHash(wallet.addr, childBytecode, currentNonce);
 
         bytes32 messageHash = keccak256(abi.encodePacked("\x19Ethereum Signed Message:\n32", txHash));
 
@@ -134,7 +167,40 @@ contract PredictiveDeployerTest is Test {
         IPredictiveDeployer(address(proxy)).deploy(invalidPrincipal, signature, childBytecode);
     }
 
-    function test_Receive(uint256 transferAmount) public {
+    function testFuzz_DeployAndUpdateNonce(uint256 pkNum) public {
+        // Setup
+        VmSafe.Wallet memory wallet = vm.createWallet(uint256(keccak256(abi.encodePacked(uint256(pkNum)))));
+
+        // Get initial nonce
+        uint256 initialNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr, hashedChildBytecode);
+
+        // Deploy
+        address actualChild = deployChild(wallet, childBytecode, hashedChildBytecode);
+
+        // Get updated nonce
+        uint256 updatedNonce = IPredictiveDeployer(address(proxy)).userNonces(wallet.addr, hashedChildBytecode);
+
+        // Assertions
+        require(actualChild != address(0), "Deployment failed");
+        require(updatedNonce == initialNonce + 1, "Nonce did not update correctly");
+    }
+
+    function testFuzz_getBytecodeHash(bytes memory _childBytecode) public {
+        if (_childBytecode.length == 0) {
+            return; // Skip the test for empty bytecode
+        }
+
+        // Act
+        bytes32 actualHash = IPredictiveDeployer(address(proxy)).getBytecodeHash(_childBytecode);
+
+        // Expected
+        bytes32 expectedHash = keccak256(_childBytecode);
+
+        // Assertions
+        assertEq(actualHash, expectedHash, "Hash mismatch");
+    }
+
+    function testFuzz_Receive(uint256 transferAmount) public {
         // Assumptions
         vm.assume(transferAmount != 0 && transferAmount <= 1000 ether);