fix: [ALC-LA2-4] revert when transferring to 0 address (#49)

alchemyplatform · Apr 5, 2024 · 90af2ae · 90af2ae
1 parent 3648cb0
commit 90af2ae
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 9 deletions.
diff --git a/src/common/BaseLightAccount.sol b/src/common/BaseLightAccount.sol
@@ -21,14 +21,10 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
         CONTRACT_WITH_ADDR
     }
 
-    /// @dev The length of the array does not match the expected length.
     error ArrayLengthMismatch();
-
-    /// @dev The signature type provided is not valid.
     error InvalidSignatureType();
-
-    /// @dev The caller is not authorized.
     error NotAuthorized(address caller);
+    error ZeroAddressNotAllowed();
 
     modifier onlyAuthorized() {
         _onlyAuthorized();
@@ -89,6 +85,9 @@ abstract contract BaseLightAccount is BaseAccount, TokenCallbackHandler, UUPSUpg
     /// @param withdrawAddress Target to send to.
     /// @param amount Amount to withdraw.
     function withdrawDepositTo(address payable withdrawAddress, uint256 amount) public onlyAuthorized {
+        if (withdrawAddress == address(0)) {
+            revert ZeroAddressNotAllowed();
+        }
         entryPoint().withdrawTo(withdrawAddress, amount);
     }
 

diff --git a/src/common/BaseLightAccountFactory.sol b/src/common/BaseLightAccountFactory.sol
@@ -11,6 +11,7 @@ abstract contract BaseLightAccountFactory is Ownable2Step {
 
     error InvalidAction();
     error TransferFailed();
+    error ZeroAddressNotAllowed();
 
     /// @notice Allow contract to receive native currency.
     receive() external payable {}
@@ -33,6 +34,9 @@ abstract contract BaseLightAccountFactory is Ownable2Step {
     /// @dev Only callable by owner.
     /// @param to Address to send native currency to.
     function withdrawStake(address payable to) external onlyOwner {
+        if (to == address(0)) {
+            revert ZeroAddressNotAllowed();
+        }
         ENTRY_POINT.withdrawStake(to);
     }
 
@@ -42,6 +46,9 @@ abstract contract BaseLightAccountFactory is Ownable2Step {
     /// @param token Address of the token to withdraw, 0 address for native currency.
     /// @param amount Amount of the token to withdraw in case of rebasing tokens.
     function withdraw(address payable to, address token, uint256 amount) external onlyOwner {
+        if (to == address(0)) {
+            revert ZeroAddressNotAllowed();
+        }
         if (token == address(0)) {
             (bool success,) = to.call{value: address(this).balance}("");
             if (!success) {

diff --git a/test/LightAccount.t.sol b/test/LightAccount.t.sol
@@ -276,6 +276,13 @@ contract LightAccountTest is Test {
         account.withdrawDepositTo(BENEFICIARY, 5);
     }
 
+    function testWithdrawDepositToZeroAddress() public {
+        account.addDeposit{value: 10}();
+        vm.prank(eoaAddress);
+        vm.expectRevert(abi.encodeWithSelector(BaseLightAccount.ZeroAddressNotAllowed.selector));
+        account.withdrawDepositTo(payable(address(0)), 5);
+    }
+
     function testOwnerCanTransferOwnership() public {
         address newOwner = address(0x100);
         vm.prank(eoaAddress);
@@ -500,7 +507,7 @@ contract LightAccountTest is Test {
                     abi.encodeCall(
                         account.upgradeToAndCall,
                         (address(newImplementation), abi.encodeCall(SimpleAccount.initialize, (address(this))))
-                        )
+                    )
                 )
             ),
             EOA_PRIVATE_KEY
@@ -549,7 +556,7 @@ contract LightAccountTest is Test {
                     bytes32(uint256(uint160(0x0000000071727De22E5E9d8BAf0edAc6f37da032)))
                 )
             ),
-            0x9861f47fa11c7176759734bb81b2c2c764ce9219f757a6aa77774550dfe37f33
+            0x8199f8205e3258d7c2f19a108ae7d87beba2ee39a352779987b9bb16c13d0763
         );
     }
 

diff --git a/test/LightAccountFactory.t.sol b/test/LightAccountFactory.t.sol
@@ -59,13 +59,25 @@ contract LightAccountFactoryTest is Test {
         assertEq(address(this).balance, 100 ether);
     }
 
+    function testWithdrawStakeToZeroAddress() public {
+        testUnlockStake();
+        vm.expectRevert(BaseLightAccountFactory.ZeroAddressNotAllowed.selector);
+        factory.withdrawStake(payable(address(0)));
+    }
+
     function testWithdraw() public {
         factory.addStake{value: 10 ether}(10 hours, 1 ether);
         assertEq(address(factory).balance, 9 ether);
         factory.withdraw(payable(address(this)), address(0), 0); // amount = balance if native currency
         assertEq(address(factory).balance, 0);
     }
 
+    function testWithdrawToZeroAddress() public {
+        factory.addStake{value: 10 ether}(10 hours, 1 ether);
+        vm.expectRevert(BaseLightAccountFactory.ZeroAddressNotAllowed.selector);
+        factory.withdraw(payable(address(0)), address(0), 0);
+    }
+
     function test2StepOwnershipTransfer() public {
         address owner1 = address(0x200);
         assertEq(factory.owner(), address(this));

diff --git a/test/MultiOwnerLightAccount.t.sol b/test/MultiOwnerLightAccount.t.sol
@@ -209,6 +209,13 @@ contract MultiOwnerLightAccountTest is Test {
         account.execute(address(lightSwitch), 0, abi.encodeCall(LightSwitch.turnOn, ()));
     }
 
+    function testWithdrawDepositToZeroAddress() public {
+        account.addDeposit{value: 10}();
+        vm.prank(eoaAddress);
+        vm.expectRevert(abi.encodeWithSelector(BaseLightAccount.ZeroAddressNotAllowed.selector));
+        account.withdrawDepositTo(payable(address(0)), 5);
+    }
+
     function testExecuteRevertingCallShouldRevertWithSameData() public {
         Reverter reverter = new Reverter();
         vm.prank(eoaAddress);
@@ -658,7 +665,7 @@ contract MultiOwnerLightAccountTest is Test {
                     abi.encodeCall(
                         account.upgradeToAndCall,
                         (address(newImplementation), abi.encodeCall(SimpleAccount.initialize, (address(this))))
-                        )
+                    )
                 )
             ),
             EOA_PRIVATE_KEY
@@ -708,7 +715,7 @@ contract MultiOwnerLightAccountTest is Test {
                     bytes32(uint256(uint160(0x0000000071727De22E5E9d8BAf0edAc6f37da032)))
                 )
             ),
-            0xff5809a60394ff57666c1f7b34605ab8a7fe83c99b833b484ed65f4fdd479c4e
+            0xbb599752b3425b84e8ae3cb828517ab8257400426f76cdb9d522785032b80568
         );
     }
 

diff --git a/test/MultiOwnerLightAccountFactory.t.sol b/test/MultiOwnerLightAccountFactory.t.sol
@@ -135,13 +135,25 @@ contract MultiOwnerLightAccountFactoryTest is Test {
         assertEq(address(this).balance, 100 ether);
     }
 
+    function testWithdrawStakeToZeroAddress() public {
+        testUnlockStake();
+        vm.expectRevert(BaseLightAccountFactory.ZeroAddressNotAllowed.selector);
+        factory.withdrawStake(payable(address(0)));
+    }
+
     function testWithdraw() public {
         factory.addStake{value: 10 ether}(10 hours, 1 ether);
         assertEq(address(factory).balance, 9 ether);
         factory.withdraw(payable(address(this)), address(0), 0); // amount = balance if native currency
         assertEq(address(factory).balance, 0);
     }
 
+    function testWithdrawToZeroAddress() public {
+        factory.addStake{value: 10 ether}(10 hours, 1 ether);
+        vm.expectRevert(BaseLightAccountFactory.ZeroAddressNotAllowed.selector);
+        factory.withdraw(payable(address(0)), address(0), 0);
+    }
+
     function test2StepOwnershipTransfer() public {
         address owner1 = address(0x200);
         assertEq(factory.owner(), address(this));