Merge pull request #17 from alchemyplatform/howy/follow-safe

chore: copy safe's signing implementation
alchemyplatform · Nov 6, 2023 · 08601d9 · 08601d9
2 parents 415ef44 + 6119eac
commit 08601d9
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 6 deletions.
diff --git a/script/Deploy_LightAccountFactory.s.sol b/script/Deploy_LightAccountFactory.s.sol
@@ -26,7 +26,7 @@ contract Deploy_LightAccountFactory is Script {
             abi.encodePacked(type(LightAccountFactory).creationCode, bytes32(uint256(uint160(address(entryPoint)))))
         );
 
-        if (initCodeHash != 0x3043a72812fec9b9987853a9b869c1a469dc6e04b0f80da3af2ecb8cf8eed209) {
+        if (initCodeHash != 0x2ad62a8bb3850247ef0c4f04e30b584e6eee7caa0e063745e90956653b90eb84) {
             revert InitCodeHashMismatch(initCodeHash);
         }
 

diff --git a/src/LightAccount.sol b/src/LightAccount.sol
@@ -58,6 +58,11 @@ contract LightAccount is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Cus
     // bytes4(keccak256("isValidSignature(bytes32,bytes)"))
     bytes4 internal constant _1271_MAGIC_VALUE = 0x1626ba7e;
     IEntryPoint private immutable _entryPoint;
+    // keccak256("EIP712Domain(string name,string version,uint256 chainId,address verifyingContract)")
+    bytes32 private constant DOMAIN_SEPARATOR_TYPEHASH =
+        0x8b73c3c69bb8fe3d512ecc4cf759cc79239f7b179b0ffacaa9a75d522b39400f;
+    // keccak256("LightAccountMessage(bytes message)");
+    bytes32 private constant LA_MSG_TYPEHASH = 0x5e3baca2936049843f06038876a12f03627b5edc98025751ecf2ac7562640199;
 
     struct LightAccountStorage {
         address owner;
@@ -224,6 +229,41 @@ contract LightAccount is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Cus
         return entryPoint().balanceOf(address(this));
     }
 
+    /**
+     * @notice Returns the domain separator for this contract, as defined in the EIP-712 standard.
+     * @return bytes32 The domain separator hash.
+     */
+    function domainSeparator() public view returns (bytes32) {
+        return keccak256(
+            abi.encode(
+                DOMAIN_SEPARATOR_TYPEHASH,
+                abi.encode("LightAccount"), // name
+                abi.encode("1"), // version
+                block.chainid, // chainId
+                address(this) // verifying contract
+            )
+        );
+    }
+
+    /**
+     * @notice Returns the pre-image of the message hash
+     * @param message Message that should be encoded.
+     * @return Encoded message.
+     */
+    function encodeMessageData(bytes memory message) public view returns (bytes memory) {
+        bytes32 messageHash = keccak256(abi.encode(LA_MSG_TYPEHASH, keccak256(message)));
+        return abi.encodePacked("\x19\x01", domainSeparator(), messageHash);
+    }
+
+    /**
+     * @notice Returns hash of a message that can be signed by owners.
+     * @param message Message that should be hashed.
+     * @return Message hash.
+     */
+    function getMessageHash(bytes memory message) public view returns (bytes32) {
+        return keccak256(encodeMessageData(message));
+    }
+
     /**
      * @dev The signature is valid if it is signed by the owner's private key
      * (if the owner is an EOA) or if it is a valid ERC-1271 signature from the
@@ -234,7 +274,9 @@ contract LightAccount is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Cus
      * @inheritdoc IERC1271
      */
     function isValidSignature(bytes32 digest, bytes memory signature) public view override returns (bytes4) {
-        if (SignatureChecker.isValidSignatureNow(owner(), digest, signature)) {
+        bytes memory messageData = encodeMessageData(abi.encode(digest));
+        bytes32 messageHash = keccak256(messageData);
+        if (SignatureChecker.isValidSignatureNow(owner(), messageHash, signature)) {
             return _1271_MAGIC_VALUE;
         }
         return 0xffffffff;

diff --git a/test/LightAccount.t.sol b/test/LightAccount.t.sol
@@ -223,20 +223,20 @@ contract LightAccountTest is Test {
 
     function testIsValidSignatureForEoaOwner() public {
         bytes32 digest = keccak256("digest");
-        bytes memory signature = _sign(EOA_PRIVATE_KEY, digest);
+        bytes memory signature = _sign(EOA_PRIVATE_KEY, account.getMessageHash(abi.encode(digest)));
         assertEq(account.isValidSignature(digest, signature), bytes4(keccak256("isValidSignature(bytes32,bytes)")));
     }
 
     function testIsValidSignatureForContractOwner() public {
         _useContractOwner();
         bytes32 digest = keccak256("digest");
-        bytes memory signature = contractOwner.sign(digest);
+        bytes memory signature = contractOwner.sign(account.getMessageHash(abi.encode(digest)));
         assertEq(account.isValidSignature(digest, signature), bytes4(keccak256("isValidSignature(bytes32,bytes)")));
     }
 
     function testIsValidSignatureRejectsInvalid() public {
         bytes32 digest = keccak256("digest");
-        bytes memory signature = _sign(123, digest);
+        bytes memory signature = _sign(123, account.getMessageHash(abi.encode(digest)));
         assertEq(account.isValidSignature(digest, signature), bytes4(0xffffffff));
     }
 
@@ -285,7 +285,7 @@ contract LightAccountTest is Test {
                     bytes32(uint256(uint160(0x5FF137D4b0FDCD49DcA30c7CF57E578a026d2789)))
                 )
             ),
-            0x3043a72812fec9b9987853a9b869c1a469dc6e04b0f80da3af2ecb8cf8eed209
+            0x2ad62a8bb3850247ef0c4f04e30b584e6eee7caa0e063745e90956653b90eb84
         );
     }