refactor: move validation and error handling out of transport layer

yarn-project/pxe/src/contract_function_simulator/oracle/interfaces.ts

@@ -70,25 +70,19 @@ export interface IUtilityExecutionOracle {
   getNoteHashMembershipWitness(
     anchorBlockHash: BlockHash,
     noteHash: Fr,
-  ): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined>;
+  ): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>>;
   getBlockHashMembershipWitness(
     anchorBlockHash: BlockHash,
     blockHash: BlockHash,
   ): Promise<MembershipWitness<typeof ARCHIVE_HEIGHT> | undefined>;
-  getNullifierMembershipWitness(
-    anchorBlockHash: BlockHash,
-    nullifier: Fr,
-  ): Promise<NullifierMembershipWitness | undefined>;
-  getPublicDataWitness(anchorBlockHash: BlockHash, leafSlot: Fr): Promise<PublicDataWitness | undefined>;
-  getLowNullifierMembershipWitness(
-    anchorBlockHash: BlockHash,
-    nullifier: Fr,
-  ): Promise<NullifierMembershipWitness | undefined>;
-  getBlockHeader(blockNumber: BlockNumber): Promise<BlockHeader | undefined>;
+  getNullifierMembershipWitness(anchorBlockHash: BlockHash, nullifier: Fr): Promise<NullifierMembershipWitness>;
+  getPublicDataWitness(anchorBlockHash: BlockHash, leafSlot: Fr): Promise<PublicDataWitness>;
+  getLowNullifierMembershipWitness(anchorBlockHash: BlockHash, nullifier: Fr): Promise<NullifierMembershipWitness>;
+  getBlockHeader(blockNumber: BlockNumber): Promise<BlockHeader>;
   getPublicKeysAndPartialAddress(
     account: AztecAddress,
   ): Promise<{ publicKeys: PublicKeys; partialAddress: PartialAddress } | undefined>;
-  getAuthWitness(messageHash: Fr): Promise<Fr[] | undefined>;
+  getAuthWitness(messageHash: Fr): Promise<Fr[]>;
   getNotes(
     owner: AztecAddress | undefined,
     storageSlot: Fr,
@@ -139,7 +133,7 @@ export interface IUtilityExecutionOracle {
     numEntries: number,
     scope: AztecAddress,
   ): Promise<void>;
-  decryptAes128(ciphertext: Buffer, iv: Buffer, symKey: Buffer): Promise<Buffer>;
+  decryptAes128(ciphertext: Buffer, iv: Buffer, symKey: Buffer): Promise<Buffer | undefined>;
   getSharedSecrets(address: AztecAddress, ephPksSlot: Fr, contractAddress: AztecAddress): Promise<Fr>;
   setContractSyncCacheInvalid(contractAddress: AztecAddress, scopes: AztecAddress[]): void;
   emitOffchainEffect(data: Fr[]): Promise<void>;

yarn-project/pxe/src/contract_function_simulator/oracle/oracle.ts

@@ -238,11 +238,6 @@ export class Oracle {
     const parsedNoteHash = Fr.fromString(noteHash);
 
     const witness = await this.handlerAsUtility().getNoteHashMembershipWitness(parsedAnchorBlockHash, parsedNoteHash);
-    if (!witness) {
-      throw new Error(
-        `Note hash ${noteHash} not found in the note hash tree at anchor block hash ${parsedAnchorBlockHash.toString()}.`,
-      );
-    }
     return witness.toNoirRepresentation();
   }
 
@@ -268,11 +263,6 @@ export class Oracle {
     const parsedNullifier = Fr.fromString(nullifier);
 
     const witness = await this.handlerAsUtility().getNullifierMembershipWitness(parsedBlockHash, parsedNullifier);
-    if (!witness) {
-      throw new Error(
-        `Nullifier witness not found for nullifier ${parsedNullifier} at block hash ${parsedBlockHash.toString()}.`,
-      );
-    }
     return witness.toNoirRepresentation();
   }
 
@@ -285,11 +275,6 @@ export class Oracle {
     const parsedNullifier = Fr.fromString(nullifier);
 
     const witness = await this.handlerAsUtility().getLowNullifierMembershipWitness(parsedBlockHash, parsedNullifier);
-    if (!witness) {
-      throw new Error(
-        `Low nullifier witness not found for nullifier ${parsedNullifier} at block hash ${parsedBlockHash.toString()}.`,
-      );
-    }
     return witness.toNoirRepresentation();
   }
 
@@ -302,11 +287,6 @@ export class Oracle {
     const parsedLeafSlot = Fr.fromString(leafSlot);
 
     const witness = await this.handlerAsUtility().getPublicDataWitness(parsedBlockHash, parsedLeafSlot);
-    if (!witness) {
-      throw new Error(
-        `Public data witness not found for slot ${parsedLeafSlot} at block hash ${parsedBlockHash.toString()}.`,
-      );
-    }
     return witness.toNoirRepresentation();
   }
 
@@ -315,19 +295,13 @@ export class Oracle {
     const parsedBlockNumber = Fr.fromString(blockNumber).toNumber();
 
     const header = await this.handlerAsUtility().getBlockHeader(BlockNumber(parsedBlockNumber));
-    if (!header) {
-      throw new Error(`Block header not found for block ${parsedBlockNumber}.`);
-    }
     return header.toFields().map(toACVMField);
   }
 
   // eslint-disable-next-line camelcase
   async aztec_utl_getAuthWitness([messageHash]: ACVMField[]): Promise<ACVMField[][]> {
     const messageHashField = Fr.fromString(messageHash);
     const witness = await this.handlerAsUtility().getAuthWitness(messageHashField);
-    if (!witness) {
-      throw new Error(`Unknown auth witness for message hash ${messageHashField}`);
-    }
     return [witness.map(toACVMField)];
   }
 
@@ -782,11 +756,11 @@ export class Oracle {
     const symKeyBuffer = fromUintArray(symKey, 8);
 
     // Noir Option<BoundedVec> is encoded as [is_some: Field, storage: Field[], length: Field].
-    try {
-      const plaintext = await this.handlerAsUtility().decryptAes128(ciphertext, ivBuffer, symKeyBuffer);
+    const plaintext = await this.handlerAsUtility().decryptAes128(ciphertext, ivBuffer, symKeyBuffer);
+    if (plaintext) {
       const [storage, length] = bufferToBoundedVec(plaintext, ciphertextBVecStorage.length);
       return [toACVMField(1), storage, length];
-    } catch {
+    } else {
       const zeroStorage = Array(ciphertextBVecStorage.length).fill(toACVMField(0));
       return [toACVMField(0), zeroStorage, toACVMField(0)];
     }

yarn-project/pxe/src/contract_function_simulator/oracle/utility_execution_oracle.ts

@@ -201,13 +201,19 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
    * @param noteHash - The note hash to find in the note hash tree.
    * @returns The membership witness containing the leaf index and sibling path
    */
-  public getNoteHashMembershipWitness(
+  public async getNoteHashMembershipWitness(
     blockHash: BlockHash,
     noteHash: Fr,
-  ): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT> | undefined> {
-    return this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
+  ): Promise<MembershipWitness<typeof NOTE_HASH_TREE_HEIGHT>> {
+    const witness = await this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
       this.aztecNode.getNoteHashMembershipWitness(blockHash, noteHash),
     );
+    if (!witness) {
+      throw new Error(
+        `Note hash ${noteHash} not found in the note hash tree at anchor block hash ${blockHash.toString()}.`,
+      );
+    }
+    return witness;
   }
 
   /**
@@ -239,13 +245,14 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
    * @param nullifier - Nullifier we try to find witness for.
    * @returns The nullifier membership witness (if found).
    */
-  public getNullifierMembershipWitness(
-    blockHash: BlockHash,
-    nullifier: Fr,
-  ): Promise<NullifierMembershipWitness | undefined> {
-    return this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
+  public async getNullifierMembershipWitness(blockHash: BlockHash, nullifier: Fr): Promise<NullifierMembershipWitness> {
+    const witness = await this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
       this.aztecNode.getNullifierMembershipWitness(blockHash, nullifier),
     );
+    if (!witness) {
+      throw new Error(`Nullifier witness not found for nullifier ${nullifier} at block hash ${blockHash.toString()}.`);
+    }
+    return witness;
   }
 
   /**
@@ -257,13 +264,19 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
    * list structure" of leaves and proving that a lower nullifier is pointing to a bigger next value than the nullifier
    * we are trying to prove non-inclusion for.
    */
-  public getLowNullifierMembershipWitness(
+  public async getLowNullifierMembershipWitness(
     blockHash: BlockHash,
     nullifier: Fr,
-  ): Promise<NullifierMembershipWitness | undefined> {
-    return this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
+  ): Promise<NullifierMembershipWitness> {
+    const witness = await this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
       this.aztecNode.getLowNullifierMembershipWitness(blockHash, nullifier),
     );
+    if (!witness) {
+      throw new Error(
+        `Low nullifier witness not found for nullifier ${nullifier} at block hash ${blockHash.toString()}.`,
+      );
+    }
+    return witness;
   }
 
   /**
@@ -272,18 +285,22 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
    * @param leafSlot - The slot of the public data tree to get the witness for.
    * @returns - The witness
    */
-  public getPublicDataWitness(blockHash: BlockHash, leafSlot: Fr): Promise<PublicDataWitness | undefined> {
-    return this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
+  public async getPublicDataWitness(blockHash: BlockHash, leafSlot: Fr): Promise<PublicDataWitness> {
+    const witness = await this.#queryWithBlockHashNotAfterAnchor(blockHash, () =>
       this.aztecNode.getPublicDataWitness(blockHash, leafSlot),
     );
+    if (!witness) {
+      throw new Error(`Public data witness not found for slot ${leafSlot} at block hash ${blockHash.toString()}.`);
+    }
+    return witness;
   }
 
   /**
    * Fetches a block header of a given block.
    * @param blockNumber - The number of a block of which to get the block header.
    * @returns Block extracted from a block with block number `blockNumber`.
    */
-  public async getBlockHeader(blockNumber: BlockNumber): Promise<BlockHeader | undefined> {
+  public async getBlockHeader(blockNumber: BlockNumber): Promise<BlockHeader> {
     const anchorBlockNumber = this.anchorBlockHeader.getBlockNumber();
     if (blockNumber > anchorBlockNumber) {
       throw new Error(`Block number ${blockNumber} is higher than current block ${anchorBlockNumber}`);
@@ -295,7 +312,10 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
     }
 
     const block = await this.aztecNode.getBlock(blockNumber);
-    return block?.header;
+    if (!block?.header) {
+      throw new Error(`Block header not found for block ${blockNumber}.`);
+    }
+    return block.header;
   }
 
   /**
@@ -342,8 +362,12 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
    * @param messageHash - Hash of the message to authenticate.
    * @returns Authentication witness for the requested message hash, or undefined if not found.
    */
-  public getAuthWitness(messageHash: Fr): Promise<Fr[] | undefined> {
-    return Promise.resolve(this.authWitnesses.find(w => w.requestHash.equals(messageHash))?.witness);
+  public getAuthWitness(messageHash: Fr): Promise<Fr[]> {
+    const witness = this.authWitnesses.find(w => w.requestHash.equals(messageHash))?.witness;
+    if (!witness) {
+      throw new Error(`Unknown auth witness for message hash ${messageHash}`);
+    }
+    return Promise.resolve(witness);
   }
 
   /**
@@ -671,9 +695,13 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra
   }
 
   // TODO(#11849): consider replacing this oracle with a pure Noir implementation of aes decryption.
-  public decryptAes128(ciphertext: Buffer, iv: Buffer, symKey: Buffer): Promise<Buffer> {
-    const aes128 = new Aes128();
-    return aes128.decryptBufferCBC(ciphertext, iv, symKey);
+  public async decryptAes128(ciphertext: Buffer, iv: Buffer, symKey: Buffer): Promise<Buffer | undefined> {
+    try {
+      const aes128 = new Aes128();
+      return await aes128.decryptBufferCBC(ciphertext, iv, symKey);
+    } catch {
+      return undefined;
+    }
   }
 
   /**

yarn-project/txe/src/constants.ts

@@ -1,3 +1,9 @@
+import { PRIVATE_LOG_CIPHERTEXT_LEN } from '@aztec/constants';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 
 export const DEFAULT_ADDRESS = AztecAddress.fromNumber(42);
+
+// Arbitrarily set at 64 because we need a bound. Nothing inherent about it.
+export const MAX_OFFCHAIN_EFFECTS_PER_TXE_QUERY = 64;
+// Must match MAX_OFFCHAIN_EFFECT_LEN in noir-projects/aztec-nr/aztec/src/test/helpers/txe_oracles.nr.
+export const MAX_OFFCHAIN_EFFECT_LEN = 2 + PRIVATE_LOG_CIPHERTEXT_LEN;

yarn-project/txe/src/oracle/interfaces.ts

@@ -38,6 +38,11 @@ export interface IAvmExecutionOracle {
   nullifierExists(siloedNullifier: Fr): Promise<boolean>;
   storageWrite(slot: Fr, value: Fr): Promise<void>;
   storageRead(slot: Fr, contractAddress: AztecAddress): Promise<Fr>;
+  returndataSize(): Promise<Fr>;
+  returndataCopy(rdOffset: number, copySize: number): Promise<Fr[]>;
+  call(l2Gas: number, daGas: number, address: AztecAddress, argsLength: number, args: Fr[]): Promise<Fr[]>;
+  staticCall(l2Gas: number, daGas: number, address: AztecAddress, argsLength: number, args: Fr[]): Promise<Fr[]>;
+  successCopy(): Promise<Fr>;
 }
 
 /**

yarn-project/txe/src/oracle/txe_oracle_public_context.ts

@@ -122,6 +122,36 @@ export class TXEOraclePublicContext implements IAvmExecutionOracle {
     return value;
   }
 
+  returndataSize(): Promise<Fr> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  returndataCopy(_rdOffset: number, _copySize: number): Promise<Fr[]> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  call(_l2Gas: number, _daGas: number, _address: AztecAddress, _argsLength: number, _args: Fr[]): Promise<Fr[]> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  staticCall(_l2Gas: number, _daGas: number, _address: AztecAddress, _argsLength: number, _args: Fr[]): Promise<Fr[]> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
+  successCopy(): Promise<Fr> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::public_context`, use `public_call` instead',
+    );
+  }
+
   async close(): Promise<L2Block> {
     this.logger.debug('Exiting Public Context, building block with collected side effects', {
       blockNumber: this.globalVariables.blockNumber,

yarn-project/txe/src/oracle/txe_oracle_top_level_context.ts

@@ -1,4 +1,8 @@
-import { CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS, NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP } from '@aztec/constants';
+import {
+  CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
+  MAX_PRIVATE_LOGS_PER_TX,
+  NUMBER_OF_L1_L2_MESSAGES_PER_ROLLUP,
+} from '@aztec/constants';
 import { BlockNumber } from '@aztec/foundation/branded-types';
 import { Schnorr } from '@aztec/foundation/crypto/schnorr';
 import { Fr } from '@aztec/foundation/curves/bn254';
@@ -175,11 +179,16 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl
 
     const txEffects = block!.body.txEffects[0];
 
+    const privateLogs = txEffects.privateLogs;
+    if (privateLogs.length > MAX_PRIVATE_LOGS_PER_TX) {
+      throw new Error(`${privateLogs.length} private logs exceed max ${MAX_PRIVATE_LOGS_PER_TX}`);
+    }
+
     return {
       txHash: txEffects.txHash,
       noteHashes: txEffects.noteHashes,
       nullifiers: txEffects.nullifiers,
-      privateLogs: txEffects.privateLogs,
+      privateLogs,
     };
   }
 

yarn-project/txe/src/oracle/txe_private_execution_oracle.ts

@@ -0,0 +1,30 @@
+import type { Fr } from '@aztec/foundation/curves/bn254';
+import { PrivateExecutionOracle } from '@aztec/pxe/simulator';
+import type { FunctionSelector } from '@aztec/stdlib/abi';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+
+/**
+ * TXE-specific subclass of PrivateExecutionOracle that forbids operations not supported in
+ * TestEnvironment::private_context. TXE uses dedicated oracle flows (e.g. private_call) instead.
+ */
+export class TXEPrivateExecutionOracle extends PrivateExecutionOracle {
+  override callPrivateFunction(
+    _targetContractAddress: AztecAddress,
+    _functionSelector: FunctionSelector,
+    _argsHash: Fr,
+    _sideEffectCounter: number,
+    _isStaticCall: boolean,
+  ): Promise<{ endSideEffectCounter: Fr; returnsHash: Fr }> {
+    throw new Error(
+      'Contract calls are forbidden inside a `TestEnvironment::private_context`, use `private_call` instead',
+    );
+  }
+
+  override assertValidPublicCalldata(_calldataHash: Fr): Promise<void> {
+    throw new Error('Enqueueing public calls is not supported in TestEnvironment::private_context');
+  }
+
+  override notifyRevertiblePhaseStart(_minRevertibleSideEffectCounter: number): Promise<void> {
+    throw new Error('Enqueueing public calls is not supported in TestEnvironment::private_context');
+  }
+}