From ec8e4385e10930c757fa9e9ba083b8422df76587 Mon Sep 17 00:00:00 2001
From: Brian Botha <brian.botha@matrix.ai>
Date: Fri, 15 Jul 2022 18:15:52 +1000
Subject: [PATCH] feat: added ability to override keypair generation with
 `privateKeyOverride` parameter for `createKeyManager`

This will skip key generation and use the provided `PrivateKey` instead. This should speed up testing by skipping the key generation.

Related #404
---
 src/PolykeyAgent.ts           |  3 ++-
 src/keys/KeyManager.ts        | 26 +++++++++++++++++++++++++-
 tests/keys/KeyManager.test.ts | 17 +++++++++++++++++
 3 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/src/PolykeyAgent.ts b/src/PolykeyAgent.ts
index e2cf14dde..a7737f913 100644
--- a/src/PolykeyAgent.ts
+++ b/src/PolykeyAgent.ts
@@ -2,7 +2,7 @@ import type { FileSystem } from './types';
 import type { PolykeyWorkerManagerInterface } from './workers/types';
 import type { ConnectionData, Host, Port } from './network/types';
 import type { SeedNodes } from './nodes/types';
-import type { KeyManagerChangeData } from './keys/types';
+import type { KeyManagerChangeData, PrivateKey } from './keys/types';
 import path from 'path';
 import process from 'process';
 import Logger from '@matrixai/logger';
@@ -108,6 +108,7 @@ class PolykeyAgent {
       rootCertDuration?: number;
       dbKeyBits?: number;
       recoveryCode?: string;
+      privateKeyOverride?: PrivateKey;
     };
     proxyConfig?: {
       authToken?: string;
diff --git a/src/keys/KeyManager.ts b/src/keys/KeyManager.ts
index 14206b67a..9c30899c2 100644
--- a/src/keys/KeyManager.ts
+++ b/src/keys/KeyManager.ts
@@ -6,6 +6,7 @@ import type {
   CertificatePemChain,
   RecoveryCode,
   KeyManagerChangeData,
+  PrivateKey,
 } from './types';
 import type { FileSystem } from '../types';
 import type { NodeId } from '../nodes/types';
@@ -40,6 +41,7 @@ class KeyManager {
     fs = require('fs'),
     logger = new Logger(this.name),
     recoveryCode,
+    privateKeyOverride,
     fresh = false,
   }: {
     keysPath: string;
@@ -51,6 +53,7 @@ class KeyManager {
     fs?: FileSystem;
     logger?: Logger;
     recoveryCode?: RecoveryCode;
+    privateKeyOverride?: PrivateKey;
     fresh?: boolean;
   }): Promise<KeyManager> {
     logger.info(`Creating ${this.name}`);
@@ -67,6 +70,7 @@ class KeyManager {
     await keyManager.start({
       password,
       recoveryCode,
+      privateKeyOverride,
       fresh,
     });
     logger.info(`Created ${this.name}`);
@@ -134,10 +138,12 @@ class KeyManager {
   public async start({
     password,
     recoveryCode,
+    privateKeyOverride,
     fresh = false,
   }: {
     password: string;
     recoveryCode?: RecoveryCode;
+    privateKeyOverride?: PrivateKey;
     fresh?: boolean;
   }): Promise<void> {
     this.logger.info(`Starting ${this.constructor.name}`);
@@ -160,6 +166,7 @@ class KeyManager {
       password,
       this.rootKeyPairBits,
       recoveryCode,
+      privateKeyOverride,
     );
     const rootCert = await this.setupRootCert(
       rootKeyPair,
@@ -561,7 +568,7 @@ class KeyManager {
     bits: number,
     recoveryCode?: RecoveryCode,
   ): Promise<KeyPair> {
-    let keyPair;
+    let keyPair: KeyPair;
     if (this.workerManager) {
       keyPair = await this.workerManager.call(async (w) => {
         let keyPair;
@@ -588,10 +595,20 @@ class KeyManager {
     return keyPair;
   }
 
+  /**
+   * Generates and writes the encrypted keypair to a the root key file.
+   * If privateKeyOverride is provided then key generation is skipped in favor of the provided key.
+   * If state already exists the privateKeyOverride is ignored.
+   * @param password
+   * @param bits - Bit-width of the generated key.
+   * @param recoveryCode - Code to generate the key from.
+   * @param privateKeyOverride - Override generation with a provided private key.
+   */
   protected async setupRootKeyPair(
     password: string,
     bits: number = 4096,
     recoveryCode: RecoveryCode | undefined,
+    privateKeyOverride: PrivateKey | undefined,
   ): Promise<[KeyPair, RecoveryCode | undefined]> {
     let rootKeyPair: KeyPair;
     let recoveryCodeNew: RecoveryCode | undefined;
@@ -610,6 +627,13 @@ class KeyManager {
       }
       return [rootKeyPair, undefined];
     } else {
+      if (privateKeyOverride != null) {
+        this.logger.info('Using provided root key pair');
+        const publicKey = keysUtils.publicKeyFromPrivateKey(privateKeyOverride);
+        rootKeyPair = { privateKey: privateKeyOverride, publicKey };
+        await this.writeRootKeyPair(rootKeyPair, password);
+        return [rootKeyPair, undefined];
+      }
       this.logger.info('Generating root key pair');
       if (recoveryCode != null) {
         // Deterministic key pair generation from recovery code
diff --git a/tests/keys/KeyManager.test.ts b/tests/keys/KeyManager.test.ts
index c1aaa345e..c2cbab188 100644
--- a/tests/keys/KeyManager.test.ts
+++ b/tests/keys/KeyManager.test.ts
@@ -161,6 +161,23 @@ describe('KeyManager', () => {
     },
     global.defaultTimeout * 2,
   );
+  test('override key generation with privateKeyOverride', async () => {
+    const keysPath = `${dataDir}/keys`;
+    const keyPair = await keysUtils.generateKeyPair(4096);
+    const mockedGenerateKeyPair = jest.spyOn(keysUtils, 'generateDeterministicKeyPair');
+    const keyManager = await KeyManager.createKeyManager({
+      keysPath,
+      password,
+      privateKeyOverride: keyPair.privateKey,
+      logger,
+    });
+    expect(mockedGenerateKeyPair).not.toHaveBeenCalled()
+    const keysPathContents = await fs.promises.readdir(keysPath);
+    expect(keysPathContents).toContain('root.pub');
+    expect(keysPathContents).toContain('root.key');
+    expect(keysUtils.publicKeyToPem(keyManager.getRootKeyPair().publicKey)).toEqual(keysUtils.publicKeyToPem(keyPair.publicKey));
+    await keyManager.stop();
+  })
   test('uses WorkerManager for generating root key pair', async () => {
     const keysPath = `${dataDir}/keys`;
     const keyManager = await KeyManager.createKeyManager({