Skip to content

Commit

Permalink
chore: remove key rotation from the key store (#8645)
Browse files Browse the repository at this point in the history 
Follow up to #8613 - the key store still had some lingering code from
back when it supported key rotation - each entry could hold multiple
keys. We no longer do this (i.e. we never append anything to an entry,
we only set it with a single key), so instead of checking if the entry
includes the hash we're looking for, we can simply check if it's the
correct (only) one.

Similarly when looking for the public key and secret key we no longer
need to iterate over the multiple entries since there's a single one.
  • Loading branch information
@nventuro
nventuro authored Sep 19, 2024
1 parent 5818018 commit d8bcb9f
Showing 1 changed file with 26 additions and 71 deletions.
97 changes: 26 additions & 71 deletions yarn-project/key-store/src/key_store.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ import { type PublicKey } from '@aztec/circuit-types';
import {
AztecAddress,
CompleteAddress,
type Fq,
Fr,
GeneratorIndex,
GrumpkinScalar,
Expand All @@ -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>;
Expand Down Expand Up @@ -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.`);
Expand Down Expand Up @@ -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);
Expand All @@ -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]);
Expand All @@ -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;
}
}

0 comments on commit d8bcb9f

Please sign in to comment.