test: added setupTestAgent for creating test agents

This is a quick and easy way to create an agent with a pre-generated key.

Related #404

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, PrivateKey } from './keys/types';
+import type { KeyManagerChangeData, PrivateKeyPem } from './keys/types';
 import path from 'path';
 import process from 'process';
 import Logger from '@matrixai/logger';
@@ -108,7 +108,7 @@ class PolykeyAgent {
       rootCertDuration?: number;
       dbKeyBits?: number;
       recoveryCode?: string;
-      privateKeyOverride?: PrivateKey;
+      privateKeyPemOverride?: PrivateKeyPem;
     };
     proxyConfig?: {
       authToken?: string;

tests/bin/agent/start.test.ts

@@ -1,19 +1,20 @@
 import type { RecoveryCode } from '@/keys/types';
 import type { StatusLive } from '@/status/types';
+import type { NodeId } from '@/nodes/types';
+import type { Host, Port } from '@/network/types';
 import path from 'path';
 import fs from 'fs';
 import readline from 'readline';
 import process from 'process';
 import * as jestMockProps from 'jest-mock-props';
 import Logger, { LogLevel, StreamHandler } from '@matrixai/logger';
-import PolykeyAgent from '@/PolykeyAgent';
 import Status from '@/status/Status';
 import * as statusErrors from '@/status/errors';
 import config from '@/config';
 import * as keysUtils from '@/keys/utils';
 import * as testBinUtils from '../utils';
-import * as testUtils from '../../utils';
 import { runDescribeIfPlatforms, runTestIfPlatforms } from '../../utils';
+import { globalRootKeyPems } from '../../globalRootKeyPems';
 
 describe('start', () => {
   const logger = new Logger('start test', LogLevel.WARN, [new StreamHandler()]);
@@ -755,10 +756,12 @@ describe('start', () => {
         logger,
       });
       const password = 'abc123';
-      // Make sure these ports are not occupied
-      const rootKeys = await keysUtils.generateKeyPair(4096);
-      const privateKeyPem = keysUtils.privateKeyToPem(rootKeys.privateKey);
-      const nodeId = keysUtils.publicKeyToNodeId(rootKeys.publicKey);
+      const privateKeyPem = globalRootKeyPems[0];
+      const nodeId = keysUtils.publicKeyToNodeId(
+        keysUtils.publicKeyFromPrivateKey(
+          keysUtils.privateKeyFromPem(privateKeyPem),
+        ),
+      );
       const agentProcess = await testBinUtils.pkSpawnSwitch(global.testCmd)(
         ['agent', 'start', '--workers', '0', '--verbose'],
         {
@@ -791,10 +794,12 @@ describe('start', () => {
         logger,
       });
       const password = 'abc123';
-      // Make sure these ports are not occupied
-      const rootKeys = await keysUtils.generateKeyPair(4096);
-      const privateKeyPem = keysUtils.privateKeyToPem(rootKeys.privateKey);
-      const nodeId = keysUtils.publicKeyToNodeId(rootKeys.publicKey);
+      const privateKeyPem = globalRootKeyPems[0];
+      const nodeId = keysUtils.publicKeyToNodeId(
+        keysUtils.publicKeyFromPrivateKey(
+          keysUtils.privateKeyFromPem(privateKeyPem),
+        ),
+      );
       const privateKeyPath = path.join(dataDir, 'private.pem');
       await fs.promises.writeFile(privateKeyPath, privateKeyPem, {
         encoding: 'utf-8',
@@ -824,42 +829,51 @@ describe('start', () => {
     },
     global.defaultTimeout * 2,
   );
-  runDescribeIfPlatforms('linux')('start with global agent', () => {
-    let globalAgentStatus: StatusLive;
-    let globalAgentClose;
+  runDescribeIfPlatforms('linux').only('start with global agent', () => {
     let agentDataDir;
-    let agent: PolykeyAgent;
-    let seedNodeId1;
-    let seedNodeHost1;
-    let seedNodePort1;
-    let seedNodeId2;
-    let seedNodeHost2;
-    let seedNodePort2;
-    beforeAll(async () => {
-      ({ globalAgentStatus, globalAgentClose } =
-        await testUtils.setupGlobalAgent(logger));
+    let agent1Status: StatusLive;
+    let agent1Stop: () => void;
+    let agent2Status: StatusLive;
+    let agent2Stop: () => void;
+    let seedNodeId1: NodeId;
+    let seedNodeHost1: Host;
+    let seedNodePort1: Port;
+    let seedNodeId2: NodeId;
+    let seedNodeHost2: Host;
+    let seedNodePort2: Port;
+    beforeEach(async () => {
       // Additional seed node
       agentDataDir = await fs.promises.mkdtemp(
         path.join(global.tmpDir, 'polykey-test-'),
       );
-      agent = await PolykeyAgent.createPolykeyAgent({
-        password: 'password',
-        nodePath: path.join(agentDataDir, 'agent'),
-        keysConfig: {
-          rootKeyPairBits: 1024,
-        },
-        logger,
-      });
-      seedNodeId1 = globalAgentStatus.data.nodeId;
-      seedNodeHost1 = globalAgentStatus.data.proxyHost;
-      seedNodePort1 = globalAgentStatus.data.proxyPort;
-      seedNodeId2 = agent.keyManager.getNodeId();
-      seedNodeHost2 = agent.grpcServerAgent.getHost();
-      seedNodePort2 = agent.grpcServerAgent.getPort();
+      const agent1Path = path.join(agentDataDir, 'agent1');
+      await fs.promises.mkdir(agent1Path);
+      ({ agentStatus: agent1Status, agentStop: agent1Stop } =
+        await testBinUtils.setupTestAgent(
+          undefined,
+          agent1Path,
+          globalRootKeyPems[0],
+          logger,
+        ));
+      const agent2Path = path.join(agentDataDir, 'agent2');
+      await fs.promises.mkdir(agent2Path);
+      ({ agentStatus: agent2Status, agentStop: agent2Stop } =
+        await testBinUtils.setupTestAgent(
+          undefined,
+          agent2Path,
+          globalRootKeyPems[1],
+          logger,
+        ));
+      seedNodeId1 = agent1Status.data.nodeId;
+      seedNodeHost1 = agent1Status.data.proxyHost;
+      seedNodePort1 = agent1Status.data.proxyPort;
+      seedNodeId2 = agent2Status.data.nodeId;
+      seedNodeHost2 = agent2Status.data.proxyHost;
+      seedNodePort2 = agent2Status.data.proxyPort;
     }, globalThis.maxTimeout);
-    afterAll(async () => {
-      await agent.stop();
-      await globalAgentClose();
+    afterEach(async () => {
+      agent1Stop();
+      agent2Stop();
       await fs.promises.rm(agentDataDir, {
         force: true,
         recursive: true,

tests/bin/utils.ts

@@ -1,5 +1,7 @@
 import type { ChildProcess } from 'child_process';
 import type ErrorPolykey from '@/ErrorPolykey';
+import type { PrivateKeyPem } from '@/keys/types';
+import type { StatusLive } from '@/status/types';
 import child_process from 'child_process';
 import os from 'os';
 import fs from 'fs';
@@ -12,6 +14,7 @@ import nexpect from 'nexpect';
 import Logger from '@matrixai/logger';
 import main from '@/bin/polykey';
 import { promise } from '@/utils';
+import * as validationUtils from '@/validation/utils';
 
 /**
  * Wrapper for execFile to make it asynchronous and non-blocking
@@ -568,6 +571,60 @@ function expectProcessError(
   }
 }
 
+/**
+ *
+ * @param cmd - Optional target command to run, usually `global.testCmd`
+ * @param agentDir - Directory to run the agent in, must exist
+ * @param privateKeyPem - Optional root key override to skip key generation
+ * @param logger
+ */
+async function setupTestAgent(
+  cmd: string | undefined,
+  agentDir: string,
+  privateKeyPem: PrivateKeyPem,
+  logger: Logger,
+): Promise<{ agentStatus: StatusLive; agentStop: () => void }> {
+  const password = 'password';
+  const agentProcess = await pkSpawnSwitch(cmd)(
+    [
+      'agent',
+      'start',
+      '--node-path',
+      agentDir,
+      '--client-host',
+      '127.0.0.1',
+      '--proxy-host',
+      '127.0.0.1',
+      '--workers',
+      '0',
+      '--format',
+      'json',
+      '--verbose',
+    ],
+    {
+      PK_PASSWORD: password,
+      PK_ROOT_KEY: privateKeyPem,
+    },
+    agentDir,
+    logger,
+  );
+  const startedProm = promise<any>();
+  agentProcess.on('error', (d) => startedProm.rejectP(d));
+  const rlOut = readline.createInterface(agentProcess.stdout!);
+  rlOut.on('line', (l) => startedProm.resolveP(JSON.parse(l.toString())));
+  const data = await startedProm.p;
+  const agentStatus: StatusLive = {
+    status: 'LIVE',
+    data: { ...data, nodeId: validationUtils.parseNodeId(data.nodeId) },
+  };
+  try {
+    return { agentStatus, agentStop: () => agentProcess.kill('SIGINT') };
+  } catch (e) {
+    agentProcess.kill('SIGINT');
+    throw e;
+  }
+}
+
 export {
   exec,
   pk,
@@ -583,4 +640,5 @@ export {
   pkExpect,
   processExit,
   expectProcessError,
+  setupTestAgent,
 };