From 4d2163e1907b43f080afe9d689a8950cb8b7187d 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 | 11 +++++++++++
 test/TWABDelegator.test.ts  | 18 ++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/contracts/TWABDelegator.sol b/contracts/TWABDelegator.sol
index 62d5dd1..b6b2d50 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);
@@ -383,6 +385,7 @@ contract TWABDelegator is ERC20, LowLevelDelegator, PermitAndMulticall {
     address _to
   ) external returns (Delegation) {
     _requireRecipientNotZeroAddress(_to);
+    _requireRecipientNotContractAddress(_to);
 
     Delegation _delegation = Delegation(_computeAddress(msg.sender, _slot));
     _transfer(_delegation, _to, _amount);
@@ -606,6 +609,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..2758179 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);
 
@@ -732,6 +744,12 @@ describe('Test Set Name', () => {
       );
     });
 
+    it('should fail to transfer tickets from a delegation if recipient is the contract address', async () => {
+      await expect(twabDelegator.transferDelegationTo(0, amount, twabDelegator.address)).to.be.revertedWith(
+        'TWABDelegator/to-not-this-addr',
+      );
+    });
+
     it('should fail to transfer tickets from an inexistent delegation', async () => {
       await expect(twabDelegator.transferDelegationTo(1, amount, owner.address)).to.be.revertedWith(
         'Transaction reverted: function call to a non-contract account',