diff --git a/.eslintignore b/.eslintignore new file mode 100644 index 0000000..b0a6f2e --- /dev/null +++ b/.eslintignore @@ -0,0 +1 @@ +wasm_bytes.ts \ No newline at end of file diff --git a/.gitignore b/.gitignore index daa0274..efede77 100644 --- a/.gitignore +++ b/.gitignore @@ -36,4 +36,6 @@ packages/prover/test_circuit/*.json packages/lib/src/circuits/ -packages/poseidon/sage/*.py \ No newline at end of file +packages/poseidon/sage/*.py + +wasm_bytes.ts \ No newline at end of file diff --git a/.prettierignore b/.prettierignore new file mode 100644 index 0000000..b0a6f2e --- /dev/null +++ b/.prettierignore @@ -0,0 +1 @@ +wasm_bytes.ts \ No newline at end of file diff --git a/packages/benchmark/node/src/node.bench_addr_membership.ts b/packages/benchmark/node/src/node.bench_addr_membership.ts index 03d3a56..dfd3f81 100644 --- a/packages/benchmark/node/src/node.bench_addr_membership.ts +++ b/packages/benchmark/node/src/node.bench_addr_membership.ts @@ -4,12 +4,10 @@ import { ecsign } from "@ethereumjs/util"; import { - SpartanWasm, Tree, Poseidon, MembershipProver, - defaultAddressMembershipConfig, - defaultWasmConfig + defaultAddressMembershipConfig } from "@personaelabs/spartan-ecdsa"; const benchAddrMembership = async () => { @@ -20,11 +18,9 @@ const benchAddrMembership = async () => { const { v, r, s } = ecsign(msgHash, privKey); const sig = `0x${r.toString("hex")}${s.toString("hex")}${v.toString(16)}`; - let wasm = new SpartanWasm(defaultWasmConfig); - // Init the Poseidon hash const poseidon = new Poseidon(); - await poseidon.initWasm(wasm); + await poseidon.initWasm(); const treeDepth = 20; const tree = new Tree(treeDepth, poseidon); @@ -57,7 +53,7 @@ const benchAddrMembership = async () => { ...defaultAddressMembershipConfig, enableProfiler: true }); - await prover.initWasm(wasm); + await prover.initWasm(); // Prove membership await prover.prove(sig, msgHash, merkleProof); diff --git a/packages/benchmark/node/src/node.bench_pubkey_membership.ts b/packages/benchmark/node/src/node.bench_pubkey_membership.ts index 6b9678a..51811a7 100644 --- a/packages/benchmark/node/src/node.bench_pubkey_membership.ts +++ b/packages/benchmark/node/src/node.bench_pubkey_membership.ts @@ -2,8 +2,6 @@ import { MembershipProver, Poseidon, Tree, - SpartanWasm, - defaultWasmConfig, defaultPubkeyMembershipConfig } from "@personaelabs/spartan-ecdsa"; import { @@ -22,11 +20,9 @@ const benchPubKeyMembership = async () => { const pubKey = ecrecover(msgHash, v, r, s); const sig = `0x${r.toString("hex")}${s.toString("hex")}${v.toString(16)}`; - let wasm = new SpartanWasm(defaultWasmConfig); - // Init the Poseidon hash const poseidon = new Poseidon(); - await poseidon.initWasm(wasm); + await poseidon.initWasm(); const treeDepth = 20; const tree = new Tree(treeDepth, poseidon); @@ -54,7 +50,7 @@ const benchPubKeyMembership = async () => { ...defaultPubkeyMembershipConfig, enableProfiler: true }); - await prover.initWasm(wasm); + await prover.initWasm(); // Prove membership await prover.prove(sig, msgHash, merkleProof); diff --git a/packages/benchmark/web/pages/index.tsx b/packages/benchmark/web/pages/index.tsx index ce8d564..a976c1f 100644 --- a/packages/benchmark/web/pages/index.tsx +++ b/packages/benchmark/web/pages/index.tsx @@ -3,9 +3,7 @@ import { MembershipProver, Tree, Poseidon, - SpartanWasm, defaultAddressMembershipConfig, - defaultWasmConfig, defaultPubkeyMembershipConfig } from "@personaelabs/spartan-ecdsa"; import { @@ -29,9 +27,8 @@ export default function Home() { const pubKey = ecrecover(msgHash, v, r, s); const sig = `0x${r.toString("hex")}${s.toString("hex")}${v.toString(16)}`; - const wasm = new SpartanWasm(defaultWasmConfig); const poseidon = new Poseidon(); - await poseidon.initWasm(wasm); + await poseidon.initWasm(); const treeDepth = 20; const pubKeyTree = new Tree(treeDepth, poseidon); @@ -55,8 +52,7 @@ export default function Home() { console.time("Full proving time"); const prover = new MembershipProver(defaultPubkeyMembershipConfig); - - prover.initWasm(wasm); + await prover.initWasm(); const { proof, publicInput } = await prover.prove( sig, @@ -81,9 +77,7 @@ export default function Home() { const { v, r, s } = ecsign(msgHash, privKey); const sig = `0x${r.toString("hex")}${s.toString("hex")}${v.toString(16)}`; - const wasm = new SpartanWasm(defaultWasmConfig); const poseidon = new Poseidon(); - await poseidon.initWasm(wasm); const treeDepth = 20; const addressTree = new Tree(treeDepth, poseidon); @@ -113,7 +107,7 @@ export default function Home() { enableProfiler: true }); - prover.initWasm(wasm); + await prover.initWasm(); const { proof, publicInput } = await prover.prove( sig, diff --git a/packages/circuits/tests/addr_membership.test.ts b/packages/circuits/tests/addr_membership.test.ts index 077a9be..9854a7e 100644 --- a/packages/circuits/tests/addr_membership.test.ts +++ b/packages/circuits/tests/addr_membership.test.ts @@ -2,12 +2,7 @@ const wasm_tester = require("circom_tester").wasm; var EC = require("elliptic").ec; import * as path from "path"; const ec = new EC("secp256k1"); -import { - Poseidon, - Tree, - SpartanWasm, - defaultWasmConfig -} from "@personaelabs/spartan-ecdsa"; +import { Poseidon, Tree } from "@personaelabs/spartan-ecdsa"; import { getEffEcdsaCircuitInput } from "./test_utils"; import { privateToAddress } from "@ethereumjs/util"; @@ -21,11 +16,9 @@ describe("membership", () => { } ); - const wasm = new SpartanWasm(defaultWasmConfig); - // Construct the tree const poseidon = new Poseidon(); - await poseidon.initWasm(wasm); + await poseidon.initWasm(); const nLevels = 10; const tree = new Tree(nLevels, poseidon); diff --git a/packages/circuits/tests/pubkey_membership.test.ts b/packages/circuits/tests/pubkey_membership.test.ts index 5202d11..01994e3 100644 --- a/packages/circuits/tests/pubkey_membership.test.ts +++ b/packages/circuits/tests/pubkey_membership.test.ts @@ -1,12 +1,7 @@ const wasm_tester = require("circom_tester").wasm; var EC = require("elliptic").ec; import * as path from "path"; -import { - Poseidon, - Tree, - SpartanWasm, - defaultWasmConfig -} from "@personaelabs/spartan-ecdsa"; +import { Poseidon, Tree } from "@personaelabs/spartan-ecdsa"; import { privateToPublic } from "@ethereumjs/util"; import { getEffEcdsaCircuitInput } from "./test_utils"; @@ -20,11 +15,9 @@ describe("pubkey_membership", () => { } ); - const wasm = new SpartanWasm(defaultWasmConfig); - // Construct the tree const poseidon = new Poseidon(); - await poseidon.initWasm(wasm); + await poseidon.initWasm(); const nLevels = 10; const tree = new Tree(nLevels, poseidon); diff --git a/packages/lib/.npmignore b/packages/lib/.npmignore index d0d957e..dda39fe 100644 --- a/packages/lib/.npmignore +++ b/packages/lib/.npmignore @@ -4,3 +4,4 @@ tsconfig.json jest.config.js copy_artifacts.sh +load_wasm.ts \ No newline at end of file diff --git a/packages/lib/load_wasm.ts b/packages/lib/load_wasm.ts new file mode 100644 index 0000000..a30763e --- /dev/null +++ b/packages/lib/load_wasm.ts @@ -0,0 +1,18 @@ +import * as fs from "fs"; + +/** + * Load the wasm file and output a typescript file with the wasm bytes embedded + */ +const loadWasm = async () => { + let wasm = fs.readFileSync("src/wasm/build/spartan_wasm_bg.wasm"); + + let bytes = new Uint8Array(wasm.buffer); + + const file = ` + export const wasmBytes = new Uint8Array([${bytes.toString()}]); + `; + + fs.writeFileSync("./src/wasm/wasm_bytes.ts", file); +}; + +loadWasm(); diff --git a/packages/lib/package.json b/packages/lib/package.json index ac2382c..d9f9e4b 100644 --- a/packages/lib/package.json +++ b/packages/lib/package.json @@ -8,8 +8,10 @@ "build/**/*" ], "scripts": { - "build": "tsc && sh ./copy_artifacts.sh", - "prepublishOnly": "yarn build" + "build": "tsc && yarn loadWasm && sh ./copy_artifacts.sh", + "prepublishOnly": "yarn build", + "prepare": "yarn loadWasm", + "loadWasm": "ts-node ./load_wasm.ts" }, "devDependencies": { "@types/jest": "^29.2.5", diff --git a/packages/lib/src/config.ts b/packages/lib/src/config.ts index 2137614..d749087 100644 --- a/packages/lib/src/config.ts +++ b/packages/lib/src/config.ts @@ -1,21 +1,11 @@ import * as path from "path"; const isWeb = typeof window !== "undefined"; -import { LeafType, ProverConfig, WasmConfig } from "./types"; - -export const defaultWasmConfig: WasmConfig = { - pathOrUrl: isWeb - ? "https://storage.googleapis.com/personae-proving-keys/spartan_wasm_bg.wasm" - : path.join(__dirname, "wasm/build/spartan_wasm_bg.wasm") -}; +import { LeafType, ProverConfig } from "./types"; // Default configs for MembershipProver // Default configs for pubkey membership proving export const defaultPubkeyMembershipConfig: ProverConfig = { - spartanWasm: isWeb - ? "https://storage.googleapis.com/personae-proving-keys/spartan_wasm_bg.wasm" - : path.join(__dirname, "wasm/build/spartan_wasm_bg.wasm"), - witnessGenWasm: isWeb ? "https://storage.googleapis.com/personae-proving-keys/membership/pubkey_membership.wasm" : path.join(__dirname, "circuits/pubkey_membership.wasm"), @@ -29,10 +19,6 @@ export const defaultPubkeyMembershipConfig: ProverConfig = { // Default configs for address membership proving export const defaultAddressMembershipConfig: ProverConfig = { - spartanWasm: isWeb - ? "https://storage.googleapis.com/personae-proving-keys/spartan_wasm_bg.wasm" - : path.join(__dirname, "wasm/build/spartan_wasm_bg.wasm"), - witnessGenWasm: isWeb ? "https://storage.googleapis.com/personae-proving-keys/membership/addr_membership.wasm" : path.join(__dirname, "circuits/addr_membership.wasm"), diff --git a/packages/lib/src/core/membership_prover.ts b/packages/lib/src/core/membership_prover.ts index 9fed71b..4839c32 100644 --- a/packages/lib/src/core/membership_prover.ts +++ b/packages/lib/src/core/membership_prover.ts @@ -1,6 +1,5 @@ import { Profiler } from "../helpers/profiler"; import { IProver, MerkleProof, NIZK, ProverConfig, LeafType } from "../types"; -import { SpartanWasm } from "../wasm"; import { bigIntToBytes, loadCircuit, @@ -11,12 +10,12 @@ import { EffEcdsaPubInput, EffEcdsaCircuitPubInput } from "../helpers/efficient_ecdsa"; +import wasm, { init } from "../wasm"; /** * ECDSA Membership Prover */ export class MembershipProver extends Profiler implements IProver { - spartanWasm!: SpartanWasm; circuit: string; witnessGenWasm: string; leafType: LeafType; @@ -29,9 +28,8 @@ export class MembershipProver extends Profiler implements IProver { this.witnessGenWasm = options.witnessGenWasm; } - async initWasm(wasm: SpartanWasm) { - this.spartanWasm = wasm; - this.spartanWasm.init(); + async initWasm() { + await init(); } // @ts-ignore @@ -40,10 +38,6 @@ export class MembershipProver extends Profiler implements IProver { msgHash: Buffer, merkleProof: MerkleProof ): Promise { - if (typeof this.spartanWasm === "undefined") { - throw new Error("wasm not initialized. Please call initWasm()."); - } - const { r, s, v } = fromSig(sig); const circuitPubInput = EffEcdsaCircuitPubInput.computeFromSig( @@ -86,11 +80,7 @@ export class MembershipProver extends Profiler implements IProver { this.timeEnd("Load circuit"); this.time("Prove"); - let proof = await this.spartanWasm.prove( - circuitBin, - witness.data, - pubInput - ); + let proof = wasm.prove(circuitBin, witness.data, pubInput); this.timeEnd("Prove"); return { proof, publicInput: effEcdsaPubInput.serialize() }; diff --git a/packages/lib/src/helpers/poseidon.ts b/packages/lib/src/helpers/poseidon.ts index 5927401..cf4b541 100644 --- a/packages/lib/src/helpers/poseidon.ts +++ b/packages/lib/src/helpers/poseidon.ts @@ -1,34 +1,22 @@ -import { SpartanWasm } from "../wasm"; import { bigIntToLeBytes, bytesLeToBigInt } from "./utils"; +import wasm, { init } from "../wasm"; export class Poseidon { - wasm!: SpartanWasm; - constructor() {} - - async initWasm(wasm: SpartanWasm) { - this.wasm = wasm; - await this.wasm.init(); - } - - assertInitWasm() { - if (typeof this.wasm === "undefined") { - throw new Error("wasm not initialized. Please call initWasm()."); - } - } - hash(inputs: bigint[]): bigint { - this.assertInitWasm(); const inputsBytes = new Uint8Array(32 * inputs.length); for (let i = 0; i < inputs.length; i++) { inputsBytes.set(bigIntToLeBytes(inputs[i], 32), i * 32); } - const result = this.wasm.poseidon(inputsBytes); + const result = wasm.poseidon(inputsBytes); return bytesLeToBigInt(result); } + async initWasm() { + await init(); + } + hashPubKey(pubKey: Buffer): bigint { - this.assertInitWasm(); const pubKeyX = BigInt("0x" + pubKey.toString("hex").slice(0, 64)); const pubKeyY = BigInt("0x" + pubKey.toString("hex").slice(64, 128)); diff --git a/packages/lib/src/types.ts b/packages/lib/src/types.ts index a23c813..881f4a3 100644 --- a/packages/lib/src/types.ts +++ b/packages/lib/src/types.ts @@ -1,8 +1,6 @@ // The same structure as MerkleProof in @zk-kit/incremental-merkle-tree. // Not directly using MerkleProof defined in @zk-kit/incremental-merkle-tree so // library users can choose whatever merkle tree management method they want. -import { SpartanWasm } from "./wasm"; - export interface MerkleProof { root: bigint; siblings: bigint[]; @@ -18,23 +16,16 @@ export interface ProverConfig { proverWasm?: string; witnessGenWasm: string; circuit: string; - spartanWasm: string; enableProfiler?: boolean; leafType: LeafType; } -export interface WasmConfig { - pathOrUrl: string; -} - export interface VerifyOptions { circuit?: string; // Path to circuit file compiled by Nova-Scotia - spartanWasm?: string; // Path to spartan wasm file enableProfiler?: boolean; } export interface IProver { - spartanWasm: SpartanWasm; circuit: string; // Path to circuit file compiled by Nova-Scotia witnessGenWasm: string; // Path to witness generator wasm file generated by Circom @@ -42,16 +33,11 @@ export interface IProver { } export interface IVerifier { - spartanWasm: SpartanWasm; // Path to spartan wasm file circuit: string; // Path to circuit file compiled by Nova-Scotia verify(proof: Uint8Array, publicInput: Uint8Array): Promise; } -export interface SpartanWasmOptions { - spartanWasm: string; -} - export enum LeafType { PubKeyHash, Address diff --git a/packages/lib/src/wasm/index.ts b/packages/lib/src/wasm/index.ts index d869020..8c33c2c 100644 --- a/packages/lib/src/wasm/index.ts +++ b/packages/lib/src/wasm/index.ts @@ -1,46 +1,11 @@ import * as wasm from "./wasm"; import _initWeb from "./wasm.js"; -import fs from "fs"; -import path from "path"; -import { WasmConfig } from "../types"; -// TODO: Rename this to just Wasm since it includes not only Spartan but also Poseidon -export class SpartanWasm { - private spartanWasmPathOrUrl: any; +import { wasmBytes } from "./wasm_bytes"; - constructor(config: WasmConfig) { - this.spartanWasmPathOrUrl = config.pathOrUrl; - } +export const init = async () => { + await wasm.initSync(wasmBytes.buffer); + wasm.init_panic_hook(); +}; - async init() { - if (typeof window === "undefined") { - await this.initNode(); - } else { - await this.initWeb(); - } - } - - prove(circuit: Uint8Array, vars: Uint8Array, public_inputs: Uint8Array) { - return wasm.prove(circuit, vars, public_inputs); - } - - verify(circuit: Uint8Array, vars: Uint8Array, public_inputs: Uint8Array) { - return wasm.verify(circuit, vars, public_inputs); - } - - poseidon(inputs: Uint8Array) { - return wasm.poseidon(inputs); - } - - private async initNode() { - const bytes = fs.readFileSync(this.spartanWasmPathOrUrl); - - await wasm.initSync(bytes); - await wasm.init_panic_hook(); - } - - private async initWeb() { - await _initWeb(this.spartanWasmPathOrUrl); - await wasm.init_panic_hook(); - } -} +export default wasm; diff --git a/packages/lib/src/wasm/wasm.js b/packages/lib/src/wasm/wasm.js index 4da191c..27ed1c3 100644 --- a/packages/lib/src/wasm/wasm.js +++ b/packages/lib/src/wasm/wasm.js @@ -400,16 +400,19 @@ function finalizeInit(instance, module) { return wasm; } -function initSync(module, maybe_memory) { +async function initSync(module, maybe_memory) { const imports = getImports(); initMemory(imports, maybe_memory); + /* if (!(module instanceof WebAssembly.Module)) { module = new WebAssembly.Module(module); } + */ + const compiled = WebAssembly.compile(module); - const instance = new WebAssembly.Instance(module, imports); + const instance = await WebAssembly.instantiate(await compiled, imports); return finalizeInit(instance, module); } diff --git a/packages/lib/tests/membership_nizk.test.ts b/packages/lib/tests/membership_nizk.test.ts index aa620f0..cade888 100644 --- a/packages/lib/tests/membership_nizk.test.ts +++ b/packages/lib/tests/membership_nizk.test.ts @@ -1,12 +1,9 @@ -import * as path from "path"; import { MembershipProver, Tree, Poseidon, defaultAddressMembershipConfig, - defaultPubkeyMembershipConfig, - SpartanWasm, - defaultWasmConfig + defaultPubkeyMembershipConfig } from "../src/lib"; import { hashPersonalMessage, @@ -36,18 +33,14 @@ describe("membership prove and verify", () => { const sig = `0x${r.toString("hex")}${s.toString("hex")}${v.toString(16)}`; let poseidon: Poseidon; - let wasm: SpartanWasm; beforeAll(async () => { - // Init Wasm - wasm = new SpartanWasm(defaultWasmConfig); - // Init Poseidon poseidon = new Poseidon(); - await poseidon.initWasm(wasm); + await poseidon.initWasm(); }); - describe("pubkey_membership prover and verify", () => { + describe.only("pubkey_membership prover and verify", () => { it("should prove and verify valid signature and merkle proof", async () => { const pubKeyTree = new Tree(treeDepth, poseidon); @@ -66,8 +59,6 @@ describe("membership prove and verify", () => { defaultPubkeyMembershipConfig ); - await pubKeyMembershipProver.initWasm(wasm); - const index = pubKeyTree.indexOf(proverPubKeyHash as bigint); const merkleProof = pubKeyTree.createProof(index); @@ -101,8 +92,6 @@ describe("membership prove and verify", () => { defaultAddressMembershipConfig ); - await addressMembershipProver.initWasm(wasm); - const index = addressTree.indexOf(proverAddress as bigint); const merkleProof = addressTree.createProof(index);