diff --git a/contracts/Rln.sol b/contracts/Rln.sol index 4df5fcc..c3e2259 100644 --- a/contracts/Rln.sol +++ b/contracts/Rln.sol @@ -8,12 +8,14 @@ contract RLN { uint256 public immutable SET_SIZE; uint256 public pubkeyIndex = 0; + // This mapping is used to keep track of the public keys that have been registered + // with the stake mapping(uint256 => uint256) public members; IPoseidonHasher public poseidonHasher; event MemberRegistered(uint256 pubkey, uint256 index); - event MemberWithdrawn(uint256 pubkey, uint256 index); + event MemberWithdrawn(uint256 pubkey); constructor( uint256 membershipDeposit, @@ -27,6 +29,7 @@ contract RLN { } function register(uint256 pubkey) external payable { + require(members[pubkey] == 0, "RLN, register: pubkey already registered"); require(pubkeyIndex < SET_SIZE, "RLN, register: set is full"); require(msg.value == MEMBERSHIP_DEPOSIT, "RLN, register: membership deposit is not satisfied"); _register(pubkey); @@ -42,54 +45,48 @@ contract RLN { } function _register(uint256 pubkey) internal { - members[pubkeyIndex] = pubkey; + // Set the pubkey to the value of the tx + members[pubkey] = msg.value; emit MemberRegistered(pubkey, pubkeyIndex); pubkeyIndex += 1; } function withdrawBatch( uint256[] calldata secrets, - uint256[] calldata pubkeyIndexes, address payable[] calldata receivers ) external { uint256 batchSize = secrets.length; require(batchSize != 0, "RLN, withdrawBatch: batch size zero"); - require(batchSize == pubkeyIndexes.length, "RLN, withdrawBatch: batch size mismatch pubkey indexes"); require(batchSize == receivers.length, "RLN, withdrawBatch: batch size mismatch receivers"); for (uint256 i = 0; i < batchSize; i++) { - _withdraw(secrets[i], pubkeyIndexes[i], receivers[i]); + _withdraw(secrets[i], receivers[i]); } } function withdraw( uint256 secret, - uint256 _pubkeyIndex, address payable receiver ) external { - _withdraw(secret, _pubkeyIndex, receiver); + _withdraw(secret, receiver); } function _withdraw( uint256 secret, - uint256 _pubkeyIndex, address payable receiver ) internal { - require(_pubkeyIndex < SET_SIZE, "RLN, _withdraw: invalid pubkey index"); - require(members[_pubkeyIndex] != 0, "RLN, _withdraw: member doesn't exist"); - require(receiver != address(0), "RLN, _withdraw: empty receiver address"); - // derive public key uint256 pubkey = hash(secret); - require(members[_pubkeyIndex] == pubkey, "RLN, _withdraw: not verified"); - - // delete member - members[_pubkeyIndex] = 0; - + require(members[pubkey] != 0, "RLN, _withdraw: member doesn't exist"); + require(receiver != address(0), "RLN, _withdraw: empty receiver address"); + // refund deposit - (bool sent, bytes memory data) = receiver.call{value: MEMBERSHIP_DEPOSIT}(""); + (bool sent, bytes memory data) = receiver.call{value: members[pubkey]}(""); require(sent, "transfer failed"); - emit MemberWithdrawn(pubkey, _pubkeyIndex); + // delete member only if refund is successful + members[pubkey] = 0; + + emit MemberWithdrawn(pubkey); } function hash(uint256 input) internal view returns (uint256) { diff --git a/test/index.ts b/test/index.ts index 59a8e04..3588506 100644 --- a/test/index.ts +++ b/test/index.ts @@ -1,4 +1,4 @@ -import { assert } from "chai"; +import { assert, expect } from "chai"; import { ethers } from "hardhat"; describe("Rln", function () { @@ -34,16 +34,41 @@ describe("Rln", function () { // We withdraw our id_commitment const receiver_address = "0x000000000000000000000000000000000000dead"; - const res_withdraw = await rln.withdraw(id_secret, reg_tree_index, receiver_address); + const res_withdraw = await rln.withdraw(id_secret, receiver_address); const txWithdrawReceipt = await res_withdraw.wait(); const wit_pubkey = txWithdrawReceipt.events[0].args.pubkey; - const wit_tree_index = txWithdrawReceipt.events[0].args.index; // We ensure the registered id_commitment is the one we passed and that the index is the same - assert(wit_pubkey.toHexString() === id_commitment, "withdraw commitment doesn't match registered commitmet"); - assert(wit_tree_index.toHexString() === reg_tree_index.toHexString(), "withdraw index doesn't match registered index"); + assert(wit_pubkey.toHexString() === id_commitment, "withdraw commitment doesn't match registered commitment"); + const pubkeyIndex = (await rln.pubkeyIndex()).toNumber(); + assert(pubkeyIndex === 1, "pubkeyIndex should be 1"); + }); + it("should not allow dupe registrations", async () => { + const PoseidonHasher = await ethers.getContractFactory("PoseidonHasher"); + const poseidonHasher = await PoseidonHasher.deploy(); + + await poseidonHasher.deployed(); + + console.log("PoseidonHasher deployed to:", poseidonHasher.address); + + const Rln = await ethers.getContractFactory("RLN"); + const rln = await Rln.deploy(1000000000000000, 20, poseidonHasher.address); + + await rln.deployed(); + + console.log("Rln deployed to:", rln.address); + + const price = await rln.MEMBERSHIP_DEPOSIT(); + + // A valid id_commitment generated in rust + const id_commitment = "0x0c3ac305f6a4fe9bfeb3eba978bc876e2a99208b8b56c80160cfb54ba8f02368" + + await rln.register(id_commitment, {value: price}); + + expect(rln.register(id_commitment, {value: price})).to.be.revertedWith("RLN, register: pubkey already registered"); + }); }); \ No newline at end of file