From 722befd2a2be0db9c0693683a9bdcb0fe1a15272 Mon Sep 17 00:00:00 2001
From: Joaquin Gonzalez <jpgonzalezra@gmail.com>
Date: Mon, 20 May 2024 14:04:13 -0300
Subject: [PATCH] test(feat): add test case for recovering account with invalid
 signature

---
 src/SimpleGuardianModule.sol | 19 +++++++++++--------
 test/SimplePlusAccount.t.sol | 14 +++++++++++++-
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/SimpleGuardianModule.sol b/src/SimpleGuardianModule.sol
index da1477e..c869d81 100644
--- a/src/SimpleGuardianModule.sol
+++ b/src/SimpleGuardianModule.sol
@@ -2,7 +2,7 @@
 pragma solidity ^0.8.25;
 
 import { ECDSA } from "@openzeppelin/contracts/utils/cryptography/ECDSA.sol";
-import { console2 } from "forge-std/src/console2.sol";
+// import { console2 } from "forge-std/src/console2.sol";
 
 abstract contract SimpleGuardianModule {
     using ECDSA for bytes32;
@@ -13,6 +13,9 @@ abstract contract SimpleGuardianModule {
     event NonceConsumed(address indexed owner, uint256 idx);
     event GuardianUpdated(address indexed previousGuardian, address indexed newGuardian);
 
+    error InvalidGuardian(address guardian);
+    error InvalidGuardianSignature();
+
     address public guardian;
     mapping(address => uint256) private _nonces;
 
@@ -46,10 +49,10 @@ abstract contract SimpleGuardianModule {
     }
 
     function _updateGuardian(address newGuardian) internal {
-        require(
-            newGuardian != address(0) && guardian != newGuardian && newGuardian != address(this),
-            "Invalid guardian address"
-        );
+        if (newGuardian == address(0) || guardian == newGuardian || newGuardian == address(this)) {
+            revert InvalidGuardian(newGuardian);
+        }
+
         address oldGuardian = guardian;
         guardian = newGuardian;
         emit GuardianUpdated(oldGuardian, newGuardian);
@@ -65,10 +68,10 @@ abstract contract SimpleGuardianModule {
         bytes32 digest = _hashTypedDataV4(structHash);
 
         address recoveredAddress = digest.recover(signature);
-        console2.log(guardian);
-
 
-        require(recoveredAddress == guardian, "Invalid guardian signature");
+        if (recoveredAddress != guardian) {
+            revert InvalidGuardianSignature();
+        }
 
         _transferOwnership(newOwner);
     }
diff --git a/test/SimplePlusAccount.t.sol b/test/SimplePlusAccount.t.sol
index 14071fc..ef4a158 100644
--- a/test/SimplePlusAccount.t.sol
+++ b/test/SimplePlusAccount.t.sol
@@ -167,11 +167,23 @@ contract SimplePlusAccountTest is AccountTest {
             EOA_PRIVATE_KEY,
             address(account)
         );
-        
+
         _executeOperation(op);
         assertEq(account.owner(), newOwner);
     }
 
+    function testRecoverAccountRevertsWithInvalidSignature() public {
+        uint256 nonce = account.getNonce(eoaAddress);
+
+        bytes32 structHash = keccak256(abi.encode(account._RECOVER_TYPEHASH(), eoaAddress, guardianAddress, nonce));
+        bytes32 digest = MessageHashUtils.toTypedDataHash(domainSeparator(address(account)), structHash);
+
+        bytes memory signature = sign(EOA_PRIVATE_KEY, digest);
+
+        vm.expectRevert(abi.encodeWithSelector(SimpleGuardianModule.InvalidGuardianSignature.selector));
+        account.recoverAccount(guardianAddress, nonce, signature);
+    }
+
     function _transferOwnership(address currentOwner, address newOwner) internal {
         vm.prank(currentOwner);
         vm.expectEmit(true, true, false, false);