Skip to content

Commit 54f8be2

Browse files
CMCDragonkaiCMCDragonkai
committed
Added tests for starting from recovery code and deriving public key bit size from existing key pair, and using testBinUtils.processExit
1 parent 037dd71 commit 54f8be2

5 files changed

Lines changed: 195 additions & 97 deletions

File tree

‎src/bin/utils/options.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ const verbose = new commander.Option('-v, --verbose', 'Log Verbose Messages')
4242
*/
4343
const fresh = new commander.Option(
4444
'--fresh',
45-
'Ignore existing state during construction'
45+
'Ignore existing state during construction',
4646
).default(false);
4747

4848
/**

‎src/keys/KeyManager.ts‎

Lines changed: 28 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -583,15 +583,13 @@ class KeyManager {
583583
let recoveryCodeNew: RecoveryCode | undefined;
584584
if (await this.existsRootKeyPair()) {
585585
if (recoveryCode != null) {
586-
// Recover the key pair with the recovery code
587-
// Check if the generated key pair matches
588-
const rootKeyPairCheck = await this.generateKeyPair(bits, recoveryCode);
589-
if (!(await this.matchRootKeyPair(rootKeyPairCheck))) {
586+
const recoveredKeyPair = await this.recoverRootKeyPair(recoveryCode);
587+
if (recoveredKeyPair == null) {
590588
throw new keysErrors.ErrorKeysRecoveryCodeIncorrect();
591589
}
592590
// Recovered key pair, write the key pair with the new password
593-
rootKeyPair = rootKeyPairCheck;
594-
await this.writeRootKeyPair(rootKeyPairCheck, password);
591+
rootKeyPair = recoveredKeyPair;
592+
await this.writeRootKeyPair(recoveredKeyPair, password);
595593
} else {
596594
// Load key pair by decrypting with password
597595
rootKeyPair = await this.readRootKeyPair(password);
@@ -698,25 +696,43 @@ class KeyManager {
698696
}
699697
}
700698

701-
protected async matchRootKeyPair(keyPair: KeyPair): Promise<boolean> {
699+
/**
700+
* Recovers root key pair with recovery code
701+
* Checks if the generated key pair public key matches
702+
* Uses the existing key pair's public key bit size
703+
* To generate the recovered key pair
704+
*/
705+
protected async recoverRootKeyPair(
706+
recoveryCode: RecoveryCode,
707+
): Promise<KeyPair | undefined> {
702708
let publicKeyPem: string;
703709
try {
704710
publicKeyPem = await this.fs.promises.readFile(this.rootPubPath, {
705711
encoding: 'utf8',
706712
});
707713
} catch (e) {
708-
if (e.code === 'ENOENT') {
709-
return false;
710-
}
711714
throw new keysErrors.ErrorRootKeysRead(e.message, {
712715
errno: e.errno,
713716
syscall: e.syscall,
714717
code: e.code,
715718
path: e.path,
716719
});
717720
}
718-
const publicKeyPemCheck = keysUtils.publicKeyToPem(keyPair.publicKey);
719-
return publicKeyPemCheck === publicKeyPem;
721+
const rootKeyPairBits = keysUtils.publicKeyBitSize(
722+
keysUtils.publicKeyFromPem(publicKeyPem),
723+
);
724+
const recoveredKeyPair = await this.generateKeyPair(
725+
rootKeyPairBits,
726+
recoveryCode,
727+
);
728+
const recoveredPublicKeyPem = keysUtils.publicKeyToPem(
729+
recoveredKeyPair.publicKey,
730+
);
731+
if (recoveredPublicKeyPem === publicKeyPem) {
732+
return recoveredKeyPair;
733+
} else {
734+
return;
735+
}
720736
}
721737

722738
protected async setupKey(

‎tests/bin/agent/start.test.ts‎

Lines changed: 134 additions & 71 deletions
Original file line numberDiff line numberDiff line change
@@ -56,13 +56,7 @@ describe('start', () => {
5656
recoveryCode.split(' ').length === 24,
5757
).toBe(true);
5858
agentProcess.kill('SIGTERM');
59-
const [exitCode, signal] = await new Promise<
60-
[number | null, NodeJS.Signals | null]
61-
>((resolve) => {
62-
agentProcess.once('exit', (code, signal) => {
63-
resolve([code, signal]);
64-
});
65-
});
59+
const [exitCode, signal] = await testBinUtils.processExit(agentProcess);
6660
expect(exitCode).toBe(null);
6761
expect(signal).toBe('SIGTERM');
6862
// Check for graceful exit
@@ -214,26 +208,18 @@ describe('start', () => {
214208
expect(stdErrLine1).toBeDefined();
215209
expect(stdErrLine1).toBe(eOutput);
216210
agentProcess2.kill('SIGQUIT');
217-
const [exitCode, signal] = await new Promise<
218-
[number | null, NodeJS.Signals | null]
219-
>((resolve) => {
220-
agentProcess2.once('exit', (code, signal) => {
221-
resolve([code, signal]);
222-
});
223-
});
211+
const [exitCode, signal] = await testBinUtils.processExit(
212+
agentProcess2,
213+
);
224214
expect(exitCode).toBe(null);
225215
expect(signal).toBe('SIGQUIT');
226216
} else if (index === 1) {
227217
expect(stdErrLine2).toBeDefined();
228218
expect(stdErrLine2).toBe(eOutput);
229219
agentProcess1.kill('SIGQUIT');
230-
const [exitCode, signal] = await new Promise<
231-
[number | null, NodeJS.Signals | null]
232-
>((resolve) => {
233-
agentProcess1.once('exit', (code, signal) => {
234-
resolve([code, signal]);
235-
});
236-
});
220+
const [exitCode, signal] = await testBinUtils.processExit(
221+
agentProcess1,
222+
);
237223
expect(exitCode).toBe(null);
238224
expect(signal).toBe('SIGQUIT');
239225
}
@@ -303,26 +289,16 @@ describe('start', () => {
303289
expect(stdErrLine1).toBeDefined();
304290
expect(stdErrLine1).toBe(eOutput);
305291
bootstrapProcess.kill('SIGTERM');
306-
const [exitCode, signal] = await new Promise<
307-
[number | null, NodeJS.Signals | null]
308-
>((resolve) => {
309-
bootstrapProcess.once('exit', (code, signal) => {
310-
resolve([code, signal]);
311-
});
312-
});
292+
const [exitCode, signal] = await testBinUtils.processExit(
293+
bootstrapProcess,
294+
);
313295
expect(exitCode).toBe(null);
314296
expect(signal).toBe('SIGTERM');
315297
} else if (index === 1) {
316298
expect(stdErrLine2).toBeDefined();
317299
expect(stdErrLine2).toBe(eOutput);
318300
agentProcess.kill('SIGTERM');
319-
const [exitCode, signal] = await new Promise<
320-
[number | null, NodeJS.Signals | null]
321-
>((resolve) => {
322-
agentProcess.once('exit', (code, signal) => {
323-
resolve([code, signal]);
324-
});
325-
});
301+
const [exitCode, signal] = await testBinUtils.processExit(agentProcess);
326302
expect(exitCode).toBe(null);
327303
expect(signal).toBe('SIGTERM');
328304
}
@@ -348,11 +324,9 @@ describe('start', () => {
348324
rlOut.once('close', reject);
349325
});
350326
agentProcess1.kill('SIGHUP');
351-
const [exitCode1, signal1] = await new Promise<[number | null, NodeJS.Signals | null]>((resolve) => {
352-
agentProcess1.once('exit', (code, signal) => {
353-
resolve([code, signal]);
354-
});
355-
});
327+
const [exitCode1, signal1] = await testBinUtils.processExit(
328+
agentProcess1,
329+
);
356330
expect(exitCode1).toBe(null);
357331
expect(signal1).toBe('SIGHUP');
358332
const agentProcess2 = await testBinUtils.pkSpawn(
@@ -364,30 +338,21 @@ describe('start', () => {
364338
dataDir,
365339
logger,
366340
);
367-
const status1 = new Status({
341+
const status = new Status({
368342
statusPath: path.join(dataDir, 'polykey', config.defaults.statusBase),
369343
fs,
370344
logger,
371345
});
372-
await status1.waitFor('LIVE');
346+
await status.waitFor('LIVE');
373347
agentProcess2.kill('SIGHUP');
374-
const [exitCode2, signal2] = await new Promise<
375-
[number | null, NodeJS.Signals | null]
376-
>((resolve) => {
377-
agentProcess2.once('exit', (code, signal) => {
378-
resolve([code, signal]);
379-
});
380-
});
348+
const [exitCode2, signal2] = await testBinUtils.processExit(
349+
agentProcess2,
350+
);
381351
expect(exitCode2).toBe(null);
382352
expect(signal2).toBe('SIGHUP');
383353
// Check for graceful exit
384-
const status2 = new Status({
385-
statusPath: path.join(dataDir, 'polykey', config.defaults.statusBase),
386-
fs,
387-
logger,
388-
});
389-
const statusInfo2 = (await status2.readStatus())!;
390-
expect(statusInfo2.status).toBe('DEAD');
354+
const statusInfo = (await status.readStatus())!;
355+
expect(statusInfo.status).toBe('DEAD');
391356
},
392357
global.defaultTimeout * 2,
393358
);
@@ -418,20 +383,21 @@ describe('start', () => {
418383
}
419384
});
420385
});
421-
const [exitCode, signal] = await new Promise<
422-
[number | null, NodeJS.Signals | null]
423-
>((resolve) => {
424-
agentProcess1.once('exit', (code, signal) => {
425-
resolve([code, signal]);
426-
});
427-
});
386+
const [exitCode, signal] = await testBinUtils.processExit(agentProcess1);
428387
expect(exitCode).toBe(null);
429388
expect(signal).toBe('SIGINT');
430389
// Unlike bootstrapping, agent start can succeed under certain compatible partial state
431390
// However in some cases, state will conflict, and the start will fail with various errors
432391
// In such cases, the `--fresh` option must be used
433392
const agentProcess2 = await testBinUtils.pkSpawn(
434-
['agent', 'start', '--root-key-pair-bits', '1024', '--fresh', '--verbose'],
393+
[
394+
'agent',
395+
'start',
396+
'--root-key-pair-bits',
397+
'1024',
398+
'--fresh',
399+
'--verbose',
400+
],
435401
{
436402
PK_NODE_PATH: path.join(dataDir, 'polykey'),
437403
PK_PASSWORD: password,
@@ -452,13 +418,7 @@ describe('start', () => {
452418
recoveryCode.split(' ').length === 24,
453419
).toBe(true);
454420
agentProcess2.kill('SIGQUIT');
455-
await new Promise<
456-
[number | null, NodeJS.Signals | null]
457-
>((resolve) => {
458-
agentProcess2.once('exit', (code, signal) => {
459-
resolve([code, signal]);
460-
});
461-
});
421+
await testBinUtils.processExit(agentProcess2);
462422
// Check for graceful exit
463423
const status = new Status({
464424
statusPath: path.join(dataDir, 'polykey', config.defaults.statusBase),
@@ -470,4 +430,107 @@ describe('start', () => {
470430
},
471431
global.defaultTimeout * 2,
472432
);
433+
test(
434+
'start from recovery code',
435+
async () => {
436+
const password1 = 'abc123';
437+
const password2 = 'new password';
438+
const status = new Status({
439+
statusPath: path.join(dataDir, 'polykey', config.defaults.statusBase),
440+
fs,
441+
logger,
442+
});
443+
const agentProcess1 = await testBinUtils.pkSpawn(
444+
[
445+
'agent',
446+
'start',
447+
'--node-path',
448+
path.join(dataDir, 'polykey'),
449+
'--root-key-pair-bits',
450+
'1024',
451+
'--verbose',
452+
],
453+
{
454+
PK_PASSWORD: password1,
455+
},
456+
dataDir,
457+
logger.getChild('agentProcess1'),
458+
);
459+
const rlOut = readline.createInterface(agentProcess1.stdout!);
460+
const recoveryCode = await new Promise<RecoveryCode>(
461+
(resolve, reject) => {
462+
rlOut.once('line', resolve);
463+
rlOut.once('close', reject);
464+
},
465+
);
466+
const statusInfo1 = (await status.readStatus())!;
467+
agentProcess1.kill('SIGTERM');
468+
await testBinUtils.processExit(agentProcess1);
469+
const recoveryCodePath = path.join(dataDir, 'recovery-code');
470+
await fs.promises.writeFile(recoveryCodePath, recoveryCode + '\n');
471+
// When recovering, having the wrong bit size is not a problem
472+
const agentProcess2 = await testBinUtils.pkSpawn(
473+
[
474+
'agent',
475+
'start',
476+
'--recovery-code-file',
477+
recoveryCodePath,
478+
'--root-key-pair-bits',
479+
'2048',
480+
'--verbose',
481+
],
482+
{
483+
PK_NODE_PATH: path.join(dataDir, 'polykey'),
484+
PK_PASSWORD: password2,
485+
},
486+
dataDir,
487+
logger.getChild('agentProcess2'),
488+
);
489+
const statusInfo2 = await status.waitFor('LIVE');
490+
expect(statusInfo2.status).toBe('LIVE');
491+
// Node Id hasn't changed
492+
expect(statusInfo1.data.nodeId).toBe(statusInfo2.data.nodeId);
493+
agentProcess2.kill('SIGTERM');
494+
await testBinUtils.processExit(agentProcess2);
495+
// Check that the password has changed
496+
const agentProcess3 = await testBinUtils.pkSpawn(
497+
['agent', 'start', '--verbose'],
498+
{
499+
PK_NODE_PATH: path.join(dataDir, 'polykey'),
500+
PK_PASSWORD: password2,
501+
},
502+
dataDir,
503+
logger.getChild('agentProcess3'),
504+
);
505+
const statusInfo3 = await status.waitFor('LIVE');
506+
expect(statusInfo3.status).toBe('LIVE');
507+
// Node ID hasn't changed
508+
expect(statusInfo1.data.nodeId).toBe(statusInfo3.data.nodeId);
509+
agentProcess3.kill('SIGTERM');
510+
await testBinUtils.processExit(agentProcess3);
511+
// Checks deterministic generation using the same recovery code
512+
// First by deleting the polykey state
513+
await fs.promises.rm(path.join(dataDir, 'polykey'), {
514+
force: true,
515+
recursive: true,
516+
});
517+
const agentProcess4 = await testBinUtils.pkSpawn(
518+
['agent', 'start', '--root-key-pair-bits', '1024', '--verbose'],
519+
{
520+
PK_NODE_PATH: path.join(dataDir, 'polykey'),
521+
PK_PASSWORD: password2,
522+
PK_RECOVERY_CODE: recoveryCode,
523+
},
524+
dataDir,
525+
logger.getChild('agentProcess4'),
526+
);
527+
const statusInfo4 = await status.waitFor('LIVE');
528+
expect(statusInfo4.status).toBe('LIVE');
529+
// Same Node ID as before
530+
expect(statusInfo1.data.nodeId).toBe(statusInfo4.data.nodeId);
531+
agentProcess4.kill('SIGTERM');
532+
await testBinUtils.processExit(agentProcess4);
533+
},
534+
global.defaultTimeout * 3,
535+
);
473536
});

0 commit comments

Comments
 (0)