diff --git a/contracts/token/TokenVesting.sol b/contracts/token/TokenVesting.sol index 3cf576f5b..6d43edff4 100644 --- a/contracts/token/TokenVesting.sol +++ b/contracts/token/TokenVesting.sol @@ -27,6 +27,7 @@ contract TokenVesting is Ownable { bool revocable; mapping (address => uint256) released; + mapping (address => bool) revoked; /** * @dev Creates a vesting contract that vests its balance of any ERC20 token to the @@ -65,17 +66,22 @@ contract TokenVesting is Ownable { } /** - * @notice Allows the owner to revoke the vesting. Tokens already vested remain in the contract. + * @notice Allows the owner to revoke the vesting. Tokens already vested + * remain in the contract, the rest are returned to the owner. * @param token ERC20 token which is being vested */ function revoke(ERC20Basic token) onlyOwner { require(revocable); + require(!revoked[token]); uint256 balance = token.balanceOf(this); uint256 vested = vestedAmount(token); + uint256 vesting = balance - vested; - token.transfer(owner, balance - vested); + revoked[token] = true; + + token.transfer(owner, vesting); Revoked(); } @@ -87,7 +93,7 @@ contract TokenVesting is Ownable { function vestedAmount(ERC20Basic token) constant returns (uint256) { if (now < cliff) { return 0; - } else if (now >= start + duration) { + } else if (now >= start + duration || revoked[token]) { return token.balanceOf(this); } else { uint256 currentBalance = token.balanceOf(this); diff --git a/test/TokenVesting.js b/test/TokenVesting.js index dee1001a7..c391c6577 100644 --- a/test/TokenVesting.js +++ b/test/TokenVesting.js @@ -80,8 +80,7 @@ contract('TokenVesting', function ([_, owner, beneficiary]) { }); it('should return the non-vested tokens when revoked by owner', async function () { - await increaseTimeTo(this.start + this.cliff + duration.weeks(1)); - await this.vesting.release(this.token.address); + await increaseTimeTo(this.start + this.cliff + duration.weeks(12)); const vested = await this.vesting.vestedAmount(this.token.address); const balance = await this.token.balanceOf(this.vesting.address); @@ -93,15 +92,25 @@ contract('TokenVesting', function ([_, owner, beneficiary]) { }); it('should keep the vested tokens when revoked by owner', async function () { - await increaseTimeTo(this.start + this.cliff + duration.weeks(1)); - await this.vesting.release(this.token.address); + await increaseTimeTo(this.start + this.cliff + duration.weeks(12)); + + const vestedPre = await this.vesting.vestedAmount(this.token.address); + + await this.vesting.revoke(this.token.address, { from: owner }); + + const vestedPost = await this.vesting.vestedAmount(this.token.address); + + vestedPre.should.bignumber.equal(vestedPost); + }); + + it('should fail to be revoked a second time', async function () { + await increaseTimeTo(this.start + this.cliff + duration.weeks(12)); const vested = await this.vesting.vestedAmount(this.token.address); await this.vesting.revoke(this.token.address, { from: owner }); - const balance = await this.token.balanceOf(this.vesting.address); - balance.should.bignumber.equal(vested); + await this.vesting.revoke(this.token.address, { from: owner }).should.be.rejectedWith(EVMThrow); }); });