Skip to content

Commit 03312a0

Browse files
jaybuidljaybuidl
committed
fix: improvements on the instant staking implementation (#1228) with feedback from @unknownunknown1 
1 parent 17502d9 commit 03312a0

File tree

4 files changed

+79
-34
lines changed

4 files changed

+79
-34
lines changed

‎contracts/src/arbitration/KlerosCore.sol

+46-17
Original file line numberDiff line numberDiff line change
@@ -9,12 +9,13 @@
99
pragma solidity 0.8.18;
1010

1111
import {IArbitrableV2, IArbitratorV2} from "./interfaces/IArbitratorV2.sol";
12-
import "./interfaces/IDisputeKit.sol";
13-
import "./interfaces/ISortitionModule.sol";
14-
import "../libraries/SafeERC20.sol";
15-
import "../libraries/Constants.sol";
16-
import "../proxy/UUPSProxiable.sol";
17-
import "../proxy/Initializable.sol";
12+
import {IDisputeKit} from "./interfaces/IDisputeKit.sol";
13+
import {ISortitionModule} from "./interfaces/ISortitionModule.sol";
14+
import {SafeERC20, IERC20} from "../libraries/SafeERC20.sol";
15+
import {Constants} from "../libraries/Constants.sol";
16+
import {OnError} from "../libraries/Types.sol";
17+
import {UUPSProxiable} from "../proxy/UUPSProxiable.sol";
18+
import {Initializable} from "../proxy/Initializable.sol";
1819

1920
/// @title KlerosCore
2021
/// Core arbitrator contract for Kleros v2.
@@ -448,17 +449,22 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
448449
/// @param _newStake The new stake.
449450
/// Note that the existing delayed stake will be nullified as non-relevant.
450451
function setStake(uint96 _courtID, uint256 _newStake) external {
451-
_setStake(msg.sender, _courtID, _newStake, false);
452+
_setStake(msg.sender, _courtID, _newStake, false, OnError.Revert);
452453
}
453454

455+
/// @dev Sets the stake of a specified account in a court, typically to apply a delayed stake or unstake inactive jurors.
456+
/// @param _account The account whose stake is being set.
457+
/// @param _courtID The ID of the court.
458+
/// @param _newStake The new stake.
459+
/// @param _alreadyTransferred Whether the PNKs have already been transferred to the contract.
454460
function setStakeBySortitionModule(
455461
address _account,
456462
uint96 _courtID,
457463
uint256 _newStake,
458464
bool _alreadyTransferred
459465
) external {
460466
if (msg.sender != address(sortitionModule)) revert SortitionModuleOnly();
461-
_setStake(_account, _courtID, _newStake, _alreadyTransferred);
467+
_setStake(_account, _courtID, _newStake, _alreadyTransferred, OnError.Return);
462468
}
463469

464470
/// @inheritdoc IArbitratorV2
@@ -994,34 +1000,57 @@ contract KlerosCore is IArbitratorV2, UUPSProxiable, Initializable {
9941000
emit DisputeKitEnabled(_courtID, _disputeKitID, _enable);
9951001
}
9961002

1003+
/// @dev If called only once then set _onError to Revert, otherwise set it to Return
1004+
/// @param _account The account to set the stake for.
1005+
/// @param _courtID The ID of the court to set the stake for.
1006+
/// @param _newStake The new stake.
1007+
/// @param _alreadyTransferred Whether the PNKs were already transferred to/from the staking contract.
1008+
/// @param _onError Whether to revert or return false on error.
1009+
/// @return Whether the stake was successfully set or not.
9971010
function _setStake(
9981011
address _account,
9991012
uint96 _courtID,
10001013
uint256 _newStake,
1001-
bool _alreadyTransferred
1002-
) internal returns (bool success) {
1014+
bool _alreadyTransferred,
1015+
OnError _onError
1016+
) internal returns (bool) {
10031017
if (_courtID == Constants.FORKING_COURT || _courtID > courts.length) {
1004-
return false; // Staking directly into the forking court is not allowed.
1018+
_stakingFailed(_onError); // Staking directly into the forking court is not allowed.
1019+
return false;
10051020
}
10061021
if (_newStake != 0 && _newStake < courts[_courtID].minStake) {
1007-
return false; // Staking less than the minimum stake is not allowed.
1022+
_stakingFailed(_onError); // Staking less than the minimum stake is not allowed.
1023+
return false;
10081024
}
10091025
(uint256 pnkDeposit, uint256 pnkWithdrawal, bool sortitionSuccess) = sortitionModule.setStake(
10101026
_account,
10111027
_courtID,
10121028
_newStake,
10131029
_alreadyTransferred
10141030
);
1015-
if (pnkDeposit > 0 && pnkWithdrawal > 0) revert StakingFailed();
1031+
if (!sortitionSuccess) {
1032+
_stakingFailed(_onError);
1033+
return false;
1034+
}
10161035
if (pnkDeposit > 0) {
1017-
// Note we don't return false after incorrect transfer because when stake is increased the transfer is done immediately, thus it can't disrupt delayed stakes' queue.
1018-
pinakion.safeTransferFrom(_account, address(this), pnkDeposit);
1019-
} else if (pnkWithdrawal > 0) {
1036+
if (!pinakion.safeTransferFrom(_account, address(this), pnkDeposit)) {
1037+
_stakingFailed(_onError);
1038+
return false;
1039+
}
1040+
}
1041+
if (pnkWithdrawal > 0) {
10201042
if (!pinakion.safeTransfer(_account, pnkWithdrawal)) {
1043+
_stakingFailed(_onError);
10211044
return false;
10221045
}
10231046
}
1024-
return sortitionSuccess;
1047+
return true;
1048+
}
1049+
1050+
/// @dev It may revert depending on the _onError parameter.
1051+
function _stakingFailed(OnError _onError) internal pure {
1052+
if (_onError == OnError.Return) return;
1053+
revert StakingFailed();
10251054
}
10261055

10271056
/// @dev Gets a court ID, the minimum number of jurors and an ID of a dispute kit from a specified extra data bytes array.

‎contracts/src/arbitration/SortitionModule.sol

+24-16
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,6 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
5353
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`.
5454
uint256 stakedPnk; // The juror's total amount of tokens staked in subcourts. Reflects actual pnk balance.
5555
uint256 lockedPnk; // The juror's total amount of tokens locked in disputes. Can reflect actual pnk balance when stakedPnk are fully withdrawn.
56-
mapping(uint96 => uint256) stakedPnkByCourt; // The amount of PNKs the juror has staked in the court in the form `stakedPnkByCourt[courtID]`.
5756
}
5857

5958
// ************************************* //
@@ -275,7 +274,7 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
275274
bool _alreadyTransferred
276275
) external override onlyByCore returns (uint256 pnkDeposit, uint256 pnkWithdrawal, bool succeeded) {
277276
Juror storage juror = jurors[_account];
278-
uint256 currentStake = juror.stakedPnkByCourt[_courtID];
277+
uint256 currentStake = stakeOf(_account, _courtID);
279278

280279
uint256 nbCourts = juror.courtIDs.length;
281280
if (_newStake == 0 && (nbCourts >= MAX_STAKE_PATHS || currentStake == 0)) {
@@ -336,22 +335,23 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
336335
if (latestIndex != 0) {
337336
DelayedStake storage delayedStake = delayedStakes[latestIndex];
338337
if (delayedStake.alreadyTransferred) {
339-
bytes32 stakePathID = _accountAndCourtIDToStakePathID(_juror, _courtID);
340338
// Sortition stake represents the stake value that was last updated during Staking phase.
341-
uint256 sortitionStake = stakeOf(bytes32(uint256(_courtID)), stakePathID);
339+
uint256 sortitionStake = stakeOf(_juror, _courtID);
340+
342341
// Withdraw the tokens that were added with the latest delayed stake.
343342
uint256 amountToWithdraw = delayedStake.stake - sortitionStake;
344343
actualAmountToWithdraw = amountToWithdraw;
345344
Juror storage juror = jurors[_juror];
346345
if (juror.stakedPnk <= actualAmountToWithdraw) {
347346
actualAmountToWithdraw = juror.stakedPnk;
348347
}
349-
// StakePnk can become lower because of penalty, thus we adjust the amount for it. stakedPnkByCourt can't be penalized so subtract the default amount.
348+
349+
// StakePnk can become lower because of penalty.
350350
juror.stakedPnk -= actualAmountToWithdraw;
351-
juror.stakedPnkByCourt[_courtID] -= amountToWithdraw;
352351
emit StakeDelayedAlreadyTransferredWithdrawn(_juror, _courtID, amountToWithdraw);
353-
// Note that if we don't delete court here it'll be duplicated after staking.
354-
if (juror.stakedPnkByCourt[_courtID] == 0) {
352+
353+
if (sortitionStake == 0) {
354+
// Delete the court otherwise it will be duplicated after staking.
355355
for (uint256 i = juror.courtIDs.length; i > 0; i--) {
356356
if (juror.courtIDs[i - 1] == _courtID) {
357357
juror.courtIDs[i - 1] = juror.courtIDs[juror.courtIDs.length - 1];
@@ -384,7 +384,6 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
384384
}
385385
// stakedPnk can become async with _currentStake (e.g. after penalty).
386386
juror.stakedPnk = (juror.stakedPnk >= _currentStake) ? juror.stakedPnk - _currentStake + _newStake : _newStake;
387-
juror.stakedPnkByCourt[_courtID] = _newStake;
388387
}
389388

390389
function _decreaseStake(
@@ -413,7 +412,6 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
413412
}
414413
// stakedPnk can become async with _currentStake (e.g. after penalty).
415414
juror.stakedPnk = (juror.stakedPnk >= _currentStake) ? juror.stakedPnk - _currentStake + _newStake : _newStake;
416-
juror.stakedPnkByCourt[_courtID] = _newStake;
417415
}
418416

419417
function lockStake(address _account, uint256 _relativeAmount) external override onlyByCore {
@@ -498,16 +496,26 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
498496
drawnAddress = _stakePathIDToAccount(tree.nodeIndexesToIDs[treeIndex]);
499497
}
500498

499+
/// @dev Get the stake of a juror in a court.
500+
/// @param _juror The address of the juror.
501+
/// @param _courtID The ID of the court.
502+
/// @return value The stake of the juror in the court.
503+
function stakeOf(address _juror, uint96 _courtID) public view returns (uint256) {
504+
bytes32 stakePathID = _accountAndCourtIDToStakePathID(_juror, _courtID);
505+
return stakeOf(bytes32(uint256(_courtID)), stakePathID);
506+
}
507+
501508
/// @dev Get the stake of a juror in a court.
502509
/// @param _key The key of the tree, corresponding to a court.
503510
/// @param _ID The stake path ID, corresponding to a juror.
504-
/// @return value The stake of the juror in the court.
505-
function stakeOf(bytes32 _key, bytes32 _ID) public view returns (uint256 value) {
511+
/// @return The stake of the juror in the court.
512+
function stakeOf(bytes32 _key, bytes32 _ID) public view returns (uint256) {
506513
SortitionSumTree storage tree = sortitionSumTrees[_key];
507514
uint treeIndex = tree.IDsToNodeIndexes[_ID];
508-
509-
if (treeIndex == 0) value = 0;
510-
else value = tree.nodes[treeIndex];
515+
if (treeIndex == 0) {
516+
return 0;
517+
}
518+
return tree.nodes[treeIndex];
511519
}
512520

513521
function getJurorBalance(
@@ -522,7 +530,7 @@ contract SortitionModule is ISortitionModule, UUPSProxiable, Initializable {
522530
Juror storage juror = jurors[_juror];
523531
totalStaked = juror.stakedPnk;
524532
totalLocked = juror.lockedPnk;
525-
stakedInCourt = juror.stakedPnkByCourt[_courtID];
533+
stakedInCourt = stakeOf(_juror, _courtID);
526534
nbCourts = juror.courtIDs.length;
527535
}
528536

‎contracts/src/libraries/Types.sol

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
// SPDX-License-Identifier: MIT
2+
3+
pragma solidity 0.8.18;
4+
5+
enum OnError {
6+
Revert,
7+
Return
8+
}

‎contracts/test/arbitration/staking.ts

+1-1
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ describe("Staking", async () => {
186186
await expect(core.setStake(2, PNK(3000)))
187187
.to.emit(sortition, "StakeDelayedAlreadyTransferred")
188188
.withArgs(deployer, 2, PNK(3000));
189-
expect(await sortition.getJurorBalance(deployer, 2)).to.be.deep.equal([PNK(5000), 0, PNK(3000), 2]); // stake has changed immediately, WARNING: this is misleading because it's not actually added to the SortitionSumTree
189+
expect(await sortition.getJurorBalance(deployer, 2)).to.be.deep.equal([PNK(5000), 0, PNK(2000), 2]); // stake has changed immediately, WARNING: this is misleading because it's not actually added to the SortitionSumTree
190190
});
191191

192192
it("Should transfer some PNK out of the juror's account", async () => {

0 commit comments

Comments
 (0)