Fix/staking bug

contracts/src/arbitration/KlerosCore.sol

@@ -64,12 +64,14 @@ contract KlerosCore is IArbitratorV2 {
         uint256 sumFeeRewardPaid; // Total sum of arbitration fees paid to coherent jurors as a reward in this round.
         uint256 sumPnkRewardPaid; // Total sum of PNK paid to coherent jurors as a reward in this round.
         IERC20 feeToken; // The token used for paying fees in this round.
+        uint256 drawIterations; // The number of iterations passed drawing the jurors for this round.
     }
 
     struct Juror {
         uint96[] courtIDs; // The IDs of courts where the juror's stake path ends. A stake path is a path from the general court to a court the juror directly staked in using `_setStake`.
-        mapping(uint96 => uint256) stakedPnk; // The amount of PNKs the juror has staked in the court in the form `stakedPnk[courtID]`.
-        mapping(uint96 => uint256) lockedPnk; // The amount of PNKs the juror has locked in the court in the form `lockedPnk[courtID]`.
+        uint256 stakedPnk; // The juror's total amount of tokens staked in subcourts. Reflects actual pnk balance.
+        uint256 lockedPnk; // The juror's total amount of tokens locked in disputes. Can reflect actual pnk balance when stakedPnk are fully withdrawn.
+        mapping(uint96 => uint256) stakedPnkByCourt; // The amount of PNKs the juror has staked in the court in the form `stakedPnkByCourt[courtID]`.
     }
 
     struct DisputeKitNode {
@@ -126,7 +128,7 @@ contract KlerosCore is IArbitratorV2 {
     // ************************************* //
 
     event StakeSet(address indexed _address, uint256 _courtID, uint256 _amount);
-    event StakeDelayed(address indexed _address, uint256 _courtID, uint256 _amount, uint256 _penalty);
+    event StakeDelayed(address indexed _address, uint256 _courtID, uint256 _amount);
     event NewPeriod(uint256 indexed _disputeID, Period _period);
     event AppealPossible(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
     event AppealDecision(uint256 indexed _disputeID, IArbitrableV2 indexed _arbitrable);
@@ -481,14 +483,14 @@ contract KlerosCore is IArbitratorV2 {
 
     /// @dev Sets the caller's stake in a court.
     /// @param _courtID The ID of the court.
-    /// @param _stake The new stake.
-    function setStake(uint96 _courtID, uint256 _stake) external {
-        if (!_setStakeForAccount(msg.sender, _courtID, _stake, 0)) revert StakingFailed();
+    /// @param _newStake The new stake.
+    function setStake(uint96 _courtID, uint256 _newStake) external {
+        if (!_setStakeForAccount(msg.sender, _courtID, _newStake)) revert StakingFailed();
     }
 
-    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake, uint256 _penalty) external {
+    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _newStake) external {
         if (msg.sender != address(sortitionModule)) revert WrongCaller();
-        _setStakeForAccount(_account, _courtID, _stake, _penalty);
+        _setStakeForAccount(_account, _courtID, _newStake);
     }
 
     /// @inheritdoc IArbitratorV2
@@ -608,21 +610,20 @@ contract KlerosCore is IArbitratorV2 {
 
         IDisputeKit disputeKit = disputeKitNodes[round.disputeKitID].disputeKit;
 
-        uint256 startIndex = round.drawnJurors.length;
-        uint256 endIndex = startIndex + _iterations <= round.nbVotes ? startIndex + _iterations : round.nbVotes;
-
-        for (uint256 i = startIndex; i < endIndex; i++) {
-            address drawnAddress = disputeKit.draw(_disputeID);
-            if (drawnAddress != address(0)) {
-                jurors[drawnAddress].lockedPnk[dispute.courtID] += round.pnkAtStakePerJuror;
-                emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
-                round.drawnJurors.push(drawnAddress);
-
-                if (round.drawnJurors.length == round.nbVotes) {
-                    sortitionModule.postDrawHook(_disputeID, currentRound);
-                }
+        uint256 startIndex = round.drawIterations;
+        for (uint256 i = startIndex; i < startIndex + _iterations && round.drawnJurors.length < round.nbVotes; i++) {
+            address drawnAddress = disputeKit.draw(_disputeID, i);
+            if (drawnAddress == address(0)) {
+                continue;
+            }
+            jurors[drawnAddress].lockedPnk += round.pnkAtStakePerJuror;
+            emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
+            round.drawnJurors.push(drawnAddress);
+            if (round.drawnJurors.length == round.nbVotes) {
+                sortitionModule.postDrawHook(_disputeID, currentRound);
             }
         }
+        round.drawIterations += _iterations;
     }
 
     /// @dev Appeals the ruling of a specified dispute.
@@ -763,16 +764,14 @@ contract KlerosCore is IArbitratorV2 {
 
         // Unlock the PNKs affected by the penalty
         address account = round.drawnJurors[_params.repartition];
-        jurors[account].lockedPnk[dispute.courtID] -= penalty;
-
-        // Apply the penalty to the staked PNKs
-        if (jurors[account].stakedPnk[dispute.courtID] >= courts[dispute.courtID].minStake + penalty) {
-            // The juror still has enough staked PNKs after penalty for this court.
-            uint256 newStake = jurors[account].stakedPnk[dispute.courtID] - penalty;
-            _setStakeForAccount(account, dispute.courtID, newStake, penalty);
-        } else if (jurors[account].stakedPnk[dispute.courtID] != 0) {
-            // The juror does not have enough staked PNKs after penalty for this court, unstake them.
-            _setStakeForAccount(account, dispute.courtID, 0, penalty);
+        jurors[account].lockedPnk -= penalty;
+
+        // Apply the penalty to the staked PNKs.
+        // Note that lockedPnk will always cover penalty while stakedPnk can become lower after manual unstaking.
+        if (jurors[account].stakedPnk >= penalty) {
+            jurors[account].stakedPnk -= penalty;
+        } else {
+            jurors[account].stakedPnk = 0;
         }
         emit TokenAndETHShift(
             account,
@@ -832,10 +831,10 @@ contract KlerosCore is IArbitratorV2 {
         uint256 pnkLocked = (round.pnkAtStakePerJuror * degreeOfCoherence) / ALPHA_DIVISOR;
 
         // Release the rest of the PNKs of the juror for this round.
-        jurors[account].lockedPnk[dispute.courtID] -= pnkLocked;
+        jurors[account].lockedPnk -= pnkLocked;
 
         // Give back the locked PNKs in case the juror fully unstaked earlier.
-        if (jurors[account].stakedPnk[dispute.courtID] == 0) {
+        if (jurors[account].stakedPnk == 0) {
             pinakion.safeTransfer(account, pnkLocked);
         }
 
@@ -1014,10 +1013,11 @@ contract KlerosCore is IArbitratorV2 {
     function getJurorBalance(
         address _juror,
         uint96 _courtID
-    ) external view returns (uint256 staked, uint256 locked, uint256 nbCourts) {
+    ) external view returns (uint256 totalStaked, uint256 totalLocked, uint256 stakedInCourt, uint256 nbCourts) {
         Juror storage juror = jurors[_juror];
-        staked = juror.stakedPnk[_courtID];
-        locked = juror.lockedPnk[_courtID];
+        totalStaked = juror.stakedPnk;
+        totalLocked = juror.lockedPnk;
+        stakedInCourt = juror.stakedPnkByCourt[_courtID];
         nbCourts = juror.courtIDs.length;
     }
 
@@ -1109,77 +1109,80 @@ contract KlerosCore is IArbitratorV2 {
     /// and `j` is the maximum number of jurors that ever staked in one of these courts simultaneously.
     /// @param _account The address of the juror.
     /// @param _courtID The ID of the court.
-    /// @param _stake The new stake.
-    /// @param _penalty Penalized amount won't be transferred back to juror when the stake is lowered.
+    /// @param _newStake The new stake.
     /// @return succeeded True if the call succeeded, false otherwise.
     function _setStakeForAccount(
         address _account,
         uint96 _courtID,
-        uint256 _stake,
-        uint256 _penalty
+        uint256 _newStake
     ) internal returns (bool succeeded) {
         if (_courtID == FORKING_COURT || _courtID > courts.length) return false;
 
         Juror storage juror = jurors[_account];
-        uint256 currentStake = juror.stakedPnk[_courtID];
+        uint256 currentStake = juror.stakedPnkByCourt[_courtID];
 
-        if (_stake != 0) {
-            // Check against locked PNKs in case the min stake was lowered.
-            if (_stake < courts[_courtID].minStake || _stake < juror.lockedPnk[_courtID]) return false;
+        if (_newStake != 0) {
+            if (_newStake < courts[_courtID].minStake) return false;
+        } else if (currentStake == 0) {
+            return false;
         }
 
-        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _stake, _penalty);
+        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _newStake);
         if (result == ISortitionModule.preStakeHookResult.failed) {
             return false;
         } else if (result == ISortitionModule.preStakeHookResult.delayed) {
-            emit StakeDelayed(_account, _courtID, _stake, _penalty);
+            emit StakeDelayed(_account, _courtID, _newStake);
             return true;
         }
 
         uint256 transferredAmount;
-        if (_stake >= currentStake) {
-            transferredAmount = _stake - currentStake;
+        if (_newStake >= currentStake) {
+            // Stake increase
+            // When stakedPnk becomes lower than lockedPnk count the locked tokens in when transferring tokens from juror.
+            // (E.g. stakedPnk = 0, lockedPnk = 150) which can happen if the juror unstaked fully while having some tokens locked.
+            uint256 previouslyLocked = (juror.lockedPnk >= juror.stakedPnk) ? juror.lockedPnk - juror.stakedPnk : 0; // underflow guard
+            transferredAmount = (_newStake >= currentStake + previouslyLocked) // underflow guard
+                ? _newStake - currentStake - previouslyLocked
+                : 0;
             if (transferredAmount > 0) {
-                if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
-                    if (currentStake == 0) {
-                        juror.courtIDs.push(_courtID);
-                    }
-                } else {
+                if (!pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
                     return false;
                 }
             }
+            if (currentStake == 0) {
+                juror.courtIDs.push(_courtID);
+            }
         } else {
-            if (_stake == 0) {
-                // Keep locked PNKs in the contract and release them after dispute is executed.
-                transferredAmount = currentStake - juror.lockedPnk[_courtID] - _penalty;
-                if (transferredAmount > 0) {
-                    if (pinakion.safeTransfer(_account, transferredAmount)) {
-                        for (uint256 i = juror.courtIDs.length; i > 0; i--) {
-                            if (juror.courtIDs[i - 1] == _courtID) {
-                                juror.courtIDs[i - 1] = juror.courtIDs[juror.courtIDs.length - 1];
-                                juror.courtIDs.pop();
-                                break;
-                            }
-                        }
-                    } else {
-                        return false;
-                    }
+            // Stake decrease: make sure locked tokens always stay in the contract. They can only be released during Execution.
+            if (juror.stakedPnk >= currentStake - _newStake + juror.lockedPnk) {
+                // We have enough pnk staked to afford withdrawal while keeping locked tokens.
+                transferredAmount = currentStake - _newStake;
+            } else if (juror.stakedPnk >= juror.lockedPnk) {
+                // Can't afford withdrawing the current stake fully. Take whatever is available while keeping locked tokens.
+                transferredAmount = juror.stakedPnk - juror.lockedPnk;
+            }
+            if (transferredAmount > 0) {
+                if (!pinakion.safeTransfer(_account, transferredAmount)) {
+                    return false;
                 }
-            } else {
-                transferredAmount = currentStake - _stake - _penalty;
-                if (transferredAmount > 0) {
-                    if (!pinakion.safeTransfer(_account, transferredAmount)) {
-                        return false;
+            }
+            if (_newStake == 0) {
+                for (uint256 i = juror.courtIDs.length; i > 0; i--) {
+                    if (juror.courtIDs[i - 1] == _courtID) {
+                        juror.courtIDs[i - 1] = juror.courtIDs[juror.courtIDs.length - 1];
+                        juror.courtIDs.pop();
+                        break;
                     }
                 }
             }
         }
 
-        // Update juror's records.
-        juror.stakedPnk[_courtID] = _stake;
+        // Note that stakedPnk can become async with currentStake (e.g. after penalty).
+        juror.stakedPnk = (juror.stakedPnk >= currentStake) ? juror.stakedPnk - currentStake + _newStake : _newStake;
+        juror.stakedPnkByCourt[_courtID] = _newStake;
 
-        sortitionModule.setStake(_account, _courtID, _stake);
-        emit StakeSet(_account, _courtID, _stake);
+        sortitionModule.setStake(_account, _courtID, _newStake);
+        emit StakeSet(_account, _courtID, _newStake);
         return true;
     }
 

contracts/src/arbitration/SortitionModule.sol

@@ -36,7 +36,6 @@ contract SortitionModule is ISortitionModule {
         address account; // The address of the juror.
         uint96 courtID; // The ID of the court.
         uint256 stake; // The new stake.
-        uint256 penalty; // Penalty value, in case the stake was set during execution.
     }
 
     // ************************************* //
@@ -185,12 +184,7 @@ contract SortitionModule is ISortitionModule {
 
         for (uint256 i = delayedStakeReadIndex; i < newDelayedStakeReadIndex; i++) {
             DelayedStake storage delayedStake = delayedStakes[i];
-            core.setStakeBySortitionModule(
-                delayedStake.account,
-                delayedStake.courtID,
-                delayedStake.stake,
-                delayedStake.penalty
-            );
+            core.setStakeBySortitionModule(delayedStake.account, delayedStake.courtID, delayedStake.stake);
             delete delayedStakes[i];
         }
         delayedStakeReadIndex = newDelayedStakeReadIndex;
@@ -199,10 +193,9 @@ contract SortitionModule is ISortitionModule {
     function preStakeHook(
         address _account,
         uint96 _courtID,
-        uint256 _stake,
-        uint256 _penalty
+        uint256 _stake
     ) external override onlyByCore returns (preStakeHookResult) {
-        (uint256 currentStake, , uint256 nbCourts) = core.getJurorBalance(_account, _courtID);
+        (, , uint256 currentStake, uint256 nbCourts) = core.getJurorBalance(_account, _courtID);
         if (currentStake == 0 && nbCourts >= MAX_STAKE_PATHS) {
             // Prevent staking beyond MAX_STAKE_PATHS but unstaking is always allowed.
             return preStakeHookResult.failed;
@@ -211,8 +204,7 @@ contract SortitionModule is ISortitionModule {
                 delayedStakes[++delayedStakeWriteIndex] = DelayedStake({
                     account: _account,
                     courtID: _courtID,
-                    stake: _stake,
-                    penalty: _penalty
+                    stake: _stake
                 });
                 return preStakeHookResult.delayed;
             }
@@ -264,7 +256,7 @@ contract SortitionModule is ISortitionModule {
     function setJurorInactive(address _account) external override onlyByCore {
         uint96[] memory courtIDs = core.getJurorCourtIDs(_account);
         for (uint256 j = courtIDs.length; j > 0; j--) {
-            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0, 0);
+            core.setStakeBySortitionModule(_account, courtIDs[j - 1], 0);
         }
     }
 
@@ -276,15 +268,15 @@ contract SortitionModule is ISortitionModule {
     /// Note that this function reverts if the sum of all values in the tree is 0.
     /// @param _key The key of the tree.
     /// @param _coreDisputeID Index of the dispute in Kleros Core.
-    /// @param _voteID ID of the voter.
+    /// @param _nonce Nonce to hash with random number.
     /// @return drawnAddress The drawn address.
     /// `O(k * log_k(n))` where
     /// `k` is the maximum number of children per node in the tree,
     ///  and `n` is the maximum number of nodes ever appended.
     function draw(
         bytes32 _key,
         uint256 _coreDisputeID,
-        uint256 _voteID
+        uint256 _nonce
     ) public view override returns (address drawnAddress) {
         require(phase == Phase.drawing, "Wrong phase.");
         SortitionSumTree storage tree = sortitionSumTrees[_key];
@@ -293,7 +285,7 @@ contract SortitionModule is ISortitionModule {
             return address(0); // No jurors staked.
         }
 
-        uint256 currentDrawnNumber = uint256(keccak256(abi.encodePacked(randomNumber, _coreDisputeID, _voteID))) %
+        uint256 currentDrawnNumber = uint256(keccak256(abi.encodePacked(randomNumber, _coreDisputeID, _nonce))) %
             tree.nodes[0];
 
         // While it still has children

contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol

@@ -181,9 +181,11 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
     /// @dev Draws the juror from the sortition tree. The drawn address is picked up by Kleros Core.
     /// Note: Access restricted to Kleros Core only.
     /// @param _coreDisputeID The ID of the dispute in Kleros Core.
+    /// @param _nonce Nonce of the drawing iteration.
     /// @return drawnAddress The drawn address.
     function draw(
-        uint256 _coreDisputeID
+        uint256 _coreDisputeID,
+        uint256 _nonce
     ) external override onlyByCore notJumped(_coreDisputeID) returns (address drawnAddress) {
         Dispute storage dispute = disputes[coreDisputeIDToLocal[_coreDisputeID]];
         Round storage round = dispute.rounds[dispute.rounds.length - 1];
@@ -193,7 +195,7 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
         bytes32 key = bytes32(uint256(courtID)); // Get the ID of the tree.
 
         // TODO: Handle the situation when no one has staked yet.
-        drawnAddress = sortitionModule.draw(key, _coreDisputeID, round.votes.length);
+        drawnAddress = sortitionModule.draw(key, _coreDisputeID, _nonce);
 
         if (_postDrawCheck(_coreDisputeID, drawnAddress)) {
             round.votes.push(Vote({account: drawnAddress, commit: bytes32(0), choice: 0, voted: false}));
@@ -558,7 +560,13 @@ contract DisputeKitClassic is BaseDisputeKit, IEvidence {
     // ************************************* //
 
     /// @inheritdoc BaseDisputeKit
-    function _postDrawCheck(uint256 /*_coreDisputeID*/, address /*_juror*/) internal pure override returns (bool) {
-        return true;
+    function _postDrawCheck(uint256 _coreDisputeID, address _juror) internal view override returns (bool) {
+        (uint96 courtID, , , , ) = core.disputes(_coreDisputeID);
+        (, uint256 lockedAmountPerJuror, , , , , , , , ) = core.getRoundInfo(
+            _coreDisputeID,
+            core.getNumberOfRounds(_coreDisputeID) - 1
+        );
+        (uint256 totalStaked, uint256 totalLocked, , ) = core.getJurorBalance(_juror, courtID);
+        return totalStaked >= totalLocked + lockedAmountPerJuror;
     }
 }