L2

p2p-org · Oct 22, 2024 · ef5da54 · ef5da54
1 parent 92e6378
commit ef5da54
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 3 deletions.
diff --git a/src/MerkleAirdrop.sol b/src/MerkleAirdrop.sol
@@ -10,6 +10,7 @@ import "./@openzeppelin/utils/Address.sol";
 error MerkleAirdrop__AlreadyClaimed();
 error MerkleAirdrop__InvalidProof();
 error MerkleAirdrop__SameRoot();
+error MerkleAirdrop__ZeroAddress();
 
 contract MerkleAirdrop is Ownable2Step {
     bytes32 public merkleRoot;
@@ -25,9 +26,7 @@ contract MerkleAirdrop is Ownable2Step {
     }
 
     function setMerkleRoot(bytes32 _merkleRoot) external onlyOwner {
-        if (_merkleRoot == merkleRoot) {
-            revert MerkleAirdrop__SameRoot();
-        }
+        require (_merkleRoot != merkleRoot, MerkleAirdrop__SameRoot());
 
         merkleRoot = _merkleRoot;
         emit MerkleAirdrop__MerkleRootSet(_merkleRoot);
@@ -52,6 +51,8 @@ contract MerkleAirdrop is Ownable2Step {
     }
 
     function recoverEther(address payable _recipient) external onlyOwner {
+        require (_recipient != address(0), MerkleAirdrop__ZeroAddress());
+
         emit MerkleAirdrop__EtherRecovered(_recipient, address(this).balance);
         Address.sendValue(_recipient, address(this).balance);
     }

diff --git a/test/TestMerkleAirdrop.sol b/test/TestMerkleAirdrop.sol
@@ -106,6 +106,32 @@ contract TestMerkleAirdrop is Test {
         merkleAirdrop.setMerkleRoot(bytes32(uint256(42)));
         vm.stopPrank();
     }
+
+    function test_RecoverEther() external {
+        vm.startPrank(deployer);
+        merkleAirdrop.transferOwnership(owner);
+        vm.stopPrank();
+
+        vm.startPrank(owner);
+        merkleAirdrop.acceptOwnership();
+
+        uint256 balanceBefore = address(merkleAirdrop).balance;
+        Address.sendValue(payable(address(merkleAirdrop)), TOTAL_ETHER);
+        uint256 balanceAfter = address(merkleAirdrop).balance;
+        assertEq(balanceAfter - balanceBefore, TOTAL_ETHER);
+
+        vm.expectRevert(MerkleAirdrop__ZeroAddress.selector);
+        merkleAirdrop.recoverEther(payable(address(0)));
+
+        uint256 balance42Before = address(42).balance;
+        merkleAirdrop.recoverEther(payable(address(42)));
+        uint256 balance42After = address(42).balance;
+
+        assertEq(address(merkleAirdrop).balance, 0);
+        assertEq(balance42After - balance42Before, TOTAL_ETHER);
+
+        vm.stopPrank();
+    }
 }
 
 struct Claim {