From 123433f1bf32ce6a109b50773bad1c98eafd4999 Mon Sep 17 00:00:00 2001 From: Hongrui Fang Date: Tue, 19 Sep 2023 19:06:11 +0800 Subject: [PATCH 1/2] implementation of simple password tree plugin --- src/pluginInteropServer/plugin.ts | 11 +++ src/pluginInteropServer/pluginServer.ts | 9 +- .../plugins/simplePasswordTree.ts | 91 +++++++++++++++++++ .../plugins/simplePreimage.ts | 13 +-- .../zkPrograms/passwordTreeProof.ts | 29 ++++++ 5 files changed, 141 insertions(+), 12 deletions(-) create mode 100644 src/pluginInteropServer/plugin.ts create mode 100644 src/pluginInteropServer/plugins/simplePasswordTree.ts create mode 100644 src/pluginInteropServer/zkPrograms/passwordTreeProof.ts diff --git a/src/pluginInteropServer/plugin.ts b/src/pluginInteropServer/plugin.ts new file mode 100644 index 0000000..9982a3e --- /dev/null +++ b/src/pluginInteropServer/plugin.ts @@ -0,0 +1,11 @@ +import { JsonProof } from "snarkyjs"; + +export type PluginType = { + compile: () => Promise; + getInputs: () => Promise; + verify: ( + jsonProof: JsonProof, + verificationKey: string, + ) => Promise<[string | boolean | undefined, string]>; + prove: (inputs: string[]) => Promise; +}; \ No newline at end of file diff --git a/src/pluginInteropServer/pluginServer.ts b/src/pluginInteropServer/pluginServer.ts index be4a6cd..775f50d 100644 --- a/src/pluginInteropServer/pluginServer.ts +++ b/src/pluginInteropServer/pluginServer.ts @@ -1,7 +1,9 @@ import express, { Request, Response } from 'express'; import bodyParser from 'body-parser'; -import { PluginType, SimplePreimage } from './plugins/simplePreimage'; +import { SimplePreimage } from './plugins/simplePreimage'; import { JsonProof } from 'snarkyjs'; +import { PluginType } from './plugin'; +import SimplePasswordTree from './plugins/simplePasswordTree'; const app = express(); const PORT = 3001; @@ -16,6 +18,10 @@ const EntryPoints: Record = { plugin: SimplePreimage, verification_key: null, }, + SimplePasswordTree: { + plugin: SimplePasswordTree, + verification_key: null, + } }; app.use(bodyParser.json()); @@ -48,6 +54,7 @@ async function prepareProofFunction( const plugin = entry.plugin; const jsonProof = await plugin.prove(data.arguments); + if (!jsonProof) throw 'entrypoint/plugin ${plugin_name} unable to generate proof'; return jsonProof; } diff --git a/src/pluginInteropServer/plugins/simplePasswordTree.ts b/src/pluginInteropServer/plugins/simplePasswordTree.ts new file mode 100644 index 0000000..108eebd --- /dev/null +++ b/src/pluginInteropServer/plugins/simplePasswordTree.ts @@ -0,0 +1,91 @@ +import { Experimental, Field, JsonProof, MerkleTree, Poseidon, verify } from "snarkyjs"; +import ProvePasswordInTreeProgram, { PASSWORD_TREE_HEIGHT, PasswordTreePublicInput, PasswordTreeWitness } from "../zkPrograms/passwordTreeProof"; +import { PluginType } from "../plugin"; + +const PasswordInTreeProverClass = Experimental.ZkProgram.Proof(ProvePasswordInTreeProgram); + +abstract class TreeStorage { + abstract getRoot(): Promise; + abstract getWitness(uid: bigint): Promise; + abstract getRole(uid: bigint): Promise; +} + +class InMemoryStorage implements TreeStorage { + roles: Map; + merkleTree: MerkleTree; + + constructor(roleMappings: Array<[bigint, Field, string]> = []) { + this.roles = new Map(); + this.merkleTree = new MerkleTree(PASSWORD_TREE_HEIGHT); + + roleMappings.forEach(([uid, password, role]) => { + this.roles.set(uid, role); + this.merkleTree.setLeaf(uid, Poseidon.hash([password])); + }) + } + + async getRoot() { return this.merkleTree.getRoot(); } + + async getWitness(uid: bigint) { + if (!this.roles.has(uid)) return undefined; + return new PasswordTreeWitness(this.merkleTree.getWitness(uid)) + } + + async getRole(uid: bigint) { return this.roles.get(uid); } +} + +const storage = new InMemoryStorage([ + [BigInt(0), Field('7555220006856562833147743033256142154591945963958408607501861037584894828141'), 'admin'], + [BigInt(1), Field('21565680844461314807147611702860246336805372493508489110556896454939225549736'), 'member'] +]); + +const compile = async (): Promise => { + console.log('Compiling SimplePasswordTree program'); + console.log(ProvePasswordInTreeProgram); + const { verificationKey } = await ProvePasswordInTreeProgram.compile(); + return verificationKey; +} + +const verifyAndGetRoleProgram = async ( + jsonProof: JsonProof, + verificationKey: string, +): Promise<[string | boolean | undefined, string]> => { + if (!verify(jsonProof, verificationKey)) { + return [false, 'proof invalid']; + } + const proof = PasswordInTreeProverClass.fromJSON(jsonProof); + const role = await storage.getRole(proof.publicInput.witness.calculateIndex().toBigInt()); + if (!role) { return [undefined, 'unknown public input']; } + return [role, 'role proved']; +} + +async function fetchPublicInput(uid: bigint): Promise { + const root = await storage.getRoot(); + const witness = await storage.getWitness(uid); + if (!witness) return undefined; + return new PasswordTreePublicInput({ root, witness }); +} + +const prove = async (inputs: string[]): Promise => { + const [uidStr, secretInput] = inputs; + const uid: bigint = BigInt(uidStr); + const publicInput = await fetchPublicInput(uid); + if (!publicInput) return undefined; + const proof = await ProvePasswordInTreeProgram.baseCase( + publicInput, Field(secretInput)); + return proof.toJSON(); +} + +// FIXME: I have no idea what this should do +const getInputs = async (): Promise => { + return Array.from(storage.roles.keys()).map(k => k.toString()); +}; + +export const SimplePasswordTree: PluginType = { + compile, + getInputs, + verify: verifyAndGetRoleProgram, + prove, +} + +export default SimplePasswordTree; \ No newline at end of file diff --git a/src/pluginInteropServer/plugins/simplePreimage.ts b/src/pluginInteropServer/plugins/simplePreimage.ts index 7bc14e8..1c9eb53 100644 --- a/src/pluginInteropServer/plugins/simplePreimage.ts +++ b/src/pluginInteropServer/plugins/simplePreimage.ts @@ -1,18 +1,9 @@ import { verify, Proof, Field, JsonProof, Experimental } from 'snarkyjs'; import { ProvePreimageProgram } from '../zkPrograms/hashPreimageProof'; +import { PluginType } from '../plugin'; const ProvePreimageProofClass = Experimental.ZkProgram.Proof(ProvePreimageProgram); -export type PluginType = { - compile: () => Promise; - getInputs: () => Promise; - verify: ( - jsonProof: JsonProof, - verificationKey: string, - ) => Promise<[string | boolean | undefined, string]>; - prove: (inputs: string[]) => Promise; -}; - const roleMapping: Record = { '7555220006856562833147743033256142154591945963958408607501861037584894828141': 'admin', @@ -42,7 +33,7 @@ const verifyAndGetRoleProgram = async ( return [role, 'role proved']; }; -const prove = async (inputs: string[]): Promise => { +const prove = async (inputs: string[]): Promise => { const [publicInput, secretInput] = inputs; console.log('simplePreimage proving for', publicInput, secretInput); const proof = await ProvePreimageProgram.baseCase( diff --git a/src/pluginInteropServer/zkPrograms/passwordTreeProof.ts b/src/pluginInteropServer/zkPrograms/passwordTreeProof.ts new file mode 100644 index 0000000..e898ccf --- /dev/null +++ b/src/pluginInteropServer/zkPrograms/passwordTreeProof.ts @@ -0,0 +1,29 @@ +import { Experimental, Field, MerkleWitness, Poseidon, Struct } from "snarkyjs"; + +export const PASSWORD_TREE_HEIGHT = 10; + +export class PasswordTreeWitness extends MerkleWitness(PASSWORD_TREE_HEIGHT) { } + +export class PasswordTreePublicInput extends Struct({ + witness: PasswordTreeWitness, + root: Field +}) { }; + +export const ProvePasswordInTreeProgram = Experimental.ZkProgram({ + publicInput: PasswordTreePublicInput, + publicOutput: Field, + + methods: { + baseCase: { + privateInputs: [Field], + method(publicInput: PasswordTreePublicInput, privateInput: Field): Field { + publicInput.witness + .calculateRoot(Poseidon.hash([privateInput])) + .assertEquals(publicInput.root); + return publicInput.witness.calculateIndex(); + } + } + } +}); + +export default ProvePasswordInTreeProgram; \ No newline at end of file From e0dcc1050472b07e861d6023aa5d175c65cd77a0 Mon Sep 17 00:00:00 2001 From: Hongrui Fang Date: Wed, 20 Sep 2023 17:11:30 +0800 Subject: [PATCH 2/2] PasswordInTreeProofClass --- src/pluginInteropServer/plugins/simplePasswordTree.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pluginInteropServer/plugins/simplePasswordTree.ts b/src/pluginInteropServer/plugins/simplePasswordTree.ts index 108eebd..7c64ba5 100644 --- a/src/pluginInteropServer/plugins/simplePasswordTree.ts +++ b/src/pluginInteropServer/plugins/simplePasswordTree.ts @@ -2,7 +2,7 @@ import { Experimental, Field, JsonProof, MerkleTree, Poseidon, verify } from "sn import ProvePasswordInTreeProgram, { PASSWORD_TREE_HEIGHT, PasswordTreePublicInput, PasswordTreeWitness } from "../zkPrograms/passwordTreeProof"; import { PluginType } from "../plugin"; -const PasswordInTreeProverClass = Experimental.ZkProgram.Proof(ProvePasswordInTreeProgram); +const PasswordInTreeProofClass = Experimental.ZkProgram.Proof(ProvePasswordInTreeProgram); abstract class TreeStorage { abstract getRoot(): Promise; @@ -53,7 +53,7 @@ const verifyAndGetRoleProgram = async ( if (!verify(jsonProof, verificationKey)) { return [false, 'proof invalid']; } - const proof = PasswordInTreeProverClass.fromJSON(jsonProof); + const proof = PasswordInTreeProofClass.fromJSON(jsonProof); const role = await storage.getRole(proof.publicInput.witness.calculateIndex().toBigInt()); if (!role) { return [undefined, 'unknown public input']; } return [role, 'role proved'];