js-waku/packages/rln/src/zerokit.ts

import * as zerokitRLN from "@waku/zerokit-rln-wasm";
import { generateSeededExtendedMembershipKey } from "@waku/zerokit-rln-wasm-utils";

import { DEFAULT_RATE_LIMIT, RATE_LIMIT_PARAMS } from "./contract/constants.js";
import { IdentityCredential } from "./identity.js";
import { WitnessCalculator } from "./resources/witness_calculator";
import { BytesUtils } from "./utils/bytes.js";
import { dateToEpoch, epochIntToBytes } from "./utils/epoch.js";
import { poseidonHash, sha256 } from "./utils/hash.js";
import { MERKLE_TREE_DEPTH } from "./utils/merkle.js";

export class Zerokit {
  public constructor(
    private readonly zkRLN: number,
    private readonly witnessCalculator: WitnessCalculator,
    private readonly _rateLimit: number = DEFAULT_RATE_LIMIT,
    private readonly rlnIdentifier: Uint8Array = new TextEncoder().encode(
      "rln/waku-rln-relay/v2.0.0"
    )
  ) {}

  public get getZkRLN(): number {
    return this.zkRLN;
  }

  public get getWitnessCalculator(): WitnessCalculator {
    return this.witnessCalculator;
  }

  public get rateLimit(): number {
    return this._rateLimit;
  }

  public generateSeededIdentityCredential(seed: string): IdentityCredential {
    const stringEncoder = new TextEncoder();
    const seedBytes = stringEncoder.encode(seed);
    const memKeys = generateSeededExtendedMembershipKey(seedBytes, true);
    return IdentityCredential.fromBytes(memKeys);
  }

  private async serializeWitness(
    idSecretHash: Uint8Array,
    pathElements: Uint8Array[],
    identityPathIndex: Uint8Array[],
    msg: Uint8Array,
    epoch: Uint8Array,
    rateLimit: number,
    messageNumberId: number
  ): Promise<Uint8Array> {
    const externalNullifier = poseidonHash(
      sha256(epoch),
      sha256(this.rlnIdentifier)
    );
    const pathElementsBytes = new Uint8Array(8 + pathElements.length * 32);
    BytesUtils.writeUintLE(pathElementsBytes, pathElements.length, 0, 8);
    for (let i = 0; i < pathElements.length; i++) {
      // We assume that the path elements are already in little-endian format
      pathElementsBytes.set(pathElements[i], 8 + i * 32);
    }
    const identityPathIndexBytes = new Uint8Array(
      8 + identityPathIndex.length * 1
    );
    BytesUtils.writeUintLE(
      identityPathIndexBytes,
      identityPathIndex.length,
      0,
      8
    );
    for (let i = 0; i < identityPathIndex.length; i++) {
      // We assume that each identity path index is already in little-endian format
      identityPathIndexBytes.set(identityPathIndex[i], 8 + i * 1);
    }
    const x = sha256(msg);
    return BytesUtils.concatenate(
      idSecretHash,
      BytesUtils.writeUintLE(new Uint8Array(32), rateLimit, 0, 32),
      BytesUtils.writeUintLE(new Uint8Array(32), messageNumberId, 0, 32),
      pathElementsBytes,
      identityPathIndexBytes,
      x,
      externalNullifier
    );
  }

  public async generateRLNProof(
    msg: Uint8Array,
    epoch: Uint8Array | Date | undefined,
    idSecretHash: Uint8Array,
    pathElements: Uint8Array[],
    identityPathIndex: Uint8Array[],
    rateLimit: number,
    messageNumberId: number
  ): Promise<Uint8Array> {
    if (epoch === undefined) {
      epoch = epochIntToBytes(dateToEpoch(new Date()));
    } else if (epoch instanceof Date) {
      epoch = epochIntToBytes(dateToEpoch(epoch));
    }

    if (epoch.length !== 32)
      throw new Error(`Epoch must be 32 bytes, got ${epoch.length}`);
    if (idSecretHash.length !== 32)
      throw new Error(
        `ID secret hash must be 32 bytes, got ${idSecretHash.length}`
      );
    if (pathElements.length !== MERKLE_TREE_DEPTH)
      throw new Error(`Path elements must be ${MERKLE_TREE_DEPTH} bytes`);
    if (identityPathIndex.length !== MERKLE_TREE_DEPTH)
      throw new Error(`Identity path index must be ${MERKLE_TREE_DEPTH} bytes`);
    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}`
      );
    }

    if (messageNumberId < 0 || messageNumberId >= rateLimit) {
      throw new Error(
        `messageNumberId must be an integer between 0 and ${rateLimit - 1}, got ${messageNumberId}`
      );
    }

    const serializedWitness = await this.serializeWitness(
      idSecretHash,
      pathElements,
      identityPathIndex,
      msg,
      epoch,
      rateLimit,
      messageNumberId
    );
    const witnessJson: Record<string, unknown> = zerokitRLN.rlnWitnessToJson(
      this.zkRLN,
      serializedWitness
    ) as Record<string, unknown>;
    const calculatedWitness: bigint[] =
      await this.witnessCalculator.calculateWitness(witnessJson);
    return zerokitRLN.generateRLNProofWithWitness(
      this.zkRLN,
      calculatedWitness,
      serializedWitness
    );
  }

  public verifyRLNProof(
    signalLength: Uint8Array,
    signal: Uint8Array,
    proof: Uint8Array,
    roots: Uint8Array[]
  ): boolean {
    if (signalLength.length !== 8)
      throw new Error("signalLength must be 8 bytes");
    if (proof.length !== 288) throw new Error("proof must be 288 bytes");
    if (roots.length == 0) throw new Error("roots array is empty");
    if (roots.find((root) => root.length !== 32)) {
      throw new Error("All roots must be 32 bytes");
    }

    return zerokitRLN.verifyWithRoots(
      this.zkRLN,
      BytesUtils.concatenate(proof, signalLength, signal),
      BytesUtils.concatenate(...roots)
    );
  }
}