diff --git a/contracts/src/arbitration/KlerosCore.sol b/contracts/src/arbitration/KlerosCore.sol index e282f4dfa..66ee11db9 100644 --- a/contracts/src/arbitration/KlerosCore.sol +++ b/contracts/src/arbitration/KlerosCore.sol @@ -127,6 +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 StakePartiallyDelayed(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); @@ -482,12 +483,28 @@ 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, 0, false)) revert StakingFailed(); } - function setStakeBySortitionModule(address _account, uint96 _courtID, uint256 _stake, uint256 _penalty) external { + function withdrawDelayedStake(uint256 _delayedStakeIndex) external { + // All necessary checks are done in sortition module. + (uint256 amountToWithdraw, uint96 courtID) = sortitionModule.removeDelayedStake(_delayedStakeIndex, msg.sender); + if (jurors[msg.sender].stakedPnk[courtID] <= amountToWithdraw) { + amountToWithdraw = jurors[msg.sender].stakedPnk[courtID]; + } + require(pinakion.safeTransfer(msg.sender, amountToWithdraw)); + jurors[msg.sender].stakedPnk[courtID] -= amountToWithdraw; + } + + function setStakeBySortitionModule( + address _account, + uint96 _courtID, + uint256 _stake, + uint256 _penalty, + bool _alreadyTransferred + ) external { if (msg.sender != address(sortitionModule)) revert WrongCaller(); - _setStakeForAccount(_account, _courtID, _stake, _penalty); + _setStakeForAccount(_account, _courtID, _stake, _penalty, _alreadyTransferred); } /// @inheritdoc IArbitratorV2 @@ -768,10 +785,11 @@ contract KlerosCore is IArbitratorV2 { 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); + // `alreadyTransferred` flag can be true only after manual stake increase, which can't happen during penalty. + _setStakeForAccount(account, dispute.courtID, newStake, penalty, false); } 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); + _setStakeForAccount(account, dispute.courtID, 0, penalty, false); } emit TokenAndETHShift( account, @@ -1110,12 +1128,14 @@ contract KlerosCore is IArbitratorV2 { /// @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 _alreadyTransferred True if the tokens were already transferred. Only relevant for delayed stake execution. /// @return succeeded True if the call succeeded, false otherwise. function _setStakeForAccount( address _account, uint96 _courtID, uint256 _stake, - uint256 _penalty + uint256 _penalty, + bool _alreadyTransferred ) internal returns (bool succeeded) { if (_courtID == FORKING_COURT || _courtID > courts.length) return false; @@ -1135,47 +1155,56 @@ contract KlerosCore is IArbitratorV2 { return true; } - uint256 transferredAmount; - if (_stake >= currentStake) { - transferredAmount = _stake - currentStake; - if (transferredAmount > 0) { - if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) { - if (currentStake == 0) { - juror.courtIDs.push(_courtID); - } - } else { - return false; - } - } - } else { - if (_stake == 0) { - // Keep locked PNKs in the contract and release them after dispute is executed. - transferredAmount = currentStake - juror.lockedPnk[_courtID] - _penalty; + // Don't transfer the tokens and only update the drawing chance if the transfer was already done. + if (!_alreadyTransferred) { + uint256 transferredAmount; + if (_stake >= currentStake) { + transferredAmount = _stake - currentStake; 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; - } + if (pinakion.safeTransferFrom(_account, address(this), transferredAmount)) { + if (currentStake == 0) { + juror.courtIDs.push(_courtID); } } else { return false; } } } else { - transferredAmount = currentStake - _stake - _penalty; - if (transferredAmount > 0) { - if (!pinakion.safeTransfer(_account, transferredAmount)) { - return false; + 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; + } + } + } else { + transferredAmount = currentStake - _stake - _penalty; + if (transferredAmount > 0) { + if (!pinakion.safeTransfer(_account, transferredAmount)) { + return false; + } } } } + + // Update juror's records. + juror.stakedPnk[_courtID] = _stake; } - // Update juror's records. - juror.stakedPnk[_courtID] = _stake; + // Transfer the tokens but don't update sortition module. + if (result == ISortitionModule.preStakeHookResult.partiallyDelayed) { + emit StakePartiallyDelayed(_account, _courtID, _stake); + return true; + } sortitionModule.setStake(_account, _courtID, _stake); emit StakeSet(_account, _courtID, _stake); diff --git a/contracts/src/arbitration/SortitionModule.sol b/contracts/src/arbitration/SortitionModule.sol index ac9eef13a..f69f8f953 100644 --- a/contracts/src/arbitration/SortitionModule.sol +++ b/contracts/src/arbitration/SortitionModule.sol @@ -37,6 +37,7 @@ contract SortitionModule is ISortitionModule { uint96 courtID; // The ID of the court. uint256 stake; // The new stake. uint256 penalty; // Penalty value, in case the stake was set during execution. + bool alreadyTransferred; // True if tokens were already transferred before delayed stake's execution. } // ************************************* // @@ -185,17 +186,38 @@ 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 - ); - delete delayedStakes[i]; + // Delayed stake could've been manually removed already. In this case simply move on to the next item. + if (delayedStake.account != address(0)) { + core.setStakeBySortitionModule( + delayedStake.account, + delayedStake.courtID, + delayedStake.stake, + delayedStake.penalty, + delayedStake.alreadyTransferred + ); + delete delayedStakes[i]; + } } delayedStakeReadIndex = newDelayedStakeReadIndex; } + /// @dev Remove the delayed stake after its partial execution in order to return the tokens. + /// @param _index Index of the delayed stake to remove. + /// @param _sender Address that attempted removal. + /// @return stake Stake amount that was discarded. + /// @return courtID ID of the court related to delayed stake. + function removeDelayedStake( + uint256 _index, + address _sender + ) external override onlyByCore returns (uint256 stake, uint96 courtID) { + DelayedStake storage delayedStake = delayedStakes[_index]; + require(delayedStake.account == _sender, "Can only remove your own stake"); + require(delayedStake.alreadyTransferred, "No tokens to return"); + stake = delayedStake.stake; + courtID = delayedStake.courtID; + delete delayedStakes[_index]; + } + function preStakeHook( address _account, uint96 _courtID, @@ -208,13 +230,18 @@ contract SortitionModule is ISortitionModule { return preStakeHookResult.failed; } else { if (phase != Phase.staking) { - delayedStakes[++delayedStakeWriteIndex] = DelayedStake({ - account: _account, - courtID: _courtID, - stake: _stake, - penalty: _penalty - }); - return preStakeHookResult.delayed; + DelayedStake storage delayedStake = delayedStakes[++delayedStakeWriteIndex]; + delayedStake.account = _account; + delayedStake.courtID = _courtID; + delayedStake.stake = _stake; + delayedStake.penalty = _penalty; + if (_stake > currentStake) { + // Actual token transfer is done right after this hook. + delayedStake.alreadyTransferred = true; + return preStakeHookResult.partiallyDelayed; + } else { + return preStakeHookResult.delayed; + } } } return preStakeHookResult.ok; @@ -264,7 +291,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, 0, false); } } diff --git a/contracts/src/arbitration/interfaces/ISortitionModule.sol b/contracts/src/arbitration/interfaces/ISortitionModule.sol index 0c5f15c6f..61b5deeb5 100644 --- a/contracts/src/arbitration/interfaces/ISortitionModule.sol +++ b/contracts/src/arbitration/interfaces/ISortitionModule.sol @@ -9,9 +9,10 @@ interface ISortitionModule { } enum preStakeHookResult { - ok, - delayed, - failed + ok, // Correct phase. All checks are passed. + partiallyDelayed, // Wrong phase but stake is increased, so transfer the tokens without updating the drawing chance. + delayed, // Wrong phase and stake is decreased. Delay the token transfer and drawing chance update. + failed // Checks didn't pass. Do no changes. } event NewPhase(Phase _phase); @@ -36,4 +37,6 @@ interface ISortitionModule { function createDisputeHook(uint256 _disputeID, uint256 _roundID) external; function postDrawHook(uint256 _disputeID, uint256 _roundID) external; + + function removeDelayedStake(uint256 _index, address _sender) external returns (uint256 amount, uint96 courtID); }