Skip to content
This repository was archived by the owner on Feb 16, 2025. It is now read-only.

Commit

Permalink
feat: rate limit
Browse files Browse the repository at this point in the history
  • Loading branch information
danisharora099 committed Jan 30, 2025
1 parent 44fc3c2 commit 20bb34e
Show file tree
Hide file tree
Showing 4 changed files with 242 additions and 7 deletions.
28 changes: 28 additions & 0 deletions src/contract/constants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
79 changes: 77 additions & 2 deletions src/contract/rln_contract.spec.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand All @@ -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;

Expand All @@ -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();

Expand Down
138 changes: 133 additions & 5 deletions src/contract/rln_contract.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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");

Expand All @@ -26,7 +27,7 @@ interface Member {
interface RLNContractOptions {
signer: ethers.Signer;
address: string;
rateLimit: number;
rateLimit?: number;
}

interface FetchMembersOptions {
Expand All @@ -50,17 +51,90 @@ 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.
*/
private constructor(
rlnInstance: RLNInstance,
options: RLNContractInitOptions
) {
const { address, signer, rateLimit, contract } = options;

if (rateLimit === undefined) {
throw new Error("rateLimit must be provided in RLNContractOptions.");
const {
address,
signer,
rateLimit = DEFAULT_RATE_LIMIT,
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.rateLimit = rateLimit;
Expand Down Expand Up @@ -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,
Expand All @@ -259,6 +337,9 @@ export class RLNContract {
);

if (!memberRegistered || !memberRegistered.args) {
log(
"Failed to register membership: No MembershipRegistered event found"
);
return undefined;
}

Expand All @@ -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();
Expand All @@ -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: {
Expand All @@ -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,
Expand All @@ -316,6 +436,9 @@ export class RLNContract {
);

if (!memberRegistered || !memberRegistered.args) {
log(
"Failed to register membership with permit: No MembershipRegistered event found"
);
return undefined;
}

Expand All @@ -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();
Expand Down
4 changes: 4 additions & 0 deletions src/rln.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down

0 comments on commit 20bb34e

Please sign in to comment.