From 9d1a83a4f8fd150a008d562347a005d13899e2c7 Mon Sep 17 00:00:00 2001 From: nicecp Date: Mon, 16 Jan 2023 17:17:21 +0800 Subject: [PATCH 1/3] lottery --- contracts/MojitoLotteryPool.sol | 174 +++++++++++++ contracts/interfaces/IMojitoSwapLottery.sol | 92 +++++++ contracts/test/MojitoLotteryMock.sol | 39 +++ test/mojito.lottery.pool.test.js | 262 ++++++++++++++++++++ 4 files changed, 567 insertions(+) create mode 100644 contracts/MojitoLotteryPool.sol create mode 100644 contracts/interfaces/IMojitoSwapLottery.sol create mode 100644 contracts/test/MojitoLotteryMock.sol create mode 100644 test/mojito.lottery.pool.test.js diff --git a/contracts/MojitoLotteryPool.sol b/contracts/MojitoLotteryPool.sol new file mode 100644 index 0000000..2ddea5b --- /dev/null +++ b/contracts/MojitoLotteryPool.sol @@ -0,0 +1,174 @@ +// SPDX-License-Identifier: MIT + +pragma solidity =0.6.12; + +import "@openzeppelin/contracts/utils/ReentrancyGuard.sol"; +import "@openzeppelin/contracts/token/ERC20/SafeERC20.sol"; +import "./interfaces/IMojitoToken.sol"; +import "./interfaces/IMojitoSwapLottery.sol"; +import "./Schedule.sol"; + +contract MojitoLotteryPool is Schedule, ReentrancyGuard { + using SafeMath for uint256; + using SafeERC20 for IERC20; + + // The Mojito token + IMojitoToken public mojito; + uint256 public startBlock; + uint256 public totalAllocPoint = 0; + + address public operator; + + struct PoolInfo { + IMojitoSwapLottery lottery; + uint256 allocPoint; + uint256 lastRewardBlock; + uint256 pendingAmount; // pending MJT + uint256 totalInject; // inject MJT + } + + PoolInfo[] public poolInfo; + mapping(uint256 => mapping(uint256 => uint256)) injectInfo; + + event AdminTokenRecovery(address token, uint256 amount); + event OperatorUpdate(address indexed from, address to); + event InjectPool(uint256 indexed pid, uint256 lotteryId, uint256 amount); + event InjectPending(uint256 indexed pid, uint256 amount); + + constructor( + IMojitoToken _mojito, + uint256 _mojitoPerBlock, + uint256 _startBlock + ) public Schedule(_mojitoPerBlock) { + mojito = _mojito; + startBlock = _startBlock; + } + + modifier onlyOwnerOrOperator() { + require(owner() == _msgSender() || operator == _msgSender(), "not owner or operator"); + _; + } + + function setOperator(address _operator) public onlyOwner { + require(_operator != address(0), "setOperator:zero address"); + address pre = operator; + operator = _operator; + emit OperatorUpdate(pre, operator); + } + + function setMojitoPerBlock(uint256 _mojitoPerBlock) public virtual override onlyOwner { + massUpdatePools(); + super.setMojitoPerBlock(_mojitoPerBlock); + } + + function add(IMojitoSwapLottery _lottery, uint256 _allocPoint, bool withUpdate) public onlyOwner { + checkPoolDuplicate(_lottery); + + if (withUpdate) { + massUpdatePools(); + } + + mojito.approve(address(_lottery), uint256(- 1)); + + totalAllocPoint = totalAllocPoint.add(_allocPoint); + uint256 lastRewardBlock = block.number > startBlock ? block.number : startBlock; + poolInfo.push(PoolInfo({ + lottery : _lottery, + allocPoint : _allocPoint, + lastRewardBlock : lastRewardBlock, + pendingAmount : 0, + totalInject : 0 + })); + } + + function set(uint256 _pid, uint256 _allocPoint) public onlyOwner { + massUpdatePools(); + totalAllocPoint = totalAllocPoint.sub(poolInfo[_pid].allocPoint).add(_allocPoint); + poolInfo[_pid].allocPoint = _allocPoint; + } + + function massUpdatePools() public { + uint256 length = poolInfo.length; + for (uint256 pid = 0; pid < length; ++pid) { + updatePool(pid); + } + } + + function updatePool(uint256 _pid) public { + PoolInfo storage pool = poolInfo[_pid]; + if (block.number <= pool.lastRewardBlock) { + return; + } + + uint256 blockReward = mintable(pool.lastRewardBlock); + uint256 mojitoReward = blockReward.mul(pool.allocPoint).div(totalAllocPoint); + mojito.mint(address(this), mojitoReward); + pool.pendingAmount = pool.pendingAmount.add(mojitoReward); + pool.lastRewardBlock = block.number; + } + + function injectPending(uint256 _pid, uint256 _amount) public { + PoolInfo storage pool = poolInfo[_pid]; + mojito.transferFrom(_msgSender(), address(this), _amount); + pool.pendingAmount = pool.pendingAmount.add(_amount); + emit InjectPending(_pid, _amount); + } + + function injectPool(uint256 _pid, bool withUpdate) public onlyOwnerOrOperator { + if (withUpdate) { + // update pendingAmount + updatePool(_pid); + } + + PoolInfo storage pool = poolInfo[_pid]; + + uint256 currentLotteryId = pool.lottery.viewCurrentLotteryId(); + pool.lottery.injectFunds(currentLotteryId, pool.pendingAmount); + + uint256 prePending = pool.pendingAmount; + pool.totalInject = pool.totalInject.add(pool.pendingAmount); + injectInfo[_pid][currentLotteryId] = injectInfo[_pid][currentLotteryId].add(pool.pendingAmount); + + pool.pendingAmount = 0; + emit InjectPool(_pid, currentLotteryId, prePending); + } + + function checkPoolDuplicate(IMojitoSwapLottery _lottery) internal view { + uint256 length = poolInfo.length; + for (uint256 pid = 0; pid < length; ++pid) { + require(poolInfo[pid].lottery != _lottery, "existing pool"); + } + } + + function poolLength() external view returns (uint256) { + return poolInfo.length; + } + + function withdrawExtraToken() public onlyOwner { + uint256 pending = totalPending(); + uint256 balance = mojito.balanceOf(address(this)); + // balance >= pending + uint256 amount = balance.sub(pending); + mojito.transfer(_msgSender(), amount); + emit AdminTokenRecovery(address(mojito), amount); + } + + function recoverWrongTokens(address _tokenAddress, uint256 _tokenAmount) external onlyOwner { + require(_tokenAddress != address(mojito), "cannot be mojito token"); + + IERC20(_tokenAddress).safeTransfer(_msgSender(), _tokenAmount); + + emit AdminTokenRecovery(_tokenAddress, _tokenAmount); + } + + function totalPending() public view returns (uint256) { + uint256 pending; + uint256 length = poolInfo.length; + for (uint256 pid = 0; pid < length; ++pid) { + pending = pending.add(poolInfo[pid].pendingAmount); + } + + return pending; + } +} + diff --git a/contracts/interfaces/IMojitoSwapLottery.sol b/contracts/interfaces/IMojitoSwapLottery.sol new file mode 100644 index 0000000..4835f5c --- /dev/null +++ b/contracts/interfaces/IMojitoSwapLottery.sol @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: MIT +pragma solidity >=0.6.0; + +interface IMojitoSwapLottery { + /** + * @notice Buy tickets for the current lottery + * @param _lotteryId: lotteryId + * @param _ticketNumbers: array of ticket numbers between 1,000,000 and 1,999,999 + * @dev Callable by users + */ + function buyTickets(uint256 _lotteryId, uint32[] calldata _ticketNumbers) external; + + /** + * @notice Claim a set of winning tickets for a lottery + * @param _lotteryId: lottery id + * @param _ticketIds: array of ticket ids + * @param _brackets: array of brackets for the ticket ids + * @dev Callable by users only, not contract! + */ + function claimTickets( + uint256 _lotteryId, + uint256[] calldata _ticketIds, + uint32[] calldata _brackets + ) external; + + /** + * @notice Close lottery + * @param _lotteryId: lottery id + * @dev Callable by operator + */ + function closeLottery(uint256 _lotteryId) external payable; + + /** + * @notice Draw the final number, calculate reward in MJT per group, and make lottery claimable + * @param _lotteryId: lottery id + * @param _autoInjection: reinjects funds into next lottery (vs. withdrawing all) + * @dev Callable by operator + */ + function drawFinalNumberAndMakeLotteryClaimable(uint256 _lotteryId, bool _autoInjection) external; + + /** + * @notice Inject funds + * @param _lotteryId: lottery id + * @param _amount: amount to inject in MJT token + * @dev Callable by operator + */ + function injectFunds(uint256 _lotteryId, uint256 _amount) external; + + /** + * @notice Start the lottery + * @dev Callable by operator + * @param _endTime: endTime of the lottery + * @param _priceTicketInMJT: price of a ticket in MJT + * @param _discountDivisor: the divisor to calculate the discount magnitude for bulks + * @param _rewardsBreakdown: breakdown of rewards per bracket (must sum to 10,000) + * @param _treasuryFee: treasury fee (10,000 = 100%, 100 = 1%) + */ + function startLottery( + uint256 _endTime, + uint256 _priceTicketInMJT, + uint256 _discountDivisor, + uint256[6] calldata _rewardsBreakdown, + uint256 _treasuryFee + ) external; + + /** + * @notice View current lottery id + */ + function viewCurrentLotteryId() external view returns (uint256); + + /** + * @notice View user ticket ids, numbers, and statuses of user for a given lottery + * @param _user: user address + * @param _lotteryId: lottery id + * @param _cursor: cursor to start where to retrieve the tickets + * @param _size: the number of tickets to retrieve + */ + function viewUserInfoForLotteryId( + address _user, + uint256 _lotteryId, + uint256 _cursor, + uint256 _size + ) + external + view + returns ( + uint256[] memory, + uint32[] memory, + bool[] memory, + uint256 + ); +} diff --git a/contracts/test/MojitoLotteryMock.sol b/contracts/test/MojitoLotteryMock.sol new file mode 100644 index 0000000..edae8cc --- /dev/null +++ b/contracts/test/MojitoLotteryMock.sol @@ -0,0 +1,39 @@ +// SPDX-License-Identifier: MIT + +pragma solidity =0.6.12; + +import "../interfaces/IMojitoToken.sol"; + +contract MojitoLotteryMock { + uint256 lotteryId; + IMojitoToken mojito; + + constructor(IMojitoToken _mojito) public { + mojito = _mojito; + } + + function viewCurrentLotteryId() external view returns (uint256) { + return lotteryId; + } + + function startLottery( + uint256 _endTime, + uint256 _priceTicketInMJT, + uint256 _discountDivisor, + uint256[6] calldata _rewardsBreakdown, + uint256 _treasuryFee + ) external { + _endTime; + _priceTicketInMJT; + _discountDivisor; + uint256 length = _rewardsBreakdown.length; + length; + _treasuryFee; + lotteryId++; + } + + function injectFunds(uint256 _lotteryId, uint256 _amount) external { + _lotteryId; + mojito.transferFrom(address(msg.sender), address(this), _amount); + } +} diff --git a/test/mojito.lottery.pool.test.js b/test/mojito.lottery.pool.test.js new file mode 100644 index 0000000..6c0d6ba --- /dev/null +++ b/test/mojito.lottery.pool.test.js @@ -0,0 +1,262 @@ +/* +* Node Version: V12.3+ +* File Name: mojito.lottery.pool.test +* Author: neo +* Date Created: 2023-01-16 +*/ + +const { + accounts, + contract, + } = require("@openzeppelin/test-environment"); +const { + BN, + expectEvent, + expectRevert, + constants, + time, + } = require("@openzeppelin/test-helpers"); +const {expect} = require("chai"); +const MojitoToken = contract.fromArtifact("MojitoToken"); +const MojitoLotteryPool = contract.fromArtifact("MojitoLotteryPool"); +const MojitoLotteryMock = contract.fromArtifact("MojitoLotteryMock"); + +describe("MojitoLotteryPool", () => { + const [caller, operator, other] = accounts; + const MojitoPerBlock = new BN("3"); + before(async () => { + this.mojito = await MojitoToken.new({from: caller}); + this.erc20 = await MojitoToken.new({from: caller}); + + this.lottery1 = await MojitoLotteryMock.new(this.mojito.address, {from: caller}); + this.lottery2 = await MojitoLotteryMock.new(this.mojito.address, {from: caller}); + + this.lotteryPool = await MojitoLotteryPool.new(this.mojito.address, MojitoPerBlock, 0, {from: caller}); + + await this.erc20.grantRole(await this.mojito.MINTER_ROLE(), caller, {from: caller}); + await this.mojito.grantRole(await this.mojito.MINTER_ROLE(), caller, {from: caller}); + await this.mojito.grantRole(await this.mojito.MINTER_ROLE(), this.lotteryPool.address, {from: caller}); + + await this.erc20.mint(other, new BN("100"), {from: caller}); + await this.mojito.mint(operator, new BN("100"), {from: caller}); + await this.mojito.mint(other, new BN("100"), {from: caller}); + + await this.mojito.approve(this.lotteryPool.address, constants.MAX_UINT256, {from: operator}); + await this.mojito.approve(this.lotteryPool.address, constants.MAX_UINT256, {from: other}); + }); + + it("setOperator(not owner)", async () => { + await expectRevert(this.lotteryPool.setOperator(operator, {from: other}), "Ownable: caller is not the owner"); + }); + + it("setOperator(zero address)", async () => { + await expectRevert(this.lotteryPool.setOperator(constants.ZERO_ADDRESS, {from: caller}), "zero address"); + }); + + it("setOperator()", async () => { + expectEvent( + await this.lotteryPool.setOperator(operator, {from: caller}), + "OperatorUpdate", + { + from: constants.ZERO_ADDRESS, + to: operator, + }, + ); + }); + + it("add(not owner)", async () => { + await expectRevert( + this.lotteryPool.add(this.lottery1.address, new BN("100"), true, {from: other}), + "Ownable: caller is not the owner", + ); + }); + + it("add()", async () => { + await this.lotteryPool.add(this.lottery1.address, new BN("10"), true, {from: caller}); + await this.lotteryPool.add(this.lottery2.address, new BN("100"), true, {from: caller}); + + expect(await this.mojito.allowance(this.lotteryPool.address, this.lottery1.address)).to.be.bignumber.equal(constants.MAX_UINT256); + expect(await this.mojito.allowance(this.lotteryPool.address, this.lottery2.address)).to.be.bignumber.equal(constants.MAX_UINT256); + + expect(await this.lotteryPool.poolLength()).to.be.bignumber.equal(new BN("2")); + expect(await this.lotteryPool.totalAllocPoint()).to.be.bignumber.equal(new BN("110")); + + const pool0 = await this.lotteryPool.poolInfo(0); + const pool1 = await this.lotteryPool.poolInfo(1); + expect(pool0.lottery).to.be.equal(this.lottery1.address); + expect(pool0.allocPoint).to.be.bignumber.equal(new BN("10")); + + expect(pool1.lottery).to.be.equal(this.lottery2.address); + expect(pool1.allocPoint).to.be.bignumber.equal(new BN("100")); + }); + + it("add(existing pool)", async () => { + await expectRevert( + this.lotteryPool.add(this.lottery1.address, new BN("100"), true, {from: caller}), + "existing pool", + ); + }); + + it("set(not owner)", async () => { + await expectRevert( + this.lotteryPool.set(0, new BN("100"), {from: other}), + "Ownable: caller is not the owner", + ); + }); + + it("set()", async () => { + const {receipt} = await this.lotteryPool.set(0, new BN("100"), {from: caller}); + + expect(await this.lotteryPool.totalAllocPoint()).to.be.bignumber.equal(new BN("200")); + const pool = await this.lotteryPool.poolInfo("0"); + expect(pool.lottery).to.be.equal(this.lottery1.address); + expect(pool.allocPoint).to.be.bignumber.equal(new BN("100")); + expect(pool.lastRewardBlock).to.be.bignumber.equal(new BN(receipt.blockNumber)); + }); + + it("poolInfo()", async () => { + const pool0 = await this.lotteryPool.poolInfo(0); + const pool1 = await this.lotteryPool.poolInfo(1); + const balance = await this.mojito.balanceOf(this.lotteryPool.address); + expect(pool0.pendingAmount.add(pool1.pendingAmount)).to.be.bignumber.equal(balance); + }); + + it("injectPending()", async () => { + const poolBefore = await this.lotteryPool.poolInfo(0); + const balanceBefore = await this.mojito.balanceOf(this.lotteryPool.address); + expectEvent( + await this.lotteryPool.injectPending(0, new BN("10"), {from: other}), + "InjectPending", + { + pid: new BN("0"), + amount: new BN("10"), + }, + ); + + const poolAfter = await this.lotteryPool.poolInfo(0); + const balanceAfter = await this.mojito.balanceOf(this.lotteryPool.address); + expect(balanceAfter.sub(balanceBefore)).to.be.bignumber.equal(new BN("10")); + expect(poolAfter.pendingAmount.sub(poolBefore.pendingAmount)).to.be.bignumber.equal(new BN("10")); + }); + + it("injectPool(not owner/operator)", async () => { + await expectRevert( + this.lotteryPool.injectPool(0, true, {from: other}), + "not owner or operator", + ); + }); + + it("injectPool(0, false)", async () => { + const pool = await this.lotteryPool.poolInfo(0); + const lotteryId = await this.lottery1.viewCurrentLotteryId(); + expectEvent( + await this.lotteryPool.injectPool(0, false, {from: caller}), + "InjectPool", + { + pid: new BN("0"), + lotteryId: lotteryId, + amount: pool.pendingAmount, + }, + ); + + const poolAfter = await this.lotteryPool.poolInfo(0); + expect(poolAfter.pendingAmount).to.be.bignumber.equal(new BN("0")); + expect(poolAfter.totalInject).to.be.bignumber.equal(pool.pendingAmount); + }); + + it("injectPool(1, true)", async () => { + const pool = await this.lotteryPool.poolInfo(1); + const lotteryId = await this.lottery2.viewCurrentLotteryId(); + + const {receipt} = await this.lotteryPool.injectPool(1, true, {from: operator}); + const blockReward = MojitoPerBlock.mul(new BN(receipt.blockNumber).sub(pool.lastRewardBlock)); + const mojitoReward = pool.allocPoint.mul(blockReward).div(await this.lotteryPool.totalAllocPoint()); + + expectEvent( + receipt, + "InjectPool", + { + pid: new BN("1"), + lotteryId: lotteryId, + amount: pool.pendingAmount.add(mojitoReward), + }, + ); + + const poolAfter = await this.lotteryPool.poolInfo(1); + expect(poolAfter.pendingAmount).to.be.bignumber.equal(new BN("0")); + expect(poolAfter.totalInject).to.be.bignumber.equal(pool.pendingAmount.add(mojitoReward)); + }); + + it("advanceBlock(+20)", async () => { + await this.mojito.transfer(this.lotteryPool.address, new BN("10"), {from: other}); + await this.erc20.transfer(this.lotteryPool.address, new BN("100"), {from: other}); + + const latest = await time.latestBlock(); + await time.advanceBlock(latest.add(new BN("20"))); + }); + + it("massUpdatePools()", async () => { + const {receipt} = await this.lotteryPool.massUpdatePools(); + + const pool0 = await this.lotteryPool.poolInfo(0); + const pool1 = await this.lotteryPool.poolInfo(1); + + expect(pool0.lastRewardBlock).to.be.bignumber.equal(new BN(receipt.blockNumber)); + expect(pool1.lastRewardBlock).to.be.bignumber.equal(new BN(receipt.blockNumber)); + }); + + it("totalPending()", async () => { + const balance = await this.mojito.balanceOf(this.lotteryPool.address); + + const pool0 = await this.lotteryPool.poolInfo(0); + const pool1 = await this.lotteryPool.poolInfo(1); + + const totalPending = pool0.pendingAmount.add(pool1.pendingAmount); + expect(totalPending).to.be.bignumber.equal(await this.lotteryPool.totalPending()); + expect(totalPending).to.be.bignumber.equal(balance.sub(new BN("10"))); + }); + + it("withdrawExtraToken(not owner)", async () => { + await expectRevert( + this.lotteryPool.withdrawExtraToken({from: operator}), + "Ownable: caller is not the owner", + ); + }); + + it("recoverWrongTokens(not owner)", async () => { + await expectRevert( + this.lotteryPool.recoverWrongTokens(this.erc20.address, 0, {from: operator}), + "Ownable: caller is not the owner", + ); + }); + + it("withdrawExtraToken()", async () => { + expectEvent( + await this.lotteryPool.withdrawExtraToken({from: caller}), + "AdminTokenRecovery", + { + token: this.mojito.address, + amount: new BN("10"), + }, + ); + + const balance = await this.mojito.balanceOf(this.lotteryPool.address); + expect(await this.lotteryPool.totalPending()).to.be.bignumber.equal(balance); + }); + + it("recoverWrongTokens", async () => { + const balance = await this.erc20.balanceOf(this.lotteryPool.address); + expectEvent( + await this.lotteryPool.recoverWrongTokens(this.erc20.address, balance, {from: caller}), + "AdminTokenRecovery", + { + token: this.erc20.address, + amount: balance, + }, + ); + + expect(balance).to.be.bignumber.equal(new BN("100")); + expect(await this.erc20.balanceOf(this.lotteryPool.address)).to.be.bignumber.equal(new BN("0")); + }); +}); + From 2ffa4b3239137ca9f82314516730fb1092519937 Mon Sep 17 00:00:00 2001 From: nicecp Date: Mon, 16 Jan 2023 17:42:39 +0800 Subject: [PATCH 2/3] inject info visible --- contracts/MojitoLotteryPool.sol | 2 +- test/mojito.lottery.pool.test.js | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/contracts/MojitoLotteryPool.sol b/contracts/MojitoLotteryPool.sol index 2ddea5b..c5dc583 100644 --- a/contracts/MojitoLotteryPool.sol +++ b/contracts/MojitoLotteryPool.sol @@ -28,7 +28,7 @@ contract MojitoLotteryPool is Schedule, ReentrancyGuard { } PoolInfo[] public poolInfo; - mapping(uint256 => mapping(uint256 => uint256)) injectInfo; + mapping(uint256 => mapping(uint256 => uint256)) public injectInfo; event AdminTokenRecovery(address token, uint256 amount); event OperatorUpdate(address indexed from, address to); diff --git a/test/mojito.lottery.pool.test.js b/test/mojito.lottery.pool.test.js index 6c0d6ba..c904d13 100644 --- a/test/mojito.lottery.pool.test.js +++ b/test/mojito.lottery.pool.test.js @@ -162,6 +162,7 @@ describe("MojitoLotteryPool", () => { const poolAfter = await this.lotteryPool.poolInfo(0); expect(poolAfter.pendingAmount).to.be.bignumber.equal(new BN("0")); expect(poolAfter.totalInject).to.be.bignumber.equal(pool.pendingAmount); + expect(await this.lotteryPool.injectInfo(0, lotteryId)).to.be.bignumber.equal(pool.pendingAmount); }); it("injectPool(1, true)", async () => { @@ -185,6 +186,7 @@ describe("MojitoLotteryPool", () => { const poolAfter = await this.lotteryPool.poolInfo(1); expect(poolAfter.pendingAmount).to.be.bignumber.equal(new BN("0")); expect(poolAfter.totalInject).to.be.bignumber.equal(pool.pendingAmount.add(mojitoReward)); + expect(await this.lotteryPool.injectInfo(1, lotteryId)).to.be.bignumber.equal(pool.pendingAmount.add(mojitoReward)); }); it("advanceBlock(+20)", async () => { From 53aba9beef5bee74b97e86eadbbfd418d1924ff6 Mon Sep 17 00:00:00 2001 From: nicecp Date: Wed, 18 Jan 2023 13:23:19 +0800 Subject: [PATCH 3/3] injector --- contracts/MojitoLotteryPool.sol | 16 +++++++++++++- test/mojito.lottery.pool.test.js | 36 ++++++++++++++++++++++++++++---- 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/contracts/MojitoLotteryPool.sol b/contracts/MojitoLotteryPool.sol index c5dc583..d5357df 100644 --- a/contracts/MojitoLotteryPool.sol +++ b/contracts/MojitoLotteryPool.sol @@ -18,6 +18,7 @@ contract MojitoLotteryPool is Schedule, ReentrancyGuard { uint256 public totalAllocPoint = 0; address public operator; + address public injector; struct PoolInfo { IMojitoSwapLottery lottery; @@ -32,6 +33,7 @@ contract MojitoLotteryPool is Schedule, ReentrancyGuard { event AdminTokenRecovery(address token, uint256 amount); event OperatorUpdate(address indexed from, address to); + event InjectorUpdate(address indexed from, address to); event InjectPool(uint256 indexed pid, uint256 lotteryId, uint256 amount); event InjectPending(uint256 indexed pid, uint256 amount); @@ -49,6 +51,11 @@ contract MojitoLotteryPool is Schedule, ReentrancyGuard { _; } + modifier onlyInjector() { + require(_msgSender() == injector, "not injector"); + _; + } + function setOperator(address _operator) public onlyOwner { require(_operator != address(0), "setOperator:zero address"); address pre = operator; @@ -56,6 +63,13 @@ contract MojitoLotteryPool is Schedule, ReentrancyGuard { emit OperatorUpdate(pre, operator); } + function setInjector(address _injector) public onlyOwner { + require(_injector != address(0), "setInjector:zero address"); + address pre = injector; + injector = _injector; + emit InjectorUpdate(pre, injector); + } + function setMojitoPerBlock(uint256 _mojitoPerBlock) public virtual override onlyOwner { massUpdatePools(); super.setMojitoPerBlock(_mojitoPerBlock); @@ -107,7 +121,7 @@ contract MojitoLotteryPool is Schedule, ReentrancyGuard { pool.lastRewardBlock = block.number; } - function injectPending(uint256 _pid, uint256 _amount) public { + function injectPending(uint256 _pid, uint256 _amount) public onlyInjector { PoolInfo storage pool = poolInfo[_pid]; mojito.transferFrom(_msgSender(), address(this), _amount); pool.pendingAmount = pool.pendingAmount.add(_amount); diff --git a/test/mojito.lottery.pool.test.js b/test/mojito.lottery.pool.test.js index c904d13..9a8d6f0 100644 --- a/test/mojito.lottery.pool.test.js +++ b/test/mojito.lottery.pool.test.js @@ -22,8 +22,8 @@ const MojitoLotteryPool = contract.fromArtifact("MojitoLotteryPool"); const MojitoLotteryMock = contract.fromArtifact("MojitoLotteryMock"); describe("MojitoLotteryPool", () => { - const [caller, operator, other] = accounts; - const MojitoPerBlock = new BN("3"); + const [caller, operator, injector, other] = accounts; + const MojitoPerBlock = new BN("3"); before(async () => { this.mojito = await MojitoToken.new({from: caller}); this.erc20 = await MojitoToken.new({from: caller}); @@ -39,9 +39,11 @@ describe("MojitoLotteryPool", () => { await this.erc20.mint(other, new BN("100"), {from: caller}); await this.mojito.mint(operator, new BN("100"), {from: caller}); + await this.mojito.mint(injector, new BN("100"), {from: caller}); await this.mojito.mint(other, new BN("100"), {from: caller}); await this.mojito.approve(this.lotteryPool.address, constants.MAX_UINT256, {from: operator}); + await this.mojito.approve(this.lotteryPool.address, constants.MAX_UINT256, {from: injector}); await this.mojito.approve(this.lotteryPool.address, constants.MAX_UINT256, {from: other}); }); @@ -50,7 +52,15 @@ describe("MojitoLotteryPool", () => { }); it("setOperator(zero address)", async () => { - await expectRevert(this.lotteryPool.setOperator(constants.ZERO_ADDRESS, {from: caller}), "zero address"); + await expectRevert(this.lotteryPool.setOperator(constants.ZERO_ADDRESS, {from: caller}), "setOperator:zero address"); + }); + + it("setInjector(not owner)", async () => { + await expectRevert(this.lotteryPool.setInjector(injector, {from: other}), "Ownable: caller is not the owner"); + }); + + it("setInjector(zero address)", async () => { + await expectRevert(this.lotteryPool.setInjector(constants.ZERO_ADDRESS, {from: caller}), "setInjector:zero address"); }); it("setOperator()", async () => { @@ -64,6 +74,17 @@ describe("MojitoLotteryPool", () => { ); }); + it("setInjector()", async () => { + expectEvent( + await this.lotteryPool.setInjector(injector, {from: caller}), + "InjectorUpdate", + { + from: constants.ZERO_ADDRESS, + to: injector, + }, + ); + }); + it("add(not owner)", async () => { await expectRevert( this.lotteryPool.add(this.lottery1.address, new BN("100"), true, {from: other}), @@ -121,11 +142,18 @@ describe("MojitoLotteryPool", () => { expect(pool0.pendingAmount.add(pool1.pendingAmount)).to.be.bignumber.equal(balance); }); + it("injectPending(not injector)", async () => { + await expectRevert( + this.lotteryPool.injectPending(0, new BN("10"), {from: other}), + "not injector", + ); + }); + it("injectPending()", async () => { const poolBefore = await this.lotteryPool.poolInfo(0); const balanceBefore = await this.mojito.balanceOf(this.lotteryPool.address); expectEvent( - await this.lotteryPool.injectPending(0, new BN("10"), {from: other}), + await this.lotteryPool.injectPending(0, new BN("10"), {from: injector}), "InjectPending", { pid: new BN("0"),