@waku/rln 0.0.6 → 0.0.8-38fdb53

package/src/encoder.ts

@@ -0,0 +1,81 @@
+import debug from "debug";
+import { proto_message, utils } from "js-waku";
+import {
+  Decoder,
+  Encoder,
+  Message,
+  ProtoMessage,
+} from "js-waku/lib/interfaces";
+
+import { MembershipKey, RLNInstance } from "./rln.js";
+
+const log = debug("waku:message:rln-encoder");
+
+export class RLNEncoder implements Encoder {
+  public contentTopic: string;
+  private readonly idKey: Uint8Array;
+
+  constructor(
+    private encoder: Encoder,
+    private rlnInstance: RLNInstance,
+    private index: number,
+    membershipKey: MembershipKey
+  ) {
+    if (index < 0) throw "invalid membership index";
+    this.idKey = membershipKey.IDKey;
+    this.contentTopic = encoder.contentTopic;
+  }
+
+  async encode(message: Message): Promise<Uint8Array | undefined> {
+    const protoMessage = await this.encodeProto(message);
+    if (!protoMessage) return;
+    return proto_message.WakuMessage.encode(protoMessage);
+  }
+
+  async encodeProto(message: Message): Promise<ProtoMessage | undefined> {
+    const protoMessage = await this.encoder.encodeProto(message);
+    if (!protoMessage) return;
+
+    const signal = toRLNSignal(message);
+
+    console.time("proof_gen_timer");
+    const proof = await this.rlnInstance.generateProof(
+      signal,
+      this.index,
+      message.timestamp,
+      this.idKey
+    );
+    console.timeEnd("proof_gen_timer");
+
+    protoMessage.rateLimitProof = proof;
+
+    return protoMessage;
+  }
+}
+
+export class RLNDecoder implements Decoder<Message> {
+  public contentTopic: string;
+
+  constructor(private decoder: Decoder<Message>) {
+    this.contentTopic = decoder.contentTopic;
+  }
+
+  decodeProto(bytes: Uint8Array): Promise<ProtoMessage | undefined> {
+    const protoMessage = proto_message.WakuMessage.decode(bytes);
+    log("Message decoded", protoMessage);
+    return Promise.resolve(protoMessage);
+  }
+
+  async decode(proto: ProtoMessage): Promise<Message | undefined> {
+    const msg = await this.decoder.decode(proto);
+    if (msg) {
+      msg.rateLimitProof = proto.rateLimitProof;
+    }
+    return msg;
+  }
+}
+
+function toRLNSignal(msg: Message): Uint8Array {
+  const contentTopicBytes = utils.utf8ToBytes(msg.contentTopic ?? "");
+  return new Uint8Array([...(msg.payload ?? []), ...contentTopicBytes]);
+}

package/src/index.ts

@@ -1,4 +1,4 @@
-import type { MembershipKey, RateLimitProof, RLNInstance } from "./rln.js";
+import type { MembershipKey, Proof, RLNInstance } from "./rln.js";
 
 // reexport the create function, dynamically imported from rln.ts
 export async function create(): Promise<RLNInstance> {
@@ -9,4 +9,4 @@ export async function create(): Promise<RLNInstance> {
   return await rlnModule.create();
 }
 
-export { RLNInstance, MembershipKey, RateLimitProof };
+export { RLNInstance, MembershipKey, Proof };

package/src/rln.ts

@@ -1,22 +1,9 @@
 import init, * as zerokitRLN from "@waku/zerokit-rln-wasm";
+import { RateLimitProof } from "js-waku/lib/interfaces";
 
-import * as resources from "./resources.js";
+import verificationKey from "./resources/verification_key.js";
 import * as wc from "./witness_calculator.js";
-
-/**
- * Convert a base64 string into uint8Array
- * @param base64
- * @returns Uint8Array
- */
-function base64ToUint8Array(base64: string): Uint8Array {
-  const binary_string = window.atob(base64);
-  const len = binary_string.length;
-  const bytes = new Uint8Array(len);
-  for (let i = 0; i < len; i++) {
-    bytes[i] = binary_string.charCodeAt(i);
-  }
-  return bytes;
-}
+import { WitnessCalculator } from "./witness_calculator.js";
 
 /**
  * Concatenate Uint8Arrays
@@ -37,10 +24,21 @@ function concatenate(...input: Uint8Array[]): Uint8Array {
   return result;
 }
 
+const stringEncoder = new TextEncoder();
+
 const DEPTH = 20;
-const VERIFICATION_KEY = base64ToUint8Array(resources.verification_key);
-const ZKEY = base64ToUint8Array(resources.zkey);
-const CIRCUIT = base64ToUint8Array(resources.circuit);
+
+async function loadWitnessCalculator(): Promise<WitnessCalculator> {
+  const url = new URL("./resources/rln.wasm", import.meta.url);
+  const response = await fetch(url);
+  return await wc.builder(new Uint8Array(await response.arrayBuffer()), false);
+}
+
+async function loadZkey(): Promise<Uint8Array> {
+  const url = new URL("./resources/rln_final.zkey", import.meta.url);
+  const response = await fetch(url);
+  return new Uint8Array(await response.arrayBuffer());
+}
 
 /**
  * Create an instance of RLN
@@ -49,19 +47,23 @@ const CIRCUIT = base64ToUint8Array(resources.circuit);
 export async function create(): Promise<RLNInstance> {
   await init();
   zerokitRLN.init_panic_hook();
-
-  const witnessCalculator = await wc.builder(CIRCUIT, false);
-  const zkRLN = zerokitRLN.newRLN(DEPTH, ZKEY, VERIFICATION_KEY);
+  const witnessCalculator = await loadWitnessCalculator();
+  const zkey = await loadZkey();
+  const vkey = stringEncoder.encode(JSON.stringify(verificationKey));
+  const zkRLN = zerokitRLN.newRLN(DEPTH, zkey, vkey);
   return new RLNInstance(zkRLN, witnessCalculator);
 }
 
 export class MembershipKey {
-  readonly IDKey: Uint8Array;
-  readonly IDCommitment: Uint8Array;
-
-  constructor(memKeys: Uint8Array) {
-    this.IDKey = memKeys.subarray(0, 32);
-    this.IDCommitment = memKeys.subarray(32);
+  constructor(
+    public readonly IDKey: Uint8Array,
+    public readonly IDCommitment: Uint8Array
+  ) {}
+
+  static fromBytes(memKeys: Uint8Array): MembershipKey {
+    const idKey = memKeys.subarray(0, 32);
+    const idCommitment = memKeys.subarray(32);
+    return new MembershipKey(idKey, idCommitment);
   }
 }
 
@@ -123,7 +125,7 @@ const shareYOffset = shareXOffset + 32;
 const nullifierOffset = shareYOffset + 32;
 const rlnIdentifierOffset = nullifierOffset + 32;
 
-export class RateLimitProof {
+export class Proof implements RateLimitProof {
   readonly proof: Uint8Array;
   readonly merkleRoot: Uint8Array;
   readonly epoch: Uint8Array;
@@ -146,26 +148,29 @@ export class RateLimitProof {
       rlnIdentifierOffset
     );
   }
+}
 
-  toBytes(): Uint8Array {
-    return concatenate(
-      this.proof,
-      this.merkleRoot,
-      this.epoch,
-      this.shareX,
-      this.shareY,
-      this.nullifier,
-      this.rlnIdentifier
-    );
-  }
+function proofToBytes(p: RateLimitProof): Uint8Array {
+  return concatenate(
+    p.proof,
+    p.merkleRoot,
+    p.epoch,
+    p.shareX,
+    p.shareY,
+    p.nullifier,
+    p.rlnIdentifier
+  );
 }
 
 export class RLNInstance {
-  constructor(private zkRLN: number, private witnessCalculator: any) {}
+  constructor(
+    private zkRLN: number,
+    private witnessCalculator: WitnessCalculator
+  ) {}
 
   generateMembershipKey(): MembershipKey {
     const memKeys = zerokitRLN.generateMembershipKey(this.zkRLN);
-    return new MembershipKey(memKeys);
+    return MembershipKey.fromBytes(memKeys);
   }
 
   insertMember(idCommitment: Uint8Array): void {
@@ -221,13 +226,16 @@ export class RLNInstance {
       rlnWitness
     );
 
-    return new RateLimitProof(proofBytes);
+    return new Proof(proofBytes);
   }
 
   verifyProof(proof: RateLimitProof | Uint8Array): boolean {
-    if (proof instanceof RateLimitProof) {
-      proof = proof.toBytes();
+    let pBytes: Uint8Array;
+    if (proof instanceof Uint8Array) {
+      pBytes = proof;
+    } else {
+      pBytes = proofToBytes(proof);
     }
-    return zerokitRLN.verifyProof(this.zkRLN, proof);
+    return zerokitRLN.verifyProof(this.zkRLN, pBytes);
   }
 }

package/src/witness_calculator.d.ts

@@ -1,4 +1,8 @@
 export async function builder(
   code: Uint8Array,
   sanityCheck: bool
-): Promise<any>;
+): Promise<WitnessCalculator>;
+
+export class WitnessCalculator {
+  calculateWitness(input, sanityCheck): Array<bigint>;
+}