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..7c64ba5 --- /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 PasswordInTreeProofClass = 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 = PasswordInTreeProofClass.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