diff --git a/yarn-project/key-store/src/key_store.ts b/yarn-project/key-store/src/key_store.ts
index f535612420a..3c6904ee805 100644
--- a/yarn-project/key-store/src/key_store.ts
+++ b/yarn-project/key-store/src/key_store.ts
@@ -2,7 +2,6 @@ import { type PublicKey } from '@aztec/circuit-types';
import {
AztecAddress,
CompleteAddress,
- type Fq,
Fr,
GeneratorIndex,
GrumpkinScalar,
@@ -21,7 +20,7 @@ import { type Bufferable, serializeToBuffer } from '@aztec/foundation/serialize'
import { type AztecKVStore, type AztecMap } from '@aztec/kv-store';
/**
- * Used for managing keys. Can hold keys of multiple accounts and allows for key rotation.
+ * Used for managing keys. Can hold keys of multiple accounts.
*/
export class KeyStore {
#keys: AztecMap<string, Buffer>;
@@ -108,53 +107,29 @@ export class KeyStore {
const [keyPrefix, account] = this.#getKeyPrefixAndAccount(pkMHash);
// Now we find the master public key for the account
- // Since each public keys buffer contains multiple public keys, we need to find the one that matches the hash.
- // Then we store the index of the key in the buffer to be able to quickly obtain the corresponding secret key.
- let pkM: PublicKey | undefined;
- let keyIndexInBuffer = 0;
- {
- const pkMsBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}pk_m`);
- if (!pkMsBuffer) {
- throw new Error(
- `Could not find ${keyPrefix}pk_m for account ${account.toString()} whose address was successfully obtained with ${keyPrefix}pk_m_hash ${pkMHash.toString()}.`,
- );
- }
+ const pkMBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}pk_m`);
+ if (!pkMBuffer) {
+ throw new Error(
+ `Could not find ${keyPrefix}pk_m for account ${account.toString()} whose address was successfully obtained with ${keyPrefix}pk_m_hash ${pkMHash.toString()}.`,
+ );
+ }
- // Now we iterate over the public keys in the buffer to find the one that matches the hash
- 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),
- );
- if (foundPkM.hash().equals(pkMHash)) {
- pkM = foundPkM;
- break;
- }
- }
+ const pkM = Point.fromBuffer(pkMBuffer);
- if (!pkM) {
- throw new Error(`Could not find ${keyPrefix}pkM for ${keyPrefix}pk_m_hash ${pkMHash.toString()}.`);
- }
+ if (!pkM.hash().equals(pkMHash)) {
+ throw new Error(`Could not find ${keyPrefix}pkM for ${keyPrefix}pk_m_hash ${pkMHash.toString()}.`);
}
// Now we find the secret key for the public key
- let skM: GrumpkinScalar | undefined;
- {
- const skMsBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}sk_m`);
- if (!skMsBuffer) {
- throw new Error(
- `Could not find ${keyPrefix}sk_m for account ${account.toString()} whose address was successfully obtained with ${keyPrefix}pk_m_hash ${pkMHash.toString()}.`,
- );
- }
-
- skM = GrumpkinScalar.fromBuffer(
- skMsBuffer.subarray(
- keyIndexInBuffer * GrumpkinScalar.SIZE_IN_BYTES,
- (keyIndexInBuffer + 1) * GrumpkinScalar.SIZE_IN_BYTES,
- ),
+ const skMBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}sk_m`);
+ if (!skMBuffer) {
+ throw new Error(
+ `Could not find ${keyPrefix}sk_m for account ${account.toString()} whose address was successfully obtained with ${keyPrefix}pk_m_hash ${pkMHash.toString()}.`,
);
}
+ const skM = GrumpkinScalar.fromBuffer(skMBuffer);
+
// We sanity check that it's possible to derive the public key from the secret key
if (!derivePublicKeyFromSecretKey(skM).equals(pkM)) {
throw new Error(`Could not derive ${keyPrefix}pkM from ${keyPrefix}skM.`);
@@ -272,30 +247,16 @@ export class KeyStore {
public getMasterSecretKey(pkM: PublicKey): Promise<GrumpkinScalar> {
const [keyPrefix, account] = this.#getKeyPrefixAndAccount(pkM);
- // We get the secret keys buffer and iterate over the values in the buffer to find the one that matches pkM
- let skM: GrumpkinScalar | undefined;
- {
- const secretKeysBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}sk_m`);
- if (!secretKeysBuffer) {
- throw new Error(
- `Could not find ${keyPrefix}sk_m for ${keyPrefix}pk_m ${pkM.toString()}. This should not happen.`,
- );
- }
-
- const numKeys = this.#calculateNumKeys(secretKeysBuffer, GrumpkinScalar);
- for (let i = 0; i < numKeys; i++) {
- const foundSkM = GrumpkinScalar.fromBuffer(
- secretKeysBuffer.subarray(i * GrumpkinScalar.SIZE_IN_BYTES, (i + 1) * GrumpkinScalar.SIZE_IN_BYTES),
- );
- if (derivePublicKeyFromSecretKey(foundSkM).equals(pkM)) {
- skM = foundSkM;
- break;
- }
- }
+ const secretKeyBuffer = this.#keys.get(`${account.toString()}-${keyPrefix}sk_m`);
+ if (!secretKeyBuffer) {
+ throw new Error(
+ `Could not find ${keyPrefix}sk_m for ${keyPrefix}pk_m ${pkM.toString()}. This should not happen.`,
+ );
+ }
- if (!skM) {
- throw new Error(`Could not find ${keyPrefix}skM for ${keyPrefix}pkM ${pkM.toString()} in secret keys buffer.`);
- }
+ const skM = GrumpkinScalar.fromBuffer(secretKeyBuffer);
+ if (!derivePublicKeyFromSecretKey(skM).equals(pkM)) {
+ throw new Error(`Could not find ${keyPrefix}skM for ${keyPrefix}pkM ${pkM.toString()} in secret keys buffer.`);
}
return Promise.resolve(skM);
@@ -310,9 +271,7 @@ export class KeyStore {
#getKeyPrefixAndAccount(value: Bufferable): [KeyPrefix, AztecAddress] {
const valueBuffer = serializeToBuffer(value);
for (const [key, val] of this.#keys.entries()) {
- // `val` can contain multiple values due to key rotation so we check if the value is in the buffer instead
- // of just calling `.equals(...)`
- if (val.includes(valueBuffer)) {
+ if (val.equals(valueBuffer)) {
for (const prefix of KEY_PREFIXES) {
if (key.includes(`-${prefix}`)) {
const account = AztecAddress.fromString(key.split('-')[0]);
@@ -323,8 +282,4 @@ export class KeyStore {
}
throw new Error(`Could not find key prefix.`);
}
-
- #calculateNumKeys(buf: Buffer, T: typeof Point | typeof Fq) {
- return buf.byteLength / T.SIZE_IN_BYTES;
- }
}