feat: rate limit

src/contract/constants.ts

@@ -6,3 +6,31 @@ export const SEPOLIA_CONTRACT = {
   address: "0xCB33Aa5B38d79E3D9Fa8B10afF38AA201399a7e3",
   abi: RLN_V2_ABI
 };
+
+/**
+ * Rate limit tiers (messages per epoch)
+ * Each membership can specify a rate limit within these bounds.
+ * @see https://github.com/waku-org/specs/blob/master/standards/core/rln-contract.md#implementation-suggestions
+ */
+export const RATE_LIMIT_TIERS = {
+  LOW: 20, // Suggested minimum rate - 20 messages per epoch
+  MEDIUM: 200,
+  HIGH: 600 // Suggested maximum rate - 600 messages per epoch
+} as const;
+
+// Default to maximum rate limit if not specified
+export const DEFAULT_RATE_LIMIT = RATE_LIMIT_TIERS.HIGH;
+
+/**
+ * Epoch length in seconds (10 minutes)
+ * This is a constant defined in the smart contract
+ */
+export const EPOCH_LENGTH = 600;
+
+// Global rate limit parameters
+export const RATE_LIMIT_PARAMS = {
+  MIN_RATE: RATE_LIMIT_TIERS.LOW,
+  MAX_RATE: RATE_LIMIT_TIERS.HIGH,
+  MAX_TOTAL_RATE: 160_000, // Maximum total rate limit across all memberships
+  EPOCH_LENGTH: EPOCH_LENGTH // Epoch length in seconds (10 minutes)
+} as const;

src/contract/rln_contract.spec.ts

@@ -1,5 +1,6 @@
 import { hexToBytes } from "@waku/utils/bytes";
-import chai from "chai";
+import chai, { expect } from "chai";
+import chaiAsPromised from "chai-as-promised";
 import spies from "chai-spies";
 import * as ethers from "ethers";
 import sinon, { SinonSandbox } from "sinon";
@@ -11,6 +12,7 @@ import { SEPOLIA_CONTRACT } from "./constants.js";
 import { RLNContract } from "./rln_contract.js";
 
 chai.use(spies);
+chai.use(chaiAsPromised);
 
 const DEFAULT_RATE_LIMIT = 10;
 
@@ -29,15 +31,88 @@ function mockRLNv2RegisteredEvent(idCommitment?: string): ethers.Event {
 
 describe("RLN Contract abstraction - RLN v2", () => {
   let sandbox: SinonSandbox;
+  let rlnInstance: any;
+  let mockedRegistryContract: any;
+  let rlnContract: RLNContract;
 
-  beforeEach(() => {
+  const mockRateLimits = {
+    minRate: 20,
+    maxRate: 600,
+    maxTotalRate: 1000,
+    currentTotalRate: 500
+  };
+
+  beforeEach(async () => {
     sandbox = sinon.createSandbox();
+    rlnInstance = await createRLN();
+    rlnInstance.zerokit.insertMember = () => undefined;
+
+    mockedRegistryContract = {
+      minMembershipRateLimit: () =>
+        Promise.resolve(ethers.BigNumber.from(mockRateLimits.minRate)),
+      maxMembershipRateLimit: () =>
+        Promise.resolve(ethers.BigNumber.from(mockRateLimits.maxRate)),
+      maxTotalRateLimit: () =>
+        Promise.resolve(ethers.BigNumber.from(mockRateLimits.maxTotalRate)),
+      currentTotalRateLimit: () =>
+        Promise.resolve(ethers.BigNumber.from(mockRateLimits.currentTotalRate)),
+      queryFilter: () => [mockRLNv2RegisteredEvent()],
+      provider: {
+        getLogs: () => [],
+        getBlockNumber: () => Promise.resolve(1000),
+        getNetwork: () => Promise.resolve({ chainId: 11155111 })
+      },
+      filters: {
+        MembershipRegistered: () => ({}),
+        MembershipRemoved: () => ({})
+      },
+      on: () => ({})
+    };
+
+    const provider = new ethers.providers.JsonRpcProvider();
+    const voidSigner = new ethers.VoidSigner(
+      SEPOLIA_CONTRACT.address,
+      provider
+    );
+    rlnContract = await RLNContract.init(rlnInstance, {
+      address: SEPOLIA_CONTRACT.address,
+      signer: voidSigner,
+      rateLimit: mockRateLimits.minRate,
+      contract: mockedRegistryContract as unknown as ethers.Contract
+    });
   });
 
   afterEach(() => {
     sandbox.restore();
   });
 
+  describe("Rate Limit Management", () => {
+    it("should get contract rate limit parameters", async () => {
+      const minRate = await rlnContract.getMinRateLimit();
+      const maxRate = await rlnContract.getMaxRateLimit();
+      const maxTotal = await rlnContract.getMaxTotalRateLimit();
+      const currentTotal = await rlnContract.getCurrentTotalRateLimit();
+
+      expect(minRate).to.equal(mockRateLimits.minRate);
+      expect(maxRate).to.equal(mockRateLimits.maxRate);
+      expect(maxTotal).to.equal(mockRateLimits.maxTotalRate);
+      expect(currentTotal).to.equal(mockRateLimits.currentTotalRate);
+    });
+
+    it("should calculate remaining total rate limit", async () => {
+      const remaining = await rlnContract.getRemainingTotalRateLimit();
+      expect(remaining).to.equal(
+        mockRateLimits.maxTotalRate - mockRateLimits.currentTotalRate
+      );
+    });
+
+    it("should set rate limit", async () => {
+      const newRate = 300; // Any value, since validation is done by contract
+      await rlnContract.setRateLimit(newRate);
+      expect(rlnContract.getRateLimit()).to.equal(newRate);
+    });
+  });
+
   it("should fetch members from events and store them in the RLN instance", async () => {
     const rlnInstance = await createRLN();
 

src/contract/rln_contract.ts

@@ -9,6 +9,7 @@ import { MerkleRootTracker } from "../root_tracker.js";
 import { zeroPadLE } from "../utils/bytes.js";
 
 import { RLN_V2_ABI } from "./abis/rlnv2.js";
+import { DEFAULT_RATE_LIMIT, RATE_LIMIT_PARAMS } from "./constants.js";
 
 const log = debug("waku:rln:contract");
 
@@ -26,7 +27,7 @@ interface Member {
 interface RLNContractOptions {
   signer: ethers.Signer;
   address: string;
-  rateLimit: number;
+  rateLimit?: number;
 }
 
 interface FetchMembersOptions {
@@ -50,6 +51,69 @@ export class RLNContract {
 
   private _members: Map<number, Member> = new Map();
 
+  /**
+   * Gets the current rate limit for this contract instance
+   */
+  public getRateLimit(): number {
+    return this.rateLimit;
+  }
+
+  /**
+   * Gets the minimum allowed rate limit from the contract
+   * @returns Promise<number> The minimum rate limit in messages per epoch
+   */
+  public async getMinRateLimit(): Promise<number> {
+    const minRate = await this.contract.minMembershipRateLimit();
+    return minRate.toNumber();
+  }
+
+  /**
+   * Gets the maximum allowed rate limit from the contract
+   * @returns Promise<number> The maximum rate limit in messages per epoch
+   */
+  public async getMaxRateLimit(): Promise<number> {
+    const maxRate = await this.contract.maxMembershipRateLimit();
+    return maxRate.toNumber();
+  }
+
+  /**
+   * Gets the maximum total rate limit across all memberships
+   * @returns Promise<number> The maximum total rate limit in messages per epoch
+   */
+  public async getMaxTotalRateLimit(): Promise<number> {
+    const maxTotalRate = await this.contract.maxTotalRateLimit();
+    return maxTotalRate.toNumber();
+  }
+
+  /**
+   * Gets the current total rate limit usage across all memberships
+   * @returns Promise<number> The current total rate limit usage in messages per epoch
+   */
+  public async getCurrentTotalRateLimit(): Promise<number> {
+    const currentTotal = await this.contract.currentTotalRateLimit();
+    return currentTotal.toNumber();
+  }
+
+  /**
+   * Gets the remaining available total rate limit that can be allocated
+   * @returns Promise<number> The remaining rate limit that can be allocated
+   */
+  public async getRemainingTotalRateLimit(): Promise<number> {
+    const [maxTotal, currentTotal] = await Promise.all([
+      this.contract.maxTotalRateLimit(),
+      this.contract.currentTotalRateLimit()
+    ]);
+    return maxTotal.sub(currentTotal).toNumber();
+  }
+
+  /**
+   * Updates the rate limit for future registrations
+   * @param newRateLimit The new rate limit to use
+   */
+  public async setRateLimit(newRateLimit: number): Promise<void> {
+    this.rateLimit = newRateLimit;
+  }
+
   /**
    * Private constructor to enforce the use of the async init method.
    */
@@ -57,10 +121,20 @@ export class RLNContract {
     rlnInstance: RLNInstance,
     options: RLNContractInitOptions
   ) {
-    const { address, signer, rateLimit, contract } = options;
+    const {
+      address,
+      signer,
+      rateLimit = DEFAULT_RATE_LIMIT,
+      contract
+    } = options;
 
-    if (rateLimit === undefined) {
-      throw new Error("rateLimit must be provided in RLNContractOptions.");
+    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.rateLimit = rateLimit;
@@ -245,6 +319,10 @@ export class RLNContract {
     identity: IdentityCredential
   ): Promise<DecryptedCredentials | undefined> {
     try {
+      log(
+        `Registering identity with rate limit: ${this.rateLimit} messages/epoch`
+      );
+
       const txRegisterResponse: ethers.ContractTransaction =
         await this.contract.register(
           identity.IDCommitmentBigInt,
@@ -259,6 +337,9 @@ export class RLNContract {
       );
 
       if (!memberRegistered || !memberRegistered.args) {
+        log(
+          "Failed to register membership: No MembershipRegistered event found"
+        );
         return undefined;
       }
 
@@ -268,6 +349,11 @@ export class RLNContract {
         index: memberRegistered.args.index
       };
 
+      log(
+        `Successfully registered membership with index ${decodedData.index} ` +
+          `and rate limit ${decodedData.rateLimit}`
+      );
+
       const network = await this.contract.provider.getNetwork();
       const address = this.contract.address;
       const membershipId = decodedData.index.toNumber();
@@ -286,6 +372,36 @@ export class RLNContract {
     }
   }
 
+  /**
+   * Helper method to get remaining messages in current epoch
+   * @param membershipId The ID of the membership to check
+   * @returns number of remaining messages allowed in current epoch
+   */
+  public async getRemainingMessages(membershipId: number): Promise<number> {
+    try {
+      const [startTime, , rateLimit] =
+        await this.contract.getMembershipInfo(membershipId);
+
+      // Calculate current epoch
+      const currentTime = Math.floor(Date.now() / 1000);
+      const epochsPassed = Math.floor(
+        (currentTime - startTime) / RATE_LIMIT_PARAMS.EPOCH_LENGTH
+      );
+      const currentEpochStart =
+        startTime + epochsPassed * RATE_LIMIT_PARAMS.EPOCH_LENGTH;
+
+      // Get message count in current epoch using contract's function
+      const messageCount = await this.contract.getMessageCount(
+        membershipId,
+        currentEpochStart
+      );
+      return Math.max(0, rateLimit.sub(messageCount).toNumber());
+    } catch (error) {
+      log(`Error getting remaining messages: ${(error as Error).message}`);
+      return 0; // Fail safe: assume no messages remaining on error
+    }
+  }
+
   public async registerWithPermitAndErase(
     identity: IdentityCredential,
     permit: {
@@ -298,6 +414,10 @@ export class RLNContract {
     idCommitmentsToErase: string[]
   ): Promise<DecryptedCredentials | undefined> {
     try {
+      log(
+        `Registering identity with permit and rate limit: ${this.rateLimit} messages/epoch`
+      );
+
       const txRegisterResponse: ethers.ContractTransaction =
         await this.contract.registerWithPermit(
           permit.owner,
@@ -316,6 +436,9 @@ export class RLNContract {
       );
 
       if (!memberRegistered || !memberRegistered.args) {
+        log(
+          "Failed to register membership with permit: No MembershipRegistered event found"
+        );
         return undefined;
       }
 
@@ -325,6 +448,11 @@ export class RLNContract {
         index: memberRegistered.args.index
       };
 
+      log(
+        `Successfully registered membership with permit. Index: ${decodedData.index}, ` +
+          `Rate limit: ${decodedData.rateLimit}, Erased ${idCommitmentsToErase.length} commitments`
+      );
+
       const network = await this.contract.provider.getNetwork();
       const address = this.contract.address;
       const membershipId = decodedData.index.toNumber();

src/rln.ts

@@ -76,6 +76,10 @@ type StartRLNOptions = {
    * If provided used for validating the network chainId and connecting to registry contract.
    */
   credentials?: EncryptedCredentials | DecryptedCredentials;
+  /**
+   * Rate limit for the member.
+   */
+  rateLimit?: number;
 };
 
 type RegisterMembershipOptions =