diff --git a/contracts/DODOToken/DODOMineV2/BaseMine.sol b/contracts/DODOToken/DODOMineV2/BaseMine.sol index 7a91342..a8d6197 100644 --- a/contracts/DODOToken/DODOMineV2/BaseMine.sol +++ b/contracts/DODOToken/DODOMineV2/BaseMine.sol @@ -183,11 +183,11 @@ contract BaseMine is InitializableOwnable { emit UpdateReward(i, newRewardPerBlock); } - function withdrawLeftOver(uint256 i) external onlyOwner { + function withdrawLeftOver(uint256 i, uint256 amount) external onlyOwner { RewardTokenInfo storage rt = rewardTokenInfos[i]; require(block.number > rt.endBlock, "DODOMineV2: MINING_NOT_FINISHED"); - IRewardVault(rt.rewardVault).withdrawLeftOver(msg.sender); + IRewardVault(rt.rewardVault).withdrawLeftOver(msg.sender,amount); emit WithdrawLeftOver(msg.sender, i); } diff --git a/contracts/DODOToken/DODOMineV2/RewardVault.sol b/contracts/DODOToken/DODOMineV2/RewardVault.sol index 84dc10d..e61dd17 100644 --- a/contracts/DODOToken/DODOMineV2/RewardVault.sol +++ b/contracts/DODOToken/DODOMineV2/RewardVault.sol @@ -14,7 +14,7 @@ import {IERC20} from "../../intf/IERC20.sol"; interface IRewardVault { function reward(address to, uint256 amount) external; - function withdrawLeftOver(address to) external; + function withdrawLeftOver(address to, uint256 amount) external; } contract RewardVault is Ownable { @@ -30,8 +30,9 @@ contract RewardVault is Ownable { IERC20(rewardToken).safeTransfer(to, amount); } - function withdrawLeftOver(address to) external onlyOwner { + function withdrawLeftOver(address to,uint256 amount) external onlyOwner { uint256 leftover = IERC20(rewardToken).balanceOf(address(this)); - IERC20(rewardToken).safeTransfer(to, leftover); + require(amount <= leftover, "VAULT_NOT_ENOUGH"); + IERC20(rewardToken).safeTransfer(to, amount); } }