feat(coordinator): add websockets support

- [x] Add websockets gateway for proof generation - [x] Update AccountSignatureGuard to work with websockets - [x] Add hooks for proof generation: progress, finish, error
privacy-scaling-explorations · Jul 9, 2024 · cbe02ac · cbe02ac
1 parent dad8c3c
commit cbe02ac
Show file tree

Hide file tree

Showing 16 changed files with 998 additions and 242 deletions.
diff --git a/contracts/tasks/helpers/ProofGenerator.ts b/contracts/tasks/helpers/ProofGenerator.ts
@@ -7,7 +7,13 @@ import { genTreeCommitment, hash3, hashLeftRight } from "maci-crypto";
 import fs from "fs";
 import path from "path";
 
-import type { ICircuitFiles, IPrepareStateParams, IProofGeneratorParams, TallyData } from "./types";
+import type {
+  ICircuitFiles,
+  IGenerateProofsOptions,
+  IPrepareStateParams,
+  IProofGeneratorParams,
+  TallyData,
+} from "./types";
 import type { Proof } from "../../ts/types";
 import type { BigNumberish } from "ethers";
 
@@ -176,7 +182,7 @@ export class ProofGenerator {
    *
    * @returns message processing proofs
    */
-  async generateMpProofs(): Promise<Proof[]> {
+  async generateMpProofs(options?: IGenerateProofsOptions): Promise<Proof[]> {
     performance.mark("mp-proofs-start");
 
     console.log(`Generating proofs of message processing...`);
@@ -188,34 +194,47 @@ export class ProofGenerator {
       totalMessageBatches += 1;
     }
 
-    const inputs: CircuitInputs[] = [];
+    try {
+      const inputs: CircuitInputs[] = [];
 
-    // while we have unprocessed messages, process them
-    while (this.poll.hasUnprocessedMessages()) {
-      // process messages in batches
-      const circuitInputs = this.poll.processMessages(
-        BigInt(this.poll.pollId),
-        this.useQuadraticVoting,
-      ) as unknown as CircuitInputs;
+      // while we have unprocessed messages, process them
+      while (this.poll.hasUnprocessedMessages()) {
+        // process messages in batches
+        const circuitInputs = this.poll.processMessages(
+          BigInt(this.poll.pollId),
+          this.useQuadraticVoting,
+        ) as unknown as CircuitInputs;
 
-      // generate the proof for this batch
-      inputs.push(circuitInputs);
+        // generate the proof for this batch
+        inputs.push(circuitInputs);
 
-      console.log(`Progress: ${this.poll.numBatchesProcessed} / ${totalMessageBatches}`);
-    }
+        console.log(`Progress: ${this.poll.numBatchesProcessed} / ${totalMessageBatches}`);
+      }
 
-    console.log("Wait until proof generation is finished");
+      console.log("Wait until proof generation is finished");
 
-    const proofs = await Promise.all(
-      inputs.map((circuitInputs, index) => this.generateProofs(circuitInputs, this.mp, `process_${index}.json`)),
-    ).then((data) => data.reduce((acc, x) => acc.concat(x), []));
+      const proofs = await Promise.all(
+        inputs.map((circuitInputs, index) =>
+          this.generateProofs(circuitInputs, this.mp, `process_${index}.json`).then((data) => {
+            options?.onBatchComplete?.({ current: index, total: totalMessageBatches, proofs: data });
+            return data;
+          }),
+        ),
+      ).then((data) => data.reduce((acc, x) => acc.concat(x), []));
 
-    console.log("Proof generation is finished");
+      console.log("Proof generation is finished");
 
-    performance.mark("mp-proofs-end");
-    performance.measure("Generate message processor proofs", "mp-proofs-start", "mp-proofs-end");
+      performance.mark("mp-proofs-end");
+      performance.measure("Generate message processor proofs", "mp-proofs-start", "mp-proofs-end");
 
-    return proofs;
+      options?.onComplete?.(proofs);
+
+      return proofs;
+    } catch (error) {
+      options?.onFail?.(error as Error);
+
+      throw error;
+    }
   }
 
   /**
@@ -224,7 +243,10 @@ export class ProofGenerator {
    * @param network - current network
    * @returns tally proofs
    */
-  async generateTallyProofs(network: Network): Promise<{ proofs: Proof[]; tallyData: TallyData }> {
+  async generateTallyProofs(
+    network: Network,
+    options?: IGenerateProofsOptions,
+  ): Promise<{ proofs: Proof[]; tallyData: TallyData }> {
     performance.mark("tally-proofs-start");
 
     console.log(`Generating proofs of vote tallying...`);
@@ -235,105 +257,118 @@ export class ProofGenerator {
       totalTallyBatches += 1;
     }
 
-    let tallyCircuitInputs: CircuitInputs;
-    const inputs: CircuitInputs[] = [];
+    try {
+      let tallyCircuitInputs: CircuitInputs;
+      const inputs: CircuitInputs[] = [];
 
-    while (this.poll.hasUntalliedBallots()) {
-      tallyCircuitInputs = (this.useQuadraticVoting
-        ? this.poll.tallyVotes()
-        : this.poll.tallyVotesNonQv()) as unknown as CircuitInputs;
+      while (this.poll.hasUntalliedBallots()) {
+        tallyCircuitInputs = (this.useQuadraticVoting
+          ? this.poll.tallyVotes()
+          : this.poll.tallyVotesNonQv()) as unknown as CircuitInputs;
 
-      inputs.push(tallyCircuitInputs);
+        inputs.push(tallyCircuitInputs);
 
-      console.log(`Progress: ${this.poll.numBatchesTallied} / ${totalTallyBatches}`);
-    }
-
-    console.log("Wait until proof generation is finished");
+        console.log(`Progress: ${this.poll.numBatchesTallied} / ${totalTallyBatches}`);
+      }
 
-    const proofs = await Promise.all(
-      inputs.map((circuitInputs, index) => this.generateProofs(circuitInputs, this.tally, `tally_${index}.json`)),
-    ).then((data) => data.reduce((acc, x) => acc.concat(x), []));
+      console.log("Wait until proof generation is finished");
 
-    console.log("Proof generation is finished");
+      const proofs = await Promise.all(
+        inputs.map((circuitInputs, index) =>
+          this.generateProofs(circuitInputs, this.tally, `tally_${index}.json`).then((data) => {
+            options?.onBatchComplete?.({ current: index, total: totalTallyBatches, proofs: data });
+            return data;
+          }),
+        ),
+      ).then((data) => data.reduce((acc, x) => acc.concat(x), []));
 
-    // verify the results
-    // Compute newResultsCommitment
-    const newResultsCommitment = genTreeCommitment(
-      this.poll.tallyResult,
-      BigInt(asHex(tallyCircuitInputs!.newResultsRootSalt as BigNumberish)),
-      this.poll.treeDepths.voteOptionTreeDepth,
-    );
+      console.log("Proof generation is finished");
 
-    // compute newSpentVoiceCreditsCommitment
-    const newSpentVoiceCreditsCommitment = hashLeftRight(
-      this.poll.totalSpentVoiceCredits,
-      BigInt(asHex(tallyCircuitInputs!.newSpentVoiceCreditSubtotalSalt as BigNumberish)),
-    );
-
-    let newPerVOSpentVoiceCreditsCommitment: bigint | undefined;
-    let newTallyCommitment: bigint;
-
-    // create the tally file data to store for verification later
-    const tallyFileData: TallyData = {
-      maci: this.maciContractAddress,
-      pollId: this.poll.pollId.toString(),
-      network: network.name,
-      chainId: network.config.chainId?.toString(),
-      isQuadratic: Boolean(this.useQuadraticVoting),
-      tallyAddress: this.tallyContractAddress,
-      newTallyCommitment: asHex(tallyCircuitInputs!.newTallyCommitment as BigNumberish),
-      results: {
-        tally: this.poll.tallyResult.map((x) => x.toString()),
-        salt: asHex(tallyCircuitInputs!.newResultsRootSalt as BigNumberish),
-        commitment: asHex(newResultsCommitment),
-      },
-      totalSpentVoiceCredits: {
-        spent: this.poll.totalSpentVoiceCredits.toString(),
-        salt: asHex(tallyCircuitInputs!.newSpentVoiceCreditSubtotalSalt as BigNumberish),
-        commitment: asHex(newSpentVoiceCreditsCommitment),
-      },
-    };
-
-    if (this.useQuadraticVoting) {
-      // Compute newPerVOSpentVoiceCreditsCommitment
-      newPerVOSpentVoiceCreditsCommitment = genTreeCommitment(
-        this.poll.perVOSpentVoiceCredits,
-        BigInt(asHex(tallyCircuitInputs!.newPerVOSpentVoiceCreditsRootSalt as BigNumberish)),
+      // verify the results
+      // Compute newResultsCommitment
+      const newResultsCommitment = genTreeCommitment(
+        this.poll.tallyResult,
+        BigInt(asHex(tallyCircuitInputs!.newResultsRootSalt as BigNumberish)),
         this.poll.treeDepths.voteOptionTreeDepth,
       );
 
-      // Compute newTallyCommitment
-      newTallyCommitment = hash3([
-        newResultsCommitment,
-        newSpentVoiceCreditsCommitment,
-        newPerVOSpentVoiceCreditsCommitment,
-      ]);
+      // compute newSpentVoiceCreditsCommitment
+      const newSpentVoiceCreditsCommitment = hashLeftRight(
+        this.poll.totalSpentVoiceCredits,
+        BigInt(asHex(tallyCircuitInputs!.newSpentVoiceCreditSubtotalSalt as BigNumberish)),
+      );
 
-      // update perVOSpentVoiceCredits in the tally file data
-      tallyFileData.perVOSpentVoiceCredits = {
-        tally: this.poll.perVOSpentVoiceCredits.map((x) => x.toString()),
-        salt: asHex(tallyCircuitInputs!.newPerVOSpentVoiceCreditsRootSalt as BigNumberish),
-        commitment: asHex(newPerVOSpentVoiceCreditsCommitment),
+      let newPerVOSpentVoiceCreditsCommitment: bigint | undefined;
+      let newTallyCommitment: bigint;
+
+      // create the tally file data to store for verification later
+      const tallyFileData: TallyData = {
+        maci: this.maciContractAddress,
+        pollId: this.poll.pollId.toString(),
+        network: network.name,
+        chainId: network.config.chainId?.toString(),
+        isQuadratic: Boolean(this.useQuadraticVoting),
+        tallyAddress: this.tallyContractAddress,
+        newTallyCommitment: asHex(tallyCircuitInputs!.newTallyCommitment as BigNumberish),
+        results: {
+          tally: this.poll.tallyResult.map((x) => x.toString()),
+          salt: asHex(tallyCircuitInputs!.newResultsRootSalt as BigNumberish),
+          commitment: asHex(newResultsCommitment),
+        },
+        totalSpentVoiceCredits: {
+          spent: this.poll.totalSpentVoiceCredits.toString(),
+          salt: asHex(tallyCircuitInputs!.newSpentVoiceCreditSubtotalSalt as BigNumberish),
+          commitment: asHex(newSpentVoiceCreditsCommitment),
+        },
       };
-    } else {
-      newTallyCommitment = hashLeftRight(newResultsCommitment, newSpentVoiceCreditsCommitment);
-    }
 
-    fs.writeFileSync(this.tallyOutputFile, JSON.stringify(tallyFileData, null, 4));
-
-    console.log(`Tally file:\n${JSON.stringify(tallyFileData, null, 4)}\n`);
-
-    // compare the commitments
-    if (asHex(newTallyCommitment) === tallyFileData.newTallyCommitment) {
-      console.log("The tally commitment is correct");
-    } else {
-      throw new Error("Error: the newTallyCommitment is invalid.");
+      if (this.useQuadraticVoting) {
+        // Compute newPerVOSpentVoiceCreditsCommitment
+        newPerVOSpentVoiceCreditsCommitment = genTreeCommitment(
+          this.poll.perVOSpentVoiceCredits,
+          BigInt(asHex(tallyCircuitInputs!.newPerVOSpentVoiceCreditsRootSalt as BigNumberish)),
+          this.poll.treeDepths.voteOptionTreeDepth,
+        );
+
+        // Compute newTallyCommitment
+        newTallyCommitment = hash3([
+          newResultsCommitment,
+          newSpentVoiceCreditsCommitment,
+          newPerVOSpentVoiceCreditsCommitment,
+        ]);
+
+        // update perVOSpentVoiceCredits in the tally file data
+        tallyFileData.perVOSpentVoiceCredits = {
+          tally: this.poll.perVOSpentVoiceCredits.map((x) => x.toString()),
+          salt: asHex(tallyCircuitInputs!.newPerVOSpentVoiceCreditsRootSalt as BigNumberish),
+          commitment: asHex(newPerVOSpentVoiceCreditsCommitment),
+        };
+      } else {
+        newTallyCommitment = hashLeftRight(newResultsCommitment, newSpentVoiceCreditsCommitment);
+      }
+
+      fs.writeFileSync(this.tallyOutputFile, JSON.stringify(tallyFileData, null, 4));
+
+      console.log(`Tally file:\n${JSON.stringify(tallyFileData, null, 4)}\n`);
+
+      // compare the commitments
+      if (asHex(newTallyCommitment) === tallyFileData.newTallyCommitment) {
+        console.log("The tally commitment is correct");
+      } else {
+        throw new Error("Error: the newTallyCommitment is invalid.");
+      }
+
+      performance.mark("tally-proofs-end");
+      performance.measure("Generate tally proofs", "tally-proofs-start", "tally-proofs-end");
+
+      options?.onComplete?.(proofs, tallyFileData);
+
+      return { proofs, tallyData: tallyFileData };
+    } catch (error) {
+      options?.onFail?.(error as Error);
+
+      throw error;
     }
-
-    performance.mark("tally-proofs-end");
-    performance.measure("Generate tally proofs", "tally-proofs-start", "tally-proofs-end");
-
-    return { proofs, tallyData: tallyFileData };
   }
 
   /**

diff --git a/contracts/tasks/helpers/types.ts b/contracts/tasks/helpers/types.ts
@@ -1,3 +1,4 @@
+import type { Proof } from "../../ts/types";
 import type { AccQueue, MACI, MessageProcessor, Poll, Tally, Verifier, VkRegistry } from "../../typechain-types";
 import type { HardhatEthersSigner } from "@nomicfoundation/hardhat-ethers/signers";
 import type {
@@ -294,6 +295,53 @@ export interface IPrepareStateParams {
   }>;
 }
 
+/**
+ * Interface that represents generate proof options
+ */
+export interface IGenerateProofsOptions {
+  /**
+   * Hook to call when batch generation is completed
+   *
+   * @param data - batch data
+   */
+  onBatchComplete?: (data: IGenerateProofsBatchData) => void;
+
+  /**
+   * Hook to call when proof generation is completed
+   *
+   * @param data - proof generated data
+   * @param tally - tally data
+   */
+  onComplete?: (data: Proof[], tally?: TallyData) => void;
+
+  /**
+   * Hook to call when generation is failed
+   *
+   * @param error - error
+   */
+  onFail?: (error: Error) => void;
+}
+
+/**
+ * Interface that represents batch data
+ */
+export interface IGenerateProofsBatchData {
+  /**
+   * Batch proofs
+   */
+  proofs: Proof[];
+
+  /**
+   * Current batch
+   */
+  current: number;
+
+  /**
+   * Total batches
+   */
+  total: number;
+}
+
 /**
  * Interface that represents prover params
  */

diff --git a/contracts/ts/index.ts b/contracts/ts/index.ts
@@ -22,7 +22,12 @@ export { Deployment } from "../tasks/helpers/Deployment";
 export { ContractStorage } from "../tasks/helpers/ContractStorage";
 export { ProofGenerator } from "../tasks/helpers/ProofGenerator";
 export { Prover } from "../tasks/helpers/Prover";
-export { EContracts } from "../tasks/helpers/types";
+export {
+  EContracts,
+  type IGenerateProofsOptions,
+  type IGenerateProofsBatchData,
+  type TallyData,
+} from "../tasks/helpers/types";
 export { linkPoseidonLibraries } from "../tasks/helpers/abi";
 
 export type { IVerifyingKeyStruct, SnarkProof, Groth16Proof, Proof } from "./types";

diff --git a/coordinator/README.md b/coordinator/README.md
@@ -9,5 +9,5 @@
 5. Run `pnpm run start` to run the service.
 6. All API calls must be called with `Authorization` header, where the value is encrypted with RSA public key you generated before. Header value contains message signature and message digest created by `COORDINATOR_ADDRESSES`. The format is `publicEncrypt({signature}:{digest})`.
    Make sure you set `COORDINATOR_ADDRESSES` env variable and sign any message with the addresses from your application (see [AccountSignatureGuard](./ts/auth/AccountSignatureGuard.service.ts)).
-7. Proofs can be generated with `POST v1/proof/generate` API method (see [dto spec](./ts/proof/dto.ts) and [controller](./ts/app.controller.ts)).
+7. Proofs can be generated with `POST v1/proof/generate` API method or with Websockets (see [dto spec](./ts/proof/dto.ts), [controller](./ts/app.controller.ts) and [wsgateway](./ts/events/events.gateway.ts)).
 8. [Swagger documentation for API methods](https://maci-coordinator.pse.dev/api)