From 1d0696872689f496597d363bfa26542ab89ca834 Mon Sep 17 00:00:00 2001
From: sklppy88 <esau@aztecprotocol.com>
Date: Wed, 23 Oct 2024 14:13:25 +0000
Subject: [PATCH] init

---
 yarn-project/key-store/src/key_store.ts | 29 +++++++++++++++++++------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/yarn-project/key-store/src/key_store.ts b/yarn-project/key-store/src/key_store.ts
index 3c6904ee8055..c6ee5e1b121f 100644
--- a/yarn-project/key-store/src/key_store.ts
+++ b/yarn-project/key-store/src/key_store.ts
@@ -10,7 +10,6 @@ import {
   KeyValidationRequest,
   type PartialAddress,
   Point,
-  computeAddress,
   computeAppSecretKey,
   deriveKeys,
   derivePublicKeyFromSecretKey,
@@ -54,8 +53,8 @@ export class KeyStore {
       publicKeys,
     } = deriveKeys(sk);
 
-    const publicKeysHash = publicKeys.hash();
-    const account = computeAddress(publicKeysHash, partialAddress);
+    const completeAddress = CompleteAddress.fromSecretKeyAndPartialAddress(sk, partialAddress);
+    const { address: account } = completeAddress;
 
     // Naming of keys is as follows ${account}-${n/iv/ov/t}${sk/pk}_m
     await this.#keys.set(`${account.toString()}-ivsk_m`, masterIncomingViewingSecretKey.toBuffer());
@@ -82,7 +81,7 @@ export class KeyStore {
     await this.#keys.set(`${account.toString()}-tpk_m_hash`, publicKeys.masterTaggingPublicKey.hash().toBuffer());
 
     // At last, we return the newly derived account address
-    return Promise.resolve(new CompleteAddress(account, publicKeys, partialAddress));
+    return Promise.resolve(completeAddress);
   }
 
   /**
@@ -104,7 +103,7 @@ export class KeyStore {
    * @returns The key validation request.
    */
   public getKeyValidationRequest(pkMHash: Fr, contractAddress: AztecAddress): Promise<KeyValidationRequest> {
-    const [keyPrefix, account] = this.#getKeyPrefixAndAccount(pkMHash);
+    const [keyPrefix, account] = this.getKeyPrefixAndAccount(pkMHash);
 
     // Now we find the master public key for the account
     const pkMBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}pk_m`);
@@ -141,6 +140,22 @@ export class KeyStore {
     return Promise.resolve(new KeyValidationRequest(pkM, skApp));
   }
 
+  /**
+   * Gets the master nullifier public key for a given account.
+   * @throws If the account does not exist in the key store.
+   * @param account - The account address for which to retrieve the master nullifier public key.
+   * @returns The master nullifier public key for the account.
+   */
+  public async getMasterNullifierPublicKey(account: AztecAddress): Promise<PublicKey> {
+    const masterNullifierPublicKeyBuffer = this.#keys.get(`${account.toString()}-npk_m`);
+    if (!masterNullifierPublicKeyBuffer) {
+      throw new Error(
+        `Account ${account.toString()} does not exist. Registered accounts: ${await this.getAccounts()}.`,
+      );
+    }
+    return Promise.resolve(Point.fromBuffer(masterNullifierPublicKeyBuffer));
+  }
+
   /**
    * Gets the master incoming viewing public key for a given account.
    * @throws If the account does not exist in the key store.
@@ -245,7 +260,7 @@ export class KeyStore {
    * @dev Used when feeding the sk_m to the kernel circuit for keys verification.
    */
   public getMasterSecretKey(pkM: PublicKey): Promise<GrumpkinScalar> {
-    const [keyPrefix, account] = this.#getKeyPrefixAndAccount(pkM);
+    const [keyPrefix, account] = this.getKeyPrefixAndAccount(pkM);
 
     const secretKeyBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}sk_m`);
     if (!secretKeyBuffer) {
@@ -268,7 +283,7 @@ export class KeyStore {
    * @dev Note that this is quite inefficient but it should not matter because there should never be too many keys
    * in the key store.
    */
-  #getKeyPrefixAndAccount(value: Bufferable): [KeyPrefix, AztecAddress] {
+  public getKeyPrefixAndAccount(value: Bufferable): [KeyPrefix, AztecAddress] {
     const valueBuffer = serializeToBuffer(value);
     for (const [key, val] of this.#keys.entries()) {
       if (val.equals(valueBuffer)) {