From d11fa4d45289fdff677382314aa1af941c9f7513 Mon Sep 17 00:00:00 2001
From: rymnc <43716372+rymnc@users.noreply.github.com>
Date: Fri, 25 Nov 2022 14:34:30 +0530
Subject: [PATCH] fix: disallow dupe registrations

---
 contracts/Rln.sol | 35 ++++++++++++++++-------------------
 test/index.ts     | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 46 insertions(+), 24 deletions(-)

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