feat: Add SMA Replay-Safe Hash (#155)

erc6900 · Aug 22, 2024 · a1510cd · a1510cd
1 parent f475848
commit a1510cd
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 15 deletions.
diff --git a/src/account/SemiModularAccount.sol b/src/account/SemiModularAccount.sol
@@ -26,6 +26,14 @@ contract SemiModularAccount is UpgradeableModularAccount {
     uint256 internal constant _SEMI_MODULAR_ACCOUNT_STORAGE_SLOT =
         0x5b9dc9aa943f8fa2653ceceda5e3798f0686455280432166ba472eca0bc17a32;
 
+    // keccak256("EIP712Domain(uint256 chainId,address verifyingContract)")
+    bytes32 private constant _DOMAIN_SEPARATOR_TYPEHASH =
+        0x47e79534a245952e8b16893a336b85a3d9ea9fa8c573f3d803afb92a79469218;
+
+    // keccak256("ReplaySafeHash(bytes32 hash)")
+    bytes32 private constant _REPLAY_SAFE_HASH_TYPEHASH =
+        0x294a8735843d4afb4f017c76faf3b7731def145ed0025fc9b1d5ce30adf113ff;
+
     ModuleEntity internal constant _FALLBACK_VALIDATION = ModuleEntity.wrap(bytes24(type(uint192).max));
 
     uint256 internal constant _SIG_VALIDATION_PASSED = 0;
@@ -88,6 +96,11 @@ contract SemiModularAccount is UpgradeableModularAccount {
         return "erc6900/reference-semi-modular-account/0.8.0";
     }
 
+    function replaySafeHash(bytes32 hash) public view virtual returns (bytes32) {
+        return
+            MessageHashUtils.toTypedDataHash({domainSeparator: _domainSeparator(), structHash: _hashStruct(hash)});
+    }
+
     function _execUserOpValidation(
         ModuleEntity userOpValidationFunction,
         PackedUserOperation memory userOp,
@@ -120,9 +133,9 @@ contract SemiModularAccount is UpgradeableModularAccount {
             if (msg.sender != fallbackSigner) {
                 revert FallbackSignerMismatch();
             }
-            return;
+        } else {
+            super._execRuntimeValidation(runtimeValidationFunction, callData, authorization);
         }
-        super._execRuntimeValidation(runtimeValidationFunction, callData, authorization);
     }
 
     function _exec1271Validation(ModuleEntity sigValidation, bytes32 hash, bytes calldata signature)
@@ -134,7 +147,7 @@ contract SemiModularAccount is UpgradeableModularAccount {
         if (sigValidation.eq(_FALLBACK_VALIDATION)) {
             address fallbackSigner = _getFallbackSigner();
 
-            if (SignatureChecker.isValidSignatureNow(fallbackSigner, hash, signature)) {
+            if (SignatureChecker.isValidSignatureNow(fallbackSigner, replaySafeHash(hash), signature)) {
                 return _1271_MAGIC_VALUE;
             }
             return _1271_INVALID;
@@ -176,11 +189,25 @@ contract SemiModularAccount is UpgradeableModularAccount {
         return address(uint160(bytes20(appendedData)));
     }
 
+    function _domainSeparator() internal view returns (bytes32) {
+        return keccak256(abi.encode(_DOMAIN_SEPARATOR_TYPEHASH, block.chainid, address(this)));
+    }
+
     function _getSemiModularAccountStorage() internal pure returns (SemiModularAccountStorage storage) {
         SemiModularAccountStorage storage _storage;
         assembly ("memory-safe") {
             _storage.slot := _SEMI_MODULAR_ACCOUNT_STORAGE_SLOT
         }
         return _storage;
     }
+
+    function _hashStruct(bytes32 hash) internal pure virtual returns (bytes32) {
+        bytes32 res;
+        assembly ("memory-safe") {
+            mstore(0x00, _REPLAY_SAFE_HASH_TYPEHASH)
+            mstore(0x20, hash)
+            res := keccak256(0, 0x40)
+        }
+        return res;
+    }
 }
diff --git a/src/modules/ReplaySafeWrapper.sol b/src/modules/ReplaySafeWrapper.sol
@@ -26,7 +26,13 @@ abstract contract ReplaySafeWrapper is ModuleEIP712 {
         });
     }
 
-    function _hashStruct(bytes32 hash) internal view virtual returns (bytes32) {
-        return keccak256(abi.encode(_REPLAY_SAFE_HASH_TYPEHASH, hash));
+    function _hashStruct(bytes32 hash) internal pure virtual returns (bytes32) {
+        bytes32 res;
+        assembly ("memory-safe") {
+            mstore(0x00, _REPLAY_SAFE_HASH_TYPEHASH)
+            mstore(0x20, hash)
+            res := keccak256(0x00, 0x40)
+        }
+        return res;
     }
 }
diff --git a/test/account/UpgradeableModularAccount.t.sol b/test/account/UpgradeableModularAccount.t.sol
@@ -402,17 +402,11 @@ contract UpgradeableModularAccountTest is AccountTestBase {
     function test_isValidSignature() public {
         bytes32 message = keccak256("hello world");
 
-        uint8 v;
-        bytes32 r;
-        bytes32 s;
+        bytes32 replaySafeHash = vm.envOr("SMA_TEST", false)
+            ? SemiModularAccount(payable(account1)).replaySafeHash(message)
+            : singleSignerValidationModule.replaySafeHash(address(account1), message);
 
-        if (vm.envOr("SMA_TEST", false)) {
-            // todo: implement replay-safe hashing for SMA
-            (v, r, s) = vm.sign(owner1Key, message);
-        } else {
-            bytes32 replaySafeHash = singleSignerValidationModule.replaySafeHash(address(account1), message);
-            (v, r, s) = vm.sign(owner1Key, replaySafeHash);
-        }
+        (uint8 v, bytes32 r, bytes32 s) = vm.sign(owner1Key, replaySafeHash);
 
         bytes memory signature = _encode1271Signature(_signerValidation, abi.encodePacked(r, s, v));