@waku/rln 0.0.2-09108d9.0 → 0.0.2-8a6571f.0

package/src/contract/rln_contract.ts

@@ -6,9 +6,10 @@ import type { IdentityCredential } from "../identity.js";
 import type { DecryptedCredentials } from "../keystore/index.js";
 import type { RLNInstance } from "../rln.js";
 import { MerkleRootTracker } from "../root_tracker.js";
-import { zeroPadLE } from "../utils/index.js";
+import { zeroPadLE } from "../utils/bytes.js";
 
-import { RLN_REGISTRY_ABI, RLN_STORAGE_ABI } from "./constants.js";
+import { RLN_ABI } from "./abi.js";
+import { DEFAULT_RATE_LIMIT, RATE_LIMIT_PARAMS } from "./constants.js";
 
 const log = new Logger("waku:rln:contract");
 
@@ -17,18 +18,21 @@ type Member = {
   index: ethers.BigNumber;
 };
 
-type Signer = ethers.Signer;
-
-type RLNContractOptions = {
-  signer: Signer;
-  registryAddress: string;
-};
+interface RLNContractOptions {
+  signer: ethers.Signer;
+  address: string;
+  rateLimit?: number;
+}
 
-type RLNStorageOptions = {
-  storageIndex?: number;
-};
+interface RLNContractInitOptions extends RLNContractOptions {
+  contract?: ethers.Contract;
+}
 
-type RLNContractInitOptions = RLNContractOptions & RLNStorageOptions;
+export interface MembershipRegisteredEvent {
+  idCommitment: string;
+  rateLimit: number;
+  index: ethers.BigNumber;
+}
 
 type FetchMembersOptions = {
   fromBlock?: number;
@@ -36,80 +40,157 @@ type FetchMembersOptions = {
   fetchChunks?: number;
 };
 
+export interface MembershipInfo {
+  index: ethers.BigNumber;
+  idCommitment: string;
+  rateLimit: number;
+  startBlock: number;
+  endBlock: number;
+  state: MembershipState;
+}
+
+export enum MembershipState {
+  Active = "Active",
+  GracePeriod = "GracePeriod",
+  Expired = "Expired",
+  ErasedAwaitsWithdrawal = "ErasedAwaitsWithdrawal"
+}
+
 export class RLNContract {
-  private registryContract: ethers.Contract;
+  public contract: ethers.Contract;
   private merkleRootTracker: MerkleRootTracker;
 
   private deployBlock: undefined | number;
-  private storageIndex: undefined | number;
-  private storageContract: undefined | ethers.Contract;
-  private _membersFilter: undefined | ethers.EventFilter;
+  private rateLimit: number;
 
   private _members: Map<number, Member> = new Map();
+  private _membersFilter: ethers.EventFilter;
+  private _membersRemovedFilter: ethers.EventFilter;
 
+  /**
+   * Asynchronous initializer for RLNContract.
+   * Allows injecting a mocked contract for testing purposes.
+   */
   public static async init(
     rlnInstance: RLNInstance,
     options: RLNContractInitOptions
   ): Promise<RLNContract> {
     const rlnContract = new RLNContract(rlnInstance, options);
 
-    await rlnContract.initStorageContract(options.signer);
     await rlnContract.fetchMembers(rlnInstance);
     rlnContract.subscribeToMembers(rlnInstance);
 
     return rlnContract;
   }
 
-  public constructor(
+  private constructor(
     rlnInstance: RLNInstance,
-    { registryAddress, signer }: RLNContractOptions
+    options: RLNContractInitOptions
   ) {
+    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;
+
     const initialRoot = rlnInstance.zerokit.getMerkleRoot();
 
-    this.registryContract = new ethers.Contract(
-      registryAddress,
-      RLN_REGISTRY_ABI,
-      signer
-    );
+    // Use the injected contract if provided; otherwise, instantiate a new one.
+    this.contract = contract || new ethers.Contract(address, RLN_ABI, signer);
     this.merkleRootTracker = new MerkleRootTracker(5, initialRoot);
+
+    // Initialize event filters for MembershipRegistered and MembershipErased
+    this._membersFilter = this.contract.filters.MembershipRegistered();
+    this._membersRemovedFilter = this.contract.filters.MembershipErased();
   }
 
-  private async initStorageContract(
-    signer: Signer,
-    options: RLNStorageOptions = {}
-  ): Promise<void> {
-    const storageIndex = options?.storageIndex
-      ? options.storageIndex
-      : await this.registryContract.usingStorageIndex();
-    const storageAddress = await this.registryContract.storages(storageIndex);
+  /**
+   * Gets the current rate limit for this contract instance
+   */
+  public getRateLimit(): number {
+    return this.rateLimit;
+  }
 
-    if (!storageAddress || storageAddress === ethers.constants.AddressZero) {
-      throw Error("No RLN Storage initialized on registry contract.");
-    }
+  /**
+   * Gets the contract address
+   */
+  public get address(): string {
+    return this.contract.address;
+  }
 
-    this.storageIndex = storageIndex;
-    this.storageContract = new ethers.Contract(
-      storageAddress,
-      RLN_STORAGE_ABI,
-      signer
-    );
-    this._membersFilter = this.storageContract.filters.MemberRegistered();
+  /**
+   * Gets the contract provider
+   */
+  public get provider(): ethers.providers.Provider {
+    return this.contract.provider;
+  }
 
-    this.deployBlock = await this.storageContract.deployedBlockNumber();
+  /**
+   * 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();
   }
 
-  public get registry(): ethers.Contract {
-    if (!this.registryContract) {
-      throw Error("Registry contract was not initialized");
-    }
-    return this.registryContract as ethers.Contract;
+  /**
+   * 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();
   }
 
-  public get contract(): ethers.Contract {
-    if (!this.storageContract) {
-      throw Error("Storage contract was not initialized");
-    }
-    return this.storageContract as ethers.Contract;
+  /**
+   * 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;
   }
 
   public get members(): Member[] {
@@ -123,7 +204,14 @@ export class RLNContract {
     if (!this._membersFilter) {
       throw Error("Members filter was not initialized.");
     }
-    return this._membersFilter as ethers.EventFilter;
+    return this._membersFilter;
+  }
+
+  private get membersRemovedFilter(): ethers.EventFilter {
+    if (!this._membersRemovedFilter) {
+      throw Error("MembersErased filter was not initialized.");
+    }
+    return this._membersRemovedFilter;
   }
 
   public async fetchMembers(
@@ -135,7 +223,14 @@ export class RLNContract {
       ...options,
       membersFilter: this.membersFilter
     });
-    this.processEvents(rlnInstance, registeredMemberEvents);
+    const removedMemberEvents = await queryFilter(this.contract, {
+      fromBlock: this.deployBlock,
+      ...options,
+      membersFilter: this.membersRemovedFilter
+    });
+
+    const events = [...registeredMemberEvents, ...removedMemberEvents];
+    this.processEvents(rlnInstance, events);
   }
 
   public processEvents(rlnInstance: RLNInstance, events: ethers.Event[]): void {
@@ -147,8 +242,8 @@ export class RLNContract {
         return;
       }
 
-      if (evt.removed) {
-        const index: ethers.BigNumber = evt.args.index;
+      if (evt.event === "MembershipErased") {
+        const index = evt.args.index as ethers.BigNumber;
         const toRemoveVal = toRemoveTable.get(evt.blockNumber);
         if (toRemoveVal != undefined) {
           toRemoveVal.push(index.toNumber());
@@ -156,7 +251,7 @@ export class RLNContract {
         } else {
           toRemoveTable.set(evt.blockNumber, [index.toNumber()]);
         }
-      } else {
+      } else if (evt.event === "MembershipRegistered") {
         let eventsPerBlock = toInsertTable.get(evt.blockNumber);
         if (eventsPerBlock == undefined) {
           eventsPerBlock = [];
@@ -177,18 +272,20 @@ export class RLNContract {
   ): void {
     toInsert.forEach((events: ethers.Event[], blockNumber: number) => {
       events.forEach((evt) => {
-        const _idCommitment = evt?.args?.idCommitment;
-        const index: ethers.BigNumber = evt?.args?.index;
+        if (!evt.args) return;
+
+        const _idCommitment = evt.args.idCommitment as string;
+        const index = evt.args.index as ethers.BigNumber;
 
         if (!_idCommitment || !index) {
           return;
         }
 
-        const idCommitment = zeroPadLE(hexToBytes(_idCommitment?._hex), 32);
+        const idCommitment = zeroPadLE(hexToBytes(_idCommitment), 32);
         rlnInstance.zerokit.insertMember(idCommitment);
         this._members.set(index.toNumber(), {
           index,
-          idCommitment: _idCommitment?._hex
+          idCommitment: _idCommitment
         });
       });
 
@@ -201,7 +298,7 @@ export class RLNContract {
     rlnInstance: RLNInstance,
     toRemove: Map<number, number[]>
   ): void {
-    const removeDescending = new Map([...toRemove].sort().reverse());
+    const removeDescending = new Map([...toRemove].reverse());
     removeDescending.forEach((indexes: number[], blockNumber: number) => {
       indexes.forEach((index) => {
         if (this._members.has(index)) {
@@ -215,63 +312,291 @@ export class RLNContract {
   }
 
   public subscribeToMembers(rlnInstance: RLNInstance): void {
-    this.contract.on(this.membersFilter, (_pubkey, _index, event) =>
-      this.processEvents(rlnInstance, [event])
+    this.contract.on(
+      this.membersFilter,
+      (
+        _idCommitment: string,
+        _rateLimit: number,
+        _index: ethers.BigNumber,
+        event: ethers.Event
+      ) => {
+        this.processEvents(rlnInstance, [event]);
+      }
+    );
+
+    this.contract.on(
+      this.membersRemovedFilter,
+      (
+        _idCommitment: string,
+        _rateLimit: number,
+        _index: ethers.BigNumber,
+        event: ethers.Event
+      ) => {
+        this.processEvents(rlnInstance, [event]);
+      }
     );
   }
 
   public async registerWithIdentity(
     identity: IdentityCredential
   ): Promise<DecryptedCredentials | undefined> {
-    if (this.storageIndex === undefined) {
-      throw Error(
-        "Cannot register credential, no storage contract index found."
+    try {
+      log.info(
+        `Registering identity with rate limit: ${this.rateLimit} messages/epoch`
       );
-    }
-    const txRegisterResponse: ethers.ContractTransaction =
-      await this.registryContract["register(uint16,uint256)"](
-        this.storageIndex,
-        identity.IDCommitmentBigInt,
-        { gasLimit: 100000 }
+
+      const txRegisterResponse: ethers.ContractTransaction =
+        await this.contract.register(
+          identity.IDCommitmentBigInt,
+          this.rateLimit,
+          [],
+          { gasLimit: 300000 }
+        );
+      const txRegisterReceipt = await txRegisterResponse.wait();
+
+      const memberRegistered = txRegisterReceipt.events?.find(
+        (event) => event.event === "MembershipRegistered"
       );
-    const txRegisterReceipt = await txRegisterResponse.wait();
 
-    // assumption: register(uint16,uint256) emits one event
-    const memberRegistered = txRegisterReceipt?.events?.[0];
+      if (!memberRegistered || !memberRegistered.args) {
+        log.error(
+          "Failed to register membership: No MembershipRegistered event found"
+        );
+        return undefined;
+      }
+
+      const decodedData: MembershipRegisteredEvent = {
+        idCommitment: memberRegistered.args.idCommitment,
+        rateLimit: memberRegistered.args.rateLimit,
+        index: memberRegistered.args.index
+      };
+
+      log.info(
+        `Successfully registered membership with index ${decodedData.index} ` +
+          `and rate limit ${decodedData.rateLimit}`
+      );
 
-    if (!memberRegistered) {
+      const network = await this.contract.provider.getNetwork();
+      const address = this.contract.address;
+      const membershipId = decodedData.index.toNumber();
+
+      return {
+        identity,
+        membership: {
+          address,
+          treeIndex: membershipId,
+          chainId: network.chainId
+        }
+      };
+    } catch (error) {
+      log.error(`Error in registerWithIdentity: ${(error as Error).message}`);
       return undefined;
     }
+  }
 
-    const decodedData = this.contract.interface.decodeEventLog(
-      "MemberRegistered",
-      memberRegistered.data
-    );
+  /**
+   * 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(
+        `Error getting remaining messages: ${(error as Error).message}`
+      );
+      return 0; // Fail safe: assume no messages remaining on error
+    }
+  }
 
-    const network = await this.registryContract.provider.getNetwork();
-    const address = this.registryContract.address;
-    const membershipId = decodedData.index.toNumber();
+  public async registerWithPermitAndErase(
+    identity: IdentityCredential,
+    permit: {
+      owner: string;
+      deadline: number;
+      v: number;
+      r: string;
+      s: string;
+    },
+    idCommitmentsToErase: string[]
+  ): Promise<DecryptedCredentials | undefined> {
+    try {
+      log.info(
+        `Registering identity with permit and rate limit: ${this.rateLimit} messages/epoch`
+      );
+
+      const txRegisterResponse: ethers.ContractTransaction =
+        await this.contract.registerWithPermit(
+          permit.owner,
+          permit.deadline,
+          permit.v,
+          permit.r,
+          permit.s,
+          identity.IDCommitmentBigInt,
+          this.rateLimit,
+          idCommitmentsToErase.map((id) => ethers.BigNumber.from(id))
+        );
+      const txRegisterReceipt = await txRegisterResponse.wait();
+
+      const memberRegistered = txRegisterReceipt.events?.find(
+        (event) => event.event === "MembershipRegistered"
+      );
 
-    return {
-      identity,
-      membership: {
-        address,
-        treeIndex: membershipId,
-        chainId: network.chainId
+      if (!memberRegistered || !memberRegistered.args) {
+        log.error(
+          "Failed to register membership with permit: No MembershipRegistered event found"
+        );
+        return undefined;
       }
-    };
+
+      const decodedData: MembershipRegisteredEvent = {
+        idCommitment: memberRegistered.args.idCommitment,
+        rateLimit: memberRegistered.args.rateLimit,
+        index: memberRegistered.args.index
+      };
+
+      log.info(
+        `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();
+
+      return {
+        identity,
+        membership: {
+          address,
+          treeIndex: membershipId,
+          chainId: network.chainId
+        }
+      };
+    } catch (error) {
+      log.error(
+        `Error in registerWithPermitAndErase: ${(error as Error).message}`
+      );
+      return undefined;
+    }
   }
 
   public roots(): Uint8Array[] {
     return this.merkleRootTracker.roots();
   }
+
+  public async withdraw(token: string, holder: string): Promise<void> {
+    try {
+      const tx = await this.contract.withdraw(token, { from: holder });
+      await tx.wait();
+    } catch (error) {
+      log.error(`Error in withdraw: ${(error as Error).message}`);
+    }
+  }
+
+  public async getMembershipInfo(
+    idCommitment: string
+  ): Promise<MembershipInfo | undefined> {
+    try {
+      const [startBlock, endBlock, rateLimit] =
+        await this.contract.getMembershipInfo(idCommitment);
+      const currentBlock = await this.contract.provider.getBlockNumber();
+
+      let state: MembershipState;
+      if (currentBlock < startBlock) {
+        state = MembershipState.Active;
+      } else if (currentBlock < endBlock) {
+        state = MembershipState.GracePeriod;
+      } else {
+        state = MembershipState.Expired;
+      }
+
+      const index = await this.getMemberIndex(idCommitment);
+      if (!index) return undefined;
+
+      return {
+        index,
+        idCommitment,
+        rateLimit: rateLimit.toNumber(),
+        startBlock: startBlock.toNumber(),
+        endBlock: endBlock.toNumber(),
+        state
+      };
+    } catch (error) {
+      return undefined;
+    }
+  }
+
+  public async extendMembership(
+    idCommitment: string
+  ): Promise<ethers.ContractTransaction> {
+    return this.contract.extendMemberships([idCommitment]);
+  }
+
+  public async eraseMembership(
+    idCommitment: string,
+    eraseFromMembershipSet: boolean = true
+  ): Promise<ethers.ContractTransaction> {
+    return this.contract.eraseMemberships(
+      [idCommitment],
+      eraseFromMembershipSet
+    );
+  }
+
+  public async registerMembership(
+    idCommitment: string,
+    rateLimit: number = DEFAULT_RATE_LIMIT
+  ): 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}`
+      );
+    }
+    return this.contract.register(idCommitment, rateLimit, []);
+  }
+
+  private async getMemberIndex(
+    idCommitment: string
+  ): Promise<ethers.BigNumber | undefined> {
+    try {
+      const events = await this.contract.queryFilter(
+        this.contract.filters.MembershipRegistered(idCommitment)
+      );
+      if (events.length === 0) return undefined;
+
+      // Get the most recent registration event
+      const event = events[events.length - 1];
+      return event.args?.index;
+    } catch (error) {
+      return undefined;
+    }
+  }
 }
 
-type CustomQueryOptions = FetchMembersOptions & {
+interface CustomQueryOptions extends FetchMembersOptions {
   membersFilter: ethers.EventFilter;
-};
+}
 
-// these value should be tested on other networks
+// These values should be tested on other networks
 const FETCH_CHUNK = 5;
 const BLOCK_RANGE = 3000;
 
@@ -286,18 +611,18 @@ async function queryFilter(
     fetchChunks = FETCH_CHUNK
   } = options;
 
-  if (!fromBlock) {
+  if (fromBlock === undefined) {
     return contract.queryFilter(membersFilter);
   }
 
-  if (!contract.signer.provider) {
-    throw Error("No provider found on the contract's signer.");
+  if (!contract.provider) {
+    throw Error("No provider found on the contract.");
   }
 
-  const toBlock = await contract.signer.provider.getBlockNumber();
+  const toBlock = await contract.provider.getBlockNumber();
 
   if (toBlock - fromBlock < fetchRange) {
-    return contract.queryFilter(membersFilter);
+    return contract.queryFilter(membersFilter, fromBlock, toBlock);
   }
 
   const events: ethers.Event[][] = [];
@@ -319,7 +644,7 @@ function splitToChunks(
   to: number,
   step: number
 ): Array<[number, number]> {
-  const chunks = [];
+  const chunks: Array<[number, number]> = [];
 
   let left = from;
   while (left < to) {
@@ -345,9 +670,18 @@ function* takeN<T>(array: T[], size: number): Iterable<T[]> {
   }
 }
 
-function ignoreErrors<T>(promise: Promise<T>, defaultValue: T): Promise<T> {
-  return promise.catch((err) => {
-    log.info(`Ignoring an error during query: ${err?.message}`);
+async function ignoreErrors<T>(
+  promise: Promise<T>,
+  defaultValue: T
+): Promise<T> {
+  try {
+    return await promise;
+  } catch (err: unknown) {
+    if (err instanceof Error) {
+      log.info(`Ignoring an error during query: ${err.message}`);
+    } else {
+      log.info(`Ignoring an unknown error during query`);
+    }
     return defaultValue;
-  });
+  }
 }