diff --git a/packages/protocol/src/state/State.ts b/packages/protocol/src/state/State.ts index acb1a42f1..2608a0c9f 100644 --- a/packages/protocol/src/state/State.ts +++ b/packages/protocol/src/state/State.ts @@ -9,6 +9,7 @@ import { StateTransition } from "../model/StateTransition"; import { StateServiceProvider } from "./StateServiceProvider"; import { RuntimeMethodExecutionContext } from "./context/RuntimeMethodExecutionContext"; +import { WitnessBlockContext } from "./WitnessBlockContext"; export class WithPath { public path?: Field; @@ -137,11 +138,17 @@ export class State extends Mixin(WithPath, WithStateServiceProvider) { this.hasPathOrFail(); - const stateTransition = StateTransition.from(this.path, option); + const { isInWitnessBlock } = container.resolve(WitnessBlockContext); - container - .resolve(RuntimeMethodExecutionContext) - .addStateTransition(stateTransition); + // If we're inside a witness block, we only want to retrieve the state + // to use as a witness but not emit an ST + if (!isInWitnessBlock) { + const stateTransition = StateTransition.from(this.path, option); + + container + .resolve(RuntimeMethodExecutionContext) + .addStateTransition(stateTransition); + } return option; } @@ -170,6 +177,12 @@ export class State extends Mixin(WithPath, WithStateServiceProvider) { toOption ); + const { isInWitnessBlock } = container.resolve(WitnessBlockContext); + + if (isInWitnessBlock) { + throw new Error("Cannot set state inside of provable block."); + } + container .resolve(RuntimeMethodExecutionContext) .addStateTransition(stateTransition); diff --git a/packages/protocol/src/state/WitnessBlockContext.ts b/packages/protocol/src/state/WitnessBlockContext.ts new file mode 100644 index 000000000..aafbfd242 --- /dev/null +++ b/packages/protocol/src/state/WitnessBlockContext.ts @@ -0,0 +1,50 @@ +import { container, singleton } from "tsyringe"; +import { Provable } from "o1js"; + +@singleton() +export class WitnessBlockContext { + public witnessBlockDepth: number = 0; + + public get isInWitnessBlock() { + return this.witnessBlockDepth > 0; + } +} + +const asyncProxyWitnessFunction = < + Ret, + F extends (...args: any[]) => Promise, +>( + originalFuncDef: F +) => { + return async (...args: Parameters) => { + const context = container.resolve(WitnessBlockContext); + context.witnessBlockDepth += 1; + const ret = await originalFuncDef(...args); + context.witnessBlockDepth -= 1; + return ret; + }; +}; + +const proxySyncWitnessFunction = < + Params extends any[], + Ret, + F extends (...args: Params) => Ret, +>( + originalFuncDef: F +) => { + return (...args: Params): Ret => { + const context = container.resolve(WitnessBlockContext); + context.witnessBlockDepth += 1; + const ret = originalFuncDef(...args); + context.witnessBlockDepth -= 1; + return ret; + }; +}; + +Provable.witnessAsync = asyncProxyWitnessFunction(Provable.witnessAsync); + +Provable.witness = proxySyncWitnessFunction(Provable.witness); + +Provable.witnessFields = proxySyncWitnessFunction(Provable.witnessFields); + +Provable.asProver = proxySyncWitnessFunction(Provable.asProver); diff --git a/packages/sdk/test/stprover-emit-sts.test.ts b/packages/sdk/test/stprover-emit-sts.test.ts new file mode 100644 index 000000000..ef9170aed --- /dev/null +++ b/packages/sdk/test/stprover-emit-sts.test.ts @@ -0,0 +1,130 @@ +import "reflect-metadata"; + +import { UInt64 } from "@proto-kit/library"; +import { runtimeMethod, runtimeModule, RuntimeModule } from "@proto-kit/module"; +import { Field, PrivateKey, Provable } from "o1js"; +import { + RuntimeMethodExecutionContext, + State, + state, +} from "@proto-kit/protocol"; +import { container } from "tsyringe"; + +import { TestingAppChain } from "../src"; + +@runtimeModule() +class StateTester extends RuntimeModule { + @state() public state1 = State.from(UInt64); + + @runtimeMethod() + public async setFail() { + await Provable.witnessAsync(Field, async () => { + await this.state1.set(UInt64.from(10)); + return Field(0); + }); + } + + @runtimeMethod() + public async setPass() { + await this.state1.set(UInt64.from(10)); + } + + @runtimeMethod() + public async getSTs() { + await this.state1.get(); + } + + @runtimeMethod() + public async getNoSTs() { + await Provable.witnessAsync(Field, async () => { + const stateReturned = await this.state1.get(); + return Field.from(stateReturned.value.toBigInt()); + }); + } +} + +describe("StateTransition", () => { + const senderKey = PrivateKey.random(); + + const appChain = TestingAppChain.fromRuntime({ + StateTester, + }); + + beforeEach(async () => { + appChain.configurePartial({ + Runtime: { + StateTester: {}, + Balances: {}, + }, + + Protocol: { + ...appChain.config.Protocol!, + }, + }); + + await appChain.start(); + appChain.setSigner(senderKey); + }); + + afterEach(async () => { + await appChain.close(); + }); + + it("should emit no sts for get", async () => { + const stateTester = appChain.runtime.resolve("StateTester"); + const context = container.resolve(RuntimeMethodExecutionContext); + + // We set the state so when we fetch it it won't error. + const tx0 = await appChain.transaction( + senderKey.toPublicKey(), + async () => { + await stateTester.setPass(); + } + ); + await tx0.sign(); + await tx0.send(); + await appChain.produceBlock(); + + const tx1 = await appChain.transaction( + senderKey.toPublicKey(), + async () => { + await stateTester.getSTs(); + } + ); + await tx1.sign(); + await tx1.send(); + const STs = context.current().result.stateTransitions; + + expect(STs.length).not.toBe(0); + + const tx2 = await appChain.transaction( + senderKey.toPublicKey(), + async () => { + await stateTester.getNoSTs(); + } + ); + await tx2.sign(); + await tx2.send(); + const STs2 = context.current().result.stateTransitions; + expect(STs2.length).toBe(0); + }); + + it("should fail outside provable code for set", async () => { + const stateTester = appChain.runtime.resolve("StateTester"); + const tx1 = await appChain.transaction( + senderKey.toPublicKey(), + async () => { + await stateTester.setPass(); + } + ); + await tx1.sign(); + await tx1.send(); + await appChain.produceBlock(); + + await expect(() => + appChain.transaction(senderKey.toPublicKey(), async () => { + await stateTester.setFail(); + }) + ).rejects.toThrow(new Error("Cannot set state inside of provable block.")); + }); +});