feat: impl of missing functionality in new key store (#5750)

Author: benesjan

noir/noir-repo/tooling/noir_js_backend_barretenberg/package.json

@@ -57,4 +57,4 @@
     "ts-node": "^10.9.1",
     "typescript": "5.4.2"
   }
-}
+}

yarn-project/circuit-types/src/keys/new_key_store.ts

@@ -1,4 +1,10 @@
-import { type AztecAddress, type Fr, type PartialAddress, type PublicKey } from '@aztec/circuits.js';
+import {
+  type AztecAddress,
+  type Fr,
+  type GrumpkinPrivateKey,
+  type PartialAddress,
+  type PublicKey,
+} from '@aztec/circuits.js';
 
 /**
  * Represents a secure storage for managing keys.
@@ -18,6 +24,12 @@ export interface NewKeyStore {
    */
   addAccount(sk: Fr, partialAddress: PartialAddress): Promise<AztecAddress>;
 
+  /**
+   * Retrieves addresses of accounts stored in the key store.
+   * @returns A Promise that resolves to an array of account addresses.
+   */
+  getAccounts(): Promise<AztecAddress[]>;
+
   /**
    * Gets the master nullifier public key for a given account.
    * @throws If the account does not exist in the key store.
@@ -76,4 +88,14 @@ export interface NewKeyStore {
    * @returns A Promise that resolves to the application outgoing viewing secret key.
    */
   getAppOutgoingViewingSecretKey(account: AztecAddress, app: AztecAddress): Promise<Fr>;
+
+  /**
+   * Retrieves the master nullifier secret key (nsk_m) corresponding to the specified master nullifier public key
+   * (Npk_m).
+   * @throws If the provided public key is not associated with any of the registered accounts.
+   * @param masterNullifierPublicKey - The master nullifier public key to get secret key for.
+   * @returns A Promise that resolves to the master nullifier secret key.
+   * @dev Used when feeding the master nullifier secret key to the kernel circuit for nullifier keys verification.
+   */
+  getMasterNullifierSecretKeyForPublicKey(masterNullifierPublicKey: PublicKey): Promise<GrumpkinPrivateKey>;
 }

yarn-project/key-store/src/new_test_key_store.test.ts

@@ -55,5 +55,17 @@ describe('NewTestKeyStore', () => {
     expect(appOutgoingViewingSecretKey.toString()).toMatchInlineSnapshot(
       `"0x2639b26510f9d30b7e173d301b263b246b7a576186be1f44cd7c86bc06773f8a"`,
     );
+
+    // Returned accounts are as expected
+    const accounts = await keyStore.getAccounts();
+    expect(accounts.toString()).toMatchInlineSnapshot(
+      `"0x0ba7834252d19c4f09d29303c269f303f40ae3d2043f921ed0bf8c0709926d4e"`,
+    );
+
+    // Manages to find master nullifer secret key for pub key
+    const masterNullifierSecretKey = await keyStore.getMasterNullifierSecretKeyForPublicKey(masterNullifierPublicKey);
+    expect(masterNullifierSecretKey.toString()).toMatchInlineSnapshot(
+      `"0x0fde74d5e504c73b58aad420dd72590fc6004571411e7f77c45378714195a52b"`,
+    );
   });
 });

yarn-project/key-store/src/new_test_key_store.ts

@@ -1,5 +1,13 @@
 import { type NewKeyStore, type PublicKey } from '@aztec/circuit-types';
-import { AztecAddress, Fr, GeneratorIndex, GrumpkinScalar, type PartialAddress, Point } from '@aztec/circuits.js';
+import {
+  AztecAddress,
+  Fr,
+  GeneratorIndex,
+  type GrumpkinPrivateKey,
+  GrumpkinScalar,
+  type PartialAddress,
+  Point,
+} from '@aztec/circuits.js';
 import { type Grumpkin } from '@aztec/circuits.js/barretenberg';
 import { poseidon2Hash, sha512ToGrumpkinScalar } from '@aztec/foundation/crypto';
 import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
@@ -75,6 +83,17 @@ export class NewTestKeyStore implements NewKeyStore {
     return Promise.resolve(accountAddress);
   }
 
+  /**
+   * Retrieves addresses of accounts stored in the key store.
+   * @returns A Promise that resolves to an array of account addresses.
+   */
+  public getAccounts(): Promise<AztecAddress[]> {
+    const allMapKeys = Array.from(this.#keys.keys());
+    // We return account addresses based on the map keys that end with '-nsk_m'
+    const accounts = allMapKeys.filter(key => key.endsWith('-nsk_m')).map(key => key.split('-')[0]);
+    return Promise.resolve(accounts.map(account => AztecAddress.fromString(account)));
+  }
+
   /**
    * Gets the master nullifier public key for a given account.
    * @throws If the account does not exist in the key store.
@@ -197,4 +216,30 @@ export class NewTestKeyStore implements NewKeyStore {
       ]),
     );
   }
+
+  /**
+   * Retrieves the master nullifier secret key (nsk_m) corresponding to the specified master nullifier public key
+   * (Npk_m).
+   * @throws If the provided public key is not associated with any of the registered accounts.
+   * @param masterNullifierPublicKey - The master nullifier public key to get secret key for.
+   * @returns A Promise that resolves to the master nullifier secret key.
+   * @dev Used when feeding the master nullifier secret key to the kernel circuit for nullifier keys verification.
+   */
+  public getMasterNullifierSecretKeyForPublicKey(masterNullifierPublicKey: PublicKey): Promise<GrumpkinPrivateKey> {
+    // We iterate over the map keys to find the account address that corresponds to the provided public key
+    for (const [key, value] of this.#keys.entries()) {
+      if (value.equals(masterNullifierPublicKey.toBuffer())) {
+        // We extract the account address from the map key
+        const accountAddress = key.split('-')[0];
+        // We fetch the secret key and return it
+        const masterNullifierSecretKeyBuffer = this.#keys.get(`${accountAddress.toString()}-nsk_m`);
+        if (!masterNullifierSecretKeyBuffer) {
+          throw new Error(`Could not find master nullifier secret key for account ${accountAddress.toString()}`);
+        }
+        return Promise.resolve(GrumpkinScalar.fromBuffer(masterNullifierSecretKeyBuffer));
+      }
+    }
+
+    throw new Error(`Could not find master nullifier secret key for public key ${masterNullifierPublicKey.toString()}`);
+  }
 }