js-rln/src/codec.ts

import debug from "debug";
import { utils } from "js-waku";
import {
  Decoder,
  Encoder,
  Message,
  ProtoMessage,
  RateLimitProof,
} from "js-waku/lib/interfaces";

import { RlnMessage } from "./message.js";
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 toWire(message: Partial<Message>): Promise<Uint8Array | undefined> {
    message.rateLimitProof = await this.generateProof(message);

    return this.encoder.toWire(message);
  }

  async toProtoObj(
    message: Partial<Message>
  ): Promise<ProtoMessage | undefined> {
    const protoMessage = await this.encoder.toProtoObj(message);
    if (!protoMessage) return;

    protoMessage.rateLimitProof = await this.generateProof(message);

    return protoMessage;
  }

  private async generateProof(
    message: Partial<Message>
  ): Promise<RateLimitProof> {
    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");
    return proof;
  }
}

export class RLNDecoder<T extends Message> implements Decoder<RlnMessage<T>> {
  constructor(private rlnInstance: RLNInstance, private decoder: Decoder<T>) {}

  get contentTopic(): string {
    return this.decoder.contentTopic;
  }

  fromWireToProtoObj(bytes: Uint8Array): Promise<ProtoMessage | undefined> {
    const protoMessage = this.decoder.fromWireToProtoObj(bytes);
    log("Message decoded", protoMessage);
    return Promise.resolve(protoMessage);
  }

  async fromProtoObj(proto: ProtoMessage): Promise<RlnMessage<T> | undefined> {
    const msg: T | undefined = await this.decoder.fromProtoObj(proto);
    if (!msg) return;
    return new RlnMessage(this.rlnInstance, msg, proto.rateLimitProof);
  }
}

function toRLNSignal(msg: Partial<Message>): Uint8Array {
  const contentTopicBytes = utils.utf8ToBytes(msg.contentTopic ?? "");
  return new Uint8Array([...(msg.payload ?? []), ...contentTopicBytes]);
}