Add oracle call to retrieve public key for a given address (#1027)

* Adds oracle call to fetch the public key for an address

* Placeholder for address validation

* Temporarily remove sanity check

circuits/cpp/src/aztec3/circuits/abis/c_bind.cpp

@@ -320,6 +320,33 @@ WASM_EXPORT void abis__compute_contract_address(uint8_t const* point_data_buf,
     NT::fr::serialize_to_buffer(contract_address, output);
 }
 
+/**
+ * @brief Compute a contract address from its partial contract address
+ * This is a WASM-export that can be called from Typescript.
+ *
+ * @details Computes a contract address by hashing the deployers public key along with the previously computed partial
+ * address Return the serialized results in the `output` buffer.
+ *
+ * @param point_data_buf point data struct as a buffer of bytes
+ * @param partial_address_data_buf partial contract address
+ * @param output buffer that will contain the output. The serialized contract address.
+ */
+WASM_EXPORT void abis__compute_contract_address_from_partial(uint8_t const* point_data_buf,
+                                                             uint8_t const* partial_address_data_buf,
+                                                             uint8_t* output)
+{
+    Point<NT> deployer_public_key;
+    NT::fr partial_address;
+
+    read(point_data_buf, deployer_public_key);
+    read(partial_address_data_buf, partial_address);
+
+    NT::fr const contract_address =
+        aztec3::circuits::compute_contract_address_from_partial(deployer_public_key, partial_address);
+
+    NT::fr::serialize_to_buffer(contract_address, output);
+}
+
 /**
  * @brief Compute a partial contract address
  * This is a WASM-export that can be called from Typescript.

circuits/cpp/src/aztec3/circuits/hash.hpp

@@ -57,21 +57,29 @@ template <typename NCT> typename NCT::fr compute_partial_contract_address(typena
     return NCT::compress(inputs, aztec3::GeneratorIndex::PARTIAL_CONTRACT_ADDRESS);
 }
 
+template <typename NCT>
+typename NCT::address compute_contract_address_from_partial(Point<NCT> const& point, typename NCT::fr partial_address)
+{
+    using fr = typename NCT::fr;
+    using address = typename NCT::address;
+
+    std::vector<fr> const inputs = {
+        point.x.fields[0], point.x.fields[1], point.y.fields[0], point.y.fields[1], partial_address,
+    };
+    return address(NCT::compress(inputs, aztec3::GeneratorIndex::CONTRACT_ADDRESS));
+}
+
 template <typename NCT> typename NCT::address compute_contract_address(Point<NCT> const& point,
                                                                        typename NCT::fr contract_address_salt,
                                                                        typename NCT::fr function_tree_root,
                                                                        typename NCT::fr constructor_hash)
 {
     using fr = typename NCT::fr;
-    using address = typename NCT::address;
 
     const fr partial_address =
         compute_partial_contract_address<NCT>(contract_address_salt, function_tree_root, constructor_hash);
 
-    std::vector<fr> const inputs = {
-        point.x.fields[0], point.x.fields[1], point.y.fields[0], point.y.fields[1], partial_address,
-    };
-    return address(NCT::compress(inputs, aztec3::GeneratorIndex::CONTRACT_ADDRESS));
+    return compute_contract_address_from_partial(point, partial_address);
 }
 
 template <typename NCT>

yarn-project/acir-simulator/src/acvm/acvm.ts

@@ -38,6 +38,7 @@ type ORACLE_NAMES =
   | 'getL1ToL2Message'
   | 'emitEncryptedLog'
   | 'emitUnencryptedLog'
+  | 'getPublicKey'
   | 'debugLog';
 
 /**

yarn-project/acir-simulator/src/client/db_oracle.ts

@@ -3,6 +3,7 @@ import { EthAddress } from '@aztec/foundation/eth-address';
 import { Fr, Point } from '@aztec/foundation/fields';
 import { FunctionAbi } from '@aztec/foundation/abi';
 import { CommitmentsDB } from '../index.js';
+import { PartialContractAddress } from '@aztec/types';
 
 /**
  * The format that noir contracts use to get notes.
@@ -57,6 +58,7 @@ export interface CommitmentDataOracleInputs {
  * The database oracle interface.
  */
 export interface DBOracle extends CommitmentsDB {
+  getPublicKey(address: AztecAddress): Promise<[Point, PartialContractAddress]>;
   getSecretKey(contractAddress: AztecAddress, pubKey: Point): Promise<Buffer>;
   getNotes(
     contractAddress: AztecAddress,

yarn-project/acir-simulator/src/client/private_execution.test.ts

@@ -10,13 +10,13 @@ import {
   PublicCallRequest,
   TxContext,
 } from '@aztec/circuits.js';
-import { computeSecretMessageHash, siloCommitment } from '@aztec/circuits.js/abis';
+import { computeContractAddressFromPartial, computeSecretMessageHash, siloCommitment } from '@aztec/circuits.js/abis';
 import { Grumpkin, pedersenPlookupCommitInputs } from '@aztec/circuits.js/barretenberg';
 import { FunctionAbi } from '@aztec/foundation/abi';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { toBufferBE } from '@aztec/foundation/bigint-buffer';
 import { EthAddress } from '@aztec/foundation/eth-address';
-import { Fr } from '@aztec/foundation/fields';
+import { Coordinate, Fr, Point } from '@aztec/foundation/fields';
 import { DebugLogger, createDebugLogger } from '@aztec/foundation/log';
 import { AppendOnlyTree, Pedersen, StandardTree, newTree } from '@aztec/merkle-tree';
 import {
@@ -27,7 +27,7 @@ import {
   ZkTokenContractAbi,
 } from '@aztec/noir-contracts/examples';
 import { PackedArguments, TxExecutionRequest } from '@aztec/types';
-import { mock } from 'jest-mock-extended';
+import { MockProxy, mock } from 'jest-mock-extended';
 import { default as levelup } from 'levelup';
 import { default as memdown, type MemDown } from 'memdown';
 import { encodeArguments } from '../abi_coder/index.js';
@@ -40,7 +40,7 @@ const createMemDown = () => (memdown as any)() as MemDown<any, any>;
 
 describe('Private Execution test suite', () => {
   let circuitsWasm: CircuitsWasm;
-  let oracle: ReturnType<typeof mock<DBOracle>>;
+  let oracle: MockProxy<DBOracle>;
   let acirSimulator: AcirSimulator;
   let historicRoots = PrivateHistoricTreeRoots.empty();
   let logger: DebugLogger;
@@ -470,4 +470,23 @@ describe('Private Execution test suite', () => {
       );
     });
   });
+
+  describe('get public key', () => {
+    it('gets the public key for an address', async () => {
+      // Tweak the contract ABI so we can extract return values
+      const abi = TestContractAbi.functions.find(f => f.name === 'getPublicKey')!;
+      abi.returnTypes = [{ kind: 'field' }, { kind: 'field' }];
+
+      // Generate a partial address, pubkey, and resulting address
+      const partialAddress = Fr.random();
+      const pubKey = new Point(new Coordinate([Fr.random(), Fr.ZERO]), new Coordinate([Fr.random(), Fr.ZERO]));
+      const wasm = await CircuitsWasm.get();
+      const address = computeContractAddressFromPartial(wasm, pubKey, partialAddress);
+      const args = [address];
+
+      oracle.getPublicKey.mockResolvedValue([pubKey, partialAddress]);
+      const result = await runSimulator({ origin: AztecAddress.random(), abi, args });
+      expect(result.returnValues).toEqual([pubKey.x.toBigInt(), pubKey.y.toBigInt()]);
+    });
+  });
 });

yarn-project/acir-simulator/src/client/private_execution.ts

@@ -81,6 +81,11 @@ export class PrivateFunctionExecution {
             ),
           ),
         ),
+      getPublicKey: async ([acvmAddress]) => {
+        const address = frToAztecAddress(fromACVMField(acvmAddress));
+        const [pubKey, partialContractAddress] = await this.context.db.getPublicKey(address);
+        return [pubKey.x.toBigInt(), pubKey.y.toBigInt(), partialContractAddress].map(toACVMField);
+      },
       getNotes: ([slot], sortBy, sortOrder, [limit], [offset], [returnSize]) =>
         this.context.getNotes(this.contractAddress, slot, sortBy, sortOrder, limit, offset, returnSize),
       getRandomField: () => Promise.resolve(toACVMField(Fr.random())),

yarn-project/acir-simulator/src/client/unconstrained_execution.ts

@@ -4,7 +4,7 @@ import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Coordinate, Fr, Point } from '@aztec/foundation/fields';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { decodeReturnValues } from '../abi_coder/decoder.js';
-import { extractReturnWitness } from '../acvm/deserialize.js';
+import { extractReturnWitness, frToAztecAddress } from '../acvm/deserialize.js';
 import { ACVMField, ZERO_ACVM_FIELD, acvm, fromACVMField, toACVMField, toACVMWitness } from '../acvm/index.js';
 import { ClientTxExecutionContext } from './client_execution_context.js';
 import { oracleDebugCallToFormattedStr } from './debug.js';
@@ -49,7 +49,11 @@ export class UnconstrainedFunctionExecution {
             ),
           ),
         ),
-
+      getPublicKey: async ([acvmAddress]) => {
+        const address = frToAztecAddress(fromACVMField(acvmAddress));
+        const [pubKey, partialContractAddress] = await this.context.db.getPublicKey(address);
+        return [pubKey.x.toBigInt(), pubKey.y.toBigInt(), partialContractAddress].map(toACVMField);
+      },
       getNotes: ([slot], sortBy, sortOrder, [limit], [offset], [returnSize]) =>
         this.context.getNotes(this.contractAddress, slot, sortBy, sortOrder, limit, offset, returnSize),
       getRandomField: () => Promise.resolve(toACVMField(Fr.random())),

yarn-project/aztec-rpc/src/account_state/account_state.ts

@@ -489,7 +489,6 @@ export class AccountState {
       contractDataOracle ?? new ContractDataOracle(this.db, this.node),
       this.db,
       keyPair,
-      this.address,
       this.node,
     );
     return new AcirSimulator(simulatorOracle);

yarn-project/aztec-rpc/src/aztec_rpc_server/aztec_rpc_server.test.ts

@@ -0,0 +1,44 @@
+import { AztecNode } from '@aztec/aztec-node';
+import { AztecAddress, CircuitsWasm, Fr } from '@aztec/circuits.js';
+import { computeContractAddressFromPartial } from '@aztec/circuits.js/abis';
+import { Grumpkin } from '@aztec/circuits.js/barretenberg';
+import { ConstantKeyPair, TestKeyStore } from '@aztec/key-store';
+import { randomBytes } from 'crypto';
+import { MockProxy, mock } from 'jest-mock-extended';
+import { MemoryDB } from '../database/memory_db.js';
+import { AztecRPCServer } from './aztec_rpc_server.js';
+
+describe('AztecRpcServer', function () {
+  let wasm: CircuitsWasm;
+  let keyStore: TestKeyStore;
+  let db: MemoryDB;
+  let node: MockProxy<AztecNode>;
+  let rpcServer: AztecRPCServer;
+
+  beforeEach(async () => {
+    keyStore = new TestKeyStore(await Grumpkin.new());
+    node = mock<AztecNode>();
+    db = new MemoryDB();
+    rpcServer = new AztecRPCServer(keyStore, node, db);
+    wasm = await CircuitsWasm.get();
+  });
+
+  it('registers a public key in the db when adding a new account', async () => {
+    const keyPair = ConstantKeyPair.random(await Grumpkin.new());
+    const pubKey = keyPair.getPublicKey();
+    const partialAddress = Fr.random();
+    const address = computeContractAddressFromPartial(wasm, pubKey, partialAddress);
+
+    await rpcServer.addAccount(await keyPair.getPrivateKey(), address, partialAddress);
+    expect(await db.getPublicKey(address)).toEqual([pubKey, partialAddress]);
+  });
+
+  // TODO(#1007)
+  it.skip('refuses to add an account with incorrect address for given partial address and pubkey', async () => {
+    const privateKey = randomBytes(32);
+    const partialAddress = Fr.random();
+    const address = AztecAddress.random();
+
+    await expect(rpcServer.addAccount(privateKey, address, partialAddress)).rejects.toThrowError(/cannot be derived/);
+  });
+});

yarn-project/aztec-rpc/src/aztec_rpc_server/aztec_rpc_server.ts

@@ -79,6 +79,16 @@ export class AztecRPCServer implements AztecRPC {
     abi = SchnorrAccountContractAbi,
   ) {
     const pubKey = this.keyStore.addAccount(privKey);
+    // TODO(#1007): ECDSA contract breaks this check, since the ecdsa public key does not match the one derived from the keystore.
+    // Once we decouple the ecdsa contract signing and encryption keys, we can re-enable this check.
+    // const wasm = await CircuitsWasm.get();
+    // const expectedAddress = computeContractAddressFromPartial(wasm, pubKey, partialContractAddress);
+    // if (!expectedAddress.equals(address)) {
+    //   throw new Error(
+    //     `Address cannot be derived from pubkey and partial address (received ${address.toString()}, derived ${expectedAddress.toString()})`,
+    //   );
+    // }
+    await this.db.addPublicKey(address, pubKey, partialContractAddress);
     await this.#initAccountState(pubKey, address, partialContractAddress, abi);
     return address;
   }

yarn-project/aztec-rpc/src/database/database.ts

@@ -1,10 +1,11 @@
-import { TxHash } from '@aztec/types';
+import { PartialContractAddress, TxHash } from '@aztec/types';
 import { ContractDatabase } from '../contract_database/index.js';
 import { NoteSpendingInfoDao } from './note_spending_info_dao.js';
 import { TxDao } from './tx_dao.js';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr, Point } from '@aztec/foundation/fields';
 import { MerkleTreeId } from '@aztec/types';
+import { PublicKey } from '@aztec/key-store';
 
 /**
  * Options for selecting items from the database.
@@ -49,4 +50,7 @@ export interface Database extends ContractDatabase {
 
   getTreeRoots(): Record<MerkleTreeId, Fr>;
   setTreeRoots(roots: Record<MerkleTreeId, Fr>): Promise<void>;
+
+  addPublicKey(address: AztecAddress, publicKey: PublicKey, partialAddress: PartialContractAddress): Promise<void>;
+  getPublicKey(address: AztecAddress): Promise<[Point, PartialContractAddress] | undefined>;
 }

yarn-project/aztec-rpc/src/database/memory_db.ts

@@ -1,11 +1,12 @@
-import { TxHash } from '@aztec/types';
+import { PartialContractAddress, TxHash } from '@aztec/types';
 import { MemoryContractDatabase } from '../contract_database/index.js';
 import { Database, GetOptions } from './database.js';
 import { NoteSpendingInfoDao } from './note_spending_info_dao.js';
 import { TxDao } from './tx_dao.js';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
 import { Fr, Point } from '@aztec/foundation/fields';
 import { MerkleTreeId } from '@aztec/types';
+import { PublicKey } from '@aztec/key-store';
 
 const sortNotes = (notes: NoteSpendingInfoDao[], sortBy: number[], sortOrder: number[]) => {
   const sortNotesLevel = (a: Fr[], b: Fr[], level = 0): number => {
@@ -36,6 +37,7 @@ export class MemoryDB extends MemoryContractDatabase implements Database {
   private txTable: TxDao[] = [];
   private noteSpendingInfoTable: NoteSpendingInfoDao[] = [];
   private treeRoots: Record<MerkleTreeId, Fr> | undefined;
+  private publicKeys: Map<bigint, [PublicKey, PartialContractAddress]> = new Map();
 
   /**
    * Retrieve a transaction from the MemoryDB using its transaction hash.
@@ -196,4 +198,13 @@ export class MemoryDB extends MemoryContractDatabase implements Database {
     this.treeRoots = roots;
     return Promise.resolve();
   }
+
+  addPublicKey(address: AztecAddress, publicKey: Point, partialAddress: PartialContractAddress): Promise<void> {
+    this.publicKeys.set(address.toBigInt(), [publicKey, partialAddress]);
+    return Promise.resolve();
+  }
+
+  getPublicKey(address: AztecAddress): Promise<[Point, Fr] | undefined> {
+    return Promise.resolve(this.publicKeys.get(address.toBigInt()));
+  }
 }

yarn-project/aztec-rpc/src/simulator_oracle/index.ts

@@ -6,7 +6,7 @@ import { FunctionAbi } from '@aztec/foundation/abi';
 import { ContractDataOracle } from '../contract_data_oracle/index.js';
 import { Database } from '../database/index.js';
 import { siloCommitment } from '@aztec/circuits.js/abis';
-import { MerkleTreeId } from '@aztec/types';
+import { MerkleTreeId, PartialContractAddress } from '@aztec/types';
 
 /**
  * A data oracle that provides information needed for simulating a transaction.
@@ -16,7 +16,6 @@ export class SimulatorOracle implements DBOracle {
     private contractDataOracle: ContractDataOracle,
     private db: Database,
     private keyPair: KeyPair,
-    private address: AztecAddress,
     private node: AztecNode,
   ) {}
 
@@ -40,6 +39,17 @@ export class SimulatorOracle implements DBOracle {
     return this.keyPair.getPrivateKey();
   }
 
+  /**
+   * Retrieve the public key associated to a given address.
+   * @param address - Address to fetch the pubkey for.
+   * @returns A public key and the corresponding partial contract address, such that the hash of the two resolves to the input address.
+   */
+  async getPublicKey(address: AztecAddress): Promise<[Point, PartialContractAddress]> {
+    const result = await this.db.getPublicKey(address);
+    if (!result) throw new Error(`Unknown public key for address ${address.toString()}`);
+    return result;
+  }
+
   /**
    * Retrieves a set of notes stored in the database for a given contract address and storage slot.
    * The query result is paginated using 'limit' and 'offset' values.

yarn-project/circuits.js/src/abis/abis.ts

@@ -170,7 +170,7 @@ export function hashConstructor(
 /**
  * Computes a contract address.
  * @param wasm - A module providing low-level wasm access.
- * @param deployerAddr - The address of the contract deployer.
+ * @param deployerPubKey - The pubkey of the contract deployer.
  * @param contractAddrSalt - The salt used as 1 one of the inputs of the contract address computation.
  * @param fnTreeRoot - The function tree root of the contract being deployed.
  * @param constructorHash - The hash of the constructor.
@@ -217,6 +217,25 @@ export function computePartialContractAddress(
   return Fr.fromBuffer(result);
 }
 
+/**
+ * Computes a contract address from its partial address and the pubkey.
+ * @param wasm - A module providing low-level wasm access.
+ * @param partial - The salt used as 1 one of the inputs of the contract address computation.
+ * @param fnTreeRoot - The function tree root of the contract being deployed.
+ * @param constructorHash - The hash of the constructor.
+ * @returns The partially constructed contract address.
+ */
+export function computeContractAddressFromPartial(wasm: IWasmModule, pubKey: Point, partialAddress: Fr): AztecAddress {
+  wasm.call('pedersen__init');
+  const result = inputBuffersToOutputBuffer(
+    wasm,
+    'abis__compute_contract_address_from_partial',
+    [pubKey.toFieldsBuffer(), partialAddress.toBuffer()],
+    32,
+  );
+  return new AztecAddress(result);
+}
+
 /**
  * Computes a siloed commitment, given the contract address and the commitment itself.
  * A siloed commitment effectively namespaces a commitment to a specific contract.