[ALC-7] add input validation to _initialize and _transferOwnership

alchemyplatform · Sep 20, 2023 · 1468da3 · 1468da3
1 parent 84716cf
commit 1468da3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/src/LightAccount.sol b/src/LightAccount.sol
@@ -82,8 +82,8 @@ contract LightAccount is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Cus
     error ArrayLengthMismatch();
 
     /**
-     * @dev The new owner is not a valid owner (e.g., `address(0)` or the
-     * account itself).
+     * @dev The new owner is not a valid owner (e.g., `address(0)`, the
+     * account itself, or the current owner).
      */
     error InvalidOwner(address owner);
 
@@ -239,6 +239,9 @@ contract LightAccount is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Cus
     }
 
     function _initialize(address anOwner) internal virtual {
+        if (anOwner == address(0)) {
+            revert InvalidOwner(address(0));
+        }
         _getStorage().owner = anOwner;
         emit LightAccountInitialized(_entryPoint, anOwner);
         emit OwnershipTransferred(address(0), anOwner);
@@ -251,6 +254,9 @@ contract LightAccount is BaseAccount, TokenCallbackHandler, UUPSUpgradeable, Cus
     function _transferOwnership(address newOwner) internal virtual {
         LightAccountStorage storage _storage = _getStorage();
         address oldOwner = _storage.owner;
+        if (newOwner == oldOwner) {
+            revert InvalidOwner(newOwner);
+        }
         _storage.owner = newOwner;
         emit OwnershipTransferred(oldOwner, newOwner);
     }

diff --git a/test/LightAccount.t.sol b/test/LightAccount.t.sol
@@ -147,6 +147,12 @@ contract LightAccountTest is Test {
         account = factory.createAccount(eoaAddress, 1);
     }
 
+    function testCannotInitializeWithZeroOwner() public {
+        LightAccountFactory factory = new LightAccountFactory(entryPoint);
+        vm.expectRevert(abi.encodeWithSelector(LightAccount.InvalidOwner.selector, (address(0))));
+        account = factory.createAccount(address(0), 1);
+    }
+
     function testAddDeposit() public {
         assertEq(account.getDeposit(), 0);
         account.addDeposit{value: 10}();
@@ -193,6 +199,12 @@ contract LightAccountTest is Test {
         account.transferOwnership(address(0x100));
     }
 
+    function testCannotTransferOwnershipToCurrentOwner() public {
+        vm.prank(eoaAddress);
+        vm.expectRevert(abi.encodeWithSelector(LightAccount.InvalidOwner.selector, (eoaAddress)));
+        account.transferOwnership(eoaAddress);
+    }
+
     function testCannotTransferOwnershipToZero() public {
         vm.prank(eoaAddress);
         vm.expectRevert(abi.encodeWithSelector(LightAccount.InvalidOwner.selector, (address(0))));