Change require statement to custom error. Add unit test

NexusMutual · Feb 3, 2023 · 40bc311 · 40bc311
1 parent 896dbcd
commit 40bc311
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 12 deletions.
diff --git a/contracts/interfaces/IStakingPool.sol b/contracts/interfaces/IStakingPool.sol
@@ -158,6 +158,7 @@ interface IStakingPool {
 
   // Staking NFTs
   error NotTokenOwnerOrApproved();
+  error TokenDoesNotBelongToPool();
 
   // Tranche & capacity
   error NewTrancheEndsBeforeInitialTranche();

diff --git a/contracts/modules/staking/StakingPool.sol b/contracts/modules/staking/StakingPool.sol
@@ -179,6 +179,14 @@ contract StakingPool is IStakingPool, Multicall {
     _;
   }
 
+  // validate token id exists and belongs to this pool
+  // stakingPoolOf() reverts for non-existent tokens
+  function requireTokenBelongsToPool(uint tokenId) internal view {
+    if (stakingNFT.stakingPoolOf(tokenId) != poolId) {
+      revert TokenDoesNotBelongToPool();
+    }
+  }
+
   constructor (
     address _stakingNFT,
     address _token,
@@ -414,12 +422,7 @@ contract StakingPool is IStakingPool, Multicall {
       address to = destination == address(0) ? msg.sender : destination;
       tokenId = stakingNFT.mint(poolId, to);
     } else {
-      // validate token id exists and belongs to this pool
-      // stakingPoolOf() reverts for non-existent tokens
-      require(
-        stakingNFT.stakingPoolOf(requestTokenId) == poolId,
-        "StakingPool: Token does not belong to this pool"
-      );
+      requireTokenBelongsToPool(requestTokenId);
       tokenId = requestTokenId;
     }
 
@@ -1086,12 +1089,7 @@ contract StakingPool is IStakingPool, Multicall {
     uint topUpAmount
   ) external whenNotPaused whenNotHalted {
 
-    // validate token id exists and belongs to this pool
-    // stakingPoolOf() reverts for non-existent tokens
-    require(
-      stakingNFT.stakingPoolOf(tokenId) == poolId,
-      "StakingPool: Token does not belong to this pool"
-    );
+    requireTokenBelongsToPool(tokenId);
 
     if (!stakingNFT.isApprovedOrOwner(msg.sender, tokenId)) {
       revert NotTokenOwnerOrApproved();

diff --git a/test/unit/StakingPool/depositTo.js b/test/unit/StakingPool/depositTo.js
@@ -644,4 +644,37 @@ describe('depositTo', function () {
       'NxmIsLockedForGovernanceVote',
     );
   });
+
+  it('should revert if trying to deposit with token from other pool', async function () {
+    const { stakingPool, nxm, stakingNFT, cover, tokenController, master } = this;
+    const [user, manager] = this.accounts.members;
+
+    const { amount, tokenId, destination } = depositToFixture;
+    const { poolId, initialPoolFee, maxPoolFee, products, ipfsDescriptionHash } = poolInitParams;
+
+    const { firstActiveTrancheId } = await getTranches(DEFAULT_PERIOD, DEFAULT_GRACE_PERIOD);
+    const otherStakingPool = await ethers.deployContract('StakingPool', [
+      stakingNFT.address,
+      nxm.address,
+      cover.address,
+      tokenController.address,
+      master.address,
+    ]);
+    await stakingPool.connect(this.coverSigner).initialize(
+      manager.address,
+      false, // isPrivatePool
+      initialPoolFee,
+      maxPoolFee,
+      products,
+      poolId + 1,
+      ipfsDescriptionHash,
+    );
+    await nxm.mint(user.address, amount.mul(2));
+    await nxm.approve(stakingPool.address, amount);
+    await nxm.approve(otherStakingPool.address, amount);
+    await stakingPool.connect(user).depositTo(amount, firstActiveTrancheId, tokenId, destination);
+    await expect(
+      otherStakingPool.connect(user).depositTo(amount, firstActiveTrancheId, 1, destination),
+    ).to.be.revertedWithCustomError(otherStakingPool, 'TokenDoesNotBelongToPool');
+  });
 });