Optimise performance of getting public keys (#213)

* Optimise performance of getting public keys

* Fix tests

* Add extra check in `deriveSecretExtension`

* Move `PUBLIC_KEY_GUARD` to separate file and add comment

* Update documentation for `publicKey` parameter

* Re-use public key for deriving CIP-3 nodes

Author: Mrtenz

src/SLIP10Node.test.ts

@@ -710,23 +710,23 @@ describe('SLIP10Node', () => {
       });
 
       // getter
-      expect(() => (node.privateKey = 'foo')).toThrow(
-        /^Cannot set property privateKey of .+ which has only a getter/iu,
-      );
+      ['privateKey', 'publicKeyBytes'].forEach((property) => {
+        expect(() => (node[property] = 'foo')).toThrow(
+          /^Cannot set property .+ of .+ which has only a getter/iu,
+        );
+      });
 
       // frozen / readonly
-      ['depth', 'privateKeyBytes', 'publicKeyBytes', 'chainCodeBytes'].forEach(
-        (property) => {
-          expect(() => (node[property] = new Uint8Array(64).fill(1))).toThrow(
-            expect.objectContaining({
-              name: 'TypeError',
-              message: expect.stringMatching(
-                `Cannot assign to read only property '${property}' of object`,
-              ),
-            }),
-          );
-        },
-      );
+      ['depth', 'privateKeyBytes', 'chainCodeBytes'].forEach((property) => {
+        expect(() => (node[property] = new Uint8Array(64).fill(1))).toThrow(
+          expect.objectContaining({
+            name: 'TypeError',
+            message: expect.stringMatching(
+              `Cannot assign to read only property '${property}' of object`,
+            ),
+          }),
+        );
+      });
     });
 
     it('throws an error if no curve is specified', async () => {
@@ -875,23 +875,23 @@ describe('SLIP10Node', () => {
       });
 
       // getter
-      expect(() => (node.privateKey = 'foo')).toThrow(
-        /^Cannot set property privateKey of .+ which has only a getter/iu,
-      );
+      ['privateKey', 'publicKeyBytes'].forEach((property) => {
+        expect(() => (node[property] = 'foo')).toThrow(
+          /^Cannot set property .+ of .+ which has only a getter/iu,
+        );
+      });
 
       // frozen / readonly
-      ['depth', 'privateKeyBytes', 'publicKeyBytes', 'chainCodeBytes'].forEach(
-        (property) => {
-          expect(() => (node[property] = new Uint8Array(64).fill(1))).toThrow(
-            expect.objectContaining({
-              name: 'TypeError',
-              message: expect.stringMatching(
-                `Cannot assign to read only property '${property}' of object`,
-              ),
-            }),
-          );
-        },
-      );
+      ['depth', 'privateKeyBytes', 'chainCodeBytes'].forEach((property) => {
+        expect(() => (node[property] = new Uint8Array(64).fill(1))).toThrow(
+          expect.objectContaining({
+            name: 'TypeError',
+            message: expect.stringMatching(
+              `Cannot assign to read only property '${property}' of object`,
+            ),
+          }),
+        );
+      });
     });
 
     it('throws an error if no curve is specified', async () => {
@@ -1107,6 +1107,30 @@ describe('SLIP10Node', () => {
     });
   });
 
+  describe('publicKeyBytes', () => {
+    it('lazily computes the public key bytes', async () => {
+      const baseNode = await SLIP10Node.fromDerivationPath({
+        derivationPath: [
+          defaultBip39NodeToken,
+          BIP44PurposeNodeToken,
+          `bip32:0'`,
+          `bip32:0'`,
+        ],
+        curve: 'secp256k1',
+      });
+
+      const { publicKey, ...json } = baseNode.toJSON();
+
+      const spy = jest.spyOn(secp256k1, 'getPublicKey');
+
+      const node = await SLIP10Node.fromExtendedKey(json);
+      expect(spy).not.toHaveBeenCalled();
+
+      expect(node.publicKeyBytes).toStrictEqual(baseNode.publicKeyBytes);
+      expect(spy).toHaveBeenCalled();
+    });
+  });
+
   describe('compressedPublicKeyBytes', () => {
     it('returns the public key in compressed form', async () => {
       const node = await SLIP10Node.fromDerivationPath({

src/SLIP10Node.ts

@@ -2,19 +2,20 @@ import { assert, bytesToHex } from '@metamask/utils';
 
 import type { BIP44CoinTypeNode } from './BIP44CoinTypeNode';
 import type { BIP44Node } from './BIP44Node';
+import { BYTES_KEY_LENGTH } from './constants';
 import type {
   Network,
   RootedSLIP10PathTuple,
   SLIP10PathTuple,
 } from './constants';
-import { BYTES_KEY_LENGTH } from './constants';
 import type { CryptographicFunctions } from './cryptography';
 import type { SupportedCurve } from './curves';
 import { getCurveByName } from './curves';
 import { deriveKeyFromPath } from './derivation';
 import { publicKeyToEthAddress } from './derivers/bip32';
 import { getDerivationPathWithSeed } from './derivers/bip39';
 import { decodeExtendedKey, encodeExtendedKey } from './extended-keys';
+import { PUBLIC_KEY_GUARD } from './guard';
 import {
   getBytes,
   getBytesUnsafe,
@@ -99,18 +100,32 @@ export type SLIP10NodeInterface = JsonSLIP10Node & {
   toJSON(): JsonSLIP10Node;
 };
 
-export type SLIP10NodeConstructorOptions = {
+type BaseSLIP10NodeConstructorOptions = {
   readonly depth: number;
   readonly masterFingerprint?: number | undefined;
   readonly parentFingerprint: number;
   readonly index: number;
   readonly network?: Network | undefined;
   readonly chainCode: Uint8Array;
-  readonly privateKey?: Uint8Array | undefined;
-  readonly publicKey: Uint8Array;
   readonly curve: SupportedCurve;
 };
 
+type SLIP10NodePrivateKeyConstructorOptions =
+  BaseSLIP10NodeConstructorOptions & {
+    readonly privateKey: Uint8Array;
+    readonly publicKey?: Uint8Array | undefined;
+  };
+
+type SLIP10NodePublicKeyConstructorOptions =
+  BaseSLIP10NodeConstructorOptions & {
+    readonly privateKey?: Uint8Array | undefined;
+    readonly publicKey: Uint8Array;
+  };
+
+export type SLIP10NodeConstructorOptions =
+  | SLIP10NodePrivateKeyConstructorOptions
+  | SLIP10NodePublicKeyConstructorOptions;
+
 export type SLIP10ExtendedKeyOptions = {
   readonly depth: number;
   readonly masterFingerprint?: number | undefined;
@@ -121,6 +136,12 @@ export type SLIP10ExtendedKeyOptions = {
   readonly privateKey?: string | Uint8Array | undefined;
   readonly publicKey?: string | Uint8Array | undefined;
   readonly curve: SupportedCurve;
+
+  /**
+   * For internal use only. This is used to ensure the public key provided to
+   * the constructor is trusted.
+   */
+  readonly guard?: typeof PUBLIC_KEY_GUARD;
 };
 
 export type SLIP10DerivationPathOptions = {
@@ -276,6 +297,7 @@ export class SLIP10Node implements SLIP10NodeInterface {
       publicKey,
       chainCode,
       curve,
+      guard,
     } = options;
 
     const chainCodeBytes = getBytes(chainCode, BYTES_KEY_LENGTH);
@@ -299,11 +321,20 @@ export class SLIP10Node implements SLIP10NodeInterface {
         privateKey,
         curveObject.privateKeyLength,
       );
+
       assert(
         curveObject.isValidPrivateKey(privateKeyBytes),
         `Invalid private key: Value is not a valid ${curve} private key.`,
       );
 
+      const trustedPublicKey =
+        guard === PUBLIC_KEY_GUARD && publicKey
+          ? // `publicKey` is typed as `string | Uint8Array`, but we know it's
+            // a `Uint8Array` because of the guard. We use `getBytes` to ensure
+            // the type is correct.
+            getBytes(publicKey, curveObject.publicKeyLength)
+          : undefined;
+
       return new SLIP10Node(
         {
           depth,
@@ -313,7 +344,7 @@ export class SLIP10Node implements SLIP10NodeInterface {
           network,
           chainCode: chainCodeBytes,
           privateKey: privateKeyBytes,
-          publicKey: await curveObject.getPublicKey(privateKeyBytes),
+          publicKey: trustedPublicKey,
           curve,
         },
         cryptographicFunctions,
@@ -478,7 +509,7 @@ export class SLIP10Node implements SLIP10NodeInterface {
 
   public readonly privateKeyBytes?: Uint8Array | undefined;
 
-  public readonly publicKeyBytes: Uint8Array;
+  #publicKeyBytes: Uint8Array | undefined;
 
   readonly #cryptographicFunctions: CryptographicFunctions;
 
@@ -503,15 +534,20 @@ export class SLIP10Node implements SLIP10NodeInterface {
       'SLIP10Node can only be constructed using `SLIP10Node.fromJSON`, `SLIP10Node.fromExtendedKey`, `SLIP10Node.fromDerivationPath`, or `SLIP10Node.fromSeed`.',
     );
 
+    assert(
+      privateKey !== undefined || publicKey !== undefined,
+      'SLIP10Node requires either a private key or a public key to be set.',
+    );
+
     this.depth = depth;
     this.masterFingerprint = masterFingerprint;
     this.parentFingerprint = parentFingerprint;
     this.index = index;
     this.network = network;
     this.chainCodeBytes = chainCode;
     this.privateKeyBytes = privateKey;
-    this.publicKeyBytes = publicKey;
     this.curve = curve;
+    this.#publicKeyBytes = publicKey;
     this.#cryptographicFunctions = cryptographicFunctions;
 
     Object.freeze(this);
@@ -533,6 +569,31 @@ export class SLIP10Node implements SLIP10NodeInterface {
     return bytesToHex(this.publicKeyBytes);
   }
 
+  /**
+   * Get the public key bytes. This will lazily derive the public key from the
+   * private key if it is not already set.
+   *
+   * @returns The public key bytes.
+   */
+  public get publicKeyBytes(): Uint8Array {
+    if (this.#publicKeyBytes !== undefined) {
+      return this.#publicKeyBytes;
+    }
+
+    // This assertion is mainly for type safety, as `SLIP10Node` requires either
+    // a private key or a public key to always be set.
+    assert(
+      this.privateKeyBytes,
+      'Either a private key or public key is required.',
+    );
+
+    this.#publicKeyBytes = getCurveByName(this.curve).getPublicKey(
+      this.privateKeyBytes,
+    );
+
+    return this.#publicKeyBytes;
+  }
+
   public get compressedPublicKeyBytes(): Uint8Array {
     return getCurveByName(this.curve).compressPublicKey(this.publicKeyBytes);
   }

src/curves/curve.ts

@@ -27,10 +27,7 @@ export type Curve = {
   curve: {
     n: bigint;
   };
-  getPublicKey: (
-    privateKey: Uint8Array,
-    compressed?: boolean,
-  ) => Uint8Array | Promise<Uint8Array>;
+  getPublicKey: (privateKey: Uint8Array, compressed?: boolean) => Uint8Array;
   isValidPrivateKey: (privateKey: Uint8Array) => boolean;
   publicAdd: (publicKey: Uint8Array, tweak: Uint8Array) => Uint8Array;
   compressPublicKey: (publicKey: Uint8Array) => Uint8Array;

src/curves/ed25519Bip32.test.ts

@@ -13,8 +13,8 @@ import fixtures from '../../test/fixtures';
 describe('getPublicKey', () => {
   fixtures.cip3.forEach((fixture) => {
     Object.values(fixture.nodes).forEach((node) => {
-      it('returns correct public key from private key', async () => {
-        const publicKey = await ed25519Bip32.getPublicKey(
+      it('returns correct public key from private key', () => {
+        const publicKey = ed25519Bip32.getPublicKey(
           hexToBytes(node.privateKey),
         );
 

src/curves/ed25519Bip32.ts

@@ -99,10 +99,10 @@ export const multiplyWithBase = (key: Uint8Array): Uint8Array => {
  * @param _compressed - Optional parameter to indicate if the public key should be compressed.
  * @returns The public key.
  */
-export const getPublicKey = async (
+export const getPublicKey = (
   privateKey: Uint8Array,
   _compressed?: boolean,
-): Promise<Uint8Array> => {
+): Uint8Array => {
   return multiplyWithBase(privateKey.slice(0, 32));
 };
 

src/derivers/bip32.ts

@@ -96,14 +96,15 @@ async function handleError(
   options: DeriveNodeArgs,
   cryptographicFunctions?: CryptographicFunctions,
 ): Promise<DeriveNodeArgs> {
-  const { childIndex, privateKey, publicKey, isHardened, curve, chainCode } =
+  const { childIndex, privateKey, publicKey, isHardened, chainCode, curve } =
     options;
 
   validateBIP32Index(childIndex + 1);
 
   if (privateKey) {
     const secretExtension = await deriveSecretExtension({
       privateKey,
+      publicKey: curve.compressPublicKey(publicKey),
       childIndex: childIndex + 1,
       isHardened,
       curve,

src/derivers/bip39.ts

@@ -15,6 +15,7 @@ import type { CryptographicFunctions } from '../cryptography';
 import { hmacSha512, pbkdf2Sha512 } from '../cryptography';
 import type { Curve, SupportedCurve } from '../curves';
 import { getCurveByName } from '../curves';
+import { PUBLIC_KEY_GUARD } from '../guard';
 import { SLIP10Node } from '../SLIP10Node';
 import { getFingerprint } from '../utils';
 
@@ -244,21 +245,24 @@ export async function createBip39KeyFromSeed(
     'Invalid private key: The private key must greater than 0 and less than the curve order.',
   );
 
+  const publicKey = curve.getPublicKey(privateKey, false);
   const masterFingerprint = getFingerprint(
-    await curve.getPublicKey(privateKey, true),
+    curve.compressPublicKey(publicKey),
     curve.compressedPublicKeyLength,
   );
 
   return SLIP10Node.fromExtendedKey(
     {
       privateKey,
+      publicKey,
       chainCode,
       masterFingerprint,
       network,
       depth: 0,
       parentFingerprint: 0,
       index: 0,
       curve: curve.name,
+      guard: PUBLIC_KEY_GUARD,
     },
     cryptographicFunctions,
   );
@@ -310,21 +314,24 @@ export async function entropyToCip3MasterNode(
 
   assert(curve.isValidPrivateKey(privateKey), 'Invalid private key.');
 
+  const publicKey = curve.getPublicKey(privateKey, false);
   const masterFingerprint = getFingerprint(
-    await curve.getPublicKey(privateKey),
+    curve.compressPublicKey(publicKey),
     curve.compressedPublicKeyLength,
   );
 
   return SLIP10Node.fromExtendedKey(
     {
       privateKey,
+      publicKey,
       chainCode,
       masterFingerprint,
       network,
       depth: 0,
       parentFingerprint: 0,
       index: 0,
       curve: curve.name,
+      guard: PUBLIC_KEY_GUARD,
     },
     cryptographicFunctions,
   );