@aztec/pxe 0.0.1-commit.b655e406 → 0.0.1-commit.b8a057fa

package/src/job_coordinator/job_coordinator.ts

@@ -0,0 +1,150 @@
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+
+/**
+ * Interface that stores must implement to support staged writes.
+ */
+export interface StagedStore {
+  /** Unique name identifying this store (used for tracking staged stores from JobCoordinator) */
+  readonly storeName: string;
+
+  /**
+   * Commits staged data to main storage.
+   * Should be called within a transaction for atomicity.
+   *
+   * @param jobId - The job identifier
+   */
+  commit(jobId: string): Promise<void>;
+
+  /**
+   * Discards staged data without committing.
+   * Called on abort.
+   *
+   * @param jobId - The job identifier
+   */
+  discardStaged(jobId: string): Promise<void>;
+}
+
+/**
+ * JobCoordinator manages job lifecycle and provides crash resilience for PXE operations.
+ *
+ * It uses a staged writes pattern:
+ * 1. When a job begins, a unique job ID is created
+ * 2. During the job, all writes go to staging (keyed by job ID)
+ * 3. On commit, staging is promoted to main storage
+ * 4. On abort, staged data is discarded
+ *
+ * Note: PXE should only rely on a single JobCoordinator instance, so it can eventually
+ * orchestrate concurrent jobs. Right now it doesn't make a difference because we're
+ * using a job queue with concurrency=1.
+ */
+export class JobCoordinator {
+  private readonly log: Logger;
+
+  /** The underlying KV store */
+  kvStore: AztecAsyncKVStore;
+
+  #currentJobId: string | undefined;
+  #stores: Map<string, StagedStore> = new Map();
+
+  constructor(kvStore: AztecAsyncKVStore, bindings?: LoggerBindings) {
+    this.kvStore = kvStore;
+    this.log = createLogger('pxe:job_coordinator', bindings);
+  }
+
+  /**
+   * Registers a staged store.
+   * Must be called during initialization for all stores that need staging support.
+   */
+  registerStore(store: StagedStore): void {
+    if (this.#stores.has(store.storeName)) {
+      throw new Error(`Store "${store.storeName}" is already registered`);
+    }
+    this.#stores.set(store.storeName, store);
+    this.log.debug(`Registered staged store: ${store.storeName}`);
+  }
+
+  /**
+   * Registers multiple staged stores.
+   */
+  registerStores(stores: StagedStore[]): void {
+    for (const store of stores) {
+      this.registerStore(store);
+    }
+  }
+
+  /**
+   * Begins a new job and returns a job ID for staged writes.
+   *
+   * @returns Job ID to pass to store operations
+   */
+  beginJob(): string {
+    if (this.#currentJobId) {
+      throw new Error(
+        `Cannot begin job: job ${this.#currentJobId} is already in progress. ` +
+          `This should not happen - ensure jobs are properly committed or aborted.`,
+      );
+    }
+
+    const jobId = randomBytes(8).toString('hex');
+    this.#currentJobId = jobId;
+
+    this.log.debug(`Started job ${jobId}`);
+    return jobId;
+  }
+
+  /**
+   * Commits a job by promoting all staged data to main storage.
+   *
+   * @param jobId - The job ID returned from beginJob
+   */
+  async commitJob(jobId: string): Promise<void> {
+    if (!this.#currentJobId || this.#currentJobId !== jobId) {
+      throw new Error(
+        `Cannot commit job ${jobId}: no matching job in progress. ` + `Current job: ${this.#currentJobId ?? 'none'}`,
+      );
+    }
+
+    this.log.debug(`Committing job ${jobId}`);
+
+    // Commit all stores atomically in a single transaction.
+    // Each store's commit is a no-op if it has no staged data (but that's up to each store to handle).
+    await this.kvStore.transactionAsync(async () => {
+      for (const store of this.#stores.values()) {
+        await store.commit(jobId);
+      }
+    });
+
+    this.#currentJobId = undefined;
+    this.log.debug(`Job ${jobId} committed successfully`);
+  }
+
+  /**
+   * Aborts a job by discarding all staged data.
+   *
+   * @param jobId - The job ID returned from beginJob
+   */
+  async abortJob(jobId: string): Promise<void> {
+    if (!this.#currentJobId || this.#currentJobId !== jobId) {
+      // Job may have already been aborted or never started properly
+      this.log.warn(`Abort called for job ${jobId} but current job is ${this.#currentJobId ?? 'none'}`);
+    }
+
+    this.log.debug(`Aborting job ${jobId}`);
+
+    for (const store of this.#stores.values()) {
+      await store.discardStaged(jobId);
+    }
+
+    this.#currentJobId = undefined;
+    this.log.debug(`Job ${jobId} aborted`);
+  }
+
+  /**
+   * Checks if there's a job currently in progress.
+   */
+  hasJobInProgress(): boolean {
+    return this.#currentJobId !== undefined;
+  }
+}

package/src/logs/log_service.ts

@@ -0,0 +1,291 @@
+import { PRIVATE_LOG_CIPHERTEXT_LEN } from '@aztec/constants';
+import type { BlockNumber } from '@aztec/foundation/branded-types';
+import type { GrumpkinScalar, Point } from '@aztec/foundation/curves/grumpkin';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
+import type { KeyStore } from '@aztec/key-store';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { BlockHash, L2TipsProvider } from '@aztec/stdlib/block';
+import type { CompleteAddress } from '@aztec/stdlib/contract';
+import type { AztecNode } from '@aztec/stdlib/interfaces/server';
+import {
+  AppTaggingSecret,
+  AppTaggingSecretKind,
+  type LogResult,
+  SiloedTag,
+  computeSharedTaggingSecret,
+} from '@aztec/stdlib/logs';
+import type { BlockHeader } from '@aztec/stdlib/tx';
+
+import {
+  type LogRetrievalRequest,
+  LogSource,
+} from '../contract_function_simulator/noir-structs/log_retrieval_request.js';
+import type { LogRetrievalResponse } from '../contract_function_simulator/noir-structs/log_retrieval_response.js';
+import type { PendingTaggedLog } from '../contract_function_simulator/noir-structs/pending_tagged_log.js';
+import { ResolvedTx } from '../contract_function_simulator/noir-structs/resolved_tx.js';
+import { AddressStore } from '../storage/address_store/address_store.js';
+import type { RecipientTaggingStore } from '../storage/tagging_store/recipient_tagging_store.js';
+import type { TaggingSecretSourcesStore } from '../storage/tagging_store/tagging_secret_sources_store.js';
+import {
+  getAllPrivateLogsByTags,
+  getAllPublicLogsByTagsFromContract,
+  syncTaggedPrivateLogs,
+} from '../tagging/index.js';
+
+/** Key used to group requests by their (fromBlock, toBlock) range so each group becomes a single node call. */
+type RangeKey = string;
+const rangeKey = (fromBlock?: BlockNumber, toBlock?: BlockNumber): RangeKey => `${fromBlock ?? ''}-${toBlock ?? ''}`;
+
+export class LogService {
+  private log: Logger;
+
+  constructor(
+    private readonly aztecNode: AztecNode,
+    private readonly anchorBlockHeader: BlockHeader,
+    private readonly l2TipsStore: L2TipsProvider,
+    private readonly keyStore: KeyStore,
+    private readonly recipientTaggingStore: RecipientTaggingStore,
+    private readonly taggingSecretSourcesStore: TaggingSecretSourcesStore,
+    private readonly addressStore: AddressStore,
+    private readonly jobId: string,
+    bindings?: LoggerBindings,
+  ) {
+    this.log = createLogger('pxe:log_service', bindings);
+  }
+
+  /** Fetches all logs matching each request's tag, returning an array of log arrays (one per request). */
+  public async fetchLogsByTag(
+    contractAddress: AztecAddress,
+    logRetrievalRequests: LogRetrievalRequest[],
+  ): Promise<LogRetrievalResponse[][]> {
+    for (const request of logRetrievalRequests) {
+      if (!contractAddress.equals(request.contractAddress)) {
+        throw new Error(`Got a log retrieval request from ${request.contractAddress}, expected ${contractAddress}`);
+      }
+    }
+
+    if (logRetrievalRequests.length === 0) {
+      return [];
+    }
+
+    const anchorBlockHash = await this.anchorBlockHeader.hash();
+
+    const [publicLogsPerRequest, privateLogsPerRequest] = await Promise.all([
+      this.#fetchPublicLogs(contractAddress, logRetrievalRequests, anchorBlockHash),
+      this.#fetchPrivateLogs(logRetrievalRequests, anchorBlockHash),
+    ]);
+
+    return logRetrievalRequests.map((_request, i) => [
+      ...publicLogsPerRequest[i].map(LogService.#toLogRetrievalResponse),
+      ...privateLogsPerRequest[i].map(LogService.#toLogRetrievalResponse),
+    ]);
+  }
+
+  async #fetchPublicLogs(
+    contractAddress: AztecAddress,
+    requests: LogRetrievalRequest[],
+    anchorBlockHash: BlockHash,
+  ): Promise<LogResult[][]> {
+    const indices = requests.flatMap((r, i) => (r.source !== LogSource.PRIVATE ? [i] : []));
+    if (indices.length === 0) {
+      return requests.map(() => []);
+    }
+
+    const resultsPerRequest: LogResult[][] = requests.map(() => []);
+    const groups = LogService.#groupByRange(indices.map(i => ({ index: i, request: requests[i] })));
+
+    await Promise.all(
+      Array.from(groups.values()).map(async group => {
+        const tags = group.entries.map(e => e.request.tag);
+        const results = await getAllPublicLogsByTagsFromContract(
+          this.aztecNode,
+          contractAddress,
+          tags,
+          anchorBlockHash,
+          { fromBlock: group.fromBlock, toBlock: group.toBlock, includeEffects: true },
+        );
+        group.entries.forEach((entry, i) => {
+          resultsPerRequest[entry.index] = results[i];
+        });
+      }),
+    );
+
+    return resultsPerRequest;
+  }
+
+  async #fetchPrivateLogs(requests: LogRetrievalRequest[], anchorBlockHash: BlockHash): Promise<LogResult[][]> {
+    const indices = requests.flatMap((r, i) => (r.source !== LogSource.PUBLIC ? [i] : []));
+    if (indices.length === 0) {
+      return requests.map(() => []);
+    }
+
+    const resultsPerRequest: LogResult[][] = requests.map(() => []);
+    const groups = LogService.#groupByRange(indices.map(i => ({ index: i, request: requests[i] })));
+
+    await Promise.all(
+      Array.from(groups.values()).map(async group => {
+        const siloedTags = await Promise.all(
+          group.entries.map(e => SiloedTag.computeFromTagAndApp(e.request.tag, e.request.contractAddress)),
+        );
+        const results = await getAllPrivateLogsByTags(this.aztecNode, siloedTags, anchorBlockHash, {
+          fromBlock: group.fromBlock,
+          toBlock: group.toBlock,
+          includeEffects: true,
+        });
+        group.entries.forEach((entry, i) => {
+          resultsPerRequest[entry.index] = results[i];
+        });
+      }),
+    );
+
+    return resultsPerRequest;
+  }
+
+  /**
+   * Groups requests by their (fromBlock, toBlock) range so each distinct range becomes a single node call with
+   * the range pushed down into the query (no in-memory filter).
+   */
+  static #groupByRange(
+    entries: Array<{ index: number; request: LogRetrievalRequest }>,
+  ): Map<RangeKey, { fromBlock?: BlockNumber; toBlock?: BlockNumber; entries: typeof entries }> {
+    const groups = new Map<RangeKey, { fromBlock?: BlockNumber; toBlock?: BlockNumber; entries: typeof entries }>();
+    for (const entry of entries) {
+      const fromBlock = entry.request.fromBlock.value;
+      const toBlock = entry.request.toBlock.value;
+      const key = rangeKey(fromBlock, toBlock);
+      const existing = groups.get(key);
+      if (existing) {
+        existing.entries.push(entry);
+      } else {
+        groups.set(key, { fromBlock, toBlock, entries: [entry] });
+      }
+    }
+    return groups;
+  }
+
+  static #toLogRetrievalResponse(log: LogResult): LogRetrievalResponse {
+    // includeEffects: true was used, so noteHashes and nullifiers are populated. Every tx has at least one nullifier
+    // (the first nullifier derived from the tx hash); empty here would indicate a buggy node.
+    const noteHashes = log.noteHashes!;
+    const nullifiers = log.nullifiers!;
+    if (nullifiers.length === 0) {
+      throw new Error(`Log for tx ${log.txHash} returned no nullifiers from the node`);
+    }
+    return {
+      // Skip the tag, and clip to the wire cap: public logs can exceed PRIVATE_LOG_CIPHERTEXT_LEN, which is the fixed
+      // size of the oracle's BoundedVec slot. A no-op for private logs, which are already within the cap.
+      logPayload: log.logData.slice(1, 1 + PRIVATE_LOG_CIPHERTEXT_LEN),
+      txHash: log.txHash,
+      uniqueNoteHashesInTx: noteHashes,
+      firstNullifierInTx: nullifiers[0],
+      blockNumber: log.blockNumber,
+      blockTimestamp: log.blockTimestamp,
+      blockHash: log.blockHash,
+    };
+  }
+
+  public async fetchTaggedLogs(
+    contractAddress: AztecAddress,
+    recipient: AztecAddress,
+    providedSecrets: AppTaggingSecret[],
+  ): Promise<PendingTaggedLog[]> {
+    this.log.verbose(
+      `Fetching tagged logs for contract ${contractAddress.toString()} and recipient ${recipient.toString()}`,
+    );
+
+    const l2Tips = await this.l2TipsStore.getL2Tips();
+
+    // Secrets PXE can enumerate for this recipient (senders via ECDH + pre-shared store secrets), plus any the app
+    // supplies explicitly for secrets PXE cannot enumerate itself (e.g. handshake-derived ones). The latter arrive
+    // already computed and are searched even when the recipient's account is unknown locally, since they need no ECDH.
+    const combinedSecrets = [...(await this.#getPointDerivedSecrets(contractAddress, recipient)), ...providedSecrets];
+
+    // These sources can overlap (a sender that is also a PXE account, or a pre-shared secret that coincides with a
+    // sender-derived one), so we deduplicate the combined set.
+    const secrets = Array.from(new Map(combinedSecrets.map(secret => [secret.toString(), secret])).values());
+
+    const logs = await syncTaggedPrivateLogs(
+      secrets,
+      this.aztecNode,
+      this.recipientTaggingStore,
+      this.anchorBlockHeader,
+      l2Tips.finalized.block.number,
+      this.jobId,
+    );
+
+    return logs.map(log => {
+      const noteHashes = log.noteHashes!;
+      const nullifiers = log.nullifiers!;
+      if (nullifiers.length === 0) {
+        throw new Error(`Log for tx ${log.txHash} returned no nullifiers from the node`);
+      }
+      return {
+        log: log.logData,
+        context: new ResolvedTx(log.txHash, noteHashes, nullifiers[0], log.blockNumber, log.blockHash.toFr()),
+      };
+    });
+  }
+
+  /**
+   * Computes the tagging secrets PXE can enumerate for a recipient: one per known sender (via ECDH) plus any
+   * pre-shared secrets registered directly for the recipient. Each registered secret is scanned under the tag
+   * streams its kind can back. Deriving the sender-based secrets requires the recipient's address preimage and keys,
+   * so returns an empty array when those are unavailable. App-supplied secrets (e.g. derived from discovered
+   * handshakes) are handled separately by the caller and do not go through here.
+   */
+  async #getPointDerivedSecrets(contractAddress: AztecAddress, recipient: AztecAddress): Promise<AppTaggingSecret[]> {
+    const recipientCompleteAddress = await this.addressStore.getCompleteAddress(recipient);
+    if (!recipientCompleteAddress || !(await this.keyStore.hasAccount(recipient))) {
+      this.log.warn(
+        `Skipping sender-derived tag retrieval for ${recipient.toString()} due to unknown address preimage`,
+      );
+      return [];
+    }
+    const recipientIvsk = await this.keyStore.getMasterIncomingViewingSecretKey(recipient);
+
+    const [senderPoints, registeredSecrets] = await Promise.all([
+      this.#getSecretsForSenders(recipientCompleteAddress, recipientIvsk),
+      this.taggingSecretSourcesStore.getSharedSecretsForRecipient(recipient),
+    ]);
+    return Promise.all([
+      ...senderPoints.map(secret => AppTaggingSecret.computeDirectional(secret, contractAddress, recipient)),
+      ...registeredSecrets.flatMap(({ kind, secret }) => {
+        switch (kind) {
+          case 'arbitrary-secret':
+            return [AppTaggingSecret.computeDirectional(secret, contractAddress, recipient)];
+          case 'handshake':
+            // A handshake-backed sender tags messages with the bare app-siloed secret, one tag domain per delivery
+            // mode.
+            return [
+              AppTaggingSecret.computeAppSiloed(secret, contractAddress, AppTaggingSecretKind.UNCONSTRAINED),
+              AppTaggingSecret.computeAppSiloed(secret, contractAddress, AppTaggingSecretKind.CONSTRAINED),
+            ];
+        }
+      }),
+    ]);
+  }
+
+  async #getSecretsForSenders(
+    recipientCompleteAddress: CompleteAddress,
+    recipientIvsk: GrumpkinScalar,
+  ): Promise<Point[]> {
+    // We implicitly add all PXE accounts as senders, this helps us decrypt tags on notes that we send to ourselves
+    // (recipient = us, sender = us).
+    const allSenders = [...(await this.taggingSecretSourcesStore.getSenders()), ...(await this.keyStore.getAccounts())];
+
+    return Promise.all(
+      allSenders.map(async sender => {
+        const taggingSecretPoint = await computeSharedTaggingSecret(recipientCompleteAddress, recipientIvsk, sender);
+
+        if (!taggingSecretPoint) {
+          // Note that all senders originate from either the TaggingSecretSourcesStore or the KeyStore.
+          throw new Error(
+            `Failed to compute a tagging secret for sender ${sender} - this implies this is an invalid address, which should not happen as they have been previously registered in PXE.`,
+          );
+        }
+
+        return taggingSecretPoint;
+      }),
+    );
+  }
+}

package/src/messages/tx_resolver_service.ts

@@ -0,0 +1,69 @@
+import { uniqueBy } from '@aztec/foundation/collection';
+import { Fr } from '@aztec/foundation/curves/bn254';
+import type { AztecNode } from '@aztec/stdlib/interfaces/server';
+import { type IndexedTxEffect, TxHash } from '@aztec/stdlib/tx';
+
+import { ResolvedTx } from '../contract_function_simulator/noir-structs/resolved_tx.js';
+
+/** Resolves transaction hashes into their on-chain context (note hashes, first nullifier, and mined block). */
+export class TxResolverService {
+  constructor(private readonly aztecNode: AztecNode) {}
+
+  /**
+   * Resolves a list of tx hashes into their on-chain context.
+   *
+   * For each tx hash, looks up the corresponding tx effect and extracts the note hashes, first nullifier, and the
+   * number and hash of the block it was mined in. Returns `null` for tx hashes that are zero, not yet available, or in
+   * blocks beyond the anchor block.
+   */
+  async resolveTxs(txHashes: Fr[], anchorBlockNumber: number): Promise<(ResolvedTx | null)[]> {
+    const nonZeroTxHashes = txHashes.filter(h => !h.isZero()).map(h => TxHash.fromField(h));
+    const uniqueTxHashes = uniqueBy(nonZeroTxHashes, h => h.toString());
+    const fetched = await Promise.all(
+      uniqueTxHashes.map(h => this.aztecNode.getTxReceipt(h, { includeTxEffect: true })),
+    );
+    const txEffects = new Map(
+      uniqueTxHashes
+        .map((h, i): [string, IndexedTxEffect | undefined] => {
+          const receipt = fetched[i];
+          if (!receipt.isMined() || !receipt.txEffect) {
+            return [h.toString(), undefined];
+          }
+          return [
+            h.toString(),
+            {
+              data: receipt.txEffect,
+              l2BlockNumber: receipt.blockNumber,
+              l2BlockHash: receipt.blockHash,
+              txIndexInBlock: receipt.txIndexInBlock,
+              slotNumber: receipt.slotNumber,
+            },
+          ];
+        })
+        .filter((entry): entry is [string, IndexedTxEffect] => entry[1] !== undefined),
+    );
+
+    return txHashes.map(txHashField => {
+      const txHash = TxHash.fromField(txHashField);
+      const txEffect = txEffects.get(txHash.toString());
+      if (!txEffect || txEffect.l2BlockNumber > anchorBlockNumber) {
+        return null;
+      }
+
+      // Every tx has at least one nullifier (the first nullifier derived from the tx hash). Hitting this condition
+      // would mean a buggy node, but since we need to access data.nullifiers[0], the defensive check does no harm.
+      const data = txEffect.data;
+      if (data.nullifiers.length === 0) {
+        throw new Error(`Tx effect for ${txHash} has no nullifiers`);
+      }
+
+      return new ResolvedTx(
+        data.txHash,
+        data.noteHashes,
+        data.nullifiers[0],
+        txEffect.l2BlockNumber,
+        txEffect.l2BlockHash.toFr(),
+      );
+    });
+  }
+}

package/src/notes/index.ts

@@ -0,0 +1 @@
+export { NoteService } from './note_service.js';