refactor: refactor key rotate and address comments from 6405

noir-projects/noir-contracts/contracts/child_contract/src/main.nr

@@ -64,6 +64,7 @@ contract Child {
     fn private_get_value(amount: Field, owner: AztecAddress) -> Field {
         let owner_npk_m_hash = get_npk_m_hash(&mut context, owner);
 
+        // TODO (#6312): This will break with key rotation. Fix this. Will not be able to find any notes after rotating key.
         let mut options = NoteGetterOptions::new();
         options = options.select(ValueNote::properties().value, amount, Option::none()).select(
             ValueNote::properties().npk_m_hash,

noir-projects/noir-contracts/contracts/docs_example_contract/src/options.nr

@@ -11,6 +11,7 @@ pub fn create_account_card_getter_options(
     account_npk_m_hash: Field,
     offset: u32
 ) -> NoteGetterOptions<CardNote, CARD_NOTE_LEN, Field> {
+    // TODO (#6312): This will break with key rotation. Fix this. Will not be able to find any notes after rotating key.
     let mut options = NoteGetterOptions::new();
     options.select(
         CardNote::properties().npk_m_hash,
@@ -26,6 +27,7 @@ pub fn create_exact_card_getter_options(
     secret: Field,
     account_npk_m_hash: Field
 ) -> NoteGetterOptions<CardNote, CARD_NOTE_LEN, Field> {
+    // TODO (#6312): This will break with key rotation. Fix this. Will not be able to find any notes after rotating key.
     let mut options = NoteGetterOptions::new();
     options.select(CardNote::properties().points, points as Field, Option::none()).select(CardNote::properties().randomness, secret, Option::none()).select(
         CardNote::properties().npk_m_hash,
@@ -54,6 +56,7 @@ pub fn filter_min_points(
 
 // docs:start:state_vars-NoteGetterOptionsFilter
 pub fn create_account_cards_with_min_points_getter_options(account_npk_m_hash: Field, min_points: u8) -> NoteGetterOptions<CardNote, CARD_NOTE_LEN, u8> {
+    // TODO (#6312): This will break with key rotation. Fix this. Will not be able to find any notes after rotating key.
     NoteGetterOptions::with_filter(filter_min_points, min_points).select(
         CardNote::properties().npk_m_hash,
         account_npk_m_hash,
@@ -64,6 +67,7 @@ pub fn create_account_cards_with_min_points_getter_options(account_npk_m_hash: F
 
 // docs:start:state_vars-NoteGetterOptionsPickOne
 pub fn create_largest_account_card_getter_options(account_npk_m_hash: Field) -> NoteGetterOptions<CardNote, CARD_NOTE_LEN, Field> {
+    // TODO (#6312): This will break with key rotation. Fix this. Will not be able to find any notes after rotating key.
     let mut options = NoteGetterOptions::new();
     options.select(
         CardNote::properties().npk_m_hash,

noir-projects/noir-contracts/contracts/ecdsa_account_contract/src/ecdsa_public_key_note.nr

@@ -24,7 +24,7 @@ impl NoteInterface<ECDSA_PUBLIC_KEY_NOTE_LEN> for EcdsaPublicKeyNote {
     // [1] = x[31]
     // [2] = y[0..31]
     // [3] = y[31]
-    // [4] = owner
+    // [4] = npk_m_hash
     fn serialize_content(self) -> [Field; ECDSA_PUBLIC_KEY_NOTE_LEN] {
         let mut x: Field = 0;
         let mut y: Field = 0;

noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr

@@ -46,7 +46,6 @@ contract StatefulTest {
     fn create_note_no_init_check(owner: AztecAddress, value: Field) {
         if (value != 0) {
             let loc = storage.notes.at(owner);
-            // TODO (#6312): This will break with key rotation. Fix this. Will not be able to spend / increment any notes after rotating key.
             increment(context, loc, value, owner);
         }
     }

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

@@ -1,5 +1,5 @@
 import { type AuthWitness, type FunctionCall, type PXE, type TxExecutionRequest } from '@aztec/circuit-types';
-import { type AztecAddress, Fr } from '@aztec/circuits.js';
+import { AztecAddress, CANONICAL_KEY_REGISTRY_ADDRESS, type Fq, Fr, type Point } from '@aztec/circuits.js';
 import { type ABIParameterVisibility, type FunctionAbi, FunctionType } from '@aztec/foundation/abi';
 
 import { type AccountInterface } from '../account/interface.js';
@@ -165,6 +165,30 @@ export class AccountWallet extends BaseWallet {
     return { isValidInPrivate, isValidInPublic };
   }
 
+  /**
+   * Returns a function interaction to rotate our master nullifer public key in the canonical key registry.
+   * @param newNpkM - The new master nullifier public key we want to use.
+   * @remarks - This does not hinder our ability to spend notes tied to a previous master nullifier public key.
+   * @returns - A function interaction.
+   */
+  public rotateNpkM(newNpkM: Point): ContractFunctionInteraction {
+    return new ContractFunctionInteraction(
+      this,
+      AztecAddress.fromBigInt(CANONICAL_KEY_REGISTRY_ADDRESS),
+      this.getRotateNpkMAbi(),
+      [this.getAddress(), newNpkM],
+    );
+  }
+
+  /**
+   * Rotates the account master nullifier secret key in our pxe / keystore
+   * @param newNskM - The new master nullifier secret key we want to use.
+   * @remarks - This does not hinder our ability to spend notes tied to a previous master nullifier public key.
+   */
+  public async rotateNskM(newNskM: Fq): Promise<void> {
+    await this.pxe.rotateNskMPxe(this.getAddress(), newNskM);
+  }
+
   /**
    * Returns a function interaction to cancel a message hash as authorized in this account.
    * @param messageHashOrIntent - The message or the caller and action to authorize/revoke
@@ -265,4 +289,38 @@ export class AccountWallet extends BaseWallet {
       returnTypes: [{ kind: 'array', length: 2, type: { kind: 'boolean' } }],
     };
   }
+
+  private getRotateNpkMAbi(): FunctionAbi {
+    return {
+      name: 'rotate_npk_m',
+      isInitializer: false,
+      functionType: FunctionType.OPEN,
+      isInternal: false,
+      parameters: [
+        {
+          name: 'address',
+          type: {
+            fields: [{ name: 'inner', type: { kind: 'field' } }],
+            kind: 'struct',
+            path: 'authwit::aztec::protocol_types::address::aztec_address::AztecAddress',
+          },
+          visibility: 'private' as ABIParameterVisibility,
+        },
+        {
+          name: 'new_npk_m',
+          type: {
+            fields: [
+              { name: 'x', type: { kind: 'field' } },
+              { name: 'y', type: { kind: 'field' } },
+            ],
+            kind: 'struct',
+            path: 'authwit::aztec::protocol_types::grumpkin_point::GrumpkinPoint',
+          },
+          visibility: 'private' as ABIParameterVisibility,
+        },
+        { name: 'nonce', type: { kind: 'field' }, visibility: 'private' as ABIParameterVisibility },
+      ],
+      returnTypes: [],
+    };
+  }
 }

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

@@ -69,8 +69,8 @@ export abstract class BaseWallet implements Wallet {
   registerAccount(secretKey: Fr, partialAddress: PartialAddress): Promise<CompleteAddress> {
     return this.pxe.registerAccount(secretKey, partialAddress);
   }
-  rotateMasterNullifierKey(account: AztecAddress, secretKey: Fq): Promise<void> {
-    return this.pxe.rotateMasterNullifierKey(account, secretKey);
+  rotateNskMPxe(address: AztecAddress, secretKey: Fq) {
+    return this.pxe.rotateNskMPxe(address, secretKey);
   }
   registerRecipient(account: CompleteAddress): Promise<void> {
     return this.pxe.registerRecipient(account);

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

@@ -61,8 +61,6 @@ export interface PXE {
    */
   registerAccount(secretKey: Fr, partialAddress: PartialAddress): Promise<CompleteAddress>;
 
-  rotateMasterNullifierKey(account: AztecAddress, secretKey: Fq): Promise<void>;
-
   /**
    * Registers a recipient in PXE. This is required when sending encrypted notes to
    * a user who hasn't deployed their account contract yet. Since their account is not deployed, their
@@ -107,6 +105,16 @@ export interface PXE {
    */
   getRecipient(address: AztecAddress): Promise<CompleteAddress | undefined>;
 
+  /**
+   * Rotates our master nullifier secret key by adding it to our keystore
+   * @param address - The address of the account we want to rotate our key for.
+   * @param newNskM - The new master nullifier secret key we want to use.
+   * @remarks - This does not hinder our ability to spend notes tied to a previous master nullifier public key.
+   * Also, one should not use this function directly without also calling the canonical key registry, and it is preferred to use the function
+   * under the AccountWallet interface, as that implementation also returns a contract interaction one could use to call the key registry with.
+   */
+  rotateNskMPxe(address: AztecAddress, newNskM: Fq): Promise<void>;
+
   /**
    * Registers a contract class in the PXE without registering any associated contract instance with it.
    *

yarn-project/end-to-end/src/e2e_key_rotation.test.ts

@@ -1,5 +1,6 @@
 import { createAccounts } from '@aztec/accounts/testing';
 import {
+  type AccountWallet,
   type AztecAddress,
   type AztecNode,
   type DebugLogger,
@@ -14,8 +15,7 @@ import {
   retryUntil,
 } from '@aztec/aztec.js';
 import { type PublicKey, derivePublicKeyFromSecretKey } from '@aztec/circuits.js';
-import { KeyRegistryContract, TestContract, TokenContract } from '@aztec/noir-contracts.js';
-import { getCanonicalKeyRegistryAddress } from '@aztec/protocol-contracts/key-registry';
+import { TestContract, TokenContract } from '@aztec/noir-contracts.js';
 
 import { jest } from '@jest/globals';
 
@@ -31,13 +31,12 @@ describe('e2e_key_rotation', () => {
   let aztecNode: AztecNode;
   let pxeA: PXE;
   let pxeB: PXE;
-  let walletA: Wallet;
-  let walletB: Wallet;
+  let walletA: AccountWallet;
+  let walletB: AccountWallet;
   let logger: DebugLogger;
   let teardownA: () => Promise<void>;
   let teardownB: () => Promise<void>;
 
-  let keyRegistryWithB: KeyRegistryContract;
   let testContract: TestContract;
   let contractWithWalletA: TokenContract;
   let contractWithWalletB: TokenContract;
@@ -58,7 +57,6 @@ describe('e2e_key_rotation', () => {
     ({ pxe: pxeB, teardown: teardownB } = await setupPXEService(aztecNode, {}, undefined, true));
 
     [walletB] = await createAccounts(pxeB, 1);
-    keyRegistryWithB = await KeyRegistryContract.at(getCanonicalKeyRegistryAddress(), walletB);
 
     // We deploy test and token contracts
     testContract = await TestContract.deploy(walletA).send().deployed();
@@ -176,9 +174,12 @@ describe('e2e_key_rotation', () => {
     {
       const newNskM = Fq.random();
       newNpkM = derivePublicKeyFromSecretKey(newNskM);
-      await pxeB.rotateMasterNullifierKey(walletB.getAddress(), newNskM);
 
-      await keyRegistryWithB.methods.rotate_npk_m(walletB.getAddress(), newNpkM, 0).send().wait();
+      // We call the registry to rotate our key here
+      await walletB.rotateNpkM(newNpkM).send().wait();
+      // We add it to our Pxe here
+      await walletB.rotateNskM(newNskM);
+
       await crossDelay();
     }
 

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

@@ -59,14 +59,14 @@ export class TestKeyStore implements KeyStore {
     // We save the keys to db associated with the account address
     await this.#keys.set(`${accountAddress.toString()}-public_keys_hash`, publicKeysHash.toBuffer());
 
-    // Naming of keys is as follows ${from}-${to}_m
+    // Naming of keys is as follows ${from}-${to}_m${_s if multiple}
     await this.#keys.set(`${accountAddress.toString()}-ivsk_m`, masterIncomingViewingSecretKey.toBuffer());
     await this.#keys.set(`${accountAddress.toString()}-ovsk_m`, masterOutgoingViewingSecretKey.toBuffer());
     await this.#keys.set(`${accountAddress.toString()}-tsk_m`, masterTaggingSecretKey.toBuffer());
     // The key of the following is different from the others because the buffer can store multiple keys
-    await this.#keys.set(`${accountAddress.toString()}-ns_keys_m`, masterNullifierSecretKey.toBuffer());
+    await this.#keys.set(`${accountAddress.toString()}-nsk_m_s`, masterNullifierSecretKey.toBuffer());
 
-    await this.#keys.set(`${accountAddress.toString()}-np_keys_m`, publicKeys.masterNullifierPublicKey.toBuffer());
+    await this.#keys.set(`${accountAddress.toString()}-npk_m_s`, publicKeys.masterNullifierPublicKey.toBuffer());
     await this.#keys.set(`${accountAddress.toString()}-ivpk_m`, publicKeys.masterIncomingViewingPublicKey.toBuffer());
     await this.#keys.set(`${accountAddress.toString()}-ovpk_m`, publicKeys.masterOutgoingViewingPublicKey.toBuffer());
     await this.#keys.set(`${accountAddress.toString()}-tpk_m`, publicKeys.masterTaggingPublicKey.toBuffer());
@@ -107,7 +107,7 @@ export class TestKeyStore implements KeyStore {
     const accountAddress = AztecAddress.fromBuffer(accountAddressBuffer);
 
     // Get the master nullifier public keys buffer for the account
-    const masterNullifierPublicKeysBuffer = this.#keys.get(`${accountAddress.toString()}-np_keys_m`);
+    const masterNullifierPublicKeysBuffer = this.#keys.get(`${accountAddress.toString()}-npk_m_s`);
     if (!masterNullifierPublicKeysBuffer) {
       throw new Error(
         `Could not find master nullifier public key for account ${accountAddress.toString()} whose address was successfully obtained with npk_m_hash ${npkMHash.toString()}.`,
@@ -120,7 +120,7 @@ export class TestKeyStore implements KeyStore {
     }
 
     // Now we iterate over the public keys in the buffer to find the one that matches the hash
-    const numKeys = masterNullifierPublicKeysBuffer.byteLength / Point.SIZE_IN_BYTES;
+    const numKeys = this.#calculateNumKeys(masterNullifierPublicKeysBuffer, Point);
     for (let i = 0; i < numKeys; i++) {
       const masterNullifierPublicKey = Point.fromBuffer(
         masterNullifierPublicKeysBuffer.subarray(i * Point.SIZE_IN_BYTES, (i + 1) * Point.SIZE_IN_BYTES),
@@ -196,7 +196,7 @@ export class TestKeyStore implements KeyStore {
 
     // Now we get the master nullifier secret keys for the account
     const masterNullifierSecretKeysBuffer = this.#keys.get(
-      `${AztecAddress.fromBuffer(accountAddressBuffer).toString()}-ns_keys_m`,
+      `${AztecAddress.fromBuffer(accountAddressBuffer).toString()}-nsk_m_s`,
     );
     if (!masterNullifierSecretKeysBuffer) {
       throw new Error(
@@ -207,7 +207,7 @@ export class TestKeyStore implements KeyStore {
     }
 
     // Now we iterate over all the secret keys to find the one that matches the hash
-    const numKeys = masterNullifierSecretKeysBuffer.byteLength / GrumpkinScalar.SIZE_IN_BYTES;
+    const numKeys = this.#calculateNumKeys(masterNullifierSecretKeysBuffer, GrumpkinScalar);
     for (let i = 0; i < numKeys; i++) {
       const secretKey = GrumpkinScalar.fromBuffer(
         masterNullifierSecretKeysBuffer.subarray(
@@ -295,7 +295,7 @@ export class TestKeyStore implements KeyStore {
     const accountAddress = AztecAddress.fromBuffer(accountAddressBuffer);
 
     // We fetch the public keys and find this specific public key's position in the buffer
-    const masterNullifierPublicKeysBuffer = this.#keys.get(`${accountAddress.toString()}-np_keys_m`);
+    const masterNullifierPublicKeysBuffer = this.#keys.get(`${accountAddress.toString()}-npk_m_s`);
     if (!masterNullifierPublicKeysBuffer) {
       throw new Error(`Could not find master nullifier public keys for account ${accountAddress.toString()}`);
     }
@@ -306,7 +306,7 @@ export class TestKeyStore implements KeyStore {
     }
 
     // Now we iterate over the public keys in the buffer to find the one that matches the hash
-    const numKeys = masterNullifierPublicKeysBuffer.byteLength / Point.SIZE_IN_BYTES;
+    const numKeys = this.#calculateNumKeys(masterNullifierPublicKeysBuffer, Point);
     let keyIndex = -1;
     for (let i = 0; i < numKeys; i++) {
       const publicKey = Point.fromBuffer(
@@ -319,7 +319,7 @@ export class TestKeyStore implements KeyStore {
     }
 
     // Now we fetch the secret keys buffer and extract the secret key at the same index
-    const masterNullifierSecretKeysBuffer = this.#keys.get(`${accountAddress.toString()}-ns_keys_m`);
+    const masterNullifierSecretKeysBuffer = this.#keys.get(`${accountAddress.toString()}-nsk_m_s`);
     if (!masterNullifierSecretKeysBuffer) {
       throw new Error(`Could not find master nullifier secret keys for account ${accountAddress.toString()}`);
     }
@@ -401,27 +401,31 @@ export class TestKeyStore implements KeyStore {
    */
   public async rotateMasterNullifierKey(account: AztecAddress, newSecretKey: Fq = Fq.random()) {
     // We append the secret key to the original secret key
-    const secretKeysBuffer = this.#keys.get(`${account.toString()}-ns_keys_m`);
+    const secretKeysBuffer = this.#keys.get(`${account.toString()}-nsk_m_s`);
     if (!secretKeysBuffer) {
       throw new Error(`Could not find nullifier secret keys for account ${account.toString()}`);
     }
 
     // We append the new secret key to the buffer of secret keys
     const newSecretKeysBuffer = Buffer.concat([secretKeysBuffer, newSecretKey.toBuffer()]);
-    await this.#keys.set(`${account.toString()}-ns_keys_m`, newSecretKeysBuffer);
+    await this.#keys.set(`${account.toString()}-nsk_m_s`, newSecretKeysBuffer);
 
     // Now we derive the public key from the new secret key and append it to the buffer of original public keys
     const newPublicKey = derivePublicKeyFromSecretKey(newSecretKey);
-    const publicKeysBuffer = this.#keys.get(`${account.toString()}-np_keys_m`);
+    const publicKeysBuffer = this.#keys.get(`${account.toString()}-npk_m_s`);
     if (!publicKeysBuffer) {
       throw new Error(`Could not find nullifier public keys for account ${account.toString()}`);
     }
 
     // We append the new public key to the buffer of public keys
     const newPublicKeysBuffer = Buffer.concat([publicKeysBuffer, newPublicKey.toBuffer()]);
-    await this.#keys.set(`${account.toString()}-np_keys_m`, newPublicKeysBuffer);
+    await this.#keys.set(`${account.toString()}-npk_m_s`, newPublicKeysBuffer);
 
     // We store a npk_m_hash-account_address map to make address easy to obtain with the hash later on
     await this.#keys.set(`${newPublicKey.hash().toString()}-npk_m_hash`, account.toBuffer());
   }
+
+  #calculateNumKeys(buf: Buffer, T: typeof Point | typeof Fq) {
+    return buf.byteLength / T.SIZE_IN_BYTES;
+  }
 }

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

@@ -160,6 +160,10 @@ export class PXEService implements PXE {
     return this.db.getAuthWitness(messageHash);
   }
 
+  async rotateNskMPxe(account: AztecAddress, secretKey: Fq = Fq.random()): Promise<void> {
+    await this.keyStore.rotateMasterNullifierKey(account, secretKey);
+  }
+
   public addCapsule(capsule: Fr[]) {
     return this.db.addCapsule(capsule);
   }
@@ -192,10 +196,6 @@ export class PXEService implements PXE {
     return accountCompleteAddress;
   }
 
-  public async rotateMasterNullifierKey(account: AztecAddress, secretKey: Fq = Fq.random()): Promise<void> {
-    await this.keyStore.rotateMasterNullifierKey(account, secretKey);
-  }
-
   public async getRegisteredAccounts(): Promise<CompleteAddress[]> {
     // Get complete addresses of both the recipients and the accounts
     const completeAddresses = await this.db.getCompleteAddresses();

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

@@ -40,9 +40,9 @@ export class Oracle {
     return unpacked.map(toACVMField);
   }
 
-  async getNullifierKeys([masterNullifierPublicKeyHash]: ACVMField[]): Promise<ACVMField[]> {
+  async getNullifierKeys([npkMHash]: ACVMField[]): Promise<ACVMField[]> {
     const { masterNullifierPublicKey, appNullifierSecretKey } = await this.typedOracle.getNullifierKeys(
-      fromACVMField(masterNullifierPublicKeyHash),
+      fromACVMField(npkMHash),
     );
 
     return [