fix: remove .transfer from loops

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
+/// Invalid receiver address, when the receiver is the contract itself or 0x0
-error InvalidWithdrawalAddress(address to);
+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 {
+    function slash(uint256 secret, address payable receiver) external {
-        _withdraw(secret, receiver);
+        _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

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

test/RLN.t.sol

@@ -167,10 +167,7 @@ contract RLNTest is Test {
         rln.registerBatch{value: badDepositAmount}(idCommitments);
     }
 
-    function test__ValidWithdraw(
+    function test__ValidSlash(uint256 idSecretHash, address payable to) public {
-        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(
+            abi.encodeWithSelector(InvalidReceiverAddress.selector, address(0))
-                InvalidWithdrawalAddress.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);
+    }
 }