perf: speed up valkeys generation (A-1361)

alexghr · alexghr · commit 2d40eee37c6b · 2026-07-02T14:33:47.000Z
diff --git a/yarn-project/cli/src/cmds/validator_keys/shared.ts b/yarn-project/cli/src/cmds/validator_keys/shared.ts
@@ -1,6 +1,6 @@
 import { prettyPrintJSON } from '@aztec/cli/utils';
 import { deriveBlsPrivateKey } from '@aztec/foundation/crypto/bls';
-import { createBn254Keystore } from '@aztec/foundation/crypto/bls/bn254_keystore';
+import { createBn254KeystoreAsync } from '@aztec/foundation/crypto/bls/bn254_keystore';
 import { computeBn254G1PublicKeyCompressed } from '@aztec/foundation/crypto/bn254';
 import type { EthAddress } from '@aztec/foundation/eth-address';
 import type { LogFn } from '@aztec/foundation/log';
@@ -91,18 +91,19 @@ export async function buildValidatorEntries(input: BuildValidatorsInput) {
     remoteSigner,
   } = input;
 
-  const summaries: ValidatorSummary[] = [];
-
-  const validators = await Promise.all(
+  const entries = await Promise.all(
     Array.from({ length: validatorCount }, async (_unused, i) => {
       const addressIndex = baseAddressIndex + i;
       const basePath = blsPath ?? defaultBlsPath;
       const perValidatorPath = withValidatorIndex(basePath, accountIndex, addressIndex);
 
       const blsPrivKey = ikm || mnemonic ? deriveBlsPrivateKey(mnemonic, ikm, perValidatorPath) : undefined;
       const blsPubCompressed = blsPrivKey ? await computeBlsPublicKeyCompressed(blsPrivKey) : undefined;
-
-      const ethAttester = deriveEthAttester(mnemonic, accountIndex, addressIndex, remoteSigner);
+      const attesterAccount = mnemonicToAccount(mnemonic, { accountIndex, addressIndex });
+      const attesterEthAddress = attesterAccount.address as unknown as string;
+      const ethAttester = remoteSigner
+        ? ({ address: attesterAccount.address as unknown as EthAddress, remoteSignerUrl: remoteSigner } as EthAccount)
+        : (('0x' + Buffer.from(attesterAccount.getHdKey().privateKey!).toString('hex')) as EthPrivateKey);
       const attester = blsPrivKey ? { eth: ethAttester, bls: blsPrivKey } : ethAttester;
 
       let publisherField: EthAccount | EthPrivateKey | (EthAccount | EthPrivateKey)[] | undefined;
@@ -126,27 +127,23 @@ export async function buildValidatorEntries(input: BuildValidatorsInput) {
         publisherField = publisherCount === 1 ? publisherAccounts[0] : publisherAccounts;
       }
 
-      const acct = mnemonicToAccount(mnemonic, {
-        accountIndex,
-        addressIndex,
-      });
-      const attesterEthAddress = acct.address as unknown as string;
-      summaries.push({
-        attesterEth: attesterEthAddress,
-        attesterBls: blsPubCompressed,
-        publisherEth: publisherAddresses.length > 0 ? publisherAddresses : undefined,
-      });
-
       return {
-        attester,
-        ...(publisherField !== undefined ? { publisher: publisherField } : {}),
-        feeRecipient,
-        coinbase: coinbase ?? attesterEthAddress,
-      } as ValidatorKeyStore;
+        validator: {
+          attester,
+          ...(publisherField !== undefined ? { publisher: publisherField } : {}),
+          feeRecipient,
+          coinbase: coinbase ?? attesterEthAddress,
+        } as ValidatorKeyStore,
+        summary: {
+          attesterEth: attesterEthAddress,
+          attesterBls: blsPubCompressed,
+          publisherEth: publisherAddresses.length > 0 ? publisherAddresses : undefined,
+        } as ValidatorSummary,
+      };
     }),
   );
 
-  return { validators, summaries };
+  return { validators: entries.map(entry => entry.validator), summaries: entries.map(entry => entry.summary) };
 }
 
 export async function resolveKeystoreOutputPath(dataDir?: string, file?: string) {
@@ -228,7 +225,7 @@ export async function writeBn254BlsKeystore(
 ): Promise<string> {
   mkdirSync(outDir, { recursive: true });
 
-  const keystore = createBn254Keystore(password, privateKeyHex, pubkeyHex, derivationPath);
+  const keystore = await createBn254KeystoreAsync(password, privateKeyHex, pubkeyHex, derivationPath);
 
   const safeBase = fileNameBase.replace(/[^a-zA-Z0-9_-]/g, '_');
   const outPath = join(outDir, `keystore-${safeBase}.json`);
@@ -241,28 +238,29 @@ export async function writeBlsBn254ToFile(
   validators: ValidatorKeyStore[],
   options: { outDir: string; password: string; blsPath?: string },
 ): Promise<void> {
-  for (let i = 0; i < validators.length; i++) {
-    const v = validators[i];
-    if (!v || typeof v !== 'object' || !('attester' in v)) {
-      continue;
-    }
-    const att = (v as any).attester;
+  await Promise.all(
+    validators.map(async (v, i) => {
+      if (!v || typeof v !== 'object' || !('attester' in v)) {
+        return;
+      }
+      const att = (v as any).attester;
 
-    // Shapes: { bls: <hex> } or { eth: <ethAccount>, bls?: <hex> } or plain EthAccount
-    const blsKey: string | undefined = typeof att === 'object' && 'bls' in att ? (att as any).bls : undefined;
-    if (!blsKey || typeof blsKey !== 'string') {
-      continue;
-    }
+      // Shapes: { bls: <hex> } or { eth: <ethAccount>, bls?: <hex> } or plain EthAccount
+      const blsKey: string | undefined = typeof att === 'object' && 'bls' in att ? (att as any).bls : undefined;
+      if (!blsKey || typeof blsKey !== 'string') {
+        return;
+      }
 
-    const pub = await computeBlsPublicKeyCompressed(blsKey);
-    const path = options.blsPath ?? defaultBlsPath;
-    const fileBase = `${String(i + 1)}_${pub.slice(2, 18)}`;
-    const keystorePath = await writeBn254BlsKeystore(options.outDir, fileBase, options.password, blsKey, pub, path);
+      const pub = await computeBlsPublicKeyCompressed(blsKey);
+      const path = options.blsPath ?? defaultBlsPath;
+      const fileBase = `${String(i + 1)}_${pub.slice(2, 18)}`;
+      const keystorePath = await writeBn254BlsKeystore(options.outDir, fileBase, options.password, blsKey, pub, path);
 
-    if (typeof att === 'object') {
-      (att as any).bls = { path: keystorePath, password: options.password };
-    }
-  }
+      if (typeof att === 'object') {
+        (att as any).bls = { path: keystorePath, password: options.password };
+      }
+    }),
+  );
 }
 
 /** Writes an Ethereum JSON V3 keystore using ethers, returns absolute path */
@@ -295,32 +293,31 @@ export async function writeEthJsonV3ToFile(
     return account;
   };
 
-  for (let i = 0; i < validators.length; i++) {
-    const v = validators[i];
-    if (!v || typeof v !== 'object') {
-      continue;
-    }
+  await Promise.all(
+    validators.map(async (v, i) => {
+      if (!v || typeof v !== 'object') {
+        return;
+      }
 
-    // attester may be string (eth), object with eth, or remote signer
-    const att = (v as any).attester;
-    if (typeof att === 'string') {
-      (v as any).attester = await maybeEncryptEth(att, `attester_${i + 1}`);
-    } else if (att && typeof att === 'object' && 'eth' in att) {
-      (att as any).eth = await maybeEncryptEth((att as any).eth, `attester_${i + 1}`);
-    }
+      // attester may be string (eth), object with eth, or remote signer
+      const att = (v as any).attester;
+      if (typeof att === 'string') {
+        (v as any).attester = await maybeEncryptEth(att, `attester_${i + 1}`);
+      } else if (att && typeof att === 'object' && 'eth' in att) {
+        (att as any).eth = await maybeEncryptEth((att as any).eth, `attester_${i + 1}`);
+      }
 
-    // publisher can be single or array
-    if ('publisher' in v) {
-      const pub = (v as any).publisher;
-      if (Array.isArray(pub)) {
-        const out: any[] = [];
-        for (let j = 0; j < pub.length; j++) {
-          out.push(await maybeEncryptEth(pub[j], `publisher_${i + 1}_${j + 1}`));
+      // publisher can be single or array
+      if ('publisher' in v) {
+        const pub = (v as any).publisher;
+        if (Array.isArray(pub)) {
+          (v as any).publisher = await Promise.all(
+            pub.map((account, j) => maybeEncryptEth(account, `publisher_${i + 1}_${j + 1}`)),
+          );
+        } else if (pub !== undefined) {
+          (v as any).publisher = await maybeEncryptEth(pub, `publisher_${i + 1}`);
         }
-        (v as any).publisher = out;
-      } else if (pub !== undefined) {
-        (v as any).publisher = await maybeEncryptEth(pub, `publisher_${i + 1}`);
       }
-    }
-  }
+    }),
+  );
 }
diff --git a/yarn-project/cli/src/cmds/validator_keys/valkeys.test.ts b/yarn-project/cli/src/cmds/validator_keys/valkeys.test.ts
@@ -181,6 +181,21 @@ describe('validator keys utilities', () => {
       expect(summaries[0].publisherEth!.length).toBe(3);
     });
 
+    it('preserves summary order for multiple validators', async () => {
+      const { summaries } = await buildValidatorEntries({
+        validatorCount: 3,
+        publisherCount: 0,
+        accountIndex: 0,
+        baseAddressIndex: 0,
+        mnemonic: TEST_MNEMONIC,
+        feeRecipient: feeRecipient,
+      });
+
+      expect(summaries.map(summary => summary.attesterEth)).toEqual(
+        [0, 1, 2].map(addressIndex => mnemonicToAccount(TEST_MNEMONIC, { accountIndex: 0, addressIndex }).address),
+      );
+    });
+
     it('derives different BLS and ETH keys when account index changes', async () => {
       // Build with account index 0
       const { validators: v1, summaries: s1 } = await buildValidatorEntries({
diff --git a/yarn-project/foundation/src/crypto/bls/bn254_keystore.ts b/yarn-project/foundation/src/crypto/bls/bn254_keystore.ts
@@ -1,7 +1,8 @@
 import { randomBytes } from '@aztec/foundation/crypto/random';
 
-import { createCipheriv, createDecipheriv, createHash, pbkdf2Sync, randomUUID } from 'crypto';
+import { createCipheriv, createDecipheriv, createHash, pbkdf2, pbkdf2Sync, randomUUID } from 'crypto';
 import { readFileSync } from 'fs';
+import { promisify } from 'util';
 import { z } from 'zod';
 
 /**
@@ -100,36 +101,23 @@ export interface Bn254KeystoreInterface {
   version: number;
 }
 
-/**
- * Creates a BN254 keystore object for a BN254 BLS private key.
- *
- * Uses PBKDF2 with SHA-256 for key derivation and AES-128-CTR for encryption,
- * following the EIP-2335 specification format.
- *
- * @param password - Password for encrypting the private key
- * @param privateKeyHex - Private key as 0x-prefixed hex string (32 bytes)
- * @param pubkeyHex - Public key as hex string (compressed or uncompressed)
- * @param derivationPath - BIP-44 style derivation path (e.g., "m/12381/3600/0/0/0")
- * @returns BN254 keystore object ready to be serialized to JSON
- * @throws Error if private key is not 32-byte hex
- */
-export function createBn254Keystore(
-  password: string,
+const pbkdf2Async = promisify(pbkdf2);
+
+function createBn254KeystoreFromDerivedKey(
   privateKeyHex: string,
   pubkeyHex: string,
   derivationPath: string,
+  salt: Buffer,
+  iv: Buffer,
+  dk: Buffer,
 ): Bn254Keystore {
   const ensureHex = (hex: string) => hex.replace(/^0x/i, '');
   const privHex = ensureHex(privateKeyHex);
   if (!/^[0-9a-fA-F]{64}$/.test(privHex)) {
     throw new Error('BLS private key must be 32-byte hex');
   }
 
-  const salt = randomBytes(32);
-  const iv = randomBytes(16);
-  const dk = pbkdf2Sync(Buffer.from(password.normalize('NFKD'), 'utf8'), salt, 262144, 32, 'sha256');
   const cipherKey = dk.subarray(0, 16);
-
   const cipher = createCipheriv('aes-128-ctr', cipherKey, iv);
   const plaintext = Buffer.from(privHex, 'hex');
   const ciphertext = Buffer.concat([cipher.update(plaintext), cipher.final()]);
@@ -166,6 +154,43 @@ export function createBn254Keystore(
   };
 }
 
+/**
+ * Creates a BN254 keystore object for a BN254 BLS private key.
+ *
+ * Uses PBKDF2 with SHA-256 for key derivation and AES-128-CTR for encryption,
+ * following the EIP-2335 specification format.
+ *
+ * @param password - Password for encrypting the private key
+ * @param privateKeyHex - Private key as 0x-prefixed hex string (32 bytes)
+ * @param pubkeyHex - Public key as hex string (compressed or uncompressed)
+ * @param derivationPath - BIP-44 style derivation path (e.g., "m/12381/3600/0/0/0")
+ * @returns BN254 keystore object ready to be serialized to JSON
+ * @throws Error if private key is not 32-byte hex
+ */
+export function createBn254Keystore(
+  password: string,
+  privateKeyHex: string,
+  pubkeyHex: string,
+  derivationPath: string,
+): Bn254Keystore {
+  const salt = randomBytes(32);
+  const iv = randomBytes(16);
+  const dk = pbkdf2Sync(Buffer.from(password.normalize('NFKD'), 'utf8'), salt, 262144, 32, 'sha256');
+  return createBn254KeystoreFromDerivedKey(privateKeyHex, pubkeyHex, derivationPath, salt, iv, dk);
+}
+
+export async function createBn254KeystoreAsync(
+  password: string,
+  privateKeyHex: string,
+  pubkeyHex: string,
+  derivationPath: string,
+): Promise<Bn254Keystore> {
+  const salt = randomBytes(32);
+  const iv = randomBytes(16);
+  const dk = await pbkdf2Async(Buffer.from(password.normalize('NFKD'), 'utf8'), salt, 262144, 32, 'sha256');
+  return createBn254KeystoreFromDerivedKey(privateKeyHex, pubkeyHex, derivationPath, salt, iv, dk);
+}
+
 /**
  * Loads and validates a BN254 keystore file.
  *
