From f123b32d7ac9383fe400a43f4970476c4d6d7a5f Mon Sep 17 00:00:00 2001
From: Daniel Kronovet <kronovet@gmail.com>
Date: Thu, 25 Feb 2021 14:28:06 -0800
Subject: [PATCH] Release locks properly in VotingToken

---
 contracts/extensions/VotingBase.sol  |  2 ++
 contracts/extensions/VotingToken.sol | 14 +++++++-------
 test/extensions/voting-token.js      | 29 ++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/contracts/extensions/VotingBase.sol b/contracts/extensions/VotingBase.sol
index 2e3861fb3c..eafcc51cac 100644
--- a/contracts/extensions/VotingBase.sol
+++ b/contracts/extensions/VotingBase.sol
@@ -238,6 +238,8 @@ abstract contract VotingBase is ColonyExtension {
       emit MotionEventSet(_motionId, REVEAL_END);
     }
 
+    postReveal(_motionId, msg.sender);
+
     tokenLocking.transfer(token, voterReward, msg.sender, true);
   }
 
diff --git a/contracts/extensions/VotingToken.sol b/contracts/extensions/VotingToken.sol
index cf05538adc..17519fdeae 100644
--- a/contracts/extensions/VotingToken.sol
+++ b/contracts/extensions/VotingToken.sol
@@ -41,7 +41,7 @@ contract VotingToken is VotingBase {
   mapping (uint256 => uint256) totalInfluences;
 
   // [motionId] => lockId
-  mapping (uint256 => uint256) locks;
+  mapping (uint256 => uint256) lockIds;
 
   // Public
 
@@ -57,15 +57,15 @@ contract VotingToken is VotingBase {
   }
 
   function postReveal(uint256 _motionId, address _user) internal override {
-    colony.unlockTokenForUser(_user, locks[_motionId]);
+    colony.unlockTokenForUser(_user, lockIds[_motionId]);
   }
 
   function postClaim(uint256 _motionId, address _user) internal override {
     uint256 lockCount = tokenLocking.getUserLock(token, _user).lockCount;
 
     // Lock may have already been released during reveal
-    if (lockCount < locks[_motionId]) {
-      colony.unlockTokenForUser(_user, locks[_motionId]);
+    if (lockCount < lockIds[_motionId]) {
+      colony.unlockTokenForUser(_user, lockIds[_motionId]);
     }
   }
 
@@ -77,7 +77,7 @@ contract VotingToken is VotingBase {
   {
     createMotion(_altTarget, _action, 1);
     motions[motionCount].maxVotes = ERC20Extended(token).totalSupply();
-    locks[motionCount] = colony.lockToken();
+    lockIds[motionCount] = colony.lockToken();
   }
 
   /// @notice Stake on a motion
@@ -109,8 +109,8 @@ contract VotingToken is VotingBase {
     internalSubmitVote(_motionId, _voteSecret);
   }
 
-  function getLock(uint256 _motionId) public view returns (uint256) {
-    return locks[_motionId];
+  function getLockId(uint256 _motionId) public view returns (uint256) {
+    return lockIds[_motionId];
   }
 
   // Internal functions
diff --git a/test/extensions/voting-token.js b/test/extensions/voting-token.js
index 3efbade8b4..67e1362e4d 100644
--- a/test/extensions/voting-token.js
+++ b/test/extensions/voting-token.js
@@ -245,6 +245,15 @@ contract("Voting Token", (accounts) => {
       expect(motion.skillId).to.eq.BN(domain1.skillId);
     });
 
+    it("can lock the token when a motion is created", async () => {
+      const action = await encodeTxData(colony, "makeTask", [1, UINT256_MAX, FAKE, 1, 0, 0]);
+      await voting.createRootMotion(ADDRESS_ZERO, action);
+      const motionId = await voting.getMotionCount();
+
+      const lockId = await voting.getLockId(motionId);
+      expect(lockId).to.not.be.zero;
+    });
+
     it("can create a motion with an alternative target", async () => {
       const action = await encodeTxData(colony, "makeTask", [1, 0, FAKE, 2, 0, 0]);
       await voting.createRootMotion(voting.address, action);
@@ -635,6 +644,18 @@ contract("Voting Token", (accounts) => {
       await voting.revealVote(motionId, SALT, NAY, { from: USER0 });
     });
 
+    it("can unlock the token once revealed", async () => {
+      await voting.submitVote(motionId, soliditySha3(SALT, NAY), { from: USER0 });
+
+      await forwardTime(SUBMIT_PERIOD, this);
+
+      await voting.revealVote(motionId, SALT, NAY, { from: USER0 });
+
+      const lockId = await voting.getLockId(motionId);
+      const { lockCount } = await tokenLocking.getUserLock(token.address, USER0);
+      expect(lockCount).to.eq.BN(lockId);
+    });
+
     it("can tally votes from two users", async () => {
       await voting.submitVote(motionId, soliditySha3(SALT, YAY), { from: USER0 });
       await voting.submitVote(motionId, soliditySha3(SALT, YAY), { from: USER1 });
@@ -910,6 +931,10 @@ contract("Voting Token", (accounts) => {
     });
 
     it("cannot take an action if there is insufficient voting power (state change actions)", async () => {
+      // Clear the locks
+      await tokenLocking.methods["deposit(address,uint256,bool)"](token.address, 0, true, { from: USER0 });
+      await tokenLocking.methods["deposit(address,uint256,bool)"](token.address, 0, true, { from: USER1 });
+
       // Set globalClaimDelay to WAD
       await colony.makeExpenditure(1, UINT256_MAX, 1);
       const expenditureId = await colony.getExpenditureCount();
@@ -955,6 +980,10 @@ contract("Voting Token", (accounts) => {
     });
 
     it("can set vote power correctly after a vote", async () => {
+      // Clear the locks
+      await tokenLocking.methods["deposit(address,uint256,bool)"](token.address, 0, true, { from: USER0 });
+      await tokenLocking.methods["deposit(address,uint256,bool)"](token.address, 0, true, { from: USER1 });
+
       await colony.makeExpenditure(1, UINT256_MAX, 1);
       const expenditureId = await colony.getExpenditureCount();