Skip to content

Commit

Permalink
feat: added ability to override keypair generation with `privateKeyOv…
Browse files Browse the repository at this point in the history
…erride` 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
  • Loading branch information
tegefaulkes committed Jul 27, 2022
1 parent 8c2a7ef commit ec8e438
Show file tree
Hide file tree
Showing 3 changed files with 44 additions and 2 deletions.
3 changes: 2 additions & 1 deletion src/PolykeyAgent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -108,6 +108,7 @@ class PolykeyAgent {
rootCertDuration?: number;
dbKeyBits?: number;
recoveryCode?: string;
privateKeyOverride?: PrivateKey;
};
proxyConfig?: {
authToken?: string;
Expand Down
26 changes: 25 additions & 1 deletion src/keys/KeyManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import type {
CertificatePemChain,
RecoveryCode,
KeyManagerChangeData,
PrivateKey,
} from './types';
import type { FileSystem } from '../types';
import type { NodeId } from '../nodes/types';
Expand Down Expand Up @@ -40,6 +41,7 @@ class KeyManager {
fs = require('fs'),
logger = new Logger(this.name),
recoveryCode,
privateKeyOverride,
fresh = false,
}: {
keysPath: string;
Expand All @@ -51,6 +53,7 @@ class KeyManager {
fs?: FileSystem;
logger?: Logger;
recoveryCode?: RecoveryCode;
privateKeyOverride?: PrivateKey;
fresh?: boolean;
}): Promise<KeyManager> {
logger.info(`Creating ${this.name}`);
Expand All @@ -67,6 +70,7 @@ class KeyManager {
await keyManager.start({
password,
recoveryCode,
privateKeyOverride,
fresh,
});
logger.info(`Created ${this.name}`);
Expand Down Expand Up @@ -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}`);
Expand All @@ -160,6 +166,7 @@ class KeyManager {
password,
this.rootKeyPairBits,
recoveryCode,
privateKeyOverride,
);
const rootCert = await this.setupRootCert(
rootKeyPair,
Expand Down Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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
Expand Down
17 changes: 17 additions & 0 deletions tests/keys/KeyManager.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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({
Expand Down

0 comments on commit ec8e438

Please sign in to comment.