o1-labs / o1js

TypeScript framework for zk-SNARKs and zkApps
https://docs.minaprotocol.com/en/zkapps/how-to-write-a-zkapp
Apache License 2.0
526 stars 119 forks source link

Using a side-loaded verification key throws an error in the SmartContract method when using ZkProgram chaining #1854

Open dfstio opened 1 month ago

dfstio commented 1 month ago

Using in the SmartContract method a proof generated by the ZkProgram that verifies the Signature throws the error:

yarn test signature.issue.sideloading.test
Error when proving Contract.setValue()
 FAIL  tests/signature.issue.sideloading.test.ts
  Side loading verification key proof
    ✓ should compile (6488 ms)
    ✓ should deploy a SmartContract (188 ms)
    ✓ should calculate the proof (12877 ms)
    ✕ should set value in SmartContract (643 ms)

  ● Side loading verification key proof › should set value in SmartContract

    thrown: Array [
      0,
      Array [
        248,
        MlBytes {
          "c": "Assert_failure",
          "l": 14,
          "t": 0,
        },
        -11,
      ],
      Array [
        0,
        MlBytes {
          "c": "src/mina/src/lib/pickles_types/plonk_types.ml",
          "l": 45,
          "t": 0,
        },
        391,
        51,
      ],
    ]

      88 |   });
      89 |
    > 90 |   it("should set value in SmartContract", async () => {
         |     ^
      91 |     const tx = await Mina.transaction({ sender, fee }, async () => {
      92 |       await contract.setValue(signatureProof, vk);
      93 |     });

      at tests/signature.issue.sideloading.test.ts:90:5
      at Object.<anonymous> (tests/signature.issue.sideloading.test.ts:52:9)

Test Suites: 1 failed, 1 total
Tests:       1 failed, 3 passed, 4 total
Snapshots:   0 total
Time:        21.391 s, estimated 43 s

The same test pass if I comment the line in the ZkProgram

signature.verify(publicKey, [data]).assertTrue();

or do not use side-loading verification key.

The code to reproduce the issue:

import { describe, expect, it } from "@jest/globals";
import {
  PublicKey,
  Mina,
  PrivateKey,
  DynamicProof,
  VerificationKey,
  ZkProgram,
  Field,
  SmartContract,
  method,
  AccountUpdate,
  state,
  State,
  Cache,
  Signature,
  verify,
  Void,
} from "o1js";

const program = ZkProgram({
  name: "SignatureIssue",
  publicInput: Field,
  publicOutput: Field,
  methods: {
    setValue: {
      privateInputs: [PublicKey, Signature],
      async method(data: Field, publicKey: PublicKey, signature: Signature) {
        // Commenting next line will allow the test to pass
        signature.verify(publicKey, [data]).assertTrue();
        return data;
      },
    },
  },
});

export class SignatureProof extends DynamicProof<Field, Field> {
  static publicInputType = Field;
  static publicOutputType = Field;
  static maxProofsVerified = 0 as const;
}

export class Contract extends SmartContract {
  @state(Field) value = State<Field>();

  @method async setValue(proof: SignatureProof, vk: VerificationKey) {
    proof.verify(vk);
    this.value.set(proof.publicOutput);
  }
}

describe("Side loading verification key proof", () => {
  let vk: VerificationKey;
  const cache: Cache = Cache.FileSystem("./cache");
  const zkAppKey = PrivateKey.randomKeypair();
  const contract = new Contract(zkAppKey.publicKey);
  let sender: Mina.TestPublicKey;
  const fee = 100_000_000;
  let signatureProof: SignatureProof;

  it("should compile", async () => {
    await Contract.compile({ cache });
    vk = (await program.compile({ cache })).verificationKey;
  });

  it("should deploy a SmartContract", async () => {
    const network = await Mina.LocalBlockchain();
    Mina.setActiveInstance(network);
    sender = network.testAccounts[0];
    const tx = await Mina.transaction({ sender, fee }, async () => {
      AccountUpdate.fundNewAccount(sender);
      await contract.deploy({});
    });
    await (await tx.sign([sender.key, zkAppKey.privateKey]).send()).wait();
  });

  it("should calculate the proof", async () => {
    const newData = Field(2);
    const owner = PrivateKey.randomKeypair();
    const signature = Signature.create(owner.privateKey, [newData]);
    const proof = await program.setValue(newData, owner.publicKey, signature);
    const ok1 = await verify(proof, vk);
    expect(ok1).toBe(true);

    signatureProof = SignatureProof.fromProof(proof);
    const ok2 = await verify(signatureProof, vk);
    expect(ok2).toBe(true);
  });

  it("should set value in SmartContract", async () => {
    const tx = await Mina.transaction({ sender, fee }, async () => {
      await contract.setValue(signatureProof, vk);
    });
    await tx.prove();
    await (await tx.sign([sender.key]).send()).wait();
  });
});
dfstio commented 1 month ago

Maybe related to https://github.com/o1-labs/o1js/issues/1561

dfstio commented 1 month ago

Chaining two ZkPrograms that use side-loading keys throws a similar error (this error does not occur if only one ZkProgram with side-loaded key is used):

Error when proving Contract.setValue()
 FAIL  tests/sideloading.issue.test.ts
  Side loading
    ✓ should compile (10078 ms)
    ✓ should prove (24360 ms)
    ✕ should deploy SmartContract and set value (748 ms)

  ● Side loading › should deploy SmartContract and set value

    thrown: Array [
      0,
      Array [
        248,
        MlBytes {
          "c": "Assert_failure",
          "l": 14,
          "t": 0,
        },
        -11,
      ],
      Array [
        0,
        MlBytes {
          "c": "src/mina/src/lib/pickles_types/vector.ml",
          "l": 40,
          "t": 0,
        },
        446,
        2,
      ],
    ]

      88 |   });
      89 |
    > 90 |   it("should deploy SmartContract and set value", async () => {
         |     ^
      91 |     const network = await Mina.LocalBlockchain();
      92 |     Mina.setActiveInstance(network);
      93 |     const sender = network.testAccounts[0];

      at tests/sideloading.issue.test.ts:90:5
      at Object.<anonymous> (tests/sideloading.issue.test.ts:64:9)

Test Suites: 1 failed, 1 total
Tests:       1 failed, 2 passed, 3 total
Snapshots:   0 total
Time:        36.347 s, estimated 52 s

Test code:

import { describe, expect, it } from "@jest/globals";
import {
  Mina,
  PrivateKey,
  DynamicProof,
  VerificationKey,
  Void,
  ZkProgram,
  Field,
  SmartContract,
  method,
  AccountUpdate,
  state,
  State,
  Cache,
} from "o1js";

class ProgramProof extends DynamicProof<Field, Void> {
  static publicInputType = Field;
  static publicOutputType = Void;
  static maxProofsVerified = 0 as const;
}

const program1 = ZkProgram({
  name: "program1",
  publicInput: Field,
  methods: {
    check: {
      privateInputs: [Field],
      async method(publicInput: Field, field: Field) {
        publicInput.assertEquals(field);
      },
    },
  },
});

const program2 = ZkProgram({
  name: "program2",
  publicInput: Field,
  methods: {
    check: {
      privateInputs: [ProgramProof, VerificationKey],
      async method(
        publicInput: Field,
        proof: ProgramProof,
        vk: VerificationKey
      ) {
        proof.verify(vk);
        proof.publicInput.assertEquals(publicInput);
      },
    },
  },
});

export class Contract extends SmartContract {
  @state(Field) value = State<Field>();

  @method async setValue(proof: ProgramProof, vk: VerificationKey) {
    proof.verify(vk);
    this.value.set(proof.publicInput);
  }
}

describe("Side loading", () => {
  let program1Vk: VerificationKey;
  let program2Vk: VerificationKey;
  let proof: ProgramProof;
  const value = Field(1);

  it("should compile", async () => {
    const cache: Cache = Cache.FileSystem("./cache");
    await Contract.compile({ cache });
    program1Vk = (await program1.compile({ cache })).verificationKey;
    program2Vk = (await program2.compile({ cache })).verificationKey;
  });

  it("should prove", async () => {
    const program1Proof = await program1.check(value, Field(1));
    const program1SideLoadedProof = ProgramProof.fromProof(program1Proof);
    const program2Proof = await program2.check(
      value,
      program1SideLoadedProof,
      program1Vk
    );
    proof = ProgramProof.fromProof(program2Proof);
    // Uncomment next line to make the test pass
    // proof = ProgramProof.fromProof(program1Proof);
  });

  it("should deploy SmartContract and set value", async () => {
    const network = await Mina.LocalBlockchain();
    Mina.setActiveInstance(network);
    const sender = network.testAccounts[0];
    const appKey = PrivateKey.randomKeypair();
    const zkApp = new Contract(appKey.publicKey);
    const tx = await Mina.transaction(
      { sender, fee: 100_000_000 },
      async () => {
        AccountUpdate.fundNewAccount(sender);
        await zkApp.deploy();
      }
    );
    await (await tx.sign([sender.key, appKey.privateKey]).send()).wait();

    const tx2 = await Mina.transaction(
      { sender, fee: 100_000_000 },
      async () => {
        // Replace program2Vk with program1Vk to make the test pass
        await zkApp.setValue(proof, program2Vk);
      }
    );
    await tx2.prove();
    await (await tx2.sign([sender.key]).send()).wait();
    expect(zkApp.value.get().toJSON()).toEqual(value.toJSON());
  });
});
dfstio commented 1 month ago

And one more similar error that occurs when I chain ZkPrograms with side-loaded keys is:

Error when proving NFTContract.updateMetadata()
    extend_exn: list too long

      at s (../../../../../../nix/store/12025gwcvnhg66dh6v6974viyq1b28ki-ocaml-base-compiler-4.14.0/lib/ocaml/stdlib.ml:29:14)
      at extend_exn (src/mina/src/lib/pickles_types/vector.ml:458:7)
      at go (src/mina/src/lib/pickles/compile.ml:56:11)
      at ../../../../../../workspace_root/src/mina/src/lib/promise/js/promise.js:38:29
      at withThreadPool (o1js/dist/node/index.cjs:3707:14)
      at prettifyStacktracePromise (o1js/dist/node/index.cjs:1951:12)
      at node_modules/o1js/dist/node/index.cjs:22487:14
      at Object.run (o1js/dist/node/index.cjs:10311:16)
      at createZkappProof (o1js/dist/node/index.cjs:22480:21)
      at addProof (o1js/dist/node/index.cjs:22473:15)
      at addMissingProofs (o1js/dist/node/index.cjs:22445:42)
      at node_modules/o1js/dist/node/index.cjs:23552:40
      at Object.<anonymous> (tests/sideloading.issue2.test.ts:237:5)

Test code:

import { describe, expect, it } from "@jest/globals";
import {
  PublicKey,
  Mina,
  PrivateKey,
  DynamicProof,
  VerificationKey,
  ZkProgram,
  Field,
  SmartContract,
  Struct,
  method,
  AccountUpdate,
  state,
  State,
  UInt32,
  DeployArgs,
  Bool,
  Cache,
  Signature,
  verify,
} from "o1js";

export class AddProof extends DynamicProof<Field, Field> {
  static publicInputType = Field;
  static publicOutputType = Field;
  static maxProofsVerified = 0 as const;
}

class NFTStateInput extends Struct({
  creator: PublicKey,
  metadata: Field,
  owner: PublicKey,
  version: UInt32,
  canChangeOwner: Bool,
}) {
  static assertEqual(a: NFTStateInput, b: NFTStateInput) {
    a.creator.assertEquals(b.creator);
    a.metadata.assertEquals(b.metadata);
    a.owner.assertEquals(b.owner);
    a.version.assertEquals(b.version);
    a.canChangeOwner.assertEquals(b.canChangeOwner);
  }
}

class NFTStateOutput extends Struct({
  metadata: Field,
  owner: PublicKey,
}) {}

const nftProgram = ZkProgram({
  name: "nftProgram",
  publicInput: NFTStateInput,
  publicOutput: NFTStateOutput,
  methods: {
    updateMetadata: {
      privateInputs: [Field, PublicKey],
      async method(
        initialState: NFTStateInput,
        metadata: Field,
        owner: PublicKey
      ) {
        initialState.owner.assertEquals(owner);
        return new NFTStateOutput({
          metadata,
          owner,
        });
      },
    },
    changeOwner: {
      privateInputs: [PublicKey], //, Signature
      async method(
        initialState: NFTStateInput,
        newOwner: PublicKey
        // https://github.com/o1-labs/o1js/issues/1854
        //signature: Signature
      ) {
        // signature
        //   .verify(initialState.owner, [
        //     ...NFTStateInput.toFields(initialState),
        //     ...newOwner.toFields(),
        //   ])
        //   .assertTrue();
        return new NFTStateOutput({
          metadata: initialState.metadata,
          owner: newOwner,
        });
      },
    },
    // https://github.com/o1-labs/o1js/issues/1854
    // Commenting the add method will make the test pass
    add: {
      privateInputs: [AddProof, VerificationKey],
      async method(
        initialState: NFTStateInput,
        proof: AddProof,
        vk: VerificationKey
      ) {
        proof.publicInput.assertEquals(initialState.metadata);
        proof.verify(vk);
        return new NFTStateOutput({
          metadata: proof.publicOutput,
          owner: initialState.owner,
        });
      },
    },
  },
});

export class NFTProof extends DynamicProof<NFTStateInput, NFTStateOutput> {
  static publicInputType = NFTStateInput;
  static publicOutputType = NFTStateOutput;
  static maxProofsVerified = 0 as const;
}

interface NFTContractDeployParams extends Exclude<DeployArgs, undefined> {
  metadata: Field;
  owner: PublicKey;
  creator: PublicKey;
  metadataVerificationKeyHash: Field;
  canChangeOwner: Bool;
}

export class NFTContract extends SmartContract {
  @state(Field) metadata = State<Field>();
  @state(PublicKey) owner = State<PublicKey>();
  @state(PublicKey) creator = State<PublicKey>();
  @state(UInt32) version = State<UInt32>();
  @state(Field) metadataVerificationKeyHash = State<Field>();
  @state(Bool) canChangeOwner = State<Bool>();

  async deploy(props: NFTContractDeployParams) {
    await super.deploy(props);
    this.metadata.set(props.metadata);
    this.owner.set(props.owner);
    this.creator.set(props.creator);
    this.metadataVerificationKeyHash.set(props.metadataVerificationKeyHash);
    this.version.set(UInt32.from(1));
    this.canChangeOwner.set(props.canChangeOwner);
  }

  @method async updateMetadata(proof: NFTProof, vk: VerificationKey) {
    this.metadataVerificationKeyHash
      .getAndRequireEquals()
      .assertEquals(vk.hash);
    proof.verify(vk);
    NFTStateInput.assertEqual(
      proof.publicInput,
      new NFTStateInput({
        creator: this.creator.getAndRequireEquals(),
        metadata: this.metadata.getAndRequireEquals(),
        owner: this.owner.getAndRequireEquals(),
        version: this.version.getAndRequireEquals(),
        canChangeOwner: this.canChangeOwner.getAndRequireEquals(),
      })
    );
    this.metadata.set(proof.publicOutput.metadata);
    this.owner.set(proof.publicOutput.owner);
  }
}

const pluginProgram = ZkProgram({
  name: "pluginProgram",
  publicInput: Field,
  publicOutput: Field,
  methods: {
    add: {
      privateInputs: [Field],
      async method(a: Field, b: Field) {
        return a.add(b);
      },
    },
  },
});

let nftProgramVk: VerificationKey;
let pluginProgramVk: VerificationKey;
const cache: Cache = Cache.FileSystem("./cache");
const owner = PrivateKey.randomKeypair();
const creator = PrivateKey.randomKeypair();
const metadata = Field(1);
const zkAppKey = PrivateKey.randomKeypair();
const nftContract = new NFTContract(zkAppKey.publicKey);
let sender: Mina.TestPublicKey;

describe("NFT with Side loading verification key", () => {
  it("should initialize a blockchain", async () => {
    const network = await Mina.LocalBlockchain();
    Mina.setActiveInstance(network);
    sender = network.testAccounts[0];
  });

  it("should compile", async () => {
    await NFTContract.compile({ cache });
    nftProgramVk = (await nftProgram.compile({ cache })).verificationKey;
    pluginProgramVk = (await pluginProgram.compile({ cache })).verificationKey;
  });

  it("should deploy a SmartContract", async () => {
    const tx = await Mina.transaction(
      { sender, fee: 100_000_000 },
      async () => {
        AccountUpdate.fundNewAccount(sender);
        await nftContract.deploy({
          creator: creator.publicKey,
          metadata,
          owner: owner.publicKey,
          metadataVerificationKeyHash: nftProgramVk.hash,
          canChangeOwner: Bool(true),
        });
      }
    );
    await (await tx.sign([sender.key, zkAppKey.privateKey]).send()).wait();
  });

  it("should update metadata", async () => {
    const newMetadata = Field(7);
    const nftInputState = new NFTStateInput({
      creator: nftContract.creator.get(),
      metadata: nftContract.metadata.get(),
      owner: nftContract.owner.get(),
      version: nftContract.version.get(),
      canChangeOwner: nftContract.canChangeOwner.get(),
    });
    const proof = await nftProgram.updateMetadata(
      nftInputState,
      newMetadata,
      nftContract.owner.get()
    );
    const metadataProof = NFTProof.fromProof(proof);
    const tx = await Mina.transaction(
      { sender, fee: 100_000_000 },
      async () => {
        await nftContract.updateMetadata(metadataProof, nftProgramVk);
      }
    );
    await tx.prove();
    await (await tx.sign([sender.key]).send()).wait();
    const metadata = nftContract.metadata.get();
    expect(metadata.toJSON()).toBe(newMetadata.toJSON());
  });
});
mitschabaude commented 1 month ago

Signature verification uses custom gates, so I think this is a matter of setting the right featureFlags on the DynamicProof

dfstio commented 1 month ago

Signature verification uses custom gates, so I think this is a matter of setting the right featureFlags on the DynamicProof

Thank you! Using FeatureFlags.allMaybe resolves the issue with the Signature.

It does not help with the other two cases that use chaining of ZkPrograms, but if I do not use side-loading verification keys at all and still chain ZkPrograms, everything works.

Trivo25 commented 1 month ago

The error when chaining multiple side loaded zkPrograms together originates in using the wrong DynamicProof setup.

In

const program1Proof = await program1.check(value, Field(1));
const program1SideLoadedProof = ProgramProof.fromProof(program1Proof);
const program2Proof = await program2.check(
  value,
  program1SideLoadedProof,
  program1Vk
);
proof = ProgramProof.fromProof(program2Proof);
// Uncomment next line to make the test pass
// proof = ProgramProof.fromProof(program1Proof);

you repurpose ProgramProof for both zkProgram proofs, although they differ in their proof structure. program1 verifies no other proof, which requires its DynamicProof class to set maxProofsVerified = 0, as done correctly. However, program2 does verify another proof - which in return requires its DynamicProof class to set maxProofsVerified = 1.

Long story short, introducing another DynamicProof class that handles programs that verify 1 proof will solve your issue


// proof class for the non-recursive zkProgram
class NonRecursiveProof extends DynamicProof<Field, Void> {
  static publicInputType = Field;
  static publicOutputType = Void;
  static maxProofsVerified = 0 as const;
  static featureFlags: FeatureFlags = FeatureFlags.allMaybe;
}

// proof class for the zkProgram that verfies the non-recursive Program
class RecusiveProof extends DynamicProof<Field, Void> {
  static publicInputType = Field;
  static publicOutputType = Void;
  static maxProofsVerified = 1 as const;
  static featureFlags: FeatureFlags = FeatureFlags.allMaybe;
}
dfstio commented 1 month ago

Long story short, introducing another DynamicProof class that handles programs that verify 1 proof will solve your issue

Thank you! This solution works.

Can I set the maxProofsVerified to the number higher than the number of proofs actually verified to be able in the future to replace the ZkProgram with another one, that maybe will verify more proofs?

According to the comments in the code,

 The `maxProofsVerified` constant is a product of the child circuit and indicates the maximum number that that circuit verifies itself.
 If you are unsure about what that is for you, you should use `2`.

https://github.com/o1-labs/o1js/blob/main/src/lib/proof-system/proof.ts#L186-L187

but when I try to set the maxProofsVerified to 2 for both DynamicProofs, I get the error

    TypeError: Cannot read properties of undefined (reading '2')

      at t2 (src/mina/src/lib/pickles_types/vector.ml:69:9)
      at func$14 (src/mina/src/lib/pickles_types/vector.ml:74:16)
      at expand_proof (src/mina/src/lib/pickles/step.ml:305:13)
      at ../../../../../../workspace_root/src/mina/src/lib/promise/js/promise.js:25:37
      at withThreadPool (o1js/dist/node/index.cjs:3707:14)
      at prettifyStacktracePromise (o1js/dist/node/index.cjs:1951:12)
      at Object.prove_ [as check] (o1js/dist/node/index.cjs:9856:18)
      at Object.<anonymous> (tests/sideloading.issue3.test.ts:83:27)

or, if I set the maxProofsVerified for RecusiveProof to 2 and for NonRecursiveProof to 0, I'm getting

Constraint unsatisfied (unreduced):
    Checked.Assert.equal
    File "src/mina/src/lib/pickles/step_verifier.ml", line 1214, characters 40-47:0
    File "src/mina/src/lib/pickles/step_verifier.ml", line 1199, characters 15-22
    File "src/mina/src/lib/pickles/step_main.ml", line 90, characters 15-22
    prevs_verified

    Constraint:
    ((basic(Equal(Var 10360)(Var 115805)))(annotation(Checked.Assert.equal)))
    Data:
    Equal 122568677290675473167057663695515585453 169212761389742147089486741819064608194

      at s (../../../../../../nix/store/12025gwcvnhg66dh6v6974viyq1b28ki-ocaml-base-compiler-4.14.0/lib/ocaml/stdlib.ml:29:14)
      at ../../../../../../nix/store/ncfn2bs6giyvk9wgxdrzskl3w7sdql4g-squashed-ocaml-dependencies/lib/ocaml/4.14.0/site-lib/base/printf.ml:6:43
      at ../../../../../../workspace_root/src/mina/src/lib/snarky/src/base/checked_runner.ml:216:13
      at equal$3 (src/mina/src/lib/snarky/src/base/snark0.ml:1087:29)
      at _iSQ_ (src/mina/src/lib/pickles/step_verifier.ml:1215:17)
      at with_label (src/mina/src/lib/snarky/src/base/snark0.ml:1253:15)
      at _iSP_ (src/mina/src/lib/pickles/step_verifier.ml:1214:24)
      at iteri$1 (../../../../../../nix/store/ncfn2bs6giyvk9wgxdrzskl3w7sdql4g-squashed-ocaml-dependencies/lib/ocaml/4.14.0/site-lib/base/array0.ml:50:18)
      at _iSN_ (src/mina/src/lib/pickles/step_verifier.ml:1204:11)
      at with_label (src/mina/src/lib/snarky/src/base/snark0.ml:1253:15)
      at verify (src/mina/src/lib/pickles/step_verifier.ml:1199:5)
      at ../../../../../../workspace_root/src/mina/src/lib/pickles/step_main.ml:108:25
      at with_label (src/mina/src/lib/snarky/src/base/snark0.ml:1253:15)
      at with_label (src/mina/src/lib/snarky/src/base/snark0.ml:1253:15)
      at _iJ1_ (src/mina/src/lib/pickles/step_main.ml:397:11)
      at ../../../../../../workspace_root/src/mina/src/lib/promise/js/promise.js:25:37
      at withThreadPool (o1js/dist/node/index.cjs:3707:14)
      at prettifyStacktracePromise (o1js/dist/node/index.cjs:1951:12)
      at node_modules/o1js/dist/node/index.cjs:22487:14
      at Object.run (o1js/dist/node/index.cjs:10311:16)
      at createZkappProof (o1js/dist/node/index.cjs:22480:21)
      at addProof (o1js/dist/node/index.cjs:22473:15)
      at addMissingProofs (o1js/dist/node/index.cjs:22445:42)
      at node_modules/o1js/dist/node/index.cjs:23552:40
      at Object.<anonymous> (tests/sideloading.issue3.test.ts:112:5)