Update isValidSignature replay protection (#620)

* Don't double hash in smart wallet isValidSignature replay protection * Update _originalMessageHash -> _hash * EIP712 typed data signature for isValidSignature
thirdweb-dev · Feb 27, 2024 · 90ca7b4 · 90ca7b4
1 parent dbe8914
commit 90ca7b4
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 19 deletions.
diff --git a/contracts/prebuilts/account/non-upgradeable/Account.sol b/contracts/prebuilts/account/non-upgradeable/Account.sol
@@ -62,13 +62,18 @@ contract Account is AccountCore, ContractMetadata, ERC1271, ERC721Holder, ERC115
             super.supportsInterface(interfaceId);
     }
 
-    /// @notice See EIP-1271
+    /**
+     *  @notice See EIP-1271
+     *
+     *  @param _hash The original message hash of the data to sign (before mixing this contract's domain separator)
+     *  @param _signature The signature produced on signing the typed data hash (result of `getMessageHash(abi.encode(rawData))`)
+     */
     function isValidSignature(
-        bytes32 _message,
+        bytes32 _hash,
         bytes memory _signature
     ) public view virtual override returns (bytes4 magicValue) {
-        bytes32 messageHash = getMessageHash(abi.encode(_message));
-        address signer = messageHash.recover(_signature);
+        bytes32 targetDigest = getMessageHash(_hash);
+        address signer = targetDigest.recover(_signature);
 
         if (isAdmin(signer)) {
             return MAGICVALUE;
@@ -89,12 +94,13 @@ contract Account is AccountCore, ContractMetadata, ERC1271, ERC721Holder, ERC115
 
     /**
      * @notice Returns the hash of message that should be signed for EIP1271 verification.
-     * @param message Message to be hashed i.e. `keccak256(abi.encode(data))`
-     * @return Hashed message
+     * @param _hash The message hash to sign for the EIP-1271 origin verifying contract.
+     * @return messageHash The digest to sign for EIP-1271 verification.
      */
-    function getMessageHash(bytes memory message) public view returns (bytes32) {
-        bytes32 messageHash = keccak256(abi.encode(MSG_TYPEHASH, keccak256(message)));
-        return keccak256(abi.encodePacked("\x19\x01", _domainSeparatorV4(), messageHash));
+    function getMessageHash(bytes32 _hash) public view returns (bytes32) {
+        bytes32 messageHash = keccak256(abi.encode(_hash));
+        bytes32 typedDataHash = keccak256(abi.encode(MSG_TYPEHASH, messageHash));
+        return keccak256(abi.encodePacked("\x19\x01", _domainSeparatorV4(), typedDataHash));
     }
 
     /*///////////////////////////////////////////////////////////////

diff --git a/contracts/prebuilts/account/utils/AccountExtension.sol b/contracts/prebuilts/account/utils/AccountExtension.sol
@@ -64,13 +64,18 @@ contract AccountExtension is ContractMetadata, ERC1271, AccountPermissions, ERC7
             super.supportsInterface(interfaceId);
     }
 
-    /// @notice See EIP-1271
+    /**
+     *  @notice See EIP-1271
+     *
+     *  @param _hash The original message hash of the data to sign (before mixing this contract's domain separator)
+     *  @param _signature The signature produced on signing the typed data hash (result of `getMessageHash(abi.encode(rawData))`)
+     */
     function isValidSignature(
-        bytes32 _message,
+        bytes32 _hash,
         bytes memory _signature
     ) public view virtual override returns (bytes4 magicValue) {
-        bytes32 messageHash = getMessageHash(abi.encode(_message));
-        address signer = messageHash.recover(_signature);
+        bytes32 targetDigest = getMessageHash(_hash);
+        address signer = targetDigest.recover(_signature);
 
         if (isAdmin(signer)) {
             return MAGICVALUE;
@@ -91,12 +96,13 @@ contract AccountExtension is ContractMetadata, ERC1271, AccountPermissions, ERC7
 
     /**
      * @notice Returns the hash of message that should be signed for EIP1271 verification.
-     * @param message Message to be hashed i.e. `keccak256(abi.encode(data))`
-     * @return Hashed message
+     * @param _hash The message hash to sign for the EIP-1271 origin verifying contract.
+     * @return messageHash The digest to sign for EIP-1271 verification.
      */
-    function getMessageHash(bytes memory message) public view returns (bytes32) {
-        bytes32 messageHash = keccak256(abi.encode(MSG_TYPEHASH, keccak256(message)));
-        return keccak256(abi.encodePacked("\x19\x01", _domainSeparatorV4(), messageHash));
+    function getMessageHash(bytes32 _hash) public view returns (bytes32) {
+        bytes32 messageHash = keccak256(abi.encode(_hash));
+        bytes32 typedDataHash = keccak256(abi.encode(MSG_TYPEHASH, messageHash));
+        return keccak256(abi.encodePacked("\x19\x01", _domainSeparatorV4(), typedDataHash));
     }
 
     /*///////////////////////////////////////////////////////////////

diff --git a/src/test/smart-wallet/AccountVulnPOC.t.sol b/src/test/smart-wallet/AccountVulnPOC.t.sol
@@ -279,7 +279,7 @@ contract SimpleAccountVulnPOCTest is BaseTest {
         // However they can bypass this by using signature verification on number contract instead
         vm.prank(accountSigner);
         bytes32 digest = keccak256(abi.encode(42));
-        bytes32 toSign = SimpleAccount(payable(account)).getMessageHash(abi.encode(digest));
+        bytes32 toSign = SimpleAccount(payable(account)).getMessageHash(digest);
         (uint8 v, bytes32 r, bytes32 s) = vm.sign(accountSignerPKey, toSign);
         bytes memory signature = abi.encodePacked(r, s, v);