diff --git a/contracts/token/TokenVesting.sol b/contracts/token/TokenVesting.sol index 3cf576f5b..2d8eabb6a 100644 --- a/contracts/token/TokenVesting.sol +++ b/contracts/token/TokenVesting.sol @@ -1,6 +1,7 @@ pragma solidity ^0.4.11; import './ERC20Basic.sol'; +import './SafeERC20.sol'; import '../ownership/Ownable.sol'; import '../math/Math.sol'; import '../math/SafeMath.sol'; @@ -13,20 +14,22 @@ import '../math/SafeMath.sol'; */ contract TokenVesting is Ownable { using SafeMath for uint256; + using SafeERC20 for ERC20Basic; event Released(uint256 amount); event Revoked(); // beneficiary of tokens after they are released - address beneficiary; + address public beneficiary; - uint256 cliff; - uint256 start; - uint256 duration; + uint256 public cliff; + uint256 public start; + uint256 public duration; - bool revocable; + bool public revocable; - mapping (address => uint256) released; + mapping (address => uint256) public released; + mapping (address => bool) public revoked; /** * @dev Creates a vesting contract that vests its balance of any ERC20 token to the @@ -52,12 +55,12 @@ contract TokenVesting is Ownable { * @notice Transfers vested tokens to beneficiary. * @param token ERC20 token which is being vested */ - function release(ERC20Basic token) { + function release(ERC20Basic token) public { uint256 vested = vestedAmount(token); require(vested > 0); - token.transfer(beneficiary, vested); + token.safeTransfer(beneficiary, vested); released[token] = released[token].add(vested); @@ -65,17 +68,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 { + function revoke(ERC20Basic token) public 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.safeTransfer(owner, vesting); Revoked(); } @@ -84,10 +92,10 @@ contract TokenVesting is Ownable { * @dev Calculates the amount that has already vested but hasn't been released yet. * @param token ERC20 token which is being vested */ - function vestedAmount(ERC20Basic token) constant returns (uint256) { + function vestedAmount(ERC20Basic token) public 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); }); });