From 20bb34e892cd83c75155acfcb87ed335e9bdeef7 Mon Sep 17 00:00:00 2001 From: Danish Arora Date: Thu, 30 Jan 2025 19:33:07 +0530 Subject: [PATCH] feat: rate limit --- src/contract/constants.ts | 28 ++++++ src/contract/rln_contract.spec.ts | 79 ++++++++++++++++- src/contract/rln_contract.ts | 136 +++++++++++++++++++++++++++++- src/rln.ts | 4 + 4 files changed, 241 insertions(+), 6 deletions(-) diff --git a/src/contract/constants.ts b/src/contract/constants.ts index 4d7900b..356d5b7 100644 --- a/src/contract/constants.ts +++ b/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; diff --git a/src/contract/rln_contract.spec.ts b/src/contract/rln_contract.spec.ts index 7e7914d..733d74c 100644 --- a/src/contract/rln_contract.spec.ts +++ b/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(); diff --git a/src/contract/rln_contract.ts b/src/contract/rln_contract.ts index 4ca8841..40b07cb 100644 --- a/src/contract/rln_contract.ts +++ b/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 = 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 The minimum rate limit in messages per epoch + */ + public async getMinRateLimit(): Promise { + const minRate = await this.contract.minMembershipRateLimit(); + return minRate.toNumber(); + } + + /** + * Gets the maximum allowed rate limit from the contract + * @returns Promise The maximum rate limit in messages per epoch + */ + public async getMaxRateLimit(): Promise { + const maxRate = await this.contract.maxMembershipRateLimit(); + return maxRate.toNumber(); + } + + /** + * Gets the maximum total rate limit across all memberships + * @returns Promise The maximum total rate limit in messages per epoch + */ + public async getMaxTotalRateLimit(): Promise { + const maxTotalRate = await this.contract.maxTotalRateLimit(); + return maxTotalRate.toNumber(); + } + + /** + * Gets the current total rate limit usage across all memberships + * @returns Promise The current total rate limit usage in messages per epoch + */ + public async getCurrentTotalRateLimit(): Promise { + const currentTotal = await this.contract.currentTotalRateLimit(); + return currentTotal.toNumber(); + } + + /** + * Gets the remaining available total rate limit that can be allocated + * @returns Promise The remaining rate limit that can be allocated + */ + public async getRemainingTotalRateLimit(): Promise { + 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 { + 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 { 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 { + 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 { 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(); diff --git a/src/rln.ts b/src/rln.ts index d5f196e..31edc98 100644 --- a/src/rln.ts +++ b/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 =