From b6d9379fc9ffab8cef67f3eede922ad74cd3158c Mon Sep 17 00:00:00 2001 From: Cameron Marshall Date: Mon, 9 Mar 2026 14:52:37 -0400 Subject: [PATCH] feat: Return proving/verifying keys from execution methods with automatic KeyStore persistence --- docs/design/key-caching.md | 215 +++++++++++++++++++++++++++ sdk/src/browser.ts | 3 +- sdk/src/keys/provider/memory.ts | 46 +++++- sdk/src/program-manager.ts | 155 ++++++++++++++++++- sdk/src/wasm.ts | 1 + sdk/tests/key-provider.test.ts | 155 +++++++++++++++++++ sdk/tests/program-manager.test.ts | 80 +++++++++- wasm/src/programs/manager/execute.rs | 154 +++++++++++++++++++ wasm/src/programs/response.rs | 70 ++++++++- 9 files changed, 865 insertions(+), 14 deletions(-) create mode 100644 docs/design/key-caching.md diff --git a/docs/design/key-caching.md b/docs/design/key-caching.md new file mode 100644 index 000000000..9d72c59d1 --- /dev/null +++ b/docs/design/key-caching.md @@ -0,0 +1,215 @@ +--- +id: key-caching +title: "Design: Execution Key Caching & Persistence" +--- + +# Execution Key Caching & Persistence + +## Problem + +SnarkVM assumes proving and verifying keys persist in a long-running process. In practice, SDK users run ephemeral +processes (serverless functions, CLI tools, web workers) where the VM is created, used, and destroyed. When a process +exits, all synthesized keys are lost. + +Key synthesis is the most expensive part of program execution -- often 30+ seconds for complex functions. Without key +persistence, every execution pays this cost, even for the same program and function. + +Users faced three problems: +1. No way to retrieve keys from execution methods for custom storage +2. No automatic persistence -- keys were discarded after each execution +3. Manual key management required understanding synthesis, caching, and locator conventions + +## Design Goals + +- **Zero-config improvement**: Users who add `cacheKeys: true` get automatic key persistence with no other changes +- **Backward compatible**: All existing code works identically; `cacheKeys` defaults to `false` +- **Layered access**: Users can use automatic persistence, manual key extraction, or both +- **No new dependencies**: Built on the existing `KeyStore` interface (PR #1207) and `AleoKeyProvider` + +## Architecture + +``` + buildExecutionTransaction({ cacheKeys: true }) + | + v + +-------------------------------------------+ + | TypeScript SDK Layer | + | ProgramManager.buildExecutionTransaction | + +-------------------------------------------+ + | | + v v + +------------------------+ +-------------------+ + | WASM Layer | | persistKeysToStore() + | buildExecutionTx- | | (auto-persist to | + | WithKeys() | | KeyStore) | + +------------------------+ +-------------------+ + | | + v v + +------------------------+ +-------------------+ + | ExecutionTransaction- | | KeyStore | + | Response | | (LocalFileKeyStore | + | { tx, pk?, vk? } | | or custom impl) | + +------------------------+ +-------------------+ + + Subsequent executions: + AleoKeyProvider.functionKeys() + 1. Memory cache --> hit? return + 2. KeyStore --> hit? cache in memory, return + 3. Network/synth --> fallback +``` + +## API Surface + +### New option: `cacheKeys` + +Added to `ExecuteOptions`, used by `buildExecutionTransaction()`, `execute()`, and `run()`. + +```typescript +const result = await programManager.buildExecutionTransaction({ + programName: "my_program.aleo", + functionName: "my_function", + inputs: ["1u32", "2u32"], + priorityFee: 0.0, + privateFee: false, + cacheKeys: true, // opt-in to key return + auto-persistence +}); + +result.transaction; // Transaction (always present) +result.provingKey; // ProvingKey | undefined +result.verifyingKey; // VerifyingKey | undefined +``` + +When `cacheKeys` is `false` (default), `buildExecutionTransaction` returns `Transaction` directly -- identical to +existing behavior. + +TypeScript function overloads enforce this at the type level: + +```typescript +// cacheKeys: true --> returns ExecutionTransactionResult +// cacheKeys: false --> returns Transaction +``` + +### Auto-persistence + +When a `KeyStore` is configured on the `AleoKeyProvider`, keys are automatically persisted after execution: + +```typescript +import { AleoKeyProvider, LocalFileKeyStore } from "@provablehq/sdk"; + +const keyStore = new LocalFileKeyStore("/path/to/keys"); +const keyProvider = new AleoKeyProvider(keyStore); +keyProvider.useCache(true); + +const programManager = new ProgramManager(networkClient, keyProvider); +programManager.setAccount(account); + +// First execution: synthesizes keys, persists to disk +await programManager.execute({ + programName: "my_program.aleo", + functionName: "my_function", + inputs: ["1u32", "2u32"], + cacheKeys: true, + keySearchParams: { cacheKey: "my_program.aleo:my_function" }, +}); + +// Second execution: loads keys from KeyStore automatically -- no re-synthesis +await programManager.execute({ + programName: "my_program.aleo", + functionName: "my_function", + inputs: ["3u32", "4u32"], + keySearchParams: { cacheKey: "my_program.aleo:my_function" }, +}); +``` + +### Key lookup chain + +`AleoKeyProvider.functionKeys()` resolves keys in this order: + +1. **In-memory cache** -- fastest, populated by `cacheKeys()` or prior lookups +2. **KeyStore** -- persistent storage (file system, database, etc.) +3. **Network fetch / synthesis** -- slowest, used as last resort + +Results from the KeyStore are automatically cached in memory for subsequent calls within the same process. + +## WASM Layer + +A new method `buildExecutionTransactionWithKeys` was added alongside the existing `buildExecutionTransaction` to avoid +breaking the `@provablehq/wasm` public API. + +```rust +// Returns ExecutionTransactionResponse { transaction, proving_key?, verifying_key? } +pub async fn execute_with_keys( + // ... same parameters as execute() ... + cache: bool, // controls whether keys are extracted +) -> Result +``` + +Key extraction happens after the `execute_program!` macro completes but before the `ProcessNative` goes out of scope. +This is the only window where keys are accessible -- once the process drops, synthesized keys are lost. + +### ExecutionTransactionResponse + +```rust +pub struct ExecutionTransactionResponse { + transaction: TransactionNative, + proving_key: Option, + verifying_key: Option, +} +``` + +Key accessors use `.take()` semantics (consuming the value). This prevents accidental double-reads and matches the +existing `ExecutionResponse` pattern. + +## Locator Convention + +Keys are persisted using dot-separated locators: + +``` +{programName}.{functionName}.prover +{programName}.{functionName}.verifier +``` + +Example: `my_program.aleo.my_function.prover` + +`LocalFileKeyStore.validateLocator()` rejects path separators (`/`, `\`), so dots are used instead of slashes. + +## Design Decisions + +### Why a new WASM method instead of modifying `execute()`? + +The `@provablehq/wasm` package has its own release cycle and public API. Adding a parameter to the existing method would +be a breaking change for any consumer that calls `execute()` positionally. A new method is additive and safe. + +### Why `.take()` semantics on key accessors? + +Proving keys can be 100+ MB. The `.take()` pattern (returning the value and setting the field to `None`) avoids cloning +large keys. It also makes ownership explicit -- once you've taken a key, you own it and the response no longer holds it. + +### Why auto-persist in the SDK layer, not WASM? + +The WASM layer has no concept of storage backends. Persistence logic belongs in the TypeScript SDK where `KeyStore` +implementations (file system, IndexedDB, custom) are available. The WASM layer's job is to make keys *available*; the +SDK layer decides what to do with them. + +### Why is `cacheKeys` opt-in (`false` by default)? + +Key extraction clones the keys from the process, temporarily doubling memory usage for that key. For users who don't +need keys returned, this cost should not be imposed. The default behavior is unchanged from before this feature. + +## Considerations + +### Memory overhead + +When `cacheKeys: true`, proving keys are cloned from the WASM process for return. For large programs, this can +temporarily double memory usage for that key (100+ MB). This matches the existing `executeFunctionOffline(cache=true)` +behavior and is expected. + +### Concurrent execution + +If multiple executions for the same program/function run concurrently, they may both persist keys simultaneously. +`LocalFileKeyStore` uses atomic writes (temp file + rename), so the last write wins safely. No data corruption occurs. + +### Persistence failures are non-blocking + +`persistKeysToStore()` catches and logs errors rather than propagating them. A failure to persist keys should not +prevent a successful execution from completing. Keys are still returned in the result object. diff --git a/sdk/src/browser.ts b/sdk/src/browser.ts index 46a9af7f5..821146d03 100644 --- a/sdk/src/browser.ts +++ b/sdk/src/browser.ts @@ -99,7 +99,7 @@ async function initializeWasm() { console.warn("initializeWasm is deprecated, you no longer need to use it"); } -export { ProgramManager, ProvingRequestOptions, ExecuteOptions, FeeAuthorizationOptions, AuthorizationOptions } from "./program-manager.js"; +export { ProgramManager, ProvingRequestOptions, ExecuteOptions, ExecutionTransactionResult, FeeAuthorizationOptions, AuthorizationOptions } from "./program-manager.js"; export { logAndThrow } from "./utils.js"; @@ -116,6 +116,7 @@ export { Execution as FunctionExecution, ExecutionRequest, ExecutionResponse, + ExecutionTransactionResponse, EncryptionToolkit, Field, GraphKey, diff --git a/sdk/src/keys/provider/memory.ts b/sdk/src/keys/provider/memory.ts index ee54378b4..c2fbb9daa 100644 --- a/sdk/src/keys/provider/memory.ts +++ b/sdk/src/keys/provider/memory.ts @@ -22,7 +22,7 @@ import { } from "../../wasm.js"; import { get } from "../../utils.js"; -import { KeyStore } from "../keystore/interface.js"; +import { KeyLocator, KeyStore } from "../keystore/interface.js"; type AleoKeyProviderInitParams = { proverUri?: string; @@ -66,6 +66,7 @@ class AleoKeyProvider implements FunctionKeyProvider { cache: Map; cacheOption: boolean; keyUris: string; + private _keyStore: KeyStore | undefined; async fetchBytes(url = "/"): Promise { try { @@ -77,14 +78,24 @@ class AleoKeyProvider implements FunctionKeyProvider { } } - constructor() { + constructor(keyStore?: KeyStore) { this.keyUris = KEY_STORE; this.cache = new Map(); this.cacheOption = false; + this._keyStore = keyStore; } async keyStore(): Promise { - return undefined; + return this._keyStore; + } + + /** + * Set the key store for persistent key storage + * + * @param {KeyStore | undefined} keyStore The key store to use for persistent storage + */ + setKeyStore(keyStore: KeyStore | undefined) { + this._keyStore = keyStore; } /** @@ -213,7 +224,34 @@ class AleoKeyProvider implements FunctionKeyProvider { } if (cacheKey) { - return this.getKeys(cacheKey); + // Check memory cache first + if (this.containsKeys(cacheKey)) { + return this.getKeys(cacheKey); + } + + // Fall back to KeyStore if configured + if (this._keyStore) { + try { + const proverLocator: KeyLocator = { locator: `${cacheKey}.prover` }; + const verifierLocator: KeyLocator = { locator: `${cacheKey}.verifier` }; + const provingKey = await this._keyStore.getProvingKey(proverLocator); + const verifyingKey = await this._keyStore.getVerifyingKey(verifierLocator); + if (provingKey && verifyingKey) { + console.debug(`Keys loaded from KeyStore for ${cacheKey}`); + // Cache in memory for faster subsequent access + if (this.cacheOption) { + this.cacheKeys(cacheKey, [provingKey, verifyingKey]); + } + return [provingKey, verifyingKey]; + } + } catch (e) { + console.debug(`KeyStore lookup failed for ${cacheKey}: ${e}`); + } + } + + throw new Error( + `Keys not found for '${cacheKey}' in memory cache or KeyStore`, + ); } } throw new Error( diff --git a/sdk/src/program-manager.ts b/sdk/src/program-manager.ts index cb5a579ee..b32e821e6 100644 --- a/sdk/src/program-manager.ts +++ b/sdk/src/program-manager.ts @@ -12,6 +12,7 @@ import { AleoKeyProvider, AleoKeyProviderParams, } from "./keys/provider/memory.js"; +import { KeyLocator } from "./keys/keystore/interface.js"; import { FunctionKeyPair @@ -21,6 +22,7 @@ import { Address, Authorization, ExecutionResponse, + ExecutionTransactionResponse, Execution as FunctionExecution, OfflineQuery, RecordPlaintext, @@ -81,6 +83,7 @@ interface DeployOptions { * @property {string | Program} [program] - Program source code to use for the transaction. * @property {ProgramImports} [imports] - Programs that the program being executed imports. * @property {number} [edition] - Edition of the program to execute the function in. + * @property {boolean} [cacheKeys] - If true, the proving and verifying keys synthesized during execution will be returned and auto-persisted to the configured KeyStore. */ interface ExecuteOptions { programName: string; @@ -97,7 +100,22 @@ interface ExecuteOptions { offlineQuery?: OfflineQuery; program?: string | Program; imports?: ProgramImports; - edition?: number, + edition?: number; + cacheKeys?: boolean; +} + +/** + * Result from building an execution transaction with key caching enabled. + * Contains the transaction and the proving/verifying keys that were synthesized during execution. + * + * @property {Transaction} transaction - The execution transaction. + * @property {ProvingKey} provingKey - The proving key synthesized during execution. + * @property {VerifyingKey} verifyingKey - The verifying key synthesized during execution. + */ +interface ExecutionTransactionResult { + transaction: Transaction; + provingKey: ProvingKey; + verifyingKey: VerifyingKey; } /** @@ -319,6 +337,43 @@ class ProgramManager { this.networkClient.headers[headerName] = value; } + /** + * Persist proving and verifying keys to the configured KeyStore, if available. + * Uses the locator convention: "{programName}.{functionName}.prover" and + * "{programName}.{functionName}.verifier" (dots instead of slashes to comply + * with LocalFileKeyStore's validateLocator). + * + * @param {string} programName The name of the program + * @param {string} functionName The name of the function + * @param {ProvingKey} [provingKey] The proving key to persist + * @param {VerifyingKey} [verifyingKey] The verifying key to persist + */ + private async persistKeysToStore( + programName: string, + functionName: string, + provingKey?: ProvingKey, + verifyingKey?: VerifyingKey, + ): Promise { + if (!provingKey || !verifyingKey) return; + + try { + const keyStore = await this.keyProvider.keyStore(); + if (!keyStore) return; + + const proverLocator: KeyLocator = { + locator: `${programName}.${functionName}.prover`, + }; + const verifierLocator: KeyLocator = { + locator: `${programName}.${functionName}.verifier`, + }; + + await keyStore.setKeys(proverLocator, verifierLocator, [provingKey, verifyingKey]); + console.debug(`Keys persisted to KeyStore for ${programName}/${functionName}`); + } catch (e) { + console.warn(`Failed to persist keys to KeyStore: ${e}`); + } + } + /** * Set the inclusion prover into the wasm memory. This should be done prior to any execution of a function with a * private record. @@ -786,9 +841,15 @@ class ProgramManager { * assert(transaction.id() === tx.id()); * }, 10000); */ + buildExecutionTransaction( + options: ExecuteOptions & { cacheKeys: true }, + ): Promise; + buildExecutionTransaction( + options: ExecuteOptions, + ): Promise; async buildExecutionTransaction( options: ExecuteOptions, - ): Promise { + ): Promise { // Destructure the options object to access the parameters const { functionName, @@ -941,6 +1002,53 @@ class ProgramManager { } // Build an execution transaction + if (options.cacheKeys) { + const response: ExecutionTransactionResponse = await WasmProgramManager.buildExecutionTransactionWithKeys( + executionPrivateKey, + program, + functionName, + inputs, + priorityFee, + feeRecord, + this.host, + imports, + provingKey, + verifyingKey, + feeProvingKey, + feeVerifyingKey, + offlineQuery, + edition, + true, + ); + + const transaction = response.getTransaction(); + + if (!response.hasKeys()) { + throw new Error("Key caching was requested but no keys were returned from execution"); + } + + const resultProvingKey = response.getProvingKey(); + const resultVerifyingKey = response.getVerifyingKey(); + + if (!resultProvingKey || !resultVerifyingKey) { + throw new Error("Key caching was requested but keys could not be extracted from execution response"); + } + + // Cache in memory for subsequent calls + const cacheKey = `${programName}.${functionName}`; + this.keyProvider.cacheKeys(cacheKey, [resultProvingKey, resultVerifyingKey]); + + // Auto-persist to KeyStore if configured + await this.persistKeysToStore( + programName, + functionName, + resultProvingKey, + resultVerifyingKey, + ); + + return { transaction, provingKey: resultProvingKey, verifyingKey: resultVerifyingKey }; + } + return await WasmProgramManager.buildExecutionTransaction( executionPrivateKey, program, @@ -1608,7 +1716,13 @@ class ProgramManager { * }, 10000); */ async execute(options: ExecuteOptions): Promise { - const tx = await this.buildExecutionTransaction(options); + let tx: Transaction; + if (options.cacheKeys) { + const result = await this.buildExecutionTransaction({ ...options, cacheKeys: true }); + tx = result.transaction; + } else { + tx = await this.buildExecutionTransaction(options); + } let feeAddress; @@ -1643,6 +1757,8 @@ class ProgramManager { * @param {VerifyingKey | undefined} verifyingKey Optional verifying key to use for the transaction * @param {PrivateKey | undefined} privateKey Optional private key to use for the transaction * @param {OfflineQuery | undefined} offlineQuery Optional offline query if creating transactions in an offline environment + * @param {number | undefined} edition Optional edition of the program to execute + * @param {boolean | undefined} cacheKeys If true, keys are extracted from the execution, persisted to the configured KeyStore, and cached in the key provider's memory. Note: this consumes the proving key from the response (response.getProvingKey() will return null afterward). * @returns {Promise} The execution response containing the outputs of the function and the proof if the program is proved. * * @example @@ -1673,7 +1789,8 @@ class ProgramManager { verifyingKey?: VerifyingKey, privateKey?: PrivateKey, offlineQuery?: OfflineQuery, - edition?: number + edition?: number, + cacheKeys?: boolean, ): Promise { // Get the private key from the account if it is not provided in the parameters let executionPrivateKey = privateKey; @@ -1705,13 +1822,13 @@ class ProgramManager { console.log("Running program offline"); console.log("Proving key: ", provingKey); console.log("Verifying key: ", verifyingKey); - return WasmProgramManager.executeFunctionOffline( + const response = await WasmProgramManager.executeFunctionOffline( executionPrivateKey, program, function_name, inputs, proveExecution, - false, + cacheKeys ?? false, imports, provingKey, verifyingKey, @@ -1719,6 +1836,30 @@ class ProgramManager { offlineQuery, edition ); + + // Auto-persist keys to KeyStore if cacheKeys is enabled. + // Note: getProvingKey() uses .take() semantics and consumes the key from the response. + // After this block, response.getProvingKey() will return null. The keys are persisted + // to the KeyStore and cached in the key provider's memory for subsequent access. + if (cacheKeys) { + try { + const programObj = Program.fromString(program); + const programName = programObj.id(); + const pk = response.getProvingKey() ?? undefined; + const vk = response.getVerifyingKey(); + if (pk && vk) { + // Cache in the key provider's memory so keys are accessible without re-synthesis + const cacheKey = `${programName}.${function_name}`; + this.keyProvider.cacheKeys(cacheKey, [pk, vk]); + + await this.persistKeysToStore(programName, function_name, pk, vk); + } + } catch (e) { + console.warn(`Could not auto-persist keys from run(): ${e}`); + } + } + + return response; } /** @@ -3729,4 +3870,4 @@ function validateTransferType(transferType: string): string { ); } -export { ProgramManager, AuthorizationOptions, FeeAuthorizationOptions, ExecuteOptions, ProvingRequestOptions }; +export { ProgramManager, AuthorizationOptions, FeeAuthorizationOptions, ExecuteOptions, ExecutionTransactionResult, ProvingRequestOptions }; diff --git a/sdk/src/wasm.ts b/sdk/src/wasm.ts index 2762cf701..4870f6712 100644 --- a/sdk/src/wasm.ts +++ b/sdk/src/wasm.ts @@ -12,6 +12,7 @@ export { ExecutionRequest, Execution, ExecutionResponse, + ExecutionTransactionResponse, Field, GraphKey, Group, diff --git a/sdk/tests/key-provider.test.ts b/sdk/tests/key-provider.test.ts index 191c8cc60..8b50303fa 100644 --- a/sdk/tests/key-provider.test.ts +++ b/sdk/tests/key-provider.test.ts @@ -1,6 +1,7 @@ import { expect } from "chai"; import { AleoKeyProvider, + AleoKeyProviderParams, CachedKeyPair, CREDITS_PROGRAM_KEYS, FunctionKeyPair, @@ -787,3 +788,157 @@ describe("Key verifier with LocalFileKeyStore (checksum verification on read)", } }); }); + +describe("AleoKeyProvider – KeyStore integration", () => { + it("constructor accepts a KeyStore and keyStore() returns it", async () => { + const tempDir = `${process.cwd()}/.keystore-test-${Date.now()}-${Math.floor(Math.random() * 1e6)}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + const provider = new AleoKeyProvider(keystore); + const retrieved = await provider.keyStore(); + expect(retrieved).to.equal(keystore); + } finally { + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + + it("constructor without KeyStore returns undefined from keyStore()", async () => { + const provider = new AleoKeyProvider(); + expect(await provider.keyStore()).to.equal(undefined); + }); + + it("setKeyStore() updates the KeyStore and keyStore() reflects the change", async () => { + const tempDir = `${process.cwd()}/.keystore-test-${Date.now()}-${Math.floor(Math.random() * 1e6)}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + const provider = new AleoKeyProvider(); + expect(await provider.keyStore()).to.equal(undefined); + + provider.setKeyStore(keystore); + expect(await provider.keyStore()).to.equal(keystore); + + provider.setKeyStore(undefined); + expect(await provider.keyStore()).to.equal(undefined); + } finally { + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + + it("functionKeys() loads keys from KeyStore when not in memory cache", async function () { + this.timeout(20000); + + const tempDir = `${process.cwd()}/.keystore-test-${Date.now()}-${Math.floor(Math.random() * 1e6)}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + // Fetch real keys and store them in the KeyStore + const fetcher = new AleoKeyProvider(); + const [prov, ver] = await fetcher.fetchCreditsKeys(CREDITS_PROGRAM_KEYS.fee_public); + const cacheKey = "my_program.aleo.my_function"; + + await keystore.setKeys( + locator(`${cacheKey}.prover`), + locator(`${cacheKey}.verifier`), + [prov, ver], + ); + + // Create a fresh provider with the KeyStore but empty memory cache + const provider = new AleoKeyProvider(keystore); + + // functionKeys should find keys in KeyStore (not in memory) + expect(provider.containsKeys(cacheKey)).to.equal(false); + const [loadedProv, loadedVer] = await provider.functionKeys( + new AleoKeyProviderParams({ cacheKey }), + ); + + expect(loadedProv).to.be.instanceOf(ProvingKey); + expect(loadedVer).to.be.instanceOf(VerifyingKey); + expect(loadedProv.checksum()).to.equal(prov.checksum()); + expect(loadedVer.checksum()).to.equal(ver.checksum()); + } finally { + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + + it("functionKeys() caches KeyStore results in memory when cacheOption is enabled", async function () { + this.timeout(20000); + + const tempDir = `${process.cwd()}/.keystore-test-${Date.now()}-${Math.floor(Math.random() * 1e6)}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + const fetcher = new AleoKeyProvider(); + const [prov, ver] = await fetcher.fetchCreditsKeys(CREDITS_PROGRAM_KEYS.fee_public); + const cacheKey = "cached_program.aleo.cached_function"; + + await keystore.setKeys( + locator(`${cacheKey}.prover`), + locator(`${cacheKey}.verifier`), + [prov, ver], + ); + + const provider = new AleoKeyProvider(keystore); + provider.useCache(true); + + // Memory cache should be empty before the call + expect(provider.containsKeys(cacheKey)).to.equal(false); + + await provider.functionKeys(new AleoKeyProviderParams({ cacheKey })); + + // After loading from KeyStore, keys should be in memory cache + expect(provider.containsKeys(cacheKey)).to.equal(true); + + // Subsequent call should work from memory cache even without KeyStore + provider.setKeyStore(undefined); + const [cachedProv, cachedVer] = provider.getKeys(cacheKey); + expect(cachedProv.checksum()).to.equal(prov.checksum()); + expect(cachedVer.checksum()).to.equal(ver.checksum()); + } finally { + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + + it("functionKeys() prefers memory cache over KeyStore", async function () { + this.timeout(60000); + + const tempDir = `${process.cwd()}/.keystore-test-${Date.now()}-${Math.floor(Math.random() * 1e6)}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + const fetcher = new AleoKeyProvider(); + const [provFee, verFee] = await fetcher.fetchCreditsKeys(CREDITS_PROGRAM_KEYS.fee_public); + const [provJoin, verJoin] = await fetcher.fetchCreditsKeys(CREDITS_PROGRAM_KEYS.join); + const cacheKey = "priority_test.aleo.my_func"; + + // Put "join" keys in KeyStore + await keystore.setKeys( + locator(`${cacheKey}.prover`), + locator(`${cacheKey}.verifier`), + [provJoin, verJoin], + ); + + // Put "fee" keys in memory cache + const provider = new AleoKeyProvider(keystore); + provider.useCache(true); + provider.cacheKeys(cacheKey, [provFee, verFee]); + + // functionKeys should return the memory-cached "fee" keys, not the KeyStore "join" keys + const [loadedProv, loadedVer] = await provider.functionKeys( + new AleoKeyProviderParams({ cacheKey }), + ); + expect(loadedProv.checksum()).to.equal(provFee.checksum()); + expect(loadedVer.checksum()).to.equal(verFee.checksum()); + } finally { + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + + it("functionKeys() throws when cacheKey not in memory or KeyStore", async () => { + const provider = new AleoKeyProvider(); + try { + await provider.functionKeys( + new AleoKeyProviderParams({ cacheKey: "nonexistent.aleo.func" }), + ); + expect.fail("should throw"); + } catch (e) { + expect(e).to.be.instanceOf(Error); + } + }); +}); diff --git a/sdk/tests/program-manager.test.ts b/sdk/tests/program-manager.test.ts index 781b561b2..ea41b6388 100644 --- a/sdk/tests/program-manager.test.ts +++ b/sdk/tests/program-manager.test.ts @@ -10,6 +10,7 @@ import { PrivateKey, Program, ProgramManager, + ProvingKey, ProvingRequest, RecordPlaintext, Transaction, @@ -17,6 +18,7 @@ import { VerifyingKey, ViewKey } from "@provablehq/sdk/%%NETWORK%%.js"; +import { LocalFileKeyStore } from "../src/node.js"; import { beaconAddressString, helloProgram, @@ -430,4 +432,80 @@ describe('Program Manager', async () => { } }); }); -}); \ No newline at end of file + + describe('Key caching integration', () => { + it('AleoKeyProvider should accept a KeyStore via constructor and expose it', async () => { + const tempDir = `${process.cwd()}/.pm-keystore-test-${Date.now()}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + const keyProviderWithStore = new AleoKeyProvider(keystore); + const pm = new ProgramManager("https://api.provable.com/v2", keyProviderWithStore); + + const store = await pm.keyProvider.keyStore(); + expect(store).to.equal(keystore); + } finally { + const $fs = await import("node:fs/promises"); + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + + // Requires full WASM execution pipeline (key synthesis + proving) + it.skip('buildExecutionTransaction with cacheKeys returns keys and transaction', async function() { + this.timeout(120000); + + const tempDir = `${process.cwd()}/.pm-keystore-test-${Date.now()}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + const keyProviderWithStore = new AleoKeyProvider(keystore); + keyProviderWithStore.useCache(true); + const pm = new ProgramManager("https://api.provable.com/v2", keyProviderWithStore); + pm.setAccount(new Account()); + + const result = await pm.buildExecutionTransaction({ + programName: "credits.aleo", + functionName: "split", + priorityFee: 0.0, + privateFee: false, + inputs: ["1u64"], + cacheKeys: true, + }); + + // Result should contain transaction and keys + expect(result.transaction).to.be.instanceOf(Transaction); + expect(result.provingKey).to.be.instanceOf(ProvingKey); + expect(result.verifyingKey).to.be.instanceOf(VerifyingKey); + } finally { + const $fs = await import("node:fs/promises"); + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + + // Requires full WASM execution pipeline (key synthesis) + it('run() with cacheKeys should auto-persist keys to KeyStore', async function() { + this.timeout(120000); + + const tempDir = `${process.cwd()}/.pm-keystore-test-${Date.now()}`; + const keystore = new LocalFileKeyStore(tempDir); + try { + const keyProviderWithStore = new AleoKeyProvider(keystore); + const pm = new ProgramManager("https://api.provable.com/v2", keyProviderWithStore); + pm.setAccount(new Account()); + + const helloWorld = "program helloworld.aleo;\n\nfunction hello:\n input r0 as u32.public;\n input r1 as u32.private;\n add r0 r1 into r2;\n output r2 as u32.private;\n"; + + // First run: keys get synthesized and auto-persisted + await pm.run(helloWorld, "hello", ["5u32", "5u32"], false, + undefined, undefined, undefined, undefined, undefined, undefined, undefined, true); + + // Keys should now be in the KeyStore + const provKey = await keystore.getProvingKey({ locator: "helloworld.aleo.hello.prover" }); + const verKey = await keystore.getVerifyingKey({ locator: "helloworld.aleo.hello.verifier" }); + expect(provKey).to.not.equal(null); + expect(verKey).to.not.equal(null); + } finally { + const $fs = await import("node:fs/promises"); + await $fs.rm(tempDir, { recursive: true, force: true }).catch(() => {}); + } + }); + }); +}); diff --git a/wasm/src/programs/manager/execute.rs b/wasm/src/programs/manager/execute.rs index af3fcc7e9..7323aff8d 100644 --- a/wasm/src/programs/manager/execute.rs +++ b/wasm/src/programs/manager/execute.rs @@ -20,6 +20,7 @@ use crate::{ Address, Authorization, ExecutionResponse, + ExecutionTransactionResponse, OfflineQuery, PrivateKey, RecordPlaintext, @@ -295,6 +296,159 @@ impl ProgramManager { Ok(Transaction::from(transaction)) } + /// Execute Aleo function and create an Aleo execution transaction, optionally caching the + /// proving and verifying keys in the response. + /// + /// This method is identical to `buildExecutionTransaction` but returns an + /// `ExecutionTransactionResponse` which contains the transaction and optionally the + /// proving/verifying keys that were synthesized during execution. + /// + /// @param private_key The private key of the sender + /// @param program The source code of the program being executed + /// @param function The name of the function to execute + /// @param inputs A javascript array of inputs to the function + /// @param priority_fee_credits The optional priority fee to be paid for the transaction + /// @param fee_record The record to spend the fee from + /// @param url The url of the Aleo network node to send the transaction to + /// @param imports (optional) Provide a list of imports to use for the function execution + /// @param proving_key (optional) Provide a proving key to use for the function execution + /// @param verifying_key (optional) Provide a verifying key to use for the function execution + /// @param fee_proving_key (optional) Provide a proving key to use for the fee execution + /// @param fee_verifying_key (optional) Provide a verifying key to use for the fee execution + /// @param offline_query An offline query object to use if building a transaction without an internet connection. + /// @param edition The edition of the program to execute. + /// @param cache If true, the proving and verifying keys will be included in the response. + /// @returns {ExecutionTransactionResponse} + #[wasm_bindgen(js_name = buildExecutionTransactionWithKeys)] + #[allow(clippy::too_many_arguments)] + pub async fn execute_with_keys( + private_key: &PrivateKey, + program: &str, + function: &str, + inputs: Array, + priority_fee_credits: f64, + fee_record: Option, + url: Option, + imports: Option, + proving_key: Option, + verifying_key: Option, + fee_proving_key: Option, + fee_verifying_key: Option, + offline_query: Option, + edition: Option, + cache: bool, + ) -> Result { + let mut process_native = ProcessNative::load_web().map_err(|err| err.to_string())?; + let process = &mut process_native; + let node_url = url.as_deref().unwrap_or(DEFAULT_URL); + + log("Check program imports are valid and add them to the process"); + let program_native = ProgramNative::from_str(program).map_err(|e| e.to_string())?; + let program_id = program_native.id().to_string(); + ProgramManager::resolve_imports(process, &program_native, imports)?; + let rng = &mut StdRng::from_entropy(); + + log(&format!("Executing function: {program_id}/{function} on-chain")); + let edition = edition.unwrap_or(1); + let (_, mut trace) = execute_program!( + process, + process_inputs!(inputs), + program, + function, + private_key, + proving_key, + verifying_key, + rng, + edition + ); + + // Extract keys from process if caching is requested (must happen before process is dropped) + let cached_keys = if cache { + let function_name = IdentifierNative::from_str(function).map_err(|err| err.to_string())?; + let pk = process + .get_proving_key(program_native.id(), &function_name) + .map_err(|_| format!("Could not find proving key for {}/{}", program_native.id(), function))?; + let vk = process + .get_verifying_key(program_native.id(), &function_name) + .map_err(|_| format!("Could not find verifying key for {}/{}", program_native.id(), function))?; + (Some(pk), Some(vk)) + } else { + (None, None) + }; + + log("Preparing inclusion proofs for execution"); + let latest_height = if let Some(offline_query) = offline_query.as_ref() { + trace.prepare_async(offline_query).await.map_err(|err| err.to_string())?; + offline_query.current_block_height().map_err(|e| e.to_string())? + } else { + let function_name = IdentifierNative::from_str(function).map_err(|err| err.to_string())?; + let view_key = + ViewKeyNative::try_from(PrivateKeyNative::from(private_key)).map_err(|err| err.to_string())?; + let query = + SnapshotQuery::try_from_inputs(node_url, &program_native, &function_name, &view_key, &inputs.to_vec()) + .await + .map_err(|err| err.to_string())?; + trace.prepare_async(&query).await.map_err(|err| err.to_string())?; + query.current_block_height().map_err(|e| e.to_string())? + }; + + log("Proving execution"); + let locator = program_native.id().to_string().add("/").add(function); + let execution = trace + .prove_execution::(&locator, VarunaVersion::V2, &mut StdRng::from_entropy()) + .map_err(|e| e.to_string())?; + + // If the function is anything other than credits.aleo/split or credits.aleo/upgrade, execute a fee. + let fee = match (program_id.as_str(), function) { + ("credits.aleo", "split") + | ("credits.aleo", "upgrade") + | ("credits.aleo", "fee_private") + | ("credits.aleo", "fee_public") => None, + _ => { + log("Calculating the minimum execution fee"); + let minimum_execution_cost = calculate_minimum_fee!(offline_query, node_url, process, &execution); + + // Check to see if the fee record has enough microcredits to pay for the deployment. + let priority_fee_microcredits = (priority_fee_credits * 1_000_000.0) as u64; + Self::validate_fee_record(&fee_record, minimum_execution_cost, priority_fee_microcredits)?; + + // Calculate the execution id. + let execution_id = execution.to_execution_id().map_err(|e| e.to_string())?; + + log("Executing fee"); + let fee = execute_fee!( + process, + private_key, + fee_record, + priority_fee_microcredits, + node_url, + fee_proving_key, + fee_verifying_key, + execution_id, + rng, + offline_query, + minimum_execution_cost + ); + Some(fee) + } + }; + + // Verify the execution + let consensus_version = + ::CONSENSUS_VERSION(latest_height).map_err(|err| err.to_string())?; + let inclusion_upgrade_height = + ::INCLUSION_UPGRADE_HEIGHT().map_err(|err| err.to_string())?; + let inclusion_version = + if latest_height >= inclusion_upgrade_height { InclusionVersion::V1 } else { InclusionVersion::V0 }; + process + .verify_execution(consensus_version, VarunaVersion::V2, inclusion_version, &execution) + .map_err(|err| err.to_string())?; + + log("Creating execution transaction"); + let transaction = TransactionNative::from_execution(execution, fee).map_err(|err| err.to_string())?; + Ok(ExecutionTransactionResponse::new(transaction, cached_keys.0, cached_keys.1)) + } + /// Execute an authorization. /// /// @param authorization The authorization to execute. diff --git a/wasm/src/programs/response.rs b/wasm/src/programs/response.rs index 46688121c..4128b97b1 100644 --- a/wasm/src/programs/response.rs +++ b/wasm/src/programs/response.rs @@ -25,7 +25,7 @@ use crate::types::native::{ VerifyingKeyNative, }; -use crate::{Execution, KeyPair, Program, ProvingKey, VerifyingKey}; +use crate::{Execution, KeyPair, Program, ProvingKey, Transaction, VerifyingKey, types::native::TransactionNative}; use std::{ops::Deref, str::FromStr}; use wasm_bindgen::{JsValue, prelude::wasm_bindgen}; @@ -148,3 +148,71 @@ impl Deref for ExecutionResponse { &self.response } } + +/// Response from building an execution transaction, containing the transaction +/// and optionally the proving/verifying keys that were synthesized or used. +#[wasm_bindgen] +pub struct ExecutionTransactionResponse { + transaction: TransactionNative, + proving_key: Option, + verifying_key: Option, +} + +#[wasm_bindgen] +impl ExecutionTransactionResponse { + pub(crate) fn new( + transaction: TransactionNative, + proving_key: Option, + verifying_key: Option, + ) -> Self { + Self { transaction, proving_key, verifying_key } + } + + /// Get the transaction from the response + /// + /// @returns {Transaction} The execution transaction + #[wasm_bindgen(js_name = "getTransaction")] + pub fn get_transaction(&self) -> Transaction { + Transaction::from(self.transaction.clone()) + } + + /// Returns true if keys are available in this response + /// + /// @returns {boolean} Whether keys were cached in this response + #[wasm_bindgen(js_name = "hasKeys")] + pub fn has_keys(&self) -> bool { + self.proving_key.is_some() && self.verifying_key.is_some() + } + + /// Returns the key pair if cached, consuming both keys. + /// Returns an error if keys were not cached. + /// Note: calling getProvingKey() or getVerifyingKey() before this method will cause it to fail. + /// + /// @returns {KeyPair} The proving and verifying key pair + #[wasm_bindgen(js_name = "getKeys")] + pub fn get_keys(&mut self) -> Result { + if let (Some(pk), Some(vk)) = (self.proving_key.take(), self.verifying_key.take()) { + Ok(KeyPair::new(ProvingKey::from(pk), VerifyingKey::from(vk))) + } else { + Err("No keys cached in this response".to_string()) + } + } + + /// Returns the proving key if available (consuming it). + /// Note the proving key is removed from the response after the first call. + /// Subsequent calls will return null. + /// + /// @returns {ProvingKey | undefined} The proving key + #[wasm_bindgen(js_name = "getProvingKey")] + pub fn get_proving_key(&mut self) -> Option { + self.proving_key.take().map(ProvingKey::from) + } + + /// Returns the verifying key if available (consuming it). + /// + /// @returns {VerifyingKey | undefined} The verifying key + #[wasm_bindgen(js_name = "getVerifyingKey")] + pub fn get_verifying_key(&mut self) -> Option { + self.verifying_key.take().map(VerifyingKey::from) + } +}