From 68bfa1efa85b3a1d3a6a0987ded6f62315a69a61 Mon Sep 17 00:00:00 2001
From: Pierrick Turelier <pierrick@turelier.com>
Date: Wed, 2 Mar 2022 10:16:22 -0600
Subject: [PATCH] fix(TWABDelegator): check contract address

---
 contracts/TWABDelegator.sol | 10 ++++++++++
 test/TWABDelegator.test.ts  | 12 ++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/contracts/TWABDelegator.sol b/contracts/TWABDelegator.sol
index 62d5dd1..54afaeb 100644
--- a/contracts/TWABDelegator.sol
+++ b/contracts/TWABDelegator.sol
@@ -187,6 +187,7 @@ contract TWABDelegator is ERC20, LowLevelDelegator, PermitAndMulticall {
    * @param _amount Amount of tickets to stake
    */
   function stake(address _to, uint256 _amount) external {
+    _requireRecipientNotContractAddress(_to);
     _requireAmountGtZero(_amount);
 
     IERC20(ticket).safeTransferFrom(msg.sender, address(this), _amount);
@@ -203,6 +204,7 @@ contract TWABDelegator is ERC20, LowLevelDelegator, PermitAndMulticall {
    */
   function unstake(address _to, uint256 _amount) external {
     _requireRecipientNotZeroAddress(_to);
+    _requireRecipientNotContractAddress(_to);
     _requireAmountGtZero(_amount);
 
     _burn(msg.sender, _amount);
@@ -606,6 +608,14 @@ contract TWABDelegator is ERC20, LowLevelDelegator, PermitAndMulticall {
     require(_amount > 0, "TWABDelegator/amount-gt-zero");
   }
 
+  /**
+   * @notice Require to verify that `_to` is not this contract address.
+   * @param _to Address to check
+   */
+  function _requireRecipientNotContractAddress(address _to) internal view {
+    require(_to != address(this), "TWABDelegator/to-not-this-addr");
+  }
+
   /**
    * @notice Require to verify that `_to` is not address zero.
    * @param _to Address to check
diff --git a/test/TWABDelegator.test.ts b/test/TWABDelegator.test.ts
index 65de819..4a1b0d5 100644
--- a/test/TWABDelegator.test.ts
+++ b/test/TWABDelegator.test.ts
@@ -127,6 +127,12 @@ describe('Test Set Name', () => {
       expect(await twabDelegator.balanceOf(stranger.address)).to.eq(amount);
     });
 
+    it('should fail to stake tickets if recipient is contract address', async () => {
+      await expect(twabDelegator.stake(twabDelegator.address, amount)).to.be.revertedWith(
+        'TWABDelegator/to-not-this-addr',
+      );
+    });
+
     it('should fail to stake tickets if recipient is address zero', async () => {
       await expect(twabDelegator.stake(AddressZero, amount)).to.be.revertedWith(
         'ERC20: mint to the zero address',
@@ -189,6 +195,12 @@ describe('Test Set Name', () => {
       );
     });
 
+    it('should fail to unstake if recipient is the contract address', async () => {
+      await expect(twabDelegator.unstake(twabDelegator.address, amount)).to.be.revertedWith(
+        'TWABDelegator/to-not-this-addr',
+      );
+    });
+
     it('should fail to unstake if amount is zero', async () => {
       await twabDelegator.stake(owner.address, amount);