From 33a7089e3ef73f124f3275af739abafe1d56e252 Mon Sep 17 00:00:00 2001 From: Jonathan Diep Date: Thu, 18 Apr 2024 11:52:14 -0700 Subject: [PATCH] fix: Keep referral amounts even after withdrawals --- contracts/Quest.sol | 2 +- test/Quest.t.sol | 14 ++++++++++---- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/contracts/Quest.sol b/contracts/Quest.sol index 3d807a11..6dfe7c49 100644 --- a/contracts/Quest.sol +++ b/contracts/Quest.sol @@ -170,7 +170,7 @@ contract Quest is ReentrancyGuardUpgradeable, PausableUpgradeable, Ownable, IQue uint256 protocolFeeForRecipient = (this.protocolFee() / 2) - referralClaimTotal; rewardToken.safeTransfer(protocolFeeRecipient, protocolFeeForRecipient); - uint256 remainingBalanceForOwner = rewardToken.balanceOf(address(this)); + uint256 remainingBalanceForOwner = rewardToken.balanceOf(address(this)) - referralClaimTotal; rewardToken.safeTransfer(owner(), remainingBalanceForOwner); questFactoryContract.withdrawCallback(questId, protocolFeeRecipient, protocolPayout, address(owner()), ownerPayout); diff --git a/test/Quest.t.sol b/test/Quest.t.sol index a29d5aea..c40964b4 100644 --- a/test/Quest.t.sol +++ b/test/Quest.t.sol @@ -401,9 +401,18 @@ contract TestQuest is Test, TestUtils, Errors, Events { ); } - vm.warp(timestamp); vm.prank(referrer); + + // verify that withdrawals can still work + quest.withdrawRemainingTokens(); + + assertEq( + SampleERC20(rewardTokenAddress).balanceOf(address(quest)), + quest.referralClaimTotal(), + "expected to have referralClaimTotal() amount left inside the contract" + ); + quest.claimReferralFees(referrer); assertEq( @@ -412,9 +421,6 @@ contract TestQuest is Test, TestUtils, Errors, Events { "referrer should claim their allocated referral rewards" ); - // verify that withdrawals can still work - quest.withdrawRemainingTokens(); - uint256 participantBalance = SampleERC20(rewardTokenAddress).balanceOf(participant); uint256 protocolFeeRecipientBalance = SampleERC20(rewardTokenAddress).balanceOf(quest.protocolFeeRecipient()); uint256 ownerBalance = SampleERC20(rewardTokenAddress).balanceOf(owner);