@aztec/pxe 0.0.1-commit.96bb3f7 → 0.0.1-commit.993d240

package/src/storage/backwards_compatibility_tests/store_spy.ts

@@ -0,0 +1,73 @@
+import type {
+  AztecAsyncArray,
+  AztecAsyncCounter,
+  AztecAsyncKVStore,
+  AztecAsyncMap,
+  AztecAsyncMultiMap,
+  AztecAsyncSet,
+  AztecAsyncSingleton,
+  Key,
+} from '@aztec/kv-store';
+
+type Value = NonNullable<any>;
+
+/** The kind of named sub-store opened on an `AztecAsyncKVStore`. */
+export type StoreKind = 'map' | 'multimap' | 'array' | 'singleton' | 'set' | 'counter';
+
+/** A single recorded `open*()` call against a spied store. */
+export interface OpenedStoreEntry {
+  name: string;
+  kind: StoreKind;
+}
+
+/** A spied `AztecAsyncKVStore` together with a function that returns the sorted log of `open*()` calls. */
+export interface StoreSpy {
+  store: AztecAsyncKVStore;
+  getOpenedStores: () => OpenedStoreEntry[];
+}
+
+/**
+ * Wraps an `AztecAsyncKVStore` so every `openMap`/`openMultiMap`/`openArray`/`openSingleton`/`openSet`/`openCounter`
+ * call is recorded before being delegated to the inner store.
+ */
+export function createStoreSpy(inner: AztecAsyncKVStore): StoreSpy {
+  const entries: OpenedStoreEntry[] = [];
+
+  const store: AztecAsyncKVStore = {
+    openMap<K extends Key, V extends Value>(name: string): AztecAsyncMap<K, V> {
+      entries.push({ name, kind: 'map' });
+      return inner.openMap<K, V>(name);
+    },
+    openMultiMap<K extends Key, V extends Value>(name: string): AztecAsyncMultiMap<K, V> {
+      entries.push({ name, kind: 'multimap' });
+      return inner.openMultiMap<K, V>(name);
+    },
+    openArray<T extends Value>(name: string): AztecAsyncArray<T> {
+      entries.push({ name, kind: 'array' });
+      return inner.openArray<T>(name);
+    },
+    openSingleton<T extends Value>(name: string): AztecAsyncSingleton<T> {
+      entries.push({ name, kind: 'singleton' });
+      return inner.openSingleton<T>(name);
+    },
+    openSet<K extends Key>(name: string): AztecAsyncSet<K> {
+      entries.push({ name, kind: 'set' });
+      return inner.openSet<K>(name);
+    },
+    openCounter<K extends Key>(name: string): AztecAsyncCounter<K> {
+      entries.push({ name, kind: 'counter' });
+      return inner.openCounter<K>(name);
+    },
+    transactionAsync: callback => inner.transactionAsync(callback),
+    clear: () => inner.clear(),
+    delete: () => inner.delete(),
+    estimateSize: () => inner.estimateSize(),
+    close: () => inner.close(),
+    backupTo: (dstPath, compact) => inner.backupTo(dstPath, compact),
+  };
+
+  const getOpenedStores = () =>
+    [...entries].sort((a, b) => (a.name === b.name ? a.kind.localeCompare(b.kind) : a.name.localeCompare(b.name)));
+
+  return { store, getOpenedStores };
+}

package/src/storage/capsule_store/capsule_service.ts

@@ -0,0 +1,90 @@
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { Capsule } from '@aztec/stdlib/tx';
+
+import type { CapsuleStore } from './capsule_store.js';
+
+/**
+ * Wraps a CapsuleStore with scope-based access control. Each operation asserts that the requested scope is in the
+ * allowed scopes list before delegating to the underlying store.
+ */
+export class CapsuleService {
+  constructor(
+    private readonly capsuleStore: CapsuleStore,
+    private readonly allowedScopes: AztecAddress[],
+  ) {}
+
+  setCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[], jobId: string, scope: AztecAddress) {
+    assertAllowedScope(scope, this.allowedScopes);
+    this.capsuleStore.setCapsule(contractAddress, slot, capsule, jobId, scope);
+  }
+
+  async getCapsule(
+    contractAddress: AztecAddress,
+    slot: Fr,
+    jobId: string,
+    scope: AztecAddress,
+    transientCapsules?: Capsule[],
+  ): Promise<Fr[] | null> {
+    assertAllowedScope(scope, this.allowedScopes);
+
+    // TODO(#12425): On the following line, the pertinent capsule gets overshadowed by the transient one. Tackle this.
+    const maybeTransientCapsule = transientCapsules?.find(
+      c =>
+        c.contractAddress.equals(contractAddress) &&
+        c.storageSlot.equals(slot) &&
+        (c.scope ?? AztecAddress.ZERO).equals(scope),
+    )?.data;
+
+    return maybeTransientCapsule ?? (await this.capsuleStore.getCapsule(contractAddress, slot, jobId, scope));
+  }
+
+  deleteCapsule(contractAddress: AztecAddress, slot: Fr, jobId: string, scope: AztecAddress) {
+    assertAllowedScope(scope, this.allowedScopes);
+    this.capsuleStore.deleteCapsule(contractAddress, slot, jobId, scope);
+  }
+
+  copyCapsule(
+    contractAddress: AztecAddress,
+    srcSlot: Fr,
+    dstSlot: Fr,
+    numEntries: number,
+    jobId: string,
+    scope: AztecAddress,
+  ): Promise<void> {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.copyCapsule(contractAddress, srcSlot, dstSlot, numEntries, jobId, scope);
+  }
+
+  appendToCapsuleArray(
+    contractAddress: AztecAddress,
+    baseSlot: Fr,
+    content: Fr[][],
+    jobId: string,
+    scope: AztecAddress,
+  ): Promise<void> {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.appendToCapsuleArray(contractAddress, baseSlot, content, jobId, scope);
+  }
+
+  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, jobId: string, scope: AztecAddress): Promise<Fr[][]> {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.readCapsuleArray(contractAddress, baseSlot, jobId, scope);
+  }
+
+  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][], jobId: string, scope: AztecAddress) {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.setCapsuleArray(contractAddress, baseSlot, content, jobId, scope);
+  }
+}
+
+function assertAllowedScope(scope: AztecAddress, allowedScopes: AztecAddress[]) {
+  if (scope.equals(AztecAddress.ZERO)) {
+    return;
+  }
+  if (!allowedScopes.some((allowed: AztecAddress) => allowed.equals(scope))) {
+    throw new Error(
+      `Scope ${scope.toString()} is not in the allowed scopes list: [${allowedScopes.map((s: AztecAddress) => s.toString()).join(', ')}]. See https://docs.aztec.network/errors/10`,
+    );
+  }
+}

package/src/storage/capsule_store/capsule_store.ts

@@ -1,14 +1,24 @@
 import { Fr } from '@aztec/foundation/curves/bn254';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore, AztecAsyncMap } from '@aztec/kv-store';
-import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+import type { StagedStore } from '../../job_coordinator/job_coordinator.js';
+
+export class CapsuleStore implements StagedStore {
+  readonly storeName = 'capsule';
 
-export class CapsuleStore {
   #store: AztecAsyncKVStore;
 
-  // Arbitrary data stored by contracts. Key is computed as `${contractAddress}:${key}`
+  // Arbitrary data stored by contracts. Key is computed as `${contractAddress}:${scope}:${key}`, using the zero
+  // address for the global scope.
   #capsules: AztecAsyncMap<string, Buffer>;
 
+  // jobId => `${contractAddress}:${scope}:${key}` => capsule data
+  // when `#stagedCapsules.get('some-job-id').get('${some-contract-address}:${some-scope}:${some-key}') === null`,
+  // it signals that the capsule was deleted during the job, so it needs to be deleted on commit
+  #stagedCapsules: Map<string, Map<string, Buffer | null>>;
+
   logger: Logger;
 
   constructor(store: AztecAsyncKVStore) {
@@ -16,21 +26,120 @@ export class CapsuleStore {
 
     this.#capsules = this.#store.openMap('capsules');
 
+    this.#stagedCapsules = new Map();
+
     this.logger = createLogger('pxe:capsule-data-provider');
   }
 
+  /**
+   * Given a job denoted by `jobId`, it returns the
+   * capsules that said job has interacted with.
+   *
+   * Capsules that haven't been committed to persistence KV storage
+   * are kept in-memory in `#stagedCapsules`, this method provides a convenient
+   * way to access that in-memory collection of data.
+   *
+   * @param jobId
+   * @returns
+   */
+  #getJobStagedCapsules(jobId: string): Map<string, Buffer | null> {
+    let jobStagedCapsules = this.#stagedCapsules.get(jobId);
+    if (!jobStagedCapsules) {
+      jobStagedCapsules = new Map();
+      this.#stagedCapsules.set(jobId, jobStagedCapsules);
+    }
+    return jobStagedCapsules;
+  }
+
+  /**
+   * Reads a capsule's slot from the staged version of the data associated to the given jobId.
+   *
+   * If it is not there, it reads it from the KV store.
+   */
+  async #getFromStage(jobId: string, dbSlotKey: string): Promise<Buffer | null | undefined> {
+    const jobStagedCapsules = this.#getJobStagedCapsules(jobId);
+    const staged: Buffer | null | undefined = jobStagedCapsules.get(dbSlotKey);
+
+    // Always issue DB read to keep IndexedDB transaction alive, even if the value is in the job staged data. This
+    // keeps IndexedDB transactions alive (they auto-commit when a new micro-task starts and there are no pending read
+    // requests). The staged value still takes precedence if it exists (including null for deletions).
+    const dbValue = await this.#loadCapsuleFromDb(dbSlotKey);
+
+    return staged !== undefined ? staged : dbValue;
+  }
+
+  /**
+   * Writes a capsule to the stage of a job.
+   */
+  #setOnStage(jobId: string, dbSlotKey: string, capsuleData: Buffer) {
+    this.#getJobStagedCapsules(jobId).set(dbSlotKey, capsuleData);
+  }
+
+  /**
+   * Deletes a capsule on the stage of a job. Note the capsule will still
+   * exist in storage until the job is committed.
+   */
+  #deleteOnStage(jobId: string, dbSlotKey: string) {
+    this.#getJobStagedCapsules(jobId).set(dbSlotKey, null);
+  }
+
+  async #loadCapsuleFromDb(dbSlotKey: string): Promise<Buffer | null> {
+    const dataBuffer = await this.#capsules.getAsync(dbSlotKey);
+    if (!dataBuffer) {
+      return null;
+    }
+
+    return dataBuffer;
+  }
+
+  /**
+   * Commits staged data to main storage.
+   * Called by JobCoordinator when a job completes successfully.
+   * Note: JobCoordinator wraps all commits in a single transaction, so we don't
+   * need our own transactionAsync here (and using one would deadlock on IndexedDB).
+   * @param jobId - The jobId identifying which staged data to commit
+   */
+  async commit(jobId: string): Promise<void> {
+    const jobStagedCapsules = this.#getJobStagedCapsules(jobId);
+
+    for (const [key, value] of jobStagedCapsules) {
+      // In the write stage, we represent deleted capsules with null
+      // (as opposed to undefined, which denotes there was never a capsule there to begin with).
+      // So we delete from actual KV store here.
+      if (value === null) {
+        await this.#capsules.delete(key);
+      } else {
+        await this.#capsules.set(key, value);
+      }
+    }
+
+    this.#stagedCapsules.delete(jobId);
+  }
+
+  /**
+   * Discards staged data without committing.
+   */
+  discardStaged(jobId: string): Promise<void> {
+    this.#stagedCapsules.delete(jobId);
+    return Promise.resolve();
+  }
+
   /**
    * Stores arbitrary information in a per-contract non-volatile database, which can later be retrieved with `loadCapsule`.
    * * If data was already stored at this slot, it is overwritten.
    * @param contractAddress - The contract address to scope the data under.
    * @param slot - The slot in the database in which to store the value. Slots need not be contiguous.
    * @param capsule - An array of field elements representing the capsule.
+   * @param jobId - The context in which this store will be visible until PXE decides to persist it to underlying KV store
    * @remarks A capsule is a "blob" of data that is passed to the contract through an oracle. It works similarly
    * to public contract storage in that it's indexed by the contract address and storage slot but instead of the global
    * network state it's backed by local PXE db.
    */
-  async storeCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[]): Promise<void> {
-    await this.#capsules.set(dbSlotToKey(contractAddress, slot), Buffer.concat(capsule.map(value => value.toBuffer())));
+  setCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[], jobId: string, scope: AztecAddress) {
+    const dbSlotKey = dbSlotToKey(contractAddress, slot, scope);
+
+    // A store overrides any pre-existing data on the slot
+    this.#setOnStage(jobId, dbSlotKey, Buffer.concat(capsule.map(value => value.toBuffer())));
   }
 
   /**
@@ -39,8 +148,18 @@ export class CapsuleStore {
    * @param slot - The slot in the database to read.
    * @returns The stored data or `null` if no data is stored under the slot.
    */
-  async loadCapsule(contractAddress: AztecAddress, slot: Fr): Promise<Fr[] | null> {
-    const dataBuffer = await this.#capsules.getAsync(dbSlotToKey(contractAddress, slot));
+  getCapsule(contractAddress: AztecAddress, slot: Fr, jobId: string, scope: AztecAddress): Promise<Fr[] | null> {
+    return this.#store.transactionAsync(() => this.#getCapsuleInternal(contractAddress, slot, jobId, scope));
+  }
+
+  /** Same as getCapsule but without its own transaction, for use inside an existing transactionAsync. */
+  async #getCapsuleInternal(
+    contractAddress: AztecAddress,
+    slot: Fr,
+    jobId: string,
+    scope: AztecAddress,
+  ): Promise<Fr[] | null> {
+    const dataBuffer = await this.#getFromStage(jobId, dbSlotToKey(contractAddress, slot, scope));
     if (!dataBuffer) {
       this.logger.trace(`Data not found for contract ${contractAddress.toString()} and slot ${slot.toString()}`);
       return null;
@@ -57,8 +176,9 @@ export class CapsuleStore {
    * @param contractAddress - The contract address under which the data is scoped.
    * @param slot - The slot in the database to delete.
    */
-  async deleteCapsule(contractAddress: AztecAddress, slot: Fr): Promise<void> {
-    await this.#capsules.delete(dbSlotToKey(contractAddress, slot));
+  deleteCapsule(contractAddress: AztecAddress, slot: Fr, jobId: string, scope: AztecAddress) {
+    // When we commit this, we will interpret null as a deletion, so we'll propagate the delete to the KV store
+    this.#deleteOnStage(jobId, dbSlotToKey(contractAddress, slot, scope));
   }
 
   /**
@@ -72,28 +192,38 @@ export class CapsuleStore {
    * @param dstSlot - The first slot to copy to.
    * @param numEntries - The number of entries to copy.
    */
-  copyCapsule(contractAddress: AztecAddress, srcSlot: Fr, dstSlot: Fr, numEntries: number): Promise<void> {
+  copyCapsule(
+    contractAddress: AztecAddress,
+    srcSlot: Fr,
+    dstSlot: Fr,
+    numEntries: number,
+    jobId: string,
+    scope: AztecAddress,
+  ): Promise<void> {
+    // This transactional context gives us "copy atomicity":
+    // there shouldn't be concurrent writes to what's being copied here.
+    // Equally important: this in practice is expected to perform thousands of DB operations
+    // and not using a transaction here would heavily impact performance.
     return this.#store.transactionAsync(async () => {
       // In order to support overlapping source and destination regions, we need to check the relative positions of source
       // and destination. If destination is ahead of source, then by the time we overwrite source elements using forward
       // indexes we'll have already read those. On the contrary, if source is ahead of destination we need to use backward
       // indexes to avoid reading elements that've been overwritten.
-
       const indexes = Array.from(Array(numEntries).keys());
       if (srcSlot.lt(dstSlot)) {
         indexes.reverse();
       }
 
       for (const i of indexes) {
-        const currentSrcSlot = dbSlotToKey(contractAddress, srcSlot.add(new Fr(i)));
-        const currentDstSlot = dbSlotToKey(contractAddress, dstSlot.add(new Fr(i)));
+        const currentSrcSlot = dbSlotToKey(contractAddress, srcSlot.add(new Fr(i)), scope);
+        const currentDstSlot = dbSlotToKey(contractAddress, dstSlot.add(new Fr(i)), scope);
 
-        const toCopy = await this.#capsules.getAsync(currentSrcSlot);
+        const toCopy = await this.#getFromStage(jobId, currentSrcSlot);
         if (!toCopy) {
           throw new Error(`Attempted to copy empty slot ${currentSrcSlot} for contract ${contractAddress.toString()}`);
         }
 
-        await this.#capsules.set(currentDstSlot, toCopy);
+        this.#setOnStage(jobId, currentDstSlot, toCopy);
       }
     });
   }
@@ -106,35 +236,51 @@ export class CapsuleStore {
    * @param baseSlot - The slot where the array length is stored
    * @param content - Array of capsule data to append
    */
-  appendToCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]): Promise<void> {
+  appendToCapsuleArray(
+    contractAddress: AztecAddress,
+    baseSlot: Fr,
+    content: Fr[][],
+    jobId: string,
+    scope: AztecAddress,
+  ): Promise<void> {
+    // We wrap this in a transaction to serialize concurrent calls from Promise.all.
+    // Without this, concurrent appends to the same array could race: both read length=0,
+    // both write at the same slots, one overwrites the other.
+    // Equally important: this in practice is expected to perform thousands of DB operations
+    // and not using a transaction here would heavily impact performance.
     return this.#store.transactionAsync(async () => {
       // Load current length, defaulting to 0 if not found
-      const lengthData = await this.loadCapsule(contractAddress, baseSlot);
+      const lengthData = await this.#getCapsuleInternal(contractAddress, baseSlot, jobId, scope);
       const currentLength = lengthData ? lengthData[0].toNumber() : 0;
 
       // Store each capsule at consecutive slots after baseSlot + 1 + currentLength
       for (let i = 0; i < content.length; i++) {
         const nextSlot = arraySlot(baseSlot, currentLength + i);
-        await this.storeCapsule(contractAddress, nextSlot, content[i]);
+        this.setCapsule(contractAddress, nextSlot, content[i], jobId, scope);
       }
 
       // Update length to include all new capsules
       const newLength = currentLength + content.length;
-      await this.storeCapsule(contractAddress, baseSlot, [new Fr(newLength)]);
+      this.setCapsule(contractAddress, baseSlot, [new Fr(newLength)], jobId, scope);
     });
   }
 
-  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr): Promise<Fr[][]> {
+  readCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, jobId: string, scope: AztecAddress): Promise<Fr[][]> {
+    // I'm leaving this transactional context here though because I'm assuming this
+    // gives us "read array atomicity": there shouldn't be concurrent writes to what's being copied
+    // here.
+    // This is one point we should revisit in the future if we want to relax the concurrency
+    // of jobs: different calls running concurrently on the same contract may cause trouble.
     return this.#store.transactionAsync(async () => {
       // Load length, defaulting to 0 if not found
-      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const maybeLength = await this.#getCapsuleInternal(contractAddress, baseSlot, jobId, scope);
       const length = maybeLength ? maybeLength[0].toBigInt() : 0n;
 
       const values: Fr[][] = [];
 
       // Read each capsule at consecutive slots after baseSlot
       for (let i = 0; i < length; i++) {
-        const currentValue = await this.loadCapsule(contractAddress, arraySlot(baseSlot, i));
+        const currentValue = await this.#getCapsuleInternal(contractAddress, arraySlot(baseSlot, i), jobId, scope);
         if (currentValue == undefined) {
           throw new Error(
             `Expected non-empty value at capsule array in base slot ${baseSlot} at index ${i} for contract ${contractAddress}`,
@@ -148,30 +294,38 @@ export class CapsuleStore {
     });
   }
 
-  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][]) {
+  setCapsuleArray(contractAddress: AztecAddress, baseSlot: Fr, content: Fr[][], jobId: string, scope: AztecAddress) {
+    // This transactional context in theory isn't so critical now because we aren't
+    // writing to DB so if there's exceptions midway and it blows up, no visible impact
+    // to persistent storage will happen.
+    // I'm leaving this transactional context here though because I'm assuming this
+    // gives us "write array atomicity": there shouldn't be concurrent writes to what's being copied
+    // here.
+    // This is one point we should revisit in the future if we want to relax the concurrency
+    // of jobs: different calls running concurrently on the same contract may cause trouble.
     return this.#store.transactionAsync(async () => {
       // Load current length, defaulting to 0 if not found
-      const maybeLength = await this.loadCapsule(contractAddress, baseSlot);
+      const maybeLength = await this.#getCapsuleInternal(contractAddress, baseSlot, jobId, scope);
       const originalLength = maybeLength ? maybeLength[0].toNumber() : 0;
 
       // Set the new length
-      await this.storeCapsule(contractAddress, baseSlot, [new Fr(content.length)]);
+      this.setCapsule(contractAddress, baseSlot, [new Fr(content.length)], jobId, scope);
 
       // Store the new content, possibly overwriting existing values
       for (let i = 0; i < content.length; i++) {
-        await this.storeCapsule(contractAddress, arraySlot(baseSlot, i), content[i]);
+        this.setCapsule(contractAddress, arraySlot(baseSlot, i), content[i], jobId, scope);
       }
 
       // Clear any stragglers
       for (let i = content.length; i < originalLength; i++) {
-        await this.deleteCapsule(contractAddress, arraySlot(baseSlot, i));
+        this.deleteCapsule(contractAddress, arraySlot(baseSlot, i), jobId, scope);
       }
     });
   }
 }
 
-function dbSlotToKey(contractAddress: AztecAddress, slot: Fr): string {
-  return `${contractAddress.toString()}:${slot.toString()}`;
+function dbSlotToKey(contractAddress: AztecAddress, slot: Fr, scope: AztecAddress): string {
+  return [contractAddress.toString(), scope.toString(), slot.toString()].join(':');
 }
 
 function arraySlot(baseSlot: Fr, index: number) {

package/src/storage/capsule_store/index.ts

@@ -1 +1,2 @@
+export { CapsuleService } from './capsule_service.js';
 export { CapsuleStore } from './capsule_store.js';