@reclaimprotocol/attestor-core 5.0.1-beta.7 → 5.0.1

package/lib/utils/zk.js

@@ -1,625 +0,0 @@
-import { concatenateUint8Arrays, crypto, generateIV } from "../utils/tls-imports.js";
-import {
-  ceilToBlockSizeMultiple,
-  CONFIG as ZK_CONFIG,
-  generateProof,
-  getBlockSizeBytes,
-  makeLocalFileFetch,
-  makeRemoteFileFetch,
-  verifyProof
-} from "@reclaimprotocol/zk-symmetric-crypto";
-import { makeGnarkOPRFOperator, makeGnarkZkOperator } from "@reclaimprotocol/zk-symmetric-crypto/gnark";
-import { makeSnarkJsZKOperator } from "@reclaimprotocol/zk-symmetric-crypto/snarkjs";
-import { makeStwoZkOperator } from "@reclaimprotocol/zk-symmetric-crypto/stwo";
-import PQueue from "p-queue";
-import {
-  DEFAULT_REMOTE_FILE_FETCH_BASE_URL,
-  DEFAULT_ZK_CONCURRENCY,
-  TOPRF_DOMAIN_SEPARATOR
-} from "../config/index.js";
-import { ZKProofEngine } from "../proto/api.js";
-import { detectEnvironment, getEnvVariable } from "../utils/env.js";
-import { AttestorError } from "../utils/error.js";
-import {
-  getPureCiphertext,
-  getRecordIV,
-  getZkAlgorithmForCipherSuite,
-  isTls13Suite,
-  strToUint8Array
-} from "../utils/generics.js";
-import { logger as LOGGER } from "../utils/logger.js";
-import { binaryHashToStr, isFullyRedacted, isRedactionCongruent, REDACTION_CHAR_CODE } from "../utils/redactions.js";
-const ZK_CONCURRENCY = +(getEnvVariable("ZK_CONCURRENCY") || DEFAULT_ZK_CONCURRENCY);
-async function makeZkProofGenerator({
-  zkOperators,
-  oprfOperators,
-  logger = LOGGER,
-  zkProofConcurrency = ZK_CONCURRENCY,
-  cipherSuite,
-  zkEngine = "snarkjs"
-}) {
-  const zkQueue = new PQueue({ concurrency: zkProofConcurrency, autoStart: true });
-  const packetsToProve = [];
-  logger = logger.child({ module: "zk", zkEngine });
-  let zkProofsToGen = 0;
-  return {
-    /**
-     * Adds the given packet to the list of packets to
-     * generate ZK proofs for.
-     *
-     * Call `generateProofs()` to finally generate the proofs
-     */
-    async addPacketToProve(packet, {
-      redactedPlaintext,
-      toprfs = [],
-      overshotToprfFromPrevBlock
-    }, onGeneratedProofs, getNextPacket) {
-      if (packet.type === "plaintext") {
-        throw new Error("Cannot generate proof for plaintext");
-      }
-      const alg = getZkAlgorithmForCipherSuite(cipherSuite);
-      const chunkSizeBytes = getChunkSizeBytes(alg);
-      const key = await crypto.exportKey(packet.encKey);
-      const iv = packet.iv;
-      const ciphertext = getPureCiphertext(packet.ciphertext, cipherSuite);
-      if (overshotToprfFromPrevBlock) {
-        redactedPlaintext.set(
-          new Uint8Array(overshotToprfFromPrevBlock.length).fill(REDACTION_CHAR_CODE)
-        );
-      }
-      const trueCiphertextLength = isTls13Suite(cipherSuite) ? ciphertext.length - 1 : ciphertext.length;
-      const packetToProve = {
-        onGeneratedProofs,
-        algorithm: alg,
-        proofsToGenerate: [],
-        toprfsToGenerate: [],
-        iv: packet.fixedIv
-      };
-      for (const toprf of toprfs) {
-        const toprfDistFromEnd = trueCiphertextLength - (toprf.dataLocation.fromIndex + toprf.dataLocation.length);
-        if (toprfDistFromEnd < 0) {
-          const nextPacket = getNextPacket();
-          if (nextPacket?.type !== "ciphertext") {
-            throw new AttestorError(
-              "ERROR_INTERNAL",
-              "TOPRF data overshoots ciphertext length, but no next ciphertext packet found"
-            );
-          }
-          if (nextPacket.encKey !== packet.encKey) {
-            throw new AttestorError(
-              "ERROR_INTERNAL",
-              "TOPRF data overshoots ciphertext length, but next packet has different encryption key"
-            );
-          }
-          const nextCiphertext = nextPacket.ciphertext.slice(0, Math.abs(toprfDistFromEnd));
-          const iv2 = nextPacket.iv;
-          toprf.overshoot = {
-            ciphertext: nextCiphertext,
-            iv: iv2,
-            recordNumber: nextPacket.recordNumber
-          };
-        }
-        const fromIndex = getIdealOffsetForToprfBlock(alg, toprf);
-        const toIndex = Math.min(fromIndex + chunkSizeBytes, ciphertext.length);
-        const slice = { fromIndex, toIndex };
-        packetToProve.toprfsToGenerate.push(getTOPRFProofGenerationParamsForSlice({
-          key,
-          iv,
-          ciphertext,
-          slice,
-          toprf: {
-            ...toprf,
-            dataLocation: {
-              ...toprf.dataLocation,
-              fromIndex: toprf.dataLocation.fromIndex - fromIndex
-            }
-          }
-        }));
-        zkProofsToGen += 1;
-        const pktToIndex = Math.min(
-          trueCiphertextLength,
-          toprf.dataLocation.fromIndex + toprf.dataLocation.length
-        );
-        const pktFromIndex = toprf.dataLocation.fromIndex;
-        for (let i = pktFromIndex; i < pktToIndex; i++) {
-          redactedPlaintext[i] = REDACTION_CHAR_CODE;
-        }
-      }
-      for (let i = 0; i < ciphertext.length; i += chunkSizeBytes) {
-        const slice = {
-          fromIndex: i,
-          toIndex: Math.min(i + chunkSizeBytes, ciphertext.length)
-        };
-        const proofParams = getProofGenerationParamsForSlice(
-          { key, iv, ciphertext, redactedPlaintext, slice }
-        );
-        if (!proofParams) {
-          continue;
-        }
-        packetToProve.proofsToGenerate.push(proofParams);
-        zkProofsToGen += 1;
-      }
-      packetsToProve.push(packetToProve);
-    },
-    getTotalChunksToProve() {
-      return zkProofsToGen;
-    },
-    async generateProofs(onChunkDone) {
-      if (!packetsToProve.length) {
-        return;
-      }
-      const start = Date.now();
-      const tasks = [];
-      for (const {
-        onGeneratedProofs,
-        algorithm,
-        proofsToGenerate,
-        toprfsToGenerate
-      } of packetsToProve) {
-        const proofs = [];
-        const toprfs = [];
-        let proofsLeft = proofsToGenerate.length + toprfsToGenerate.length;
-        for (const proofToGen of proofsToGenerate) {
-          tasks.push(zkQueue.add(async () => {
-            const proof = await generateZkProofForChunk(algorithm, proofToGen);
-            onChunkDone?.();
-            proofs.push(proof);
-            proofsLeft -= 1;
-            if (proofsLeft === 0) {
-              onGeneratedProofs(proofs, toprfs);
-            }
-          }, { throwOnTimeout: true }));
-        }
-        for (const toprfToGen of toprfsToGenerate) {
-          tasks.push(zkQueue.add(async () => {
-            const toprf = await generateOprfProofForChunk(algorithm, toprfToGen);
-            onChunkDone?.();
-            toprfs.push(toprf);
-            proofsLeft -= 1;
-            if (proofsLeft === 0) {
-              onGeneratedProofs(proofs, toprfs);
-            }
-          }, { throwOnTimeout: true }));
-        }
-      }
-      await Promise.all(tasks);
-      logger?.info(
-        { durationMs: Date.now() - start, zkProofsToGen },
-        "generated ZK proofs"
-      );
-      packetsToProve.splice(0, packetsToProve.length);
-      zkProofsToGen = 0;
-      const alg = getZkAlgorithmForCipherSuite(cipherSuite);
-      const zkOperator = await getZkOperatorForAlgorithm(alg);
-      zkOperator.release?.();
-    }
-  };
-  async function generateZkProofForChunk(algorithm, {
-    startIdx,
-    redactedPlaintext,
-    privateInput,
-    publicInput
-  }) {
-    const operator = getZkOperatorForAlgorithm(algorithm);
-    const proof = await generateProof(
-      { algorithm, privateInput, publicInput, operator, logger }
-    );
-    logger?.debug({ startIdx }, "generated proof for chunk");
-    return {
-      proofData: typeof proof.proofData === "string" ? strToUint8Array(proof.proofData) : proof.proofData,
-      decryptedRedactedCiphertext: proof.plaintext || new Uint8Array(),
-      redactedPlaintext,
-      startIdx
-    };
-  }
-  async function generateOprfProofForChunk(algorithm, { startIdx, privateInput, publicInput, toprf }) {
-    const operator = getOprfOperatorForAlgorithm(algorithm);
-    const toprfLocations = [];
-    if (toprf?.overshoot) {
-      const { dataLocation, overshoot: { ciphertext } } = toprf;
-      toprfLocations.push(
-        {
-          pos: dataLocation.fromIndex,
-          len: dataLocation.length - ciphertext.length
-        },
-        {
-          pos: ceilToBlockSizeMultiple(
-            dataLocation.fromIndex + dataLocation.length,
-            algorithm
-          ),
-          len: ciphertext.length
-        }
-      );
-    } else if (toprf) {
-      toprfLocations.push({
-        pos: toprf.dataLocation.fromIndex,
-        len: toprf.dataLocation.length
-      });
-    }
-    const proof = await generateProof(
-      {
-        algorithm,
-        privateInput,
-        publicInput,
-        operator,
-        logger,
-        ...toprf ? {
-          toprf: {
-            locations: toprfLocations,
-            output: toprf.nullifier,
-            responses: toprf.responses,
-            domainSeparator: TOPRF_DOMAIN_SEPARATOR
-          },
-          mask: toprf.mask
-        } : {}
-      }
-    );
-    logger?.debug({ toprfLocations }, "generated TOPRF proof for chunk");
-    return {
-      startIdx,
-      proofData: typeof proof.proofData === "string" ? strToUint8Array(proof.proofData) : proof.proofData,
-      payload: toprf
-    };
-  }
-  function getZkOperatorForAlgorithm(algorithm) {
-    return zkOperators?.[algorithm] || makeDefaultZkOperator(algorithm, zkEngine, logger);
-  }
-  function getOprfOperatorForAlgorithm(algorithm) {
-    return oprfOperators?.[algorithm] || makeDefaultOPRFOperator(algorithm, zkEngine, logger);
-  }
-}
-async function verifyZkPacket({
-  cipherSuite,
-  ciphertext,
-  zkReveal,
-  zkOperators,
-  oprfOperators,
-  logger = LOGGER,
-  zkEngine = "snarkjs",
-  iv,
-  recordNumber,
-  toprfOvershotNullifier,
-  getNextPacket
-}) {
-  const { proofs, toprfs, oprfRawMarkers } = zkReveal;
-  const algorithm = getZkAlgorithmForCipherSuite(cipherSuite);
-  const recordIV = getRecordIV(ciphertext, cipherSuite);
-  ciphertext = new Uint8Array(getPureCiphertext(ciphertext, cipherSuite));
-  const realRedactedPlaintext = new Uint8Array(ciphertext.length).fill(REDACTION_CHAR_CODE);
-  const replacements = await Promise.all(toprfs.map(async (toprf, i) => {
-    try {
-      return await verifyToprfProofPacket(toprf);
-    } catch (e) {
-      e.message += ` (TOPRF proof ${i}, from ${toprf.payload?.dataLocation?.fromIndex}, record ${recordNumber})`;
-      throw e;
-    }
-  }));
-  await Promise.all(proofs.map(async (proof, i) => {
-    try {
-      await verifyZkProofPacket(proof);
-    } catch (e) {
-      e.message += ` (ZK proof ${i}, startIdx ${proof.startIdx}, record ${recordNumber})`;
-      throw e;
-    }
-  }));
-  for (const { set, startIdx } of replacements) {
-    realRedactedPlaintext.set(set, startIdx);
-  }
-  if (toprfOvershotNullifier) {
-    realRedactedPlaintext.set(toprfOvershotNullifier);
-  }
-  return { redactedPlaintext: realRedactedPlaintext, oprfRawMarkers };
-  async function verifyZkProofPacket({
-    proofData,
-    decryptedRedactedCiphertext,
-    redactedPlaintext,
-    startIdx
-  }) {
-    const ciphertextChunkEnd = startIdx + redactedPlaintext.length;
-    const ciphertextChunk = ciphertext.slice(startIdx, ciphertextChunkEnd);
-    for (let i = 0; i < ciphertextChunk.length; i++) {
-      if (redactedPlaintext[i] === REDACTION_CHAR_CODE) {
-        ciphertextChunk[i] = REDACTION_CHAR_CODE;
-      }
-    }
-    let nonce = concatenateUint8Arrays([iv, recordIV]);
-    if (!recordIV.length) {
-      nonce = generateIV(nonce, recordNumber);
-    }
-    const ciphertextInput = {
-      ciphertext: ciphertextChunk,
-      iv: nonce,
-      offsetBytes: startIdx
-    };
-    if (!isRedactionCongruent(redactedPlaintext, decryptedRedactedCiphertext)) {
-      throw new Error("redacted ciphertext not congruent");
-    }
-    await verifyProof(
-      {
-        proof: {
-          algorithm,
-          proofData,
-          plaintext: decryptedRedactedCiphertext
-        },
-        publicInput: ciphertextInput,
-        logger,
-        operator: getZkOperator()
-      }
-    );
-    logger?.debug(
-      { startIdx, endIdx: startIdx + redactedPlaintext.length },
-      "verified proof"
-    );
-    realRedactedPlaintext.set(redactedPlaintext, startIdx);
-  }
-  async function verifyToprfProofPacket({ startIdx, proofData, payload: toprf }) {
-    if (!toprf?.dataLocation || !toprf.responses || !toprf.nullifier) {
-      throw new Error("invalid TOPRF proof payload");
-    }
-    const { dataLocation, nullifier } = toprf;
-    const ciphertextChunkEnd = Math.min(ciphertext.length, getChunkSizeBytes(algorithm) + startIdx);
-    const isLastChunk = ciphertextChunkEnd >= ciphertext.length;
-    const ciphertextChunk = ciphertext.slice(startIdx, ciphertextChunkEnd);
-    let nonce = concatenateUint8Arrays([iv, recordIV]);
-    if (!recordIV.length) {
-      nonce = generateIV(nonce, recordNumber);
-    }
-    const ciphertextInput = {
-      ciphertext: ciphertextChunk,
-      iv: nonce,
-      offsetBytes: startIdx
-    };
-    let pubInput = ciphertextInput;
-    const nulliferStr = binaryHashToStr(nullifier, dataLocation.length);
-    const locations = [];
-    const toprfEndIdx = dataLocation.fromIndex + dataLocation.length;
-    const trueCiphLen = isLastChunk && isTls13Suite(cipherSuite) ? ciphertextChunk.length - 1 : ciphertextChunk.length;
-    const overshoot = toprfEndIdx - trueCiphLen;
-    if (overshoot > 0) {
-      const nextPkt = getNextPacket(
-        strToUint8Array(nulliferStr.slice(dataLocation.length - overshoot))
-      );
-      if (!nextPkt) {
-        throw new Error("OPRF data overshot, but no next packet found");
-      }
-      const nextRecordIV = getRecordIV(ciphertext, cipherSuite);
-      let nextNonce = concatenateUint8Arrays([iv, nextRecordIV]);
-      if (!nextRecordIV.length) {
-        nextNonce = generateIV(nextNonce, recordNumber + 1);
-      }
-      pubInput = [
-        ciphertextInput,
-        {
-          ciphertext: nextPkt.slice(0, overshoot),
-          iv: nextNonce,
-          offsetBytes: 0
-        }
-      ];
-      locations.push(
-        {
-          pos: dataLocation.fromIndex,
-          len: dataLocation.length - overshoot
-        },
-        {
-          pos: ceilToBlockSizeMultiple(
-            dataLocation.fromIndex + dataLocation.length,
-            algorithm
-          ),
-          len: overshoot
-        }
-      );
-    } else {
-      locations.push({
-        pos: dataLocation.fromIndex,
-        len: dataLocation.length
-      });
-    }
-    await verifyProof(
-      {
-        proof: { algorithm, proofData, plaintext: void 0 },
-        publicInput: pubInput,
-        logger,
-        operator: getOprfOperator(),
-        toprf: {
-          locations,
-          domainSeparator: TOPRF_DOMAIN_SEPARATOR,
-          output: nullifier,
-          responses: toprf.responses
-        }
-      }
-    );
-    logger?.debug({ locations }, "verified TOPRF proof");
-    return {
-      set: strToUint8Array(
-        nulliferStr.slice(0, locations[0].len)
-      ),
-      startIdx: locations[0].pos + startIdx
-    };
-  }
-  function getZkOperator() {
-    return zkOperators?.[algorithm] || makeDefaultZkOperator(algorithm, zkEngine, logger);
-  }
-  function getOprfOperator() {
-    return oprfOperators?.[algorithm] || makeDefaultOPRFOperator(algorithm, zkEngine, logger);
-  }
-}
-function getChunkSizeBytes(alg) {
-  const { chunkSize, bitsPerWord } = ZK_CONFIG[alg];
-  return chunkSize * bitsPerWord / 8;
-}
-const zkEngines = {};
-const oprfEngines = {};
-const operatorMakers = {
-  "snarkjs": makeSnarkJsZKOperator,
-  "gnark": makeGnarkZkOperator,
-  "stwo": makeStwoZkOperator
-};
-const OPRF_OPERATOR_MAKERS = {
-  "gnark": makeGnarkOPRFOperator
-};
-function makeDefaultZkOperator(algorithm, zkEngine, logger) {
-  let zkOperators = zkEngines[zkEngine];
-  if (!zkOperators) {
-    zkEngines[zkEngine] = {};
-    zkOperators = zkEngines[zkEngine];
-  }
-  if (!zkOperators[algorithm]) {
-    const opType = getOperatorType();
-    const zkBaseUrl = opType === "remote" ? getZkResourcesBaseUrl() : void 0;
-    logger?.info({ type: opType, algorithm, zkBaseUrl }, "fetching zk operator");
-    const fetcher = opType === "local" ? makeLocalFileFetch() : makeRemoteFileFetch({ baseUrl: zkBaseUrl, logger });
-    const maker = operatorMakers[zkEngine];
-    if (!maker) {
-      throw new Error(`No ZK operator maker for ${zkEngine}`);
-    }
-    zkOperators[algorithm] = maker({ algorithm, fetcher });
-  }
-  return zkOperators[algorithm];
-}
-function getOperatorType() {
-  const envop = getEnvVariable("ZK_OPERATOR_TYPE");
-  if (envop === "local" || envop === "remote") {
-    return envop;
-  }
-  return detectEnvironment() === "node" ? "local" : "remote";
-}
-function makeDefaultOPRFOperator(algorithm, zkEngine, logger) {
-  let operators = oprfEngines[zkEngine];
-  if (!operators) {
-    oprfEngines[zkEngine] = {};
-    operators = oprfEngines[zkEngine];
-  }
-  if (!operators[algorithm]) {
-    const type = getOperatorType();
-    const zkBaseUrl = type === "remote" ? getZkResourcesBaseUrl() : void 0;
-    logger?.info({ type, algorithm, zkBaseUrl }, "fetching oprf operator");
-    const fetcher = type === "local" ? makeLocalFileFetch() : makeRemoteFileFetch({ baseUrl: zkBaseUrl, logger });
-    const maker = OPRF_OPERATOR_MAKERS[zkEngine];
-    if (!maker) {
-      throw new Error(`No OPRF operator maker for ${zkEngine}`);
-    }
-    operators[algorithm] = maker({ algorithm, fetcher });
-  }
-  return operators[algorithm];
-}
-function getEngineString(engine) {
-  if (engine === ZKProofEngine.ZK_ENGINE_GNARK) {
-    return "gnark";
-  }
-  if (engine === ZKProofEngine.ZK_ENGINE_SNARKJS) {
-    return "snarkjs";
-  }
-  if (engine === ZKProofEngine.ZK_ENGINE_STWO) {
-    return "stwo";
-  }
-  throw new Error(`Unknown ZK engine: ${engine}`);
-}
-function getEngineProto(engine) {
-  if (engine === "gnark") {
-    return ZKProofEngine.ZK_ENGINE_GNARK;
-  }
-  if (engine === "snarkjs") {
-    return ZKProofEngine.ZK_ENGINE_SNARKJS;
-  }
-  if (engine === "stwo") {
-    return ZKProofEngine.ZK_ENGINE_STWO;
-  }
-  throw new Error(`Unknown ZK engine: ${engine}`);
-}
-function getProofGenerationParamsForSlice({
-  key,
-  iv,
-  ciphertext,
-  redactedPlaintext,
-  slice: { fromIndex, toIndex }
-}) {
-  const ciphertextChunk = ciphertext.slice(fromIndex, toIndex);
-  const plaintextChunk = redactedPlaintext.slice(fromIndex, toIndex);
-  if (isFullyRedacted(plaintextChunk)) {
-    return;
-  }
-  for (let i = 0; i < ciphertextChunk.length; i++) {
-    if (plaintextChunk[i] === REDACTION_CHAR_CODE) {
-      ciphertextChunk[i] = REDACTION_CHAR_CODE;
-    }
-  }
-  return {
-    startIdx: fromIndex,
-    redactedPlaintext: plaintextChunk,
-    privateInput: { key },
-    publicInput: { ciphertext: ciphertextChunk, iv, offsetBytes: fromIndex }
-  };
-}
-function getTOPRFProofGenerationParamsForSlice({
-  key,
-  iv,
-  ciphertext,
-  slice: { fromIndex, toIndex },
-  toprf
-}) {
-  const ciphertextChunk = ciphertext.slice(fromIndex, toIndex);
-  if (toprf?.overshoot) {
-    const {
-      overshoot: { ciphertext: overshootCiphertext, iv: overshootIv }
-    } = toprf;
-    return {
-      privateInput: { key },
-      publicInput: [
-        {
-          ciphertext: ciphertextChunk,
-          iv,
-          offsetBytes: fromIndex
-        },
-        { ciphertext: overshootCiphertext, iv: overshootIv }
-      ],
-      toprf,
-      startIdx: fromIndex
-    };
-  }
-  return {
-    privateInput: { key },
-    publicInput: { ciphertext: ciphertextChunk, iv, offsetBytes: fromIndex },
-    toprf,
-    startIdx: fromIndex
-  };
-}
-function getIdealOffsetForToprfBlock(alg, { dataLocation, overshoot }) {
-  const chunkSizeBytes = getChunkSizeBytes(alg);
-  const blockSizeBytes = getBlockSizeBytes(alg);
-  const offsetChunks = Math.floor(dataLocation.fromIndex / chunkSizeBytes);
-  const endOffsetChunks = Math.floor((dataLocation.fromIndex + dataLocation.length) / chunkSizeBytes);
-  if (endOffsetChunks === offsetChunks) {
-    const start = offsetChunks * chunkSizeBytes;
-    if (overshoot) {
-      const overshootBlocks = Math.ceil(overshoot.ciphertext.length / blockSizeBytes);
-      return start + overshootBlocks * blockSizeBytes;
-    }
-    return start;
-  }
-  const offsetBytes = Math.floor(dataLocation.fromIndex / blockSizeBytes) * blockSizeBytes;
-  const endOffsetBytes = Math.ceil((dataLocation.fromIndex + dataLocation.length) / blockSizeBytes);
-  if (endOffsetBytes - offsetBytes > chunkSizeBytes) {
-    throw new AttestorError(
-      "ERROR_BAD_REQUEST",
-      "OPRF data cannot fit into a single chunk"
-    );
-  }
-  return offsetBytes;
-}
-function getZkResourcesBaseUrl() {
-  if (typeof ATTESTOR_BASE_URL !== "string") {
-    return DEFAULT_REMOTE_FILE_FETCH_BASE_URL;
-  }
-  return new URL(
-    DEFAULT_REMOTE_FILE_FETCH_BASE_URL,
-    ATTESTOR_BASE_URL
-  ).toString();
-}
-export {
-  getEngineProto,
-  getEngineString,
-  makeDefaultOPRFOperator,
-  makeDefaultZkOperator,
-  makeZkProofGenerator,
-  verifyZkPacket
-};