Merge pull request #492 from frangio/fix-vesting-revoke

TokenVesting bugfix

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);

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);
   });
 
 });