From f219a5882c91d04c048ebb79ec703b0fb0e0d903 Mon Sep 17 00:00:00 2001
From: Kingter <83567446+kingster-will@users.noreply.github.com>
Date: Tue, 22 Oct 2024 23:31:49 -0700
Subject: [PATCH] chore(wip): prevent transfer to zero or wip (#285)

prevent send wip to zero address or wip contract as well self approval.

issue: resolves #283
---
 contracts/src/token/WIP.sol    | 37 +++++++++++++++++
 contracts/test/token/WIP.t.sol | 75 ++++++++++++++++++++++++++++++++++
 2 files changed, 112 insertions(+)

diff --git a/contracts/src/token/WIP.sol b/contracts/src/token/WIP.sol
index 9cebd47e..718c59b4 100644
--- a/contracts/src/token/WIP.sol
+++ b/contracts/src/token/WIP.sol
@@ -11,6 +11,10 @@ contract WIP is ERC20 {
     event Withdrawal(address indexed to, uint amount);
     /// @notice emitted when a transfer of IP fails
     error IPTransferFailed();
+    /// @notice emitted when an invalid transfer recipient is detected
+    error ERC20InvalidReceiver(address receiver);
+    /// @notice emitted when an invalid transfer spender is detected
+    error ERC20InvalidSpender(address spender);
 
     /// @notice triggered when IP is deposited in exchange for WIP
     receive() external payable {
@@ -46,6 +50,39 @@ contract WIP is ERC20 {
         return "WIP";
     }
 
+    /// @notice approves `spender` to spend `amount` of WIP
+    function approve(address spender, uint256 amount) public override returns (bool) {
+        if (spender == msg.sender) {
+            revert ERC20InvalidSpender(msg.sender);
+        }
+
+        return super.approve(spender, amount);
+    }
+
+    /// @notice transfers `amount` of WIP to a recipient `to`
+    function transfer(address to, uint256 amount) public override returns (bool) {
+        if (to == address(0)) {
+            revert ERC20InvalidReceiver(address(0));
+        }
+        if (to == address(this)) {
+            revert ERC20InvalidReceiver(address(this));
+        }
+
+        return super.transfer(to, amount);
+    }
+
+    /// @notice transfers `amount` of WIP from `from` to a recipient `to`
+    function transferFrom(address from, address to, uint256 amount) public override returns (bool) {
+        if (to == address(0)) {
+            revert ERC20InvalidReceiver(address(0));
+        }
+        if (to == address(this)) {
+            revert ERC20InvalidReceiver(address(this));
+        }
+
+        return super.transferFrom(from, to, amount);
+    }
+
     /// @dev Sets Permit2 contract's allowance to infinity.
     function _givePermit2InfiniteAllowance() internal pure override returns (bool) {
         return true;
diff --git a/contracts/test/token/WIP.t.sol b/contracts/test/token/WIP.t.sol
index a352df85..70830ef4 100644
--- a/contracts/test/token/WIP.t.sol
+++ b/contracts/test/token/WIP.t.sol
@@ -119,5 +119,80 @@ contract WIPTest is Test {
         assertEq(wip.totalSupply(), depositAmount - withdrawAmount);
     }
 
+    function testTransferToZeroAddressReverts() public {
+        address owner = address(0x123);
+
+        vm.deal(owner, 1 ether);
+
+        vm.prank(owner);
+        wip.deposit{ value: 1 ether }();
+
+        assertEq(wip.balanceOf(owner), 1 ether);
+
+        vm.expectRevert(abi.encodeWithSelector(WIP.ERC20InvalidReceiver.selector, address(0)));
+        vm.prank(owner);
+        wip.transfer(address(0), 1 ether);
+    }
+
+    function testTransferToWIPContractReverts() public {
+        address owner = address(0x123);
+
+        vm.deal(owner, 1 ether);
+
+        vm.prank(owner);
+        wip.deposit{ value: 1 ether }();
+
+        assertEq(wip.balanceOf(owner), 1 ether);
+
+        vm.expectRevert(abi.encodeWithSelector(WIP.ERC20InvalidReceiver.selector, address(wip)));
+        vm.prank(owner);
+        wip.transfer(address(wip), 1 ether);
+    }
+
+    function testTransferFromReceiverIsWIPContractReverts() public {
+        address owner = address(0x123);
+
+        vm.deal(owner, 1 ether);
+
+        vm.prank(owner);
+        wip.deposit{ value: 1 ether }();
+
+        assertEq(wip.balanceOf(owner), 1 ether);
+
+        vm.expectRevert(abi.encodeWithSelector(WIP.ERC20InvalidReceiver.selector, address(wip)));
+        vm.prank(owner);
+        wip.transferFrom(owner, address(wip), 1 ether);
+    }
+
+    function testTransferFromReceiverIsZeroAddressReverts() public {
+        address owner = address(0x123);
+
+        vm.deal(owner, 1 ether);
+
+        vm.prank(owner);
+        wip.deposit{ value: 1 ether }();
+
+        assertEq(wip.balanceOf(owner), 1 ether);
+
+        vm.expectRevert(abi.encodeWithSelector(WIP.ERC20InvalidReceiver.selector, address(0)));
+        vm.prank(owner);
+        wip.transferFrom(owner, address(0), 1 ether);
+    }
+
+    function testApprovalToSelfReverts() public {
+        address owner = address(0x123);
+
+        vm.deal(owner, 1 ether);
+
+        vm.prank(owner);
+        wip.deposit{ value: 1 ether }();
+
+        assertEq(wip.balanceOf(owner), 1 ether);
+
+        vm.expectRevert(abi.encodeWithSelector(WIP.ERC20InvalidSpender.selector, owner));
+        vm.prank(owner);
+        wip.approve(owner, 1 ether);
+    }
+
     receive() external payable {}
 }