Merge pull request #1 from mlabs-haskell/connor/simple-password-tree

Simple password tree plugin

Author: chfanghr

src/pluginInteropServer/plugin.ts

@@ -0,0 +1,11 @@
+import { JsonProof } from "snarkyjs";
+
+export type PluginType = {
+  compile: () => Promise<string>;
+  getInputs: () => Promise<string[]>;
+  verify: (
+      jsonProof: JsonProof,
+      verificationKey: string,
+  ) => Promise<[string | boolean | undefined, string]>;
+  prove: (inputs: string[]) => Promise<undefined | JsonProof>;
+};

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<string, EntryPoint> = {
         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;
 }
 

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<Field>;
+  abstract getWitness(uid: bigint): Promise<undefined | PasswordTreeWitness>;
+  abstract getRole(uid: bigint): Promise<undefined | string>;
+}
+
+class InMemoryStorage implements TreeStorage {
+  roles: Map<bigint, string>;
+  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<string> => {
+  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<undefined | PasswordTreePublicInput> {
+  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<undefined | JsonProof> => {
+  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<string[]> => {
+  return Array.from(storage.roles.keys()).map(k => k.toString());
+};
+
+export const SimplePasswordTree: PluginType = {
+  compile,
+  getInputs,
+  verify: verifyAndGetRoleProgram,
+  prove,
+}
+
+export default SimplePasswordTree;

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<string>;
-    getInputs: () => Promise<string[]>;
-    verify: (
-        jsonProof: JsonProof,
-        verificationKey: string,
-    ) => Promise<[string | boolean | undefined, string]>;
-    prove: (inputs: string[]) => Promise<JsonProof>;
-};
-
 const roleMapping: Record<string, string> = {
     '7555220006856562833147743033256142154591945963958408607501861037584894828141':
         'admin',
@@ -42,7 +33,7 @@ const verifyAndGetRoleProgram = async (
     return [role, 'role proved'];
 };
 
-const prove = async (inputs: string[]): Promise<JsonProof> => {
+const prove = async (inputs: string[]): Promise<undefined | JsonProof> => {
     const [publicInput, secretInput] = inputs;
     console.log('simplePreimage proving for', publicInput, secretInput);
     const proof = await ProvePreimageProgram.baseCase(

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;