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/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(loc, value, owner);
         }
     }

yarn-project/aztec.js/src/account/interface.ts

@@ -1,6 +1,6 @@
 import { type AuthWitness, type CompleteAddress, type FunctionCall } from '@aztec/circuit-types';
 import { type AztecAddress } from '@aztec/circuits.js';
-import { type Fr } from '@aztec/foundation/fields';
+import { type Fq, type Fr } from '@aztec/foundation/fields';
 
 import { type ContractFunctionInteraction } from '../contract/contract_function_interaction.js';
 import { type EntrypointInterface } from '../entrypoint/entrypoint.js';
@@ -51,4 +51,20 @@ export interface AccountInterface extends AuthWitnessProvider, EntrypointInterfa
   /** Returns the rollup version for this account */
   getVersion(): Fr;
 }
+
+/**
+ * Handler for interfacing with an account's ability to rotate its keys.
+ */
+export interface AccountKeyRotationInterface {
+  /**
+   * Rotates the account master nullifier key pair.
+   * @param newNskM - The new master nullifier secret key we want to use.
+   * @remarks - This function also calls the canonical key registry with the account's new derived master nullifier public key.
+   * We are doing it this way to avoid user error, in the case that a user rotates their keys in the key registry,
+   * but fails to do so in the key store. This leads to unspendable notes.
+   *
+   * This does not hinder our ability to spend notes tied to a previous master nullifier public key, provided we have the master nullifier secret key for it.
+   */
+  rotateNullifierKeys(newNskM: Fq): Promise<void>;
+}
 // docs:end:account-interface

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

@@ -1,8 +1,8 @@
 import { type PXE } from '@aztec/circuit-types';
 
-import { type AccountInterface } from './interface.js';
+import { type AccountInterface, type AccountKeyRotationInterface } from './interface.js';
 
 /**
  * The wallet interface.
  */
-export type Wallet = AccountInterface & PXE;
+export type Wallet = AccountInterface & PXE & AccountKeyRotationInterface;

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, Fq, Fr, derivePublicKeyFromSecretKey } 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 };
   }
 
+  /**
+   * Rotates the account master nullifier key pair.
+   * @param newNskM - The new master nullifier secret key we want to use.
+   * @remarks - This function also calls the canonical key registry with the account's new derived master nullifier public key.
+   * We are doing it this way to avoid user error, in the case that a user rotates their keys in the key registry,
+   * but fails to do so in the key store. This leads to unspendable notes.
+   *
+   * This does not hinder our ability to spend notes tied to a previous master nullifier public key, provided we have the master nullifier secret key for it.
+   */
+  public async rotateNullifierKeys(newNskM: Fq = Fq.random()): Promise<void> {
+    // We rotate our secret key in the keystore first, because if the subsequent interaction fails, there are no bad side-effects.
+    // If vice versa (the key registry is called first), but the call to the PXE fails, we will end up in a situation with unspendable notes, as we have not committed our
+    // nullifier secret key to our wallet.
+    await this.pxe.rotateNskM(this.getAddress(), newNskM);
+    const interaction = new ContractFunctionInteraction(
+      this,
+      AztecAddress.fromBigInt(CANONICAL_KEY_REGISTRY_ADDRESS),
+      this.getRotateNpkMAbi(),
+      [this.getAddress(), derivePublicKeyFromSecretKey(newNskM), Fr.ZERO],
+    );
+
+    await interaction.send().wait();
+  }
+
   /**
    * 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
@@ -268,4 +292,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

@@ -54,6 +54,8 @@ export abstract class BaseWallet implements Wallet {
         },
   ): Promise<AuthWitness>;
 
+  abstract rotateNullifierKeys(newNskM: Fq): Promise<void>;
+
   getAddress() {
     return this.getCompleteAddress().address;
   }
@@ -69,8 +71,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);
+  rotateNskM(address: AztecAddress, secretKey: Fq) {
+    return this.pxe.rotateNskM(address, secretKey);
   }
   registerRecipient(account: CompleteAddress): Promise<void> {
     return this.pxe.registerRecipient(account);

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

@@ -1,5 +1,5 @@
 import { type AuthWitness, type PXE, type TxExecutionRequest } from '@aztec/circuit-types';
-import { type CompleteAddress, type Fr } from '@aztec/circuits.js';
+import { type CompleteAddress, type Fq, type Fr } from '@aztec/circuits.js';
 
 import { DefaultEntrypoint } from '../entrypoint/default_entrypoint.js';
 import { type EntrypointInterface, type ExecutionRequestInit } from '../entrypoint/entrypoint.js';
@@ -42,4 +42,8 @@ export class SignerlessWallet extends BaseWallet {
   createAuthWit(_messageHash: Fr): Promise<AuthWitness> {
     throw new Error('Method not implemented.');
   }
+
+  rotateNullifierKeys(_newNskM: Fq): Promise<void> {
+    throw new Error('Method not implemented.');
+  }
 }

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,18 @@ export interface PXE {
    */
   getRecipient(address: AztecAddress): Promise<CompleteAddress | undefined>;
 
+  /**
+   * Rotates master nullifier keys.
+   * @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 - One should not use this function directly without also calling the canonical key registry to rotate
+   * the new master nullifier secret key's derived master nullifier public key.
+   * Therefore, it is preferred to use rotateNullifierKeys on AccountWallet, as that handles the call to the Key Registry as well.
+   *
+   * This does not hinder our ability to spend notes tied to a previous master nullifier public key, provided we have the master nullifier secret key for it.
+   */
+  rotateNskM(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

@@ -14,8 +14,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';
 
@@ -37,7 +36,6 @@ describe('e2e_key_rotation', () => {
   let teardownA: () => Promise<void>;
   let teardownB: () => Promise<void>;
 
-  let keyRegistryWithB: KeyRegistryContract;
   let testContract: TestContract;
   let contractWithWalletA: TokenContract;
   let contractWithWalletB: TokenContract;
@@ -58,7 +56,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 +173,11 @@ 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();
+      // This function saves the new nullifier secret key for the account in our PXE,
+      // and calls the key registry with the derived nullifier public key.
+      await walletB.rotateNullifierKeys(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;
+  }
 }