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

sklppy88 · benesjan · web-flow · commit 6f3dab87f1cc · 2024-05-21T09:12:46.000-05:00
Addresses comments from #6405, refactors key rotate --------- Co-authored-by: Jan Beneš <janbenes1234@gmail.com>
diff --git a/docs/docs/guides/js_apps/rotate_keys.md b/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
diff --git a/noir-projects/noir-contracts/contracts/child_contract/src/main.nr b/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,
diff --git a/noir-projects/noir-contracts/contracts/docs_example_contract/src/options.nr b/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,
diff --git a/noir-projects/noir-contracts/contracts/stateful_test_contract/src/main.nr b/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);
         }
     }
diff --git a/yarn-project/aztec.js/src/account/interface.ts b/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
diff --git a/yarn-project/aztec.js/src/account/wallet.ts b/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;
diff --git a/yarn-project/aztec.js/src/wallet/account_wallet.ts b/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: [],
+    };
+  }
 }
diff --git a/yarn-project/aztec.js/src/wallet/base_wallet.ts b/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);
diff --git a/yarn-project/aztec.js/src/wallet/signerless_wallet.ts b/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.');
+  }
 }
diff --git a/yarn-project/circuit-types/src/interfaces/pxe.ts b/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.
    *
diff --git a/yarn-project/end-to-end/src/e2e_key_rotation.test.ts b/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();
     }
 
diff --git a/yarn-project/key-store/src/key_store.ts b/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;
+  }
 }
diff --git a/yarn-project/pxe/src/pxe_service/pxe_service.ts b/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();
diff --git a/yarn-project/simulator/src/client/private_execution.test.ts b/yarn-project/simulator/src/client/private_execution.test.ts

Original file line number	Diff line number	Diff line change
`@@ -46,7 +46,6 @@ contract StatefulTest {`
`46`	`46`	`fn create_note_no_init_check(owner: AztecAddress, value: Field) {`
`47`	`47`	`if (value != 0) {`
`48`	`48`	`let loc = storage.notes.at(owner);`
`49`		`- // TODO (#6312): This will break with key rotation. Fix this. Will not be able to spend / increment any notes after rotating key.`
`50`	`49`	`increment(loc, value, owner);`
`51`	`50`	`}`
`52`	`51`	`}`
Original file line number	Diff line number	Diff line change
`@@ -54,6 +54,8 @@ export abstract class BaseWallet implements Wallet {`
`54`	`54`	`},`
`55`	`55`	`): Promise<AuthWitness>;`
`56`	`56`
	`57`	`+ abstract rotateNullifierKeys(newNskM: Fq): Promise<void>;`
	`58`	`+`
`57`	`59`	`getAddress() {`
`58`	`60`	`return this.getCompleteAddress().address;`
`59`	`61`	`}`
`@@ -69,8 +71,8 @@ export abstract class BaseWallet implements Wallet {`
`69`	`71`	`registerAccount(secretKey: Fr, partialAddress: PartialAddress): Promise<CompleteAddress> {`
`70`	`72`	`return this.pxe.registerAccount(secretKey, partialAddress);`
`71`	`73`	`}`
`72`		`- rotateMasterNullifierKey(account: AztecAddress, secretKey: Fq): Promise<void> {`
`73`		`- return this.pxe.rotateMasterNullifierKey(account, secretKey);`
	`74`	`+ rotateNskM(address: AztecAddress, secretKey: Fq) {`
	`75`	`+ return this.pxe.rotateNskM(address, secretKey);`
`74`	`76`	`}`
`75`	`77`	`registerRecipient(account: CompleteAddress): Promise<void> {`
`76`	`78`	`return this.pxe.registerRecipient(account);`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,5 @@`
`1`	`1`	`import { type AuthWitness, type PXE, type TxExecutionRequest } from '@aztec/circuit-types';`
`2`		`-import { type CompleteAddress, type Fr } from '@aztec/circuits.js';`
	`2`	`+import { type CompleteAddress, type Fq, type Fr } from '@aztec/circuits.js';`
`3`	`3`
`4`	`4`	`import { DefaultEntrypoint } from '../entrypoint/default_entrypoint.js';`
`5`	`5`	`import { type EntrypointInterface, type ExecutionRequestInit } from '../entrypoint/entrypoint.js';`
`@@ -42,4 +42,8 @@ export class SignerlessWallet extends BaseWallet {`
`42`	`42`	`createAuthWit(_messageHash: Fr): Promise<AuthWitness> {`
`43`	`43`	`throw new Error('Method not implemented.');`
`44`	`44`	`}`
	`45`	`+`
	`46`	`+ rotateNullifierKeys(_newNskM: Fq): Promise<void> {`
	`47`	`+ throw new Error('Method not implemented.');`
	`48`	`+ }`
`45`	`49`	`}`