From 99c6625b40a93c2cfd5fefae9232659348e8bb92 Mon Sep 17 00:00:00 2001 From: rymnc <43716372+rymnc@users.noreply.github.com> Date: Thu, 30 Mar 2023 20:23:08 +0530 Subject: [PATCH] fix: remove .transfer from loops --- contracts/Rln.sol | 45 +++++++++++++++++++------ docs/index.md | 50 +++++++++++++++++++++++----- test/RLN.t.sol | 85 +++++++++++++++++++++++++++++++++-------------- 3 files changed, 135 insertions(+), 45 deletions(-) diff --git a/contracts/Rln.sol b/contracts/Rln.sol index 2ec11ab..1b2b07a 100644 --- a/contracts/Rln.sol +++ b/contracts/Rln.sol @@ -23,8 +23,8 @@ error DuplicateIdCommitment(); /// @param givenReceiversLen The length of the receivers array error MismatchedBatchSize(uint256 givenSecretsLen, uint256 givenReceiversLen); -/// Invalid withdrawal address, when the receiver is the contract itself or 0x0 -error InvalidWithdrawalAddress(address to); +/// Invalid receiver address, when the receiver is the contract itself or 0x0 +error InvalidReceiverAddress(address to); /// Member is not registered error MemberNotRegistered(uint256 idCommitment); @@ -32,6 +32,12 @@ error MemberNotRegistered(uint256 idCommitment); /// Member has no stake error MemberHasNoStake(uint256 idCommitment); +/// User has insufficient balance to withdraw +error InsufficientWithdrawalBalance(); + +/// Contract has insufficient balance to return +error InsufficientContractBalance(); + contract RLN { /// @notice The deposit amount required to register as a member uint256 public immutable MEMBERSHIP_DEPOSIT; @@ -51,8 +57,11 @@ contract RLN { /// @notice The membership status of each member mapping(uint256 => bool) public members; + /// @notice The balance of each user that can be withdrawn + mapping(address => uint256) public withdrawalBalance; + /// @notice The Poseidon hasher contract - IPoseidonHasher public poseidonHasher; + IPoseidonHasher public immutable poseidonHasher; /// Emitted when a new member is added to the set /// @param idCommitment The idCommitment of the member @@ -91,8 +100,9 @@ contract RLN { uint256 requiredDeposit = MEMBERSHIP_DEPOSIT * idCommitmentlen; if (msg.value != requiredDeposit) revert InsufficientDeposit(requiredDeposit, msg.value); + for (uint256 i = 0; i < idCommitmentlen; i++) { - _register(idCommitments[i], msg.value / idCommitmentlen); + _register(idCommitments[i], MEMBERSHIP_DEPOSIT); } } @@ -113,7 +123,7 @@ contract RLN { /// Allows a user to slash a batch of members /// @param secrets array of idSecretHashes /// @param receivers array of addresses to receive the funds - function withdrawBatch( + function slashBatch( uint256[] calldata secrets, address payable[] calldata receivers ) external { @@ -122,23 +132,23 @@ contract RLN { if (batchSize != receivers.length) revert MismatchedBatchSize(secrets.length, receivers.length); for (uint256 i = 0; i < batchSize; i++) { - _withdraw(secrets[i], receivers[i]); + _slash(secrets[i], receivers[i]); } } /// Allows a user to slash a member /// @param secret The idSecretHash of the member - function withdraw(uint256 secret, address payable receiver) external { - _withdraw(secret, receiver); + function slash(uint256 secret, address payable receiver) external { + _slash(secret, receiver); } /// Slashes a member by removing them from the set, and transferring their /// stake to the receiver /// @param secret The idSecretHash of the member /// @param receiver The address to receive the funds - function _withdraw(uint256 secret, address payable receiver) internal { + function _slash(uint256 secret, address payable receiver) internal { if (receiver == address(this) || receiver == address(0)) - revert InvalidWithdrawalAddress(receiver); + revert InvalidReceiverAddress(receiver); // derive idCommitment uint256 idCommitment = hash(secret); @@ -154,11 +164,24 @@ contract RLN { stakedAmounts[idCommitment] = 0; // refund deposit - receiver.transfer(amountToTransfer); + withdrawalBalance[receiver] += amountToTransfer; emit MemberWithdrawn(idCommitment); } + /// Allows a user to withdraw funds allocated to them upon slashing a member + function withdraw() external { + uint256 amount = withdrawalBalance[msg.sender]; + + if (amount == 0) revert InsufficientWithdrawalBalance(); + if (amount > address(this).balance) + revert InsufficientContractBalance(); + + withdrawalBalance[msg.sender] = 0; + + payable(msg.sender).transfer(amount); + } + /// Hashes a value using the Poseidon hasher /// NOTE: The variant of Poseidon we use accepts only 1 input, assume n=2, and the second input is 0 /// @param input The value to hash diff --git a/docs/index.md b/docs/index.md index aa200b2..a8566af 100644 --- a/docs/index.md +++ b/docs/index.md @@ -890,13 +890,13 @@ Batch size mismatch, when the length of secrets and receivers are not equal | givenSecretsLen | uint256 | The length of the secrets array | | givenReceiversLen | uint256 | The length of the receivers array | -## InvalidWithdrawalAddress +## InvalidReceiverAddress ```solidity -error InvalidWithdrawalAddress(address to) +error InvalidReceiverAddress(address to) ``` -Invalid withdrawal address, when the receiver is the contract itself or 0x0 +Invalid receiver address, when the receiver is the contract itself or 0x0 ## MemberNotRegistered @@ -914,6 +914,22 @@ error MemberHasNoStake(uint256 idCommitment) Member has no stake +## InsufficientWithdrawalBalance + +```solidity +error InsufficientWithdrawalBalance() +``` + +User has insufficient balance to withdraw + +## InsufficientContractBalance + +```solidity +error InsufficientContractBalance() +``` + +Contract has insufficient balance to return + ## RLN ### MEMBERSHIP_DEPOSIT @@ -964,6 +980,14 @@ mapping(uint256 => bool) members The membership status of each member +### withdrawalBalance + +```solidity +mapping(address => uint256) withdrawalBalance +``` + +The balance of each user that can be withdrawn + ### poseidonHasher ```solidity @@ -1050,10 +1074,10 @@ Registers a member | idCommitment | uint256 | The idCommitment of the member | | stake | uint256 | The amount of eth staked by the member | -### withdrawBatch +### slashBatch ```solidity -function withdrawBatch(uint256[] secrets, address payable[] receivers) external +function slashBatch(uint256[] secrets, address payable[] receivers) external ``` Allows a user to slash a batch of members @@ -1065,10 +1089,10 @@ Allows a user to slash a batch of members | secrets | uint256[] | array of idSecretHashes | | receivers | address payable[] | array of addresses to receive the funds | -### withdraw +### slash ```solidity -function withdraw(uint256 secret, address payable receiver) external +function slash(uint256 secret, address payable receiver) external ``` Allows a user to slash a member @@ -1080,10 +1104,10 @@ Allows a user to slash a member | secret | uint256 | The idSecretHash of the member | | receiver | address payable | | -### \_withdraw +### \_slash ```solidity -function _withdraw(uint256 secret, address payable receiver) internal +function _slash(uint256 secret, address payable receiver) internal ``` Slashes a member by removing them from the set, and transferring their @@ -1096,6 +1120,14 @@ stake to the receiver | secret | uint256 | The idSecretHash of the member | | receiver | address payable | The address to receive the funds | +### withdraw + +```solidity +function withdraw() external +``` + +Allows a user to withdraw funds allocated to them upon slashing a member + ### hash ```solidity diff --git a/test/RLN.t.sol b/test/RLN.t.sol index fe9fcce..b56e8eb 100644 --- a/test/RLN.t.sol +++ b/test/RLN.t.sol @@ -167,10 +167,7 @@ contract RLNTest is Test { rln.registerBatch{value: badDepositAmount}(idCommitments); } - function test__ValidWithdraw( - uint256 idSecretHash, - address payable to - ) public { + function test__ValidSlash(uint256 idSecretHash, address payable to) public { // avoid precompiles, etc // TODO: wrap both of these in a single function assumePayable(to); @@ -182,52 +179,51 @@ contract RLNTest is Test { assertEq(rln.stakedAmounts(idCommitment), MEMBERSHIP_DEPOSIT); uint256 balanceBefore = to.balance; - rln.withdraw(idSecretHash, to); + rln.slash(idSecretHash, to); + vm.prank(to); + rln.withdraw(); assertEq(rln.stakedAmounts(idCommitment), 0); assertEq(rln.members(idCommitment), false); assertEq(to.balance, balanceBefore + MEMBERSHIP_DEPOSIT); } - function test__InvalidWithdraw__ToZeroAddress() public { + function test__InvalidSlash__ToZeroAddress() public { uint256 idSecretHash = 19014214495641488759237505126948346942972912379615652741039992445865937985820; uint256 idCommitment = poseidon.hash(idSecretHash); rln.register{value: MEMBERSHIP_DEPOSIT}(idCommitment); assertEq(rln.stakedAmounts(idCommitment), MEMBERSHIP_DEPOSIT); vm.expectRevert( - abi.encodeWithSelector( - InvalidWithdrawalAddress.selector, - address(0) - ) + abi.encodeWithSelector(InvalidReceiverAddress.selector, address(0)) ); - rln.withdraw(idSecretHash, payable(address(0))); + rln.slash(idSecretHash, payable(address(0))); } - function test__InvalidWithdraw__ToRlnAddress() public { + function test__InvalidSlash__ToRlnAddress() public { uint256 idSecretHash = 19014214495641488759237505126948346942972912379615652741039992445865937985820; uint256 idCommitment = poseidon.hash(idSecretHash); rln.register{value: MEMBERSHIP_DEPOSIT}(idCommitment); assertEq(rln.stakedAmounts(idCommitment), MEMBERSHIP_DEPOSIT); vm.expectRevert( abi.encodeWithSelector( - InvalidWithdrawalAddress.selector, + InvalidReceiverAddress.selector, address(rln) ) ); - rln.withdraw(idSecretHash, payable(address(rln))); + rln.slash(idSecretHash, payable(address(rln))); } - function test__InvalidWithdraw__InvalidIdCommitment( + function test__InvalidSlash__InvalidIdCommitment( uint256 idSecretHash ) public { uint256 idCommitment = poseidon.hash(idSecretHash); vm.expectRevert( abi.encodeWithSelector(MemberNotRegistered.selector, idCommitment) ); - rln.withdraw(idSecretHash, payable(address(this))); + rln.slash(idSecretHash, payable(address(this))); } // this shouldn't be possible, but just in case - function test__InvalidWithdraw__NoStake( + function test__InvalidSlash__NoStake( uint256 idSecretHash, address payable to ) public { @@ -240,7 +236,7 @@ contract RLNTest is Test { rln.register{value: MEMBERSHIP_DEPOSIT}(idCommitment); assertEq(rln.stakedAmounts(idCommitment), MEMBERSHIP_DEPOSIT); - rln.withdraw(idSecretHash, to); + rln.slash(idSecretHash, to); assertEq(rln.stakedAmounts(idCommitment), 0); assertEq(rln.members(idCommitment), false); @@ -255,10 +251,10 @@ contract RLNTest is Test { vm.expectRevert( abi.encodeWithSelector(MemberHasNoStake.selector, idCommitment) ); - rln.withdraw(idSecretHash, to); + rln.slash(idSecretHash, to); } - function test__ValidBatchWithdraw( + function test__ValidBatchSlash( uint256[] calldata idSecretHashes, address payable to ) public { @@ -281,7 +277,7 @@ contract RLNTest is Test { } uint256 balanceBefore = to.balance; - rln.withdrawBatch( + rln.slashBatch( idSecretHashes, repeatElementIntoArray(idSecretHashes.length, to) ); @@ -289,20 +285,22 @@ contract RLNTest is Test { assertEq(rln.stakedAmounts(idCommitments[i]), 0); assertEq(rln.members(idCommitments[i]), false); } + vm.prank(to); + rln.withdraw(); assertEq( to.balance, balanceBefore + MEMBERSHIP_DEPOSIT * idCommitmentlen ); } - function test__InvalidBatchWithdraw__EmptyBatch() public { + function test__InvalidBatchSlash__EmptyBatch() public { uint256[] memory idSecretHashes = new uint256[](0); address payable[] memory to = new address payable[](0); vm.expectRevert(EmptyBatch.selector); - rln.withdrawBatch(idSecretHashes, to); + rln.slashBatch(idSecretHashes, to); } - function test__InvalidBatchWithdraw__MismatchInputSize( + function test__InvalidBatchSlash__MismatchInputSize( uint256[] calldata idSecretHashes, address payable to ) public { @@ -319,9 +317,46 @@ contract RLNTest is Test { numberOfReceivers ) ); - rln.withdrawBatch( + rln.slashBatch( idSecretHashes, repeatElementIntoArray(numberOfReceivers, to) ); } + + function test__InvalidWithdraw__InsufficientWithdrawalBalance() public { + vm.expectRevert(InsufficientWithdrawalBalance.selector); + rln.withdraw(); + } + + function test__InvalidWithdraw__InsufficientContractBalance() public { + uint256 idSecretHash = 19014214495641488759237505126948346942972912379615652741039992445865937985820; + uint256 idCommitment = poseidon.hash(idSecretHash); + rln.register{value: MEMBERSHIP_DEPOSIT}(idCommitment); + assertEq(rln.stakedAmounts(idCommitment), MEMBERSHIP_DEPOSIT); + rln.slash(idSecretHash, payable(address(this))); + assertEq(rln.stakedAmounts(idCommitment), 0); + assertEq(rln.members(idCommitment), false); + + vm.deal(address(rln), 0); + vm.expectRevert(InsufficientContractBalance.selector); + rln.withdraw(); + } + + function test__ValidWithdraw(address payable to) public { + assumePayable(to); + assumeNoPrecompiles(to); + + uint256 idSecretHash = 19014214495641488759237505126948346942972912379615652741039992445865937985820; + uint256 idCommitment = poseidon.hash(idSecretHash); + + rln.register{value: MEMBERSHIP_DEPOSIT}(idCommitment); + assertEq(rln.stakedAmounts(idCommitment), MEMBERSHIP_DEPOSIT); + rln.slash(idSecretHash, to); + assertEq(rln.stakedAmounts(idCommitment), 0); + assertEq(rln.members(idCommitment), false); + + vm.prank(to); + rln.withdraw(); + assertEq(rln.withdrawalBalance(to), 0); + } }