fix: Add error handling for invalid new owner

jpgonzalezra · May 20, 2024 · cc78cf7 · cc78cf7
1 parent 722befd
commit cc78cf7
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/src/SimpleGuardianModule.sol b/src/SimpleGuardianModule.sol
@@ -13,6 +13,7 @@ abstract contract SimpleGuardianModule {
     event NonceConsumed(address indexed owner, uint256 idx);
     event GuardianUpdated(address indexed previousGuardian, address indexed newGuardian);
 
+    error InvalidNewOwner(address owner);
     error InvalidGuardian(address guardian);
     error InvalidGuardianSignature();
 
@@ -26,8 +27,6 @@ abstract contract SimpleGuardianModule {
 
     /**
      * @notice Retuns a nonce for a given address.
-     * @param   from  Address.
-     * @return  uint256 Nonce Value.
      */
     function getNonce(address from) external view virtual returns (uint256) {
         return _nonces[from];
@@ -59,9 +58,9 @@ abstract contract SimpleGuardianModule {
     }
 
     function recoverAccount(address newOwner, uint256 nonce, bytes calldata signature) external {
-        require(
-            newOwner != address(0) && _owner() != newOwner && newOwner != address(this), "Invalid new owner address"
-        );
+        if (newOwner == address(0) || _owner() == newOwner || newOwner == address(this)) {
+            revert InvalidNewOwner(newOwner);
+        }
 
         _verifyAndConsumeNonce(newOwner, nonce);
         bytes32 structHash = keccak256(abi.encode(_RECOVER_TYPEHASH, _owner(), newOwner, nonce));

diff --git a/test/SimplePlusAccount.t.sol b/test/SimplePlusAccount.t.sol
@@ -32,6 +32,7 @@ contract SimplePlusAccountTest is AccountTest {
         SimplePlusAccountFactory factory = new SimplePlusAccountFactory(entryPoint);
         account = factory.createAccount(eoaAddress, 1);
         vm.deal(address(account), 1 << 128);
+        emit GuardianUpdated(address(0), guardianAddress);
         account.initGuardian(guardianAddress);
     }
 
@@ -110,10 +111,7 @@ contract SimplePlusAccountTest is AccountTest {
      * Guardian test cases
      */
     function testGuardianCanTransferOwnership() public {
-        vm.prank(guardianAddress);
-        emit GuardianUpdated(address(0), guardianAddress);
         uint256 nonce = account.getNonce(eoaAddress);
-
         address newOwner = address(0x100);
         bytes32 structHash = keccak256(abi.encode(account._RECOVER_TYPEHASH(), eoaAddress, newOwner, nonce));
         bytes32 digest = MessageHashUtils.toTypedDataHash(domainSeparator(address(account)), structHash);
@@ -126,9 +124,19 @@ contract SimplePlusAccountTest is AccountTest {
         assertEq(account.owner(), newOwner);
     }
 
-    function testEntryPointExecutesRecoverAccountByGuardian() public {
-        vm.prank(guardianAddress);
+    function testGuardianCannotTransferOwnershipPerInvalidNewOwner() public {
+        uint256 nonce = account.getNonce(eoaAddress);
+
+        bytes32 structHash = keccak256(abi.encode(account._RECOVER_TYPEHASH(), eoaAddress, eoaAddress, nonce));
+        bytes32 digest = MessageHashUtils.toTypedDataHash(domainSeparator(address(account)), structHash);
+
+        bytes memory signature = sign(GUARDIAN_PRIVATE_KEY, digest);
 
+        vm.expectRevert(abi.encodeWithSelector(SimpleGuardianModule.InvalidNewOwner.selector, eoaAddress));
+        account.recoverAccount(eoaAddress, nonce, signature);
+    }
+
+    function testEntryPointExecutesRecoverAccountByGuardian() public {
         uint256 nonce = account.getNonce(eoaAddress);
         address newOwner = address(0x100);