@aztec/pxe 0.0.1-commit.3100065 → 0.0.1-commit.330febf

package/src/private_kernel/private_kernel_execution_prover.ts

@@ -1,6 +1,7 @@
+import { CircuitKind } from '@aztec/bb.js';
 import { MAX_APPS_PER_KERNEL } from '@aztec/constants';
 import { uniqueBy } from '@aztec/foundation/collection';
-import { vkAsFieldsMegaHonk } from '@aztec/foundation/crypto/keys';
+import { vkAsFields } from '@aztec/foundation/crypto/keys';
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { pushTestData } from '@aztec/foundation/testing';
@@ -20,9 +21,13 @@ import {
   type PrivateKernelExecutionProofOutput,
   PrivateKernelInit2CircuitPrivateInputs,
   PrivateKernelInit3CircuitPrivateInputs,
+  PrivateKernelInit4CircuitPrivateInputs,
+  PrivateKernelInit5CircuitPrivateInputs,
   PrivateKernelInitCircuitPrivateInputs,
   PrivateKernelInner2CircuitPrivateInputs,
   PrivateKernelInner3CircuitPrivateInputs,
+  PrivateKernelInner4CircuitPrivateInputs,
+  PrivateKernelInner5CircuitPrivateInputs,
   PrivateKernelInnerCircuitPrivateInputs,
   PrivateKernelResetTailCircuitPrivateInputs,
   type PrivateKernelSimulateOutput,
@@ -145,6 +150,7 @@ export class PrivateKernelExecutionProver {
             bytecode: output.bytecode,
             witness: output.outputWitness,
             vk: output.verificationKey.keyAsBytes,
+            kind: CircuitKind.Kernel,
             timings: {
               witgen: witgenTimer.ms(),
             },
@@ -230,6 +236,7 @@ export class PrivateKernelExecutionProver {
       bytecode: tailOutput.bytecode,
       witness: tailOutput.outputWitness,
       vk: tailOutput.verificationKey.keyAsBytes,
+      kind: CircuitKind.Kernel,
       timings: {
         witgen: witgenTimer.ms(),
       },
@@ -261,6 +268,7 @@ export class PrivateKernelExecutionProver {
         bytecode: hidingOutput.bytecode,
         witness: hidingOutput.outputWitness,
         vk: hidingOutput.verificationKey.keyAsBytes,
+        kind: CircuitKind.HidingKernel,
         timings: {
           witgen: witgenTimer.ms(),
         },
@@ -270,7 +278,11 @@ export class PrivateKernelExecutionProver {
     if (profileMode == 'gates' || profileMode == 'full') {
       for (const entry of executionSteps) {
         const gateCountTimer = new Timer();
-        const gateCount = await this.proofCreator.computeGateCountForCircuit(entry.bytecode, entry.functionName);
+        const gateCount = await this.proofCreator.computeGateCountForCircuit(
+          entry.bytecode,
+          entry.functionName,
+          entry.kind,
+        );
         entry.gateCount = gateCount;
         entry.timings.gateCount = gateCountTimer.ms();
       }
@@ -366,6 +378,7 @@ export class PrivateKernelExecutionProver {
       bytecode: next.acir,
       witness: next.partialWitness,
       vk: next.vk,
+      kind: CircuitKind.App,
       timings: {
         witgen: next.profileResult?.timings.witgen ?? 0,
         oracles: next.profileResult?.timings.oracles,
@@ -399,8 +412,8 @@ export class PrivateKernelExecutionProver {
   ) {
     const { contractAddress, functionSelector } = publicInputs.callContext;
 
-    const vkAsFields = await vkAsFieldsMegaHonk(vkAsBuffer);
-    const vk = await VerificationKeyAsFields.fromKey(vkAsFields);
+    const vkFields = await vkAsFields(vkAsBuffer, CircuitKind.App);
+    const vk = await VerificationKeyAsFields.fromKey(vkFields);
 
     const { currentContractClassId, publicKeys, saltedInitializationHash } =
       await this.oracle.getContractAddressPreimage(contractAddress);
@@ -514,6 +527,47 @@ export class PrivateKernelExecutionProver {
           functionName = 'private_kernel_init_3';
           break;
         }
+        case 4: {
+          const proofInput = new PrivateKernelInit4CircuitPrivateInputs(
+            txRequest,
+            vkTreeRoot,
+            ProtocolContractsList,
+            apps[0],
+            apps[1],
+            apps[2],
+            apps[3],
+            isPrivateOnlyTx,
+            minRevertibleSideEffectCounter,
+          );
+          this.log.debug(`Calling private kernel init_4 with isPrivateOnly ${isPrivateOnlyTx}`);
+          pushTestData('private-kernel-inputs-init-4', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInit4Output(proofInput)
+            : await this.proofCreator.simulateInit4(proofInput);
+          functionName = 'private_kernel_init_4';
+          break;
+        }
+        case 5: {
+          const proofInput = new PrivateKernelInit5CircuitPrivateInputs(
+            txRequest,
+            vkTreeRoot,
+            ProtocolContractsList,
+            apps[0],
+            apps[1],
+            apps[2],
+            apps[3],
+            apps[4],
+            isPrivateOnlyTx,
+            minRevertibleSideEffectCounter,
+          );
+          this.log.debug(`Calling private kernel init_5 with isPrivateOnly ${isPrivateOnlyTx}`);
+          pushTestData('private-kernel-inputs-init-5', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInit5Output(proofInput)
+            : await this.proofCreator.simulateInit5(proofInput);
+          functionName = 'private_kernel_init_5';
+          break;
+        }
         default:
           throw new Error(`Unsupported init kernel batch size: ${apps.length}`);
       }
@@ -548,6 +602,37 @@ export class PrivateKernelExecutionProver {
           functionName = 'private_kernel_inner_3';
           break;
         }
+        case 4: {
+          const proofInput = new PrivateKernelInner4CircuitPrivateInputs(
+            previousKernelData,
+            apps[0],
+            apps[1],
+            apps[2],
+            apps[3],
+          );
+          pushTestData('private-kernel-inputs-inner-4', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInner4Output(proofInput)
+            : await this.proofCreator.simulateInner4(proofInput);
+          functionName = 'private_kernel_inner_4';
+          break;
+        }
+        case 5: {
+          const proofInput = new PrivateKernelInner5CircuitPrivateInputs(
+            previousKernelData,
+            apps[0],
+            apps[1],
+            apps[2],
+            apps[3],
+            apps[4],
+          );
+          pushTestData('private-kernel-inputs-inner-5', proofInput);
+          output = generateWitnesses
+            ? await this.proofCreator.generateInner5Output(proofInput)
+            : await this.proofCreator.simulateInner5(proofInput);
+          functionName = 'private_kernel_inner_5';
+          break;
+        }
         default:
           throw new Error(`Unsupported inner kernel batch size: ${apps.length}`);
       }
@@ -558,6 +643,7 @@ export class PrivateKernelExecutionProver {
       bytecode: output.bytecode,
       witness: output.outputWitness,
       vk: output.verificationKey.keyAsBytes,
+      kind: CircuitKind.Kernel,
       timings: {
         witgen: witgenTimer.ms(),
       },

package/src/pxe.ts

@@ -37,7 +37,7 @@ import type {
   PrivateKernelTailCircuitPublicInputs,
 } from '@aztec/stdlib/kernel';
 import { computeAddressSecret } from '@aztec/stdlib/keys';
-import { deriveEcdhSharedSecretPoint } from '@aztec/stdlib/logs';
+import { deriveEcdhSharedSecretPoint, protectFromForgery } from '@aztec/stdlib/logs';
 import {
   BlockHeader,
   type ContractOverrides,
@@ -62,7 +62,6 @@ import { BlockSynchronizer } from './block_synchronizer/index.js';
 import type { PXEConfig } from './config/index.js';
 import { ContractClassService } from './contract/contract_class_service.js';
 import { ContractSyncService } from './contract/contract_sync_service.js';
-import { BenchmarkedNodeFactory } from './contract_function_simulator/benchmarked_node.js';
 import {
   ContractFunctionSimulator,
   generateSimulatedProvingResult,
@@ -74,6 +73,7 @@ import { PrivateEventFilterValidator } from './events/private_event_filter_valid
 import type { ExecutionHooks } from './hooks/index.js';
 import { JobCoordinator } from './job_coordinator/job_coordinator.js';
 import { TxResolverService } from './messages/tx_resolver_service.js';
+import { type CachingAztecNode, withCache } from './node/caching_aztec_node.js';
 import {
   PrivateKernelExecutionProver,
   type PrivateKernelExecutionProverConfig,
@@ -230,7 +230,7 @@ export type RegisteredTaggingSecretSource =
  */
 export class PXE {
   private constructor(
-    private node: AztecNode,
+    private node: CachingAztecNode,
     private nodeDebug: AztecNodeDebug | undefined,
     private db: AztecAsyncKVStore,
     private blockStateSynchronizer: BlockSynchronizer,
@@ -315,18 +315,22 @@ export class PXE {
       l2TipsStore,
       factStore,
     } = openPxeStores(store, initialBlockHash);
-    const contractClassService = new ContractClassService(node, contractStore);
+    // Every PXE consumer reads through this one wrapper, so a read cached by one is served to the rest. Only
+    // immutable, hash-pinned reads are cached (the rule lives on `withCache`), which makes it safe regardless of
+    // the consumer's anchor block; the block synchronizer wipes it on anchor updates to bound memory.
+    const readCachedNode = withCache(node);
+    const contractClassService = new ContractClassService(readCachedNode, contractStore);
     const contractSyncService = new ContractSyncService(
-      node,
+      readCachedNode,
       contractStore,
       contractClassService,
       noteStore,
       createLogger('pxe:contract_sync', bindings),
     );
-    const txResolver = new TxResolverService(node);
+    const txResolver = new TxResolverService(readCachedNode);
 
     const synchronizer = new BlockSynchronizer(
-      node,
+      readCachedNode,
       store,
       anchorBlockStore,
       noteStore,
@@ -334,7 +338,6 @@ export class PXE {
       factStore,
       l2TipsStore,
       contractSyncService,
-      contractClassService,
       config,
       bindings,
     );
@@ -355,7 +358,7 @@ export class PXE {
     const jobQueue = new SerialQueue();
 
     const pxe = new PXE(
-      node,
+      readCachedNode,
       nodeDebug,
       store,
       synchronizer,
@@ -410,7 +413,7 @@ export class PXE {
       noteStore: this.noteStore,
       keyStore: this.keyStore,
       addressStore: this.addressStore,
-      aztecNode: BenchmarkedNodeFactory.create(this.node),
+      aztecNode: this.node,
       l2TipsStore: this.l2TipsStore,
       senderTaggingStore: this.senderTaggingStore,
       recipientTaggingStore: this.recipientTaggingStore,
@@ -890,9 +893,11 @@ export class PXE {
   }
 
   /**
-   * Derives the shared secret of a handshake from its ephemeral public key: `S = addressSecret(recipient) * ephPk`,
-   * the same derivation the recipient-side scan uses for handshakes discovered onchain. The ephemeral key is
-   * reconstructed from its x-coordinate with a positive y, matching how the protocol generates it.
+   * Derives the shared secret of a handshake from its ephemeral public key: the raw ECDH secret
+   * `S = addressSecret(recipient) * ephPk` with the forgery protection applied,
+   * matching the protected secret the registry stores and the derivation the recipient-side scan uses for
+   * handshakes discovered onchain. The ephemeral key is reconstructed from its x-coordinate with a positive y,
+   * matching how the protocol generates it.
    *
    * @throws If `ephPk` is not the x-coordinate of a Grumpkin curve point, or if `recipient` is not an account of
    * this PXE, since the derivation needs its keys.
@@ -914,7 +919,8 @@ export class PXE {
 
     const ivskM = await this.keyStore.getMasterIncomingViewingSecretKey(recipient);
     const addressSecret = await computeAddressSecret(await completeAddress.getPreaddress(), ivskM);
-    return deriveEcdhSharedSecretPoint(addressSecret, ephPkPoint);
+    const rawSecret = await deriveEcdhSharedSecretPoint(addressSecret, ephPkPoint);
+    return protectFromForgery(rawSecret, ephPkPoint, await recipient.toAddressPoint());
   }
 
   /** Registers a resolved shared secret scoped to a recipient. Returns whether it was newly added. */
@@ -1013,6 +1019,7 @@ export class PXE {
     // computationally demanding that it'd be rare for someone to try to do it concurrently regardless.
     return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
+      const recording = this.node.startRecording();
       try {
         const syncTimer = new Timer();
         await this.#maybeSync();
@@ -1060,7 +1067,7 @@ export class PXE {
 
         const txProvingResult = new TxProvingResult(privateExecutionResult, publicInputs, chonkProof!, {
           timings,
-          nodeRPCCalls: contractFunctionSimulator?.getStats().nodeRPCCalls,
+          nodeRPCCalls: recording.stop(),
         });
 
         // We keep track of which tagging indices we've used in this tx so that we don't repeat them in future txs
@@ -1081,6 +1088,9 @@ export class PXE {
         return txProvingResult;
       } catch (err: any) {
         throw this.#contextualizeError(err, inspect(txRequest), inspect(privateExecutionResult));
+      } finally {
+        // Idempotent cleanup for the error and early-exit paths. The success path already stopped the recording.
+        recording.stop();
       }
     });
   }
@@ -1098,6 +1108,7 @@ export class PXE {
     // We disable concurrent profiles for consistency with simulateTx.
     return this.#putInJobQueue(async jobId => {
       const totalTimer = new Timer();
+      const recording = this.node.startRecording();
       try {
         const txInfo = {
           origin: txRequest.origin,
@@ -1163,10 +1174,12 @@ export class PXE {
             total - ((syncTime ?? 0) + (proving ?? 0) + perFunction.reduce((acc, { time }) => acc + time, 0)),
         };
 
-        const simulatorStats = contractFunctionSimulator.getStats();
-        return new TxProfileResult(executionSteps, { timings, nodeRPCCalls: simulatorStats.nodeRPCCalls });
+        return new TxProfileResult(executionSteps, { timings, nodeRPCCalls: recording.stop() });
       } catch (err: any) {
         throw this.#contextualizeError(err, inspect(txRequest), `profileMode=${profileMode}`);
+      } finally {
+        // Idempotent cleanup for the error and early-exit paths. The success path already stopped the recording.
+        recording.stop();
       }
     });
   }
@@ -1205,6 +1218,7 @@ export class PXE {
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another simulation is currently modifying).
     return this.#putInJobQueue(async jobId => {
+      const recording = this.node.startRecording();
       try {
         const totalTimer = new Timer();
         const txInfo = {
@@ -1325,10 +1339,9 @@ export class PXE {
             : {}),
         });
 
-        const simulatorStats = contractFunctionSimulator.getStats();
         return TxSimulationResult.fromPrivateSimulationResultAndPublicOutput(privateSimulationResult, publicOutput, {
           timings,
-          nodeRPCCalls: simulatorStats.nodeRPCCalls,
+          nodeRPCCalls: recording.stop(),
         });
       } catch (err: any) {
         throw this.#contextualizeError(
@@ -1338,6 +1351,9 @@ export class PXE {
           `skipTxValidation=${skipTxValidation}`,
           `scopes=${scopes.map(s => s.toString()).join(', ')}`,
         );
+      } finally {
+        // Idempotent cleanup for the error and early-exit paths. The success path already stopped the recording.
+        recording.stop();
       }
     });
   }
@@ -1354,6 +1370,7 @@ export class PXE {
     // to the capsules), and we need to prevent concurrent runs from interfering with one another (e.g. attempting to
     // delete the same read value, or reading values that another execution is currently modifying).
     return this.#putInJobQueue(async jobId => {
+      const recording = this.node.startRecording();
       try {
         const totalTimer = new Timer();
         const syncTimer = new Timer();
@@ -1393,12 +1410,11 @@ export class PXE {
           unaccounted: totalTime - (syncTime + perFunction.reduce((acc, { time }) => acc + time, 0)),
         };
 
-        const simulationStats = contractFunctionSimulator.getStats();
         return {
           result: executionResult,
           offchainEffects,
           anchorBlockTimestamp: anchorBlockHeader.globalVariables.timestamp,
-          stats: { timings, nodeRPCCalls: simulationStats.nodeRPCCalls },
+          stats: { timings, nodeRPCCalls: recording.stop() },
         };
       } catch (err: any) {
         const { to, name, args } = call;
@@ -1408,6 +1424,9 @@ export class PXE {
           `executeUtility ${to}:${name}(${stringifiedArgs})`,
           `scopes=${scopes.map(s => s.toString()).join(', ')}`,
         );
+      } finally {
+        // Idempotent cleanup for the error and early-exit paths. The success path already stopped the recording.
+        recording.stop();
       }
     });
   }

package/src/storage/contract_store/contract_store.ts

@@ -17,6 +17,7 @@ import {
   findFunctionAbiBySelector,
   findFunctionArtifactBySelector,
   getFunctionDebugMetadata,
+  getFunctionReturnType,
 } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import {
@@ -371,7 +372,7 @@ export class ContractStore {
       hideMsgSender: false,
       isStatic: functionDao.isStatic,
       args: encodeArguments(functionDao, args),
-      returnTypes: functionDao.returnTypes,
+      returnType: getFunctionReturnType(functionDao),
     });
   }
 }

package/src/storage/fact_store/fact_store.ts

@@ -82,7 +82,7 @@ export class FactStore implements StagedStore {
   }
 
   /**
-   * Records a fact in a collection, visible under the given scope.
+   * Records a fact in a collection.
    *
    * The collection is created implicitly on the first fact recorded for its key: recording into an existing collection
    * just adds to it.

package/src/storage/fact_store/index.ts

@@ -5,6 +5,7 @@ export type { Fact } from './stored_fact.js';
 export {
   OriginBlockState,
   classifyOriginBlockState,
+  originBlockStateFromNumber,
   anchoredTipBlockNumbers,
   toFactWithOriginState,
   type TipBlockNumbers,

package/src/storage/fact_store/origin_state.ts

@@ -12,7 +12,7 @@ import type { Fact } from './stored_fact.js';
  * - `Finalized`: L1-finalized.
  *
  * The numeric discriminants must stay in sync with the Noir `OriginBlockState` in
- * `noir-projects/aztec-nr/aztec/src/facts/origin_state.nr`: PXE serializes this value into the `Fact` oracle response
+ * `noir-projects/labs/aztec-nr/aztec/src/facts/origin_state.nr`: PXE serializes this value into the `Fact` oracle response
  * and Noir deserializes it via `from_u8`, which rejects any value outside this set.
  */
 export enum OriginBlockState {
@@ -21,6 +21,15 @@ export enum OriginBlockState {
   Finalized = 3,
 }
 
+/** Parses a numeric origin-block-state discriminant, rejecting unknown values. */
+export function originBlockStateFromNumber(value: number): OriginBlockState {
+  if (!(value in OriginBlockState)) {
+    const validNames = Object.keys(OriginBlockState).filter(k => isNaN(Number(k)));
+    throw new Error(`Invalid OriginBlockState value ${value}, expected one of ${validNames.join(', ')}`);
+  }
+  return value as OriginBlockState;
+}
+
 /** The two chain-tip block numbers needed to classify an origin block (`finalized <= proven` always holds). */
 export type TipBlockNumbers = { provenBlockNumber: number; finalizedBlockNumber: number };
 

package/src/storage/store_identity.ts

@@ -0,0 +1,72 @@
+import type { EthAddress } from '@aztec/foundation/eth-address';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+import { DatabaseVersion } from '@aztec/stdlib/database-version/version';
+
+/** The triple that determine which physical store a logical store name maps to. */
+export type StoreIdentity = {
+  /** Chain ID of the L1 the rollup is deployed to. */
+  l1ChainId: number;
+  /** Address of the rollup contract the store's data pertains to. */
+  rollupAddress: EthAddress;
+  /** Schema version of the data held in the store. */
+  schemaVersion: number;
+};
+
+/**
+ * Composes the store-name discriminator for a store identity. Two identities map to the same physical store iff
+ * their slugs are equal, so the format must stay stable: `<l1ChainId>-<rollupAddress>-v<schemaVersion>` — changing
+ * it orphans every existing store.
+ */
+export function storeIdentitySlug({ l1ChainId, rollupAddress, schemaVersion }: StoreIdentity): string {
+  return `${l1ChainId}-${rollupAddress.toString()}-v${schemaVersion}`;
+}
+
+/** Composes the physical store name for a logical store name and identity. */
+export function effectiveStoreName(name: string, identity: StoreIdentity): string {
+  return `${name}_${storeIdentitySlug(identity)}`;
+}
+
+/**
+ * Invariant check for stores whose physical name carries their identity: the recorded version can
+ * only disagree with the expected one if there is a store-naming bug.
+ */
+export async function assertStoreIdentity(
+  store: AztecAsyncKVStore,
+  storeName: string,
+  identity: StoreIdentity,
+): Promise<void> {
+  const expected = new DatabaseVersion(identity.schemaVersion, identity.rollupAddress);
+  const singleton = store.openSingleton<string>('dbVersion');
+  const stored = await singleton.getAsync();
+  if (stored === undefined) {
+    await singleton.set(expected.toBuffer().toString('utf-8'));
+    return;
+  }
+  let storedVersion: DatabaseVersion;
+  try {
+    storedVersion = DatabaseVersion.fromBuffer(Buffer.from(stored, 'utf-8'));
+  } catch {
+    throw new StoreIdentityMismatchError(storeName, expected.toString(), stored);
+  }
+  if (!storedVersion.equals(expected)) {
+    throw new StoreIdentityMismatchError(storeName, expected.toString(), storedVersion.toString());
+  }
+}
+
+/**
+ * Thrown when a store's recorded identity does not match the identity it was opened under. Since the identity is
+ * part of the physical store name, this can only indicate a store-naming bug; the store is left untouched.
+ */
+export class StoreIdentityMismatchError extends Error {
+  constructor(
+    public readonly storeName: string,
+    public readonly expected: string,
+    public readonly actual: string,
+  ) {
+    super(
+      `Store '${storeName}' records identity ${actual} but was opened as ${expected}. ` +
+        `Refusing to open; data was NOT modified.`,
+    );
+    this.name = 'StoreIdentityMismatchError';
+  }
+}