Fix self-updating verification key

CHANGELOG.md

@@ -17,7 +17,13 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## [Unreleased](https://github.com/o1-labs/snarkyjs/compare/21de489...HEAD)
 
-> No unreleased changes yet
+### Breaking changes
+
+- Change type of verification key returned by `SmartContract.compile()` to match `VerificationKey` https://github.com/o1-labs/snarkyjs/pull/812
+
+### Fixed
+
+- Update the zkApp verification key from within one of its own methods, via proof https://github.com/o1-labs/snarkyjs/pull/812
 
 ## [0.9.4](https://github.com/o1-labs/snarkyjs/compare/9acec55...21de489)
 

src/examples/vk_regression.ts

@@ -11,7 +11,7 @@ await isReady;
 let dump = process.argv[4] === '--dump';
 let jsonPath = process.argv[dump ? 5 : 4];
 
-const Contracts: typeof SmartContract[] = [
+const Contracts: (typeof SmartContract)[] = [
   Voting_,
   Membership_,
   HelloWorld,
@@ -54,7 +54,7 @@ async function checkVk(contracts: typeof Contracts) {
       verificationKey: { data, hash },
     } = await c.compile();
 
-    if (data !== vk.data || hash !== vk.hash) {
+    if (data !== vk.data || hash.toString() !== vk.hash) {
       errorStack += `\n\nRegression test for contract ${
         c.name
       } failed, because of a verification key mismatch.
@@ -83,7 +83,10 @@ async function dumpVk(contracts: typeof Contracts) {
   for await (const c of contracts) {
     let { verificationKey } = await c.compile();
     newEntries[c.name] = {
-      verificationKey,
+      verificationKey: {
+        data: verificationKey.data,
+        hash: verificationKey.hash.toString(),
+      },
     };
   }
 

src/examples/zkapps/zkapp-self-update.ts

@@ -0,0 +1,78 @@
+/**
+ * This example deploys a zkApp and then updates its verification key via proof, self-replacing the zkApp
+ */
+import {
+  SmartContract,
+  VerificationKey,
+  method,
+  Permissions,
+  isReady,
+  PrivateKey,
+  Mina,
+  AccountUpdate,
+  Circuit,
+} from 'snarkyjs';
+
+class Foo extends SmartContract {
+  init() {
+    super.init();
+    this.account.permissions.set({
+      ...Permissions.default(),
+      setVerificationKey: Permissions.proof(),
+    });
+  }
+
+  @method replaceVerificationKey(verificationKey: VerificationKey) {
+    this.account.verificationKey.set(verificationKey);
+  }
+}
+
+class Bar extends SmartContract {
+  @method call() {
+    Circuit.log('Bar');
+  }
+}
+
+// setup
+
+await isReady;
+
+const Local = Mina.LocalBlockchain({ proofsEnabled: true });
+Mina.setActiveInstance(Local);
+
+const zkAppPrivateKey = PrivateKey.random();
+const zkAppAddress = zkAppPrivateKey.toPublicKey();
+const zkApp = new Foo(zkAppAddress);
+
+const { privateKey: deployerKey, publicKey: deployerAccount } =
+  Local.testAccounts[0];
+
+// deploy first verification key
+
+await Foo.compile();
+
+const tx = await Mina.transaction(deployerAccount, () => {
+  AccountUpdate.fundNewAccount(deployerAccount);
+  zkApp.deploy();
+});
+await tx.prove();
+await tx.sign([deployerKey, zkAppPrivateKey]).send();
+
+const fooVerificationKey = Mina.getAccount(zkAppAddress).zkapp?.verificationKey;
+Circuit.log('original verification key', fooVerificationKey);
+
+// update verification key
+
+const { verificationKey: barVerificationKey } = await Bar.compile();
+
+const tx2 = await Mina.transaction(deployerAccount, () => {
+  zkApp.replaceVerificationKey(barVerificationKey);
+});
+await tx2.prove();
+await tx2.sign([deployerKey]).send();
+
+const updatedVerificationKey =
+  Mina.getAccount(zkAppAddress).zkapp?.verificationKey;
+
+// should be different from Foo
+Circuit.log('updated verification key', updatedVerificationKey);

src/lib/account_update.ts

@@ -21,6 +21,7 @@ import * as Encoding from './encoding.js';
 import { hashWithPrefix, packToFields } from './hash.js';
 import { prefixes } from '../js_crypto/constants.js';
 import { Context } from './global-context.js';
+import { assert } from './errors.js';
 
 // external API
 export { AccountUpdate, Permissions, ZkappPublicInput };
@@ -1732,13 +1733,23 @@ const Authorization = {
     accountUpdate.lazyAuthorization = { ...signature, kind: 'lazy-signature' };
   },
   setProofAuthorizationKind(
-    { body }: AccountUpdate,
+    { body, id }: AccountUpdate,
     priorAccountUpdates?: AccountUpdate[]
   ) {
     body.authorizationKind.isSigned = Bool(false);
     body.authorizationKind.isProved = Bool(true);
     let hash = Circuit.witness(Field, () => {
-      priorAccountUpdates ??= zkAppProver.getData().transaction.accountUpdates;
+      let proverData = zkAppProver.getData();
+      let isProver = proverData !== undefined;
+      assert(
+        isProver || priorAccountUpdates !== undefined,
+        'Called `setProofAuthorizationKind()` outside the prover without passing in `priorAccountUpdates`.'
+      );
+      let myAccountUpdateId = isProver ? proverData.accountUpdate.id : id;
+      priorAccountUpdates ??= proverData.transaction.accountUpdates;
+      priorAccountUpdates = priorAccountUpdates.filter(
+        (a) => a.id !== myAccountUpdateId
+      );
       let priorAccountUpdatesFlat = CallForest.toFlatList(
         priorAccountUpdates,
         false

src/lib/errors.ts

@@ -2,7 +2,7 @@ import { Types } from '../provable/types.js';
 import { TokenId } from './account_update.js';
 import { Int64 } from './int.js';
 
-export { invalidTransactionError };
+export { invalidTransactionError, Bug, assert };
 
 const ErrorHandlers = {
   Invalid_fee_excess({
@@ -101,3 +101,18 @@ function invalidTransactionError(
   // fallback if we don't have a good error message yet
   return rawErrors;
 }
+
+/**
+ * An error that was assumed cannot happen, and communicates to users that it's not their fault but an internal bug.
+ */
+function Bug(message: string) {
+  return Error(
+    `${message}\nThis shouldn't have happened and indicates an internal bug.`
+  );
+}
+/**
+ * Make an assertion. When failing, this will communicate to users it's not their fault but indicates an internal bug.
+ */
+function assert(condition: boolean, message = 'Failed assertion.') {
+  if (!condition) throw Bug(message);
+}

src/lib/zkapp.ts

@@ -694,12 +694,13 @@ class SmartContract {
       this
     );
 
-    let verificationKey = getVerificationKeyArtifact();
+    let verificationKey_ = getVerificationKeyArtifact();
+    let verificationKey = {
+      data: verificationKey_.data,
+      hash: Field(verificationKey_.hash),
+    } satisfies VerificationKey;
     this._provers = provers;
-    this._verificationKey = {
-      data: verificationKey.data,
-      hash: Field(verificationKey.hash),
-    };
+    this._verificationKey = verificationKey;
     // TODO: instead of returning provers, return an artifact from which provers can be recovered
     return { verificationKey, provers, verify };
   }