feat: Verify registered artifact matches instance class id (#5297)

When registering a new contract instance with an artifact, validate that
the computed artifact class id matches the class id from the instance
preimage.

yarn-project/aztec.js/src/index.ts

@@ -82,7 +82,6 @@ export {
   AztecNode,
   Body,
   CompleteAddress,
-  ContractWithArtifact,
   ExtendedNote,
   FunctionCall,
   GrumpkinPrivateKey,

yarn-project/aztec.js/src/wallet/base_wallet.ts

@@ -1,6 +1,5 @@
 import {
   AuthWitness,
-  ContractWithArtifact,
   ExtendedNote,
   FunctionCall,
   GetUnencryptedLogsResponse,
@@ -76,7 +75,10 @@ export abstract class BaseWallet implements Wallet {
   getRecipient(address: AztecAddress): Promise<CompleteAddress | undefined> {
     return this.pxe.getRecipient(address);
   }
-  registerContract(contract: ContractWithArtifact): Promise<void> {
+  registerContract(contract: {
+    /** Instance */ instance: ContractInstanceWithAddress;
+    /** Associated artifact */ artifact?: ContractArtifact;
+  }): Promise<void> {
     return this.pxe.registerContract(contract);
   }
   registerContractClass(artifact: ContractArtifact): Promise<void> {

yarn-project/circuit-types/src/interfaces/index.ts

@@ -1,6 +1,5 @@
 export * from './aztec-node.js';
 export * from './pxe.js';
-export * from './contract-with-artifact.js';
 export * from './sync-status.js';
 export * from './configs.js';
 export * from './nullifier_tree.js';

yarn-project/circuit-types/src/interfaces/pxe.ts

@@ -11,7 +11,6 @@ import { NoteFilter } from '../notes/note_filter.js';
 import { Tx, TxHash, TxReceipt } from '../tx/index.js';
 import { TxEffect } from '../tx_effect.js';
 import { TxExecutionRequest } from '../tx_execution_request.js';
-import { ContractWithArtifact } from './contract-with-artifact.js';
 import { SyncStatus } from './sync-status.js';
 
 // docs:start:pxe-interface
@@ -112,9 +111,9 @@ export interface PXE {
    * deploying a contract. Dapps that wish to interact with contracts already deployed should register
    * these contracts in their users' PXE Service through this method.
    *
-   * @param contract - An object containing contract artifact and instance.
+   * @param contract - A contract instance to register, with an optional artifact which can be omitted if the contract class has already been registered.
    */
-  registerContract(contract: ContractWithArtifact): Promise<void>;
+  registerContract(contract: { instance: ContractInstanceWithAddress; artifact?: ContractArtifact }): Promise<void>;
 
   /**
    * Retrieves the addresses of contracts added to this PXE Service.

yarn-project/circuit-types/src/mocks.ts

@@ -4,6 +4,8 @@ import {
   MAX_PUBLIC_CALL_STACK_LENGTH_PER_TX,
   MAX_REVERTIBLE_PUBLIC_CALL_STACK_LENGTH_PER_TX,
   Proof,
+  computeContractClassId,
+  getContractClassFromArtifact,
 } from '@aztec/circuits.js';
 import { ContractArtifact } from '@aztec/foundation/abi';
 import { makeTuple } from '@aztec/foundation/array';
@@ -13,7 +15,6 @@ import { Fr } from '@aztec/foundation/fields';
 import { Tuple } from '@aztec/foundation/serialize';
 import { ContractInstanceWithAddress, SerializableContractInstance } from '@aztec/types/contracts';
 
-import { ContractWithArtifact } from './interfaces/index.js';
 import { FunctionL2Logs, Note, TxL2Logs } from './logs/index.js';
 import { makePrivateKernelTailCircuitPublicInputs, makePublicCallRequest } from './mocks_to_purge.js';
 import { ExtendedNote } from './notes/index.js';
@@ -62,10 +63,11 @@ export const randomContractArtifact = (): ContractArtifact => ({
 export const randomContractInstanceWithAddress = (opts: { contractClassId?: Fr } = {}): ContractInstanceWithAddress =>
   SerializableContractInstance.random(opts).withAddress(AztecAddress.random());
 
-export const randomDeployedContract = (): ContractWithArtifact => ({
-  artifact: randomContractArtifact(),
-  instance: randomContractInstanceWithAddress(),
-});
+export const randomDeployedContract = () => {
+  const artifact = randomContractArtifact();
+  const contractClassId = computeContractClassId(getContractClassFromArtifact(artifact));
+  return { artifact, instance: randomContractInstanceWithAddress({ contractClassId }) };
+};
 
 export const randomExtendedNote = ({
   note = Note.random(),

yarn-project/pxe/src/pxe_service/pxe_service.ts

@@ -1,7 +1,6 @@
 import {
   AuthWitness,
   AztecNode,
-  ContractWithArtifact,
   ExtendedNote,
   FunctionCall,
   GetUnencryptedLogsResponse,
@@ -32,7 +31,6 @@ import {
   PartialAddress,
   PrivateKernelTailCircuitPublicInputs,
   PublicCallRequest,
-  computeArtifactHash,
   computeContractClassId,
   getContractClassFromArtifact,
 } from '@aztec/circuits.js';
@@ -43,7 +41,6 @@ import { Fr } from '@aztec/foundation/fields';
 import { SerialQueue } from '@aztec/foundation/fifo';
 import { DebugLogger, createDebugLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
-import { required } from '@aztec/foundation/validation';
 import {
   AcirSimulator,
   ExecutionResult,
@@ -222,22 +219,31 @@ export class PXEService implements PXE {
     this.log.info(`Added contract class ${artifact.name} with id ${contractClassId}`);
   }
 
-  public async registerContract(contract: ContractWithArtifact) {
-    const instance = contract.instance;
-    const artifact =
-      'artifact' in contract
-        ? contract.artifact
-        : required(
-            await this.db.getContractArtifact(contract.contractClassId),
-            `Unknown artifact for class id ${contract.contractClassId} when registering ${instance.address}`,
-          );
-    const artifactHash = computeArtifactHash(artifact);
-    const contractClassId = computeContractClassId(getContractClassFromArtifact({ ...artifact, artifactHash }));
-    await this.db.addContractArtifact(contractClassId, artifact);
+  public async registerContract(contract: { instance: ContractInstanceWithAddress; artifact?: ContractArtifact }) {
+    const { instance } = contract;
+    let { artifact } = contract;
+
+    if (artifact) {
+      // If the user provides an artifact, validate it against the expected class id and register it
+      const contractClassId = computeContractClassId(getContractClassFromArtifact(artifact));
+      if (!contractClassId.equals(instance.contractClassId)) {
+        throw new Error(
+          `Artifact does not match expected class id (computed ${contractClassId} but instance refers to ${instance.contractClassId})`,
+        );
+      }
+      await this.db.addContractArtifact(contractClassId, artifact);
+    } else {
+      // Otherwise, make sure there is an artifact already registered for that class id
+      artifact = await this.db.getContractArtifact(instance.contractClassId);
+      if (!artifact) {
+        throw new Error(
+          `Missing contract artifact for class id ${instance.contractClassId} for contract ${instance.address}`,
+        );
+      }
+    }
+
+    this.log.info(`Added contract ${artifact.name} at ${instance.address.toString()}`);
     await this.db.addContractInstance(instance);
-    const hasPortal = instance.portalContractAddress && !instance.portalContractAddress.isZero();
-    const portalInfo = hasPortal ? ` with portal ${instance.portalContractAddress.toChecksumString()}` : '';
-    this.log.info(`Added contract ${artifact.name} at ${instance.address.toString()}${portalInfo}`);
     await this.synchronizer.reprocessDeferredNotesForContract(instance.address);
   }
 

yarn-project/pxe/src/pxe_service/test/pxe_test_suite.ts

@@ -1,5 +1,4 @@
 import {
-  ContractWithArtifact,
   PXE,
   TxExecutionRequest,
   randomContractArtifact,
@@ -90,7 +89,7 @@ export const pxeTestSuite = (testName: string, pxeSetup: () => Promise<PXE>) =>
     });
 
     it('successfully adds a contract', async () => {
-      const contracts: ContractWithArtifact[] = [randomDeployedContract(), randomDeployedContract()];
+      const contracts = [randomDeployedContract(), randomDeployedContract()];
       for (const contract of contracts) {
         await pxe.registerContract(contract);
       }
@@ -109,15 +108,19 @@ export const pxeTestSuite = (testName: string, pxeSetup: () => Promise<PXE>) =>
       await pxe.registerContractClass(artifact);
       expect(await pxe.getContractClass(contractClassId)).toEqual(contractClass);
 
-      await pxe.registerContract({ contractClassId, instance });
+      await pxe.registerContract({ instance });
       expect(await pxe.getContractInstance(instance.address)).toEqual(instance);
     });
 
     it('refuses to register a contract with a class that has not been registered', async () => {
       const instance = randomContractInstanceWithAddress();
-      await expect(pxe.registerContract({ contractClassId: Fr.random(), instance })).rejects.toThrow(
-        /Unknown artifact/i,
-      );
+      await expect(pxe.registerContract({ instance })).rejects.toThrow(/Missing contract artifact/i);
+    });
+
+    it('refuses to register a contract with an artifact with mismatching class id', async () => {
+      const artifact = randomContractArtifact();
+      const instance = randomContractInstanceWithAddress();
+      await expect(pxe.registerContract({ instance, artifact })).rejects.toThrow(/Artifact does not match/i);
     });
 
     it('throws when simulating a tx targeting public entrypoint', async () => {