fix(KlerosCore): staking logic fix

Author: unknownunknown1

Diff for: contracts/src/arbitration/KlerosCore.sol

@@ -68,8 +68,9 @@ contract KlerosCore is IArbitratorV2 {
 
     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 +127,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);
@@ -483,12 +484,12 @@ contract KlerosCore is IArbitratorV2 {
     /// @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();
+        if (!_setStakeForAccount(msg.sender, _courtID, _stake)) revert StakingFailed();
     }
 
-    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake, uint256 _penalty) external {
+    function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake) external {
         if (msg.sender != address(sortitionModule)) revert WrongCaller();
-        _setStakeForAccount(_account, _courtID, _stake, _penalty);
+        _setStakeForAccount(_account, _courtID, _stake);
     }
 
     /// @inheritdoc IArbitratorV2
@@ -614,7 +615,7 @@ contract KlerosCore is IArbitratorV2 {
         for (uint256 i = startIndex; i < endIndex; i++) {
             address drawnAddress = disputeKit.draw(_disputeID);
             if (drawnAddress != address(0)) {
-                jurors[drawnAddress].lockedPnk[dispute.courtID] += round.pnkAtStakePerJuror;
+                jurors[drawnAddress].lockedPnk += round.pnkAtStakePerJuror;
                 emit Draw(drawnAddress, _disputeID, currentRound, round.drawnJurors.length);
                 round.drawnJurors.push(drawnAddress);
 
@@ -763,16 +764,12 @@ 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 if there ara any.
+        // Note that lockedPnk will always cover penalty while stakedPnk can become lower after manual unstaking.
+        if (jurors[account].stakedPnk >= penalty) {
+            jurors[account].stakedPnk -= penalty;
         }
         emit TokenAndETHShift(
             account,
@@ -832,10 +829,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 +1011,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;
     }
 
@@ -1110,35 +1108,33 @@ contract KlerosCore is IArbitratorV2 {
     /// @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.
     /// @return succeeded True if the call succeeded, false otherwise.
-    function _setStakeForAccount(
-        address _account,
-        uint96 _courtID,
-        uint256 _stake,
-        uint256 _penalty
-    ) internal returns (bool succeeded) {
+    function _setStakeForAccount(address _account, uint96 _courtID, uint256 _stake) 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 (_stake < courts[_courtID].minStake) return false;
         }
 
-        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _stake, _penalty);
+        ISortitionModule.preStakeHookResult result = sortitionModule.preStakeHook(_account, _courtID, _stake);
         if (result == ISortitionModule.preStakeHookResult.failed) {
             return false;
         } else if (result == ISortitionModule.preStakeHookResult.delayed) {
-            emit StakeDelayed(_account, _courtID, _stake, _penalty);
+            emit StakeDelayed(_account, _courtID, _stake);
             return true;
         }
 
         uint256 transferredAmount;
         if (_stake >= currentStake) {
-            transferredAmount = _stake - currentStake;
+            // 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;
+            transferredAmount = (_stake >= currentStake + previouslyLocked)
+                ? _stake - currentStake - previouslyLocked
+                : 0;
             if (transferredAmount > 0) {
                 if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) {
                     if (currentStake == 0) {
@@ -1150,8 +1146,14 @@ contract KlerosCore is IArbitratorV2 {
             }
         } else {
             if (_stake == 0) {
-                // Keep locked PNKs in the contract and release them after dispute is executed.
-                transferredAmount = currentStake - juror.lockedPnk[_courtID] - _penalty;
+                // Make sure locked tokens always stay in the contract. They can only be released during Execution.
+                if (juror.stakedPnk >= currentStake + juror.lockedPnk) {
+                    // We have enough pnk staked to afford withdrawal of the current stake while keeping locked tokens.
+                    transferredAmount = currentStake;
+                } 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)) {
                         for (uint256 i = juror.courtIDs.length; i > 0; i--) {
@@ -1166,7 +1168,13 @@ contract KlerosCore is IArbitratorV2 {
                     }
                 }
             } else {
-                transferredAmount = currentStake - _stake - _penalty;
+                if (juror.stakedPnk >= currentStake - _stake + juror.lockedPnk) {
+                    // We have enough pnk staked to afford withdrawal while keeping locked tokens.
+                    transferredAmount = currentStake - _stake;
+                } 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;
@@ -1175,8 +1183,9 @@ contract KlerosCore is IArbitratorV2 {
             }
         }
 
-        // 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 + _stake : _stake;
+        juror.stakedPnkByCourt[_courtID] = _stake;
 
         sortitionModule.setStake(_account, _courtID, _stake);
         emit StakeSet(_account, _courtID, _stake);

Diff for: 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);
         }
     }
 

Diff for: contracts/src/arbitration/dispute-kits/DisputeKitClassic.sol

@@ -558,7 +558,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;
     }
 }

Diff for: contracts/src/arbitration/dispute-kits/DisputeKitSybilResistant.sol

@@ -587,9 +587,8 @@ contract DisputeKitSybilResistant is BaseDisputeKit, IEvidence {
             _coreDisputeID,
             core.getNumberOfRounds(_coreDisputeID) - 1
         );
-        (uint256 staked, uint256 locked, ) = core.getJurorBalance(_juror, courtID);
-        (, , uint256 minStake, , , , ) = core.courts(courtID);
-        if (staked < locked + lockedAmountPerJuror || staked < minStake) {
+        (uint256 totalStaked, uint256 totalLocked, , ) = core.getJurorBalance(_juror, courtID);
+        if (totalStaked < totalLocked + lockedAmountPerJuror) {
             return false;
         } else {
             return _proofOfHumanity(_juror);

Diff for: contracts/src/arbitration/interfaces/ISortitionModule.sol

@@ -26,12 +26,7 @@ interface ISortitionModule {
 
     function draw(bytes32 _court, uint256 _coreDisputeID, uint256 _voteID) external view returns (address);
 
-    function preStakeHook(
-        address _account,
-        uint96 _courtID,
-        uint256 _stake,
-        uint256 _penalty
-    ) external returns (preStakeHookResult);
+    function preStakeHook(address _account, uint96 _courtID, uint256 _stake) external returns (preStakeHookResult);
 
     function createDisputeHook(uint256 _disputeID, uint256 _roundID) external;