@aztec-labs/pxe 6.0.0-nightly.20260829

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-labs/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/base_staging_store.ts

@@ -0,0 +1,227 @@
+import { Semaphore } from '@aztec-labs/foundation/queue';
+import type { AztecAsyncKVStore } from '@aztec-labs/kv-store';
+
+import type { Rollbackable } from './rollbackable.js';
+import type { ChangeSetId, StagedStore } from './staged_write_coordinator.js';
+
+/**
+ * Base class for stores that stage their writes per change set, flushing them on commit and dropping them on abort.
+ *
+ * A change set's staged data exists only between {@link beginChangeSet} and {@link commitChangeSet}/
+ * {@link discardChangeSet}. An operation arriving on behalf of a change set that already ended (e.g. one discarded on
+ * abort) finds no matching change set and throws, so it cannot stage new data for a dead one.
+ *
+ * The staged data and the store's kv handles live in this base class and are only reachable through
+ * {@link withChangeSet} (change set operations, staged data only), {@link withChangeSetAndDb} (change set operations
+ * that also read the DB), {@link flushChangeSet} (the commit-time write-back) and {@link applyRollback} (the reorg
+ * truncation).
+ *
+ * The class is thread safe: an internal lock serializes the operations run through {@link withChangeSet} and
+ * {@link withChangeSetAndDb}, so two of them issued concurrently (e.g. under `Promise.all`) never interleave across
+ * awaits. Subclasses can therefore read staged data and write it back without handling atomicity themselves.
+ *
+ * @typeParam TChangeSet - The in-memory buffer a change set accumulates its writes in, built empty by `buildChangeSet`
+ * on every {@link beginChangeSet} and dropped when the change set ends. A store keyed by id might stage
+ * `{ items: Map<string, Item> }`.
+ * @typeParam TDb - The store's kv handles, opened once at construction by `buildDb` and shared by every change set.
+ * That same store would open `{ items: AztecAsyncMap<string, Buffer>; itemsByBlockNumber: AztecAsyncMultiMap }`.
+ */
+export abstract class BaseStagingStore<TChangeSet, TDb> implements StagedStore, Rollbackable {
+  public readonly storeName: string;
+
+  readonly #store: AztecAsyncKVStore;
+  readonly #db: TDb;
+  readonly #buildChangeSet: () => TChangeSet;
+  readonly #lock = new Semaphore(1);
+
+  /** The change set currently open, if any. */
+  #current: OpenChangeSet<TChangeSet> | undefined;
+
+  protected constructor({
+    storeName,
+    store,
+    buildChangeSet,
+    buildDb,
+  }: {
+    /** Unique name identifying this store, used in error messages and for registration with StagedWriteCoordinator. */
+    storeName: string;
+    /** The backing kv store; the runners open their transactions on it. */
+    store: AztecAsyncKVStore;
+    /** Creates the empty staged data a change set starts with. */
+    buildChangeSet: () => TChangeSet;
+    /** Opens the store's kv handles. */
+    buildDb: (store: AztecAsyncKVStore) => TDb;
+  }) {
+    this.storeName = storeName;
+    this.#store = store;
+    this.#buildChangeSet = buildChangeSet;
+    this.#db = buildDb(store);
+  }
+
+  /**
+   * Opens the change set, so its operations are accepted until {@link commitChangeSet} or {@link discardChangeSet}.
+   *
+   * @throws If a change set is already open.
+   */
+  beginChangeSet(changeSetId: ChangeSetId): void {
+    if (this.#current !== undefined) {
+      throw new Error(
+        `Store "${this.storeName}": cannot open change set "${changeSetId}" because change set ` +
+          `"${this.#current.changeSetId}" is already open`,
+      );
+    }
+    this.#current = { changeSetId, changeSet: this.#buildChangeSet(), inFlight: 0 };
+  }
+
+  /**
+   * Commits the change set's staged data: flushes it via {@link flushChangeSet}, then closes the change set. Runs
+   * inside the transaction owned by the caller.
+   *
+   * Not meant to be overridden: subclasses implement {@link flushChangeSet}.
+   *
+   * @throws If the change set is not open, or still has operations in flight.
+   */
+  async commitChangeSet(changeSetId: ChangeSetId): Promise<void> {
+    const current = this.#currentOrThrow(changeSetId);
+    if (current.inFlight > 0) {
+      throw new Error(
+        `Store "${this.storeName}": cannot commit change set "${changeSetId}" while ${current.inFlight} of its ` +
+          `operations are in flight`,
+      );
+    }
+    await this.flushChangeSet(current.changeSet, this.#db);
+    this.#closeChangeSet(changeSetId);
+  }
+
+  /** Closes the change set, discarding any staged data without committing. A no-op if it is not open. */
+  discardChangeSet(changeSetId: ChangeSetId): void {
+    this.#closeChangeSet(changeSetId);
+  }
+
+  /**
+   * Rolls the store back to `toBlock` via {@link applyRollback}.
+   *
+   * Must be called inside a transaction owned by the caller, since it opens none of its own.
+   *
+   * Not meant to be overridden: subclasses implement {@link applyRollback}.
+   *
+   * @throws If a change set is open, since its staged writes could later be committed anchored to deleted blocks.
+   */
+  async rollbackToBlock(toBlock: number): Promise<void> {
+    if (this.#current !== undefined) {
+      throw new Error(
+        `Store "${this.storeName}": cannot roll back while change set "${this.#current.changeSetId}" is open`,
+      );
+    }
+    await this.applyRollback(toBlock, this.#db);
+  }
+
+  /**
+   * Writes the change set's staged data to persistent storage. Runs inside the caller's transaction: it must not
+   * open a transaction of its own or call {@link withChangeSetAndDb}.
+   */
+  protected abstract flushChangeSet(changeSet: TChangeSet, db: TDb): Promise<void>;
+
+  /**
+   * Deletes the state originating from blocks strictly above `toBlock`. Runs inside the transaction owned by
+   * {@link rollbackToBlock}'s caller: it must not open a transaction of its own or call {@link withChangeSetAndDb}.
+   */
+  protected abstract applyRollback(toBlock: number, db: TDb): Promise<void>;
+
+  /**
+   * Runs a change set operation that reads the DB. Takes the store's lock, opens a transaction, and calls `fn` with
+   * the change set's staged data and a read-only view of the DB (writes are staged in memory until
+   * {@link flushChangeSet} runs on commit).
+   *
+   * Prefer {@link withChangeSet} unless `fn` actually reads the DB.
+   *
+   * The lock makes the store thread safe: two operations issued concurrently (e.g. under `Promise.all`) cannot
+   * interleave across awaits, so `fn` can read staged data and write it back without handling atomicity itself.
+   *
+   * @throws If the change set is not open.
+   */
+  protected withChangeSetAndDb<R>(
+    changeSetId: ChangeSetId,
+    fn: (changeSet: TChangeSet, db: ReadonlyDb<TDb>) => Promise<R>,
+  ): Promise<R> {
+    return this.#runLocked(changeSetId, () =>
+      this.#store.transactionAsync(() => {
+        // Re-resolve after the wait: the change set may have been discarded while this operation queued on the lock
+        // or on the store's transaction queue.
+        const current = this.#currentOrThrow(changeSetId);
+        return fn(current.changeSet, this.#db);
+      }),
+    );
+  }
+
+  /**
+   * Runs a change set operation over the staged data alone: takes the store's lock and calls `fn` with the change
+   * set's staged data, opening no db transaction.
+   *
+   * Use {@link withChangeSetAndDb} instead when `fn` reads the DB.
+   *
+   * @throws If the change set is not open.
+   */
+  protected withChangeSet<R>(changeSetId: ChangeSetId, fn: (changeSet: TChangeSet) => R | Promise<R>): Promise<R> {
+    return this.#runLocked(changeSetId, async () => {
+      // Re-resolve after the wait: the change set may have been discarded while this operation queued on the lock.
+      const current = this.#currentOrThrow(changeSetId);
+      return await fn(current.changeSet);
+    });
+  }
+
+  async #runLocked<R>(changeSetId: ChangeSetId, run: () => Promise<R>): Promise<R> {
+    const entered = this.#currentOrThrow(changeSetId);
+    entered.inFlight++;
+    try {
+      await this.#lock.acquire();
+      try {
+        return await run();
+      } finally {
+        this.#lock.release();
+      }
+    } finally {
+      // Count down on the change set this operation entered, which may no longer be the open one.
+      entered.inFlight--;
+    }
+  }
+
+  #currentOrThrow(changeSetId: ChangeSetId): OpenChangeSet<TChangeSet> {
+    if (this.#current?.changeSetId !== changeSetId) {
+      throw new Error(`Store "${this.storeName}": change set "${changeSetId}" is not open`);
+    }
+    return this.#current;
+  }
+
+  #closeChangeSet(changeSetId: ChangeSetId): void {
+    if (this.#current?.changeSetId === changeSetId) {
+      this.#current = undefined;
+    }
+  }
+}
+
+/**
+ * View of a store's kv handles restricted to the kv interfaces' read methods. Change set operations receive this view:
+ * while a change set is open the DB is read-only, since all writes are staged in memory until
+ * {@link BaseStagingStore.flushChangeSet} runs on commit.
+ */
+export type ReadonlyDb<T> = {
+  readonly [K in keyof T]: Pick<T[K], Extract<keyof T[K], ReadMethod>>;
+};
+
+/** The kv read surface staging stores use: what {@link ReadonlyDb} exposes while a change set is open. */
+type ReadMethod =
+  | 'getAsync'
+  | 'hasAsync'
+  | 'entriesAsync'
+  | 'valuesAsync'
+  | 'keysAsync'
+  | 'sizeAsync'
+  | 'getValuesAsync'
+  | 'getValueCountAsync';
+
+/**
+ * The change set in progress: its id, its staged data (created empty when it opens), and how many of its operations
+ * are still running.
+ */
+type OpenChangeSet<T> = { changeSetId: ChangeSetId; changeSet: T; inFlight: number };

package/src/storage/capsule_store/capsule_service.ts

@@ -0,0 +1,97 @@
+import type { Fr } from '@aztec-labs/foundation/curves/bn254';
+import { AztecAddress } from '@aztec-labs/stdlib/aztec-address';
+import type { Capsule } from '@aztec-labs/stdlib/tx';
+
+import { assertAllowedScope } from '../allowed_scopes.js';
+import type { ChangeSetId } from '../staged_write_coordinator.js';
+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 {
+  private readonly allowedScopes: AztecAddress[];
+
+  constructor(
+    private readonly capsuleStore: CapsuleStore,
+    allowedScopes: AztecAddress[],
+  ) {
+    // The zero address denotes the global capsule scope, which is always permitted.
+    this.allowedScopes = [...allowedScopes, AztecAddress.ZERO];
+  }
+
+  setCapsule(contractAddress: AztecAddress, slot: Fr, capsule: Fr[], changeSetId: ChangeSetId, scope: AztecAddress) {
+    assertAllowedScope(scope, this.allowedScopes);
+    this.capsuleStore.setCapsule(contractAddress, slot, capsule, changeSetId, scope);
+  }
+
+  async getCapsule(
+    contractAddress: AztecAddress,
+    slot: Fr,
+    changeSetId: ChangeSetId,
+    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, changeSetId, scope));
+  }
+
+  deleteCapsule(contractAddress: AztecAddress, slot: Fr, changeSetId: ChangeSetId, scope: AztecAddress) {
+    assertAllowedScope(scope, this.allowedScopes);
+    this.capsuleStore.deleteCapsule(contractAddress, slot, changeSetId, scope);
+  }
+
+  copyCapsule(
+    contractAddress: AztecAddress,
+    srcSlot: Fr,
+    dstSlot: Fr,
+    numEntries: number,
+    changeSetId: ChangeSetId,
+    scope: AztecAddress,
+  ): Promise<void> {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.copyCapsule(contractAddress, srcSlot, dstSlot, numEntries, changeSetId, scope);
+  }
+
+  appendToCapsuleArray(
+    contractAddress: AztecAddress,
+    baseSlot: Fr,
+    content: Fr[][],
+    changeSetId: ChangeSetId,
+    scope: AztecAddress,
+  ): Promise<void> {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.appendToCapsuleArray(contractAddress, baseSlot, content, changeSetId, scope);
+  }
+
+  readCapsuleArray(
+    contractAddress: AztecAddress,
+    baseSlot: Fr,
+    changeSetId: ChangeSetId,
+    scope: AztecAddress,
+  ): Promise<Fr[][]> {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.readCapsuleArray(contractAddress, baseSlot, changeSetId, scope);
+  }
+
+  setCapsuleArray(
+    contractAddress: AztecAddress,
+    baseSlot: Fr,
+    content: Fr[][],
+    changeSetId: ChangeSetId,
+    scope: AztecAddress,
+  ) {
+    assertAllowedScope(scope, this.allowedScopes);
+    return this.capsuleStore.setCapsuleArray(contractAddress, baseSlot, content, changeSetId, scope);
+  }
+}