From 76f86de9c754d48f6a6c8d7b1dfe91e4d3080f1e Mon Sep 17 00:00:00 2001
From: Danish Arora <danisharora099@gmail.com>
Date: Thu, 27 Mar 2025 10:54:51 +0530
Subject: [PATCH] chore: add rateLimit to MembershipInfo

---
 packages/rln/src/contract/rln_contract.ts | 43 ++++++++++++-----------
 packages/rln/src/keystore/keystore.ts     |  3 +-
 packages/rln/src/keystore/types.ts        |  1 +
 3 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/packages/rln/src/contract/rln_contract.ts b/packages/rln/src/contract/rln_contract.ts
index e0fded9628..863a9ee12c 100644
--- a/packages/rln/src/contract/rln_contract.ts
+++ b/packages/rln/src/contract/rln_contract.ts
@@ -95,15 +95,7 @@ export class RLNContract {
       contract
     } = options;
 
-    if (
-      rateLimit < RATE_LIMIT_PARAMS.MIN_RATE ||
-      rateLimit > RATE_LIMIT_PARAMS.MAX_RATE
-    ) {
-      throw new Error(
-        `Rate limit must be between ${RATE_LIMIT_PARAMS.MIN_RATE} and ${RATE_LIMIT_PARAMS.MAX_RATE} messages per epoch`
-      );
-    }
-
+    this.validateRateLimit(rateLimit);
     this.rateLimit = rateLimit;
 
     const initialRoot = rlnInstance.zerokit.getMerkleRoot();
@@ -118,6 +110,21 @@ export class RLNContract {
     this._membersExpiredFilter = this.contract.filters.MembershipExpired();
   }
 
+  /**
+   * Validates that the rate limit is within the allowed range
+   * @throws Error if the rate limit is outside the allowed range
+   */
+  private validateRateLimit(rateLimit: number): void {
+    if (
+      rateLimit < RATE_LIMIT_PARAMS.MIN_RATE ||
+      rateLimit > RATE_LIMIT_PARAMS.MAX_RATE
+    ) {
+      throw new Error(
+        `Rate limit must be between ${RATE_LIMIT_PARAMS.MIN_RATE} and ${RATE_LIMIT_PARAMS.MAX_RATE} messages per epoch`
+      );
+    }
+  }
+
   /**
    * Gets the current rate limit for this contract instance
    */
@@ -192,6 +199,7 @@ export class RLNContract {
    * @param newRateLimit The new rate limit to use
    */
   public async setRateLimit(newRateLimit: number): Promise<void> {
+    this.validateRateLimit(newRateLimit);
     this.rateLimit = newRateLimit;
   }
 
@@ -469,7 +477,8 @@ export class RLNContract {
         membership: {
           address,
           treeIndex: membershipId,
-          chainId: network.chainId
+          chainId: network.chainId,
+          rateLimit: decodedData.membershipRateLimit.toNumber()
         }
       };
     } catch (error) {
@@ -594,7 +603,8 @@ export class RLNContract {
         membership: {
           address,
           treeIndex: membershipId,
-          chainId: network.chainId
+          chainId: network.chainId,
+          rateLimit: decodedData.membershipRateLimit.toNumber()
         }
       };
     } catch (error) {
@@ -669,16 +679,9 @@ export class RLNContract {
 
   public async registerMembership(
     idCommitment: string,
-    rateLimit: number = DEFAULT_RATE_LIMIT
+    rateLimit: number = this.rateLimit
   ): Promise<ethers.ContractTransaction> {
-    if (
-      rateLimit < RATE_LIMIT_PARAMS.MIN_RATE ||
-      rateLimit > RATE_LIMIT_PARAMS.MAX_RATE
-    ) {
-      throw new Error(
-        `Rate limit must be between ${RATE_LIMIT_PARAMS.MIN_RATE} and ${RATE_LIMIT_PARAMS.MAX_RATE}`
-      );
-    }
+    this.validateRateLimit(rateLimit);
     return this.contract.register(idCommitment, rateLimit, []);
   }
 
diff --git a/packages/rln/src/keystore/keystore.ts b/packages/rln/src/keystore/keystore.ts
index 84836186bf..f39bfaf8a0 100644
--- a/packages/rln/src/keystore/keystore.ts
+++ b/packages/rln/src/keystore/keystore.ts
@@ -274,7 +274,8 @@ export class Keystore {
         membership: {
           treeIndex: _.get(obj, "treeIndex"),
           chainId: _.get(obj, "membershipContract.chainId"),
-          address: _.get(obj, "membershipContract.address")
+          address: _.get(obj, "membershipContract.address"),
+          rateLimit: _.get(obj, "membershipContract.rateLimit")
         }
       };
     } catch (err) {
diff --git a/packages/rln/src/keystore/types.ts b/packages/rln/src/keystore/types.ts
index 9d8e86a891..a14867c1e1 100644
--- a/packages/rln/src/keystore/types.ts
+++ b/packages/rln/src/keystore/types.ts
@@ -11,6 +11,7 @@ export type MembershipInfo = {
   chainId: number;
   address: string;
   treeIndex: number;
+  rateLimit: number;
 };
 
 export type KeystoreEntity = {