diff --git a/packages/rln/src/contract/rln_contract.ts b/packages/rln/src/contract/rln_contract.ts index 12b733f6e1..f81a418135 100644 --- a/packages/rln/src/contract/rln_contract.ts +++ b/packages/rln/src/contract/rln_contract.ts @@ -1,11 +1,11 @@ import { Logger } from "@waku/utils"; import { hexToBytes } from "@waku/utils/bytes"; -import { ethers } from "ethers"; import type { RLNInstance } from "../rln.js"; import { MerkleRootTracker } from "../root_tracker.js"; import { zeroPadLE } from "../utils/bytes.js"; +import { ContractStateError } from "./errors.js"; import { RLNBaseContract } from "./rln_base_contract.js"; import { RLNContractInitOptions } from "./types.js"; @@ -14,6 +14,7 @@ const log = new Logger("waku:rln:contract"); export class RLNContract extends RLNBaseContract { private instance: RLNInstance; private merkleRootTracker: MerkleRootTracker; + private lastSyncedBlock: number = 0; /** * Asynchronous initializer for RLNContract. @@ -24,121 +25,145 @@ export class RLNContract extends RLNBaseContract { options: RLNContractInitOptions ): Promise { const rlnContract = new RLNContract(rlnInstance, options); - + await rlnContract.syncState(); return rlnContract; } + /** + * Override base contract method to keep Merkle tree in sync + * Registers a new membership with the given commitment and rate limit + */ + public override async registerMembership( + idCommitment: string, + rateLimit: number = this.getRateLimit() + ): Promise { + await super.registerMembership(idCommitment, rateLimit); + await this.syncState(); + } + + /** + * Override base contract method to keep Merkle tree in sync + * Erases an existing membership from the contract + */ + public override async eraseMembership( + idCommitment: string, + eraseFromMembershipSet: boolean = true + ): Promise { + await super.eraseMembership(idCommitment, eraseFromMembershipSet); + await this.syncState(); + } + + /** + * Gets the current Merkle root + * Returns the latest valid root or empty array if no roots exist + */ + public async getMerkleRoot(): Promise { + await this.syncState(); + const roots = this.merkleRootTracker.roots(); + return roots.length > 0 ? roots[0] : new Uint8Array(); + } + private constructor( rlnInstance: RLNInstance, options: RLNContractInitOptions ) { super(options); - this.instance = rlnInstance; - const initialRoot = rlnInstance.zerokit.getMerkleRoot(); this.merkleRootTracker = new MerkleRootTracker(5, initialRoot); } - public override processEvents(events: ethers.Event[]): void { - const toRemoveTable = new Map(); - const toInsertTable = new Map(); + /** + * Syncs the local Merkle tree with the current contract state + */ + private async syncState(): Promise { + try { + const currentBlock = await this.provider.getBlockNumber(); - events.forEach((evt) => { - if (!evt.args) { + // If we're already synced, just get new members + if (this.lastSyncedBlock > 0) { + await this.syncNewMembers(this.lastSyncedBlock, currentBlock); + this.lastSyncedBlock = currentBlock; return; } - if ( - evt.event === "MembershipErased" || - evt.event === "MembershipExpired" - ) { - let index = evt.args.index; + // First time sync - get all members + const nextIndex = await this.contract.nextFreeIndex(); + const members = await this.getMembersInRange(0, nextIndex.toNumber()); - if (!index) { - return; + // Clear existing members by deleting them one by one + // This effectively resets the tree without needing resetTree() + for (let i = 0; i < nextIndex.toNumber(); i++) { + try { + this.instance.zerokit.deleteMember(i); + } catch (error) { + // Ignore errors for non-existent members + continue; } - - if (typeof index === "number" || typeof index === "string") { - index = ethers.BigNumber.from(index); - } else { - log.error("Index is not a number or string", { - index, - event: evt - }); - return; - } - - const toRemoveVal = toRemoveTable.get(evt.blockNumber); - if (toRemoveVal != undefined) { - toRemoveVal.push(index.toNumber()); - toRemoveTable.set(evt.blockNumber, toRemoveVal); - } else { - toRemoveTable.set(evt.blockNumber, [index.toNumber()]); - } - } else if (evt.event === "MembershipRegistered") { - let eventsPerBlock = toInsertTable.get(evt.blockNumber); - if (eventsPerBlock == undefined) { - eventsPerBlock = []; - } - - eventsPerBlock.push(evt); - toInsertTable.set(evt.blockNumber, eventsPerBlock); } - }); - this.removeMembers(this.instance, toRemoveTable); - this.insertMembers(this.instance, toInsertTable); + // Insert all members + for (const member of members) { + const idCommitment = zeroPadLE(hexToBytes(member.idCommitment), 32); + this.instance.zerokit.insertMember(idCommitment); + } + + // Update root tracker + const currentRoot = this.instance.zerokit.getMerkleRoot(); + this.merkleRootTracker.pushRoot(currentBlock, currentRoot); + this.lastSyncedBlock = currentBlock; + + log.info( + `Synced ${members.length} members to current block ${currentBlock}` + ); + } catch (error) { + log.error("Failed to sync state", error); + throw new ContractStateError("Failed to sync contract state"); + } } - private insertMembers( - rlnInstance: RLNInstance, - toInsert: Map - ): void { - toInsert.forEach((events: ethers.Event[], blockNumber: number) => { - events.forEach((evt) => { - if (!evt.args) return; + /** + * Syncs new members added between fromBlock and toBlock + */ + private async syncNewMembers( + fromBlock: number, + toBlock: number + ): Promise { + // Get members that were added + const filter = this.contract.filters.MembershipRegistered(); + const addEvents = await this.contract.queryFilter( + filter, + fromBlock, + toBlock + ); - const _idCommitment = evt.args.idCommitment as string; - let index = evt.args.index; + // Get members that were removed + const removeFilter = this.contract.filters.MembershipErased(); + const removeEvents = await this.contract.queryFilter( + removeFilter, + fromBlock, + toBlock + ); - if (!_idCommitment || !index) { - return; - } + // Process removals first (in reverse block order) + for (const evt of removeEvents.sort( + (a, b) => b.blockNumber - a.blockNumber + )) { + if (!evt.args) continue; + const index = evt.args.index.toNumber(); + this.instance.zerokit.deleteMember(index); + this.merkleRootTracker.backFill(evt.blockNumber); + } - if (typeof index === "number" || typeof index === "string") { - index = ethers.BigNumber.from(index); - } - - const idCommitment = zeroPadLE(hexToBytes(_idCommitment), 32); - rlnInstance.zerokit.insertMember(idCommitment); - - const numericIndex = index.toNumber(); - this._members.set(numericIndex, { - index, - idCommitment: _idCommitment - }); - }); - - const currentRoot = rlnInstance.zerokit.getMerkleRoot(); - this.merkleRootTracker.pushRoot(blockNumber, currentRoot); - }); - } - - private removeMembers( - rlnInstance: RLNInstance, - toRemove: Map - ): void { - const removeDescending = new Map([...toRemove].reverse()); - removeDescending.forEach((indexes: number[], blockNumber: number) => { - indexes.forEach((index) => { - if (this._members.has(index)) { - this._members.delete(index); - rlnInstance.zerokit.deleteMember(index); - } - }); - - this.merkleRootTracker.backFill(blockNumber); - }); + // Then process additions + for (const evt of addEvents) { + if (!evt.args) continue; + const idCommitment = zeroPadLE(hexToBytes(evt.args.idCommitment), 32); + this.instance.zerokit.insertMember(idCommitment); + this.merkleRootTracker.pushRoot( + evt.blockNumber, + this.instance.zerokit.getMerkleRoot() + ); + } } }