Add ability to accept ether payment

EthSign · Apr 28, 2024 · f820548 · f820548
1 parent 1bf727b
commit f820548
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 3 deletions.
diff --git a/contracts/core/TTUFeeCollector.sol b/contracts/core/TTUFeeCollector.sol
@@ -18,8 +18,16 @@ contract TTUFeeCollector is ITTUFeeCollector, Ownable {
 
     constructor() Ownable(_msgSender()) {}
 
+    receive() external payable {}
+
     function withdrawFee(IERC20 token, uint256 amount) external onlyOwner {
-        token.safeTransfer(owner(), amount);
+        if (address(token) == address(0)) {
+            (bool success, bytes memory data) = owner().call{value: amount}("");
+            // solhint-disable-next-line custom-errors
+            require(success, string(data));
+        } else {
+            token.safeTransfer(owner(), amount);
+        }
     }
 
     function setDefaultFeeBips(uint256 bips) external onlyOwner {

diff --git a/test/TTUFeeCollector.t.sol b/test/TTUFeeCollector.t.sol
@@ -23,6 +23,8 @@ contract TTUFeeCollectorTest is Test {
     // Ownable
     error OwnableUnauthorizedAccount(address account);
 
+    receive() external payable {}
+
     function setUp() public {
         instance = new TTUFeeCollector();
         mockErc20 = IERC20(address(new MockERC20()));
@@ -130,14 +132,21 @@ contract TTUFeeCollectorTest is Test {
         }
     }
 
-    function testFuzz_withdrawFee_success(uint256 amount) public {
+    function testFuzz_withdrawFee_success_erc20(uint256 amount) public {
         MockERC20(address(mockErc20)).mint(address(this), amount);
         mockErc20.safeTransfer(address(instance), amount);
         uint256 balanceBefore = mockErc20.balanceOf(address(this));
         instance.withdrawFee(IERC20(mockErc20), amount);
         assertEq(mockErc20.balanceOf(address(this)) - balanceBefore, amount);
     }
 
+    function testFuzz_withdrawFee_success_ether(uint128 amount) public {
+        vm.deal(address(instance), amount);
+        uint256 balanceBefore = address(this).balance;
+        instance.withdrawFee(IERC20(address(0)), amount);
+        assertEq(address(this).balance - balanceBefore, amount);
+    }
+
     function testFuzz_withdrawFee_fail_notOwner(
         address notOwner,
         uint256 amount
@@ -159,7 +168,9 @@ contract TTUFeeCollectorTest is Test {
         address notMockErc20,
         uint256 amount
     ) public {
-        vm.assume(notMockErc20 != address(mockErc20));
+        vm.assume(
+            notMockErc20 != address(mockErc20) && notMockErc20 != address(0)
+        );
         MockERC20(address(mockErc20)).mint(address(this), amount);
         mockErc20.safeTransfer(address(instance), amount);
         vm.expectRevert();