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
piplabs · Oct 23, 2024 · f219a58 · f219a58
1 parent fe0144d
commit f219a58
Show file tree

Hide file tree

Showing 2 changed files with 112 additions and 0 deletions.
diff --git 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
@@ -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 {}
 }