refactor: refactor key rotate and address comments from 6405 (#6450)

Addresses comments from #6405, refactors key rotate

---------

Co-authored-by: Jan Beneš <janbenes1234@gmail.com>

docs/docs/guides/js_apps/rotate_keys.md

@@ -26,16 +26,6 @@ You will need to import these from Aztec.js:
 
 ## Rotate nullifier secret and public key
 
-Call `rotateMasterNullifierKey` on the PXE to rotate the secret key. 
+Call `rotateNullifierKeys` on the AccountWallet to rotate the secret key in the PXE and call the key registry with the new derived public key.
 
-#include_code rotateMasterNullifierKey yarn-project/end-to-end/src/e2e_key_rotation.test.ts typescript
-
-## Rotate public key 
-
-Connect to the key registry contract with your wallet. 
-
-#include_code keyRegistryWithB yarn-project/end-to-end/src/e2e_key_rotation.test.ts typescript
-
-Then `rotate_npk_m` on the key registry contract to rotate the public key:
-
-#include_code rotate_npk_m yarn-project/end-to-end/src/e2e_key_rotation.test.ts typescript
+#include_code rotateNullifierKeys yarn-project/end-to-end/src/e2e_key_rotation.test.ts typescript

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,39 @@ 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.PUBLIC,
+      isInternal: false,
+      isStatic: 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

@@ -15,11 +15,9 @@ import {
 } from '@aztec/aztec.js';
 // docs:start:imports
 import { type PublicKey, derivePublicKeyFromSecretKey } from '@aztec/circuits.js';
-import { KeyRegistryContract } from '@aztec/noir-contracts.js';
-// docs:end:imports
 import { TestContract, TokenContract } from '@aztec/noir-contracts.js';
-import { getCanonicalKeyRegistryAddress } from '@aztec/protocol-contracts/key-registry';
 
+// docs:end:imports
 import { jest } from '@jest/globals';
 
 import { expectsNumOfNoteEncryptedLogsInTheLastBlockToBe, setup, setupPXEService } from './fixtures/utils.js';
@@ -40,7 +38,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;
@@ -59,10 +56,8 @@ describe('e2e_key_rotation', () => {
     } = await setup(1));
 
     ({ pxe: pxeB, teardown: teardownB } = await setupPXEService(aztecNode, {}, undefined, true));
-    // docs:start:keyRegistryWithB
     [walletB] = await createAccounts(pxeB, 1);
-    keyRegistryWithB = await KeyRegistryContract.at(getCanonicalKeyRegistryAddress(), walletB);
-    // docs:end:keyRegistryWithB
+
     // We deploy test and token contracts
     testContract = await TestContract.deploy(walletA).send().deployed();
     const tokenInstance = await deployTokenContract(initialBalance, walletA.getAddress(), pxeA);
@@ -181,12 +176,12 @@ describe('e2e_key_rotation', () => {
       const newNskM = Fq.random();
       newNpkM = derivePublicKeyFromSecretKey(newNskM);
       // docs:end:create_keys
-      // docs:start:rotateMasterNullifierKey
-      await pxeB.rotateMasterNullifierKey(walletB.getAddress(), newNskM);
-      // docs:end:rotateMasterNullifierKey
-      // docs:start:rotate_npk_m
-      await keyRegistryWithB.methods.rotate_npk_m(walletB.getAddress(), newNpkM, 0).send().wait();
-      // docs:end:rotate_npk_m
+
+      // docs:start:rotateNullifierKeys
+      // 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);
+      // docs:end:rotateNullifierKeys
       await crossDelay();
     }
 

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

@@ -124,7 +124,7 @@ export class KeyStore {
       }
 
       // Now we iterate over the public keys in the buffer to find the one that matches the hash
-      const numKeys = pkMsBuffer.byteLength / Point.SIZE_IN_BYTES;
+      const numKeys = this.#calculateNumKeys(pkMsBuffer, Point);
       for (; keyIndexInBuffer < numKeys; keyIndexInBuffer++) {
         const foundPkM = Point.fromBuffer(
           pkMsBuffer.subarray(keyIndexInBuffer * Point.SIZE_IN_BYTES, (keyIndexInBuffer + 1) * Point.SIZE_IN_BYTES),
@@ -289,7 +289,7 @@ export class KeyStore {
         );
       }
 
-      const numKeys = secretKeysBuffer.byteLength / GrumpkinScalar.SIZE_IN_BYTES;
+      const numKeys = this.#calculateNumKeys(secretKeysBuffer, GrumpkinScalar);
       for (let i = 0; i < numKeys; i++) {
         const foundSk = GrumpkinScalar.fromBuffer(
           secretKeysBuffer.subarray(i * GrumpkinScalar.SIZE_IN_BYTES, (i + 1) * GrumpkinScalar.SIZE_IN_BYTES),
@@ -396,4 +396,8 @@ export class KeyStore {
 
     await this.#keys.set(key, serializeToBuffer([currentValue, value]));
   }
+
+  #calculateNumKeys(buf: Buffer, T: typeof Point | typeof Fq) {
+    return buf.byteLength / T.SIZE_IN_BYTES;
+  }
 }

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

@@ -37,7 +37,7 @@ import {
 import { computeNoteHashNonce, siloNullifier } from '@aztec/circuits.js/hash';
 import { type ContractArtifact, type DecodedReturn, FunctionSelector, encodeArguments } from '@aztec/foundation/abi';
 import { arrayNonEmptyLength, padArrayEnd } from '@aztec/foundation/collection';
-import { Fq, Fr } from '@aztec/foundation/fields';
+import { type Fq, Fr } from '@aztec/foundation/fields';
 import { SerialQueue } from '@aztec/foundation/fifo';
 import { type DebugLogger, createDebugLogger } from '@aztec/foundation/log';
 import { Timer } from '@aztec/foundation/timer';
@@ -161,6 +161,10 @@ export class PXEService implements PXE {
     return this.db.getAuthWitness(messageHash);
   }
 
+  async rotateNskM(account: AztecAddress, secretKey: Fq): Promise<void> {
+    await this.keyStore.rotateMasterNullifierKey(account, secretKey);
+  }
+
   public addCapsule(capsule: Fr[]) {
     return this.db.addCapsule(capsule);
   }
@@ -193,10 +197,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();