Merge pull request #296 from MatrixAI/bin-test-fixes

Session Bin Test Fixes
MatrixAI · Dec 13, 2021 · e630109 · e630109
2 parents 8804b28 + ea8493f
commit e630109
Show file tree

Hide file tree

Showing 102 changed files with 2,143 additions and 2,088 deletions.
diff --git a/package-lock.json b/package-lock.json
diff --git a/package.json b/package.json
@@ -65,7 +65,7 @@
     "lintfix": "eslint '{src,tests}/**/*.{js,ts}' --fix",
     "docs": "rm -r ./docs || true; typedoc --gitRevision master --tsconfig ./tsconfig.build.json --out ./docs src && touch ./docs/.nojekyll",
     "proto-generate": "scripts/proto-generate.sh",
-    "polykey": "ts-node --require tsconfig-paths/register --transpile-only src/bin/polykey.ts"
+    "polykey": "ts-node --require tsconfig-paths/register --compiler typescript-cached-transpile --transpile-only src/bin/polykey.ts"
   },
   "dependencies": {
     "@grpc/grpc-js": "1.3.7",
@@ -95,6 +95,7 @@
     "pako": "^1.0.11",
     "prompts": "^2.4.1",
     "readable-stream": "^3.6.0",
+    "resource-counter": "^1.2.4",
     "threads": "^1.6.5",
     "ts-custom-error": "^3.2.0",
     "utp-native": "^2.5.3",
@@ -130,6 +131,7 @@
     "ts-node": "^10.4.0",
     "tsconfig-paths": "^3.9.0",
     "typedoc": "^0.21.5",
-    "typescript": "^4.1.3"
+    "typescript": "^4.1.3",
+    "typescript-cached-transpile": "0.0.6"
   }
 }
diff --git a/src/agent/GRPCClientAgent.ts b/src/agent/GRPCClientAgent.ts
@@ -38,16 +38,17 @@ class GRPCClientAgent extends GRPCClient<AgentServiceClient> {
     logger?: Logger;
   }): Promise<GRPCClientAgent> {
     logger.info(`Creating ${this.name}`);
-    const { client, serverCertChain } = await super.createClient({
-      clientConstructor: AgentServiceClient,
-      nodeId,
-      host,
-      port,
-      tlsConfig,
-      proxyConfig,
-      timeout,
-      logger,
-    });
+    const { client, serverCertChain, flowCountInterceptor } =
+      await super.createClient({
+        clientConstructor: AgentServiceClient,
+        nodeId,
+        host,
+        port,
+        tlsConfig,
+        proxyConfig,
+        timeout,
+        logger,
+      });
     const grpcClientAgent = new GRPCClientAgent({
       client,
       nodeId,
@@ -56,6 +57,7 @@ class GRPCClientAgent extends GRPCClient<AgentServiceClient> {
       tlsConfig,
       proxyConfig,
       serverCertChain,
+      flowCountInterceptor,
       logger,
     });
     logger.info(`Created ${this.name}`);

diff --git a/src/bin/agent/CommandLock.ts b/src/bin/agent/CommandLock.ts
@@ -9,7 +9,6 @@ class CommandLock extends CommandPolykey {
     this.description('Lock the Client and Clear the Existing Token');
     this.action(async (options) => {
       const { default: Session } = await import('../../sessions/Session');
-      // Just delete the session token
       const session = new Session({
         sessionTokenPath: path.join(
           options.nodePath,
@@ -18,6 +17,7 @@ class CommandLock extends CommandPolykey {
         fs: this.fs,
         logger: this.logger.getChild(Session.name),
       });
+      // Destroy local session
       await session.destroy();
     });
   }

diff --git a/src/bin/agent/CommandLockAll.ts b/src/bin/agent/CommandLockAll.ts
@@ -1,8 +1,10 @@
 import type PolykeyClient from '../../PolykeyClient';
+import path from 'path';
 import CommandPolykey from '../CommandPolykey';
 import * as binUtils from '../utils';
 import * as binOptions from '../utils/options';
 import * as binProcessors from '../utils/processors';
+import config from '../../config';
 
 class CommandLockAll extends CommandPolykey {
   constructor(...args: ConstructorParameters<typeof CommandPolykey>) {
@@ -14,7 +16,9 @@ class CommandLockAll extends CommandPolykey {
     this.addOption(binOptions.clientPort);
     this.action(async (options) => {
       const { default: PolykeyClient } = await import('../../PolykeyClient');
+      const { default: Session } = await import('../../sessions/Session');
       const utilsPB = await import('../../proto/js/polykey/v1/utils/utils_pb');
+
       const clientOptions = await binProcessors.processClientOptions(
         options.nodePath,
         options.nodeId,
@@ -23,7 +27,19 @@ class CommandLockAll extends CommandPolykey {
         this.fs,
         this.logger.getChild(binProcessors.processClientOptions.name),
       );
-      let pkClient: PolykeyClient | undefined;
+      const meta = await binProcessors.processAuthentication(
+        options.passwordFile,
+        this.fs,
+      );
+      const session = new Session({
+        sessionTokenPath: path.join(
+          options.nodePath,
+          config.defaults.tokenBase,
+        ),
+        fs: this.fs,
+        logger: this.logger.getChild(Session.name),
+      });
+      let pkClient: PolykeyClient;
       this.exitHandlers.handlers.push(async () => {
         if (pkClient != null) await pkClient.stop();
       });
@@ -35,18 +51,15 @@ class CommandLockAll extends CommandPolykey {
           nodePath: options.nodePath,
           logger: this.logger.getChild(PolykeyClient.name),
         });
-        const meta = await binProcessors.processAuthentication(
-          options.passwordFile,
-          this.fs,
-        );
-        const grpcClient = pkClient.grpcClient;
         const emptyMessage = new utilsPB.EmptyMessage();
         await binUtils.retryAuthentication(
-          (auth) => grpcClient.sessionsLockAll(emptyMessage, auth),
+          (auth) => pkClient.grpcClient.sessionsLockAll(emptyMessage, auth),
           meta,
         );
+        // Destroy local session
+        await session.destroy();
       } finally {
-        if (pkClient != null) await pkClient.stop();
+        if (pkClient! != null) await pkClient.stop();
       }
     });
   }

diff --git a/src/bin/agent/CommandStart.ts b/src/bin/agent/CommandStart.ts
@@ -9,7 +9,7 @@ import CommandPolykey from '../CommandPolykey';
 import * as binOptions from '../utils/options';
 import * as binProcessors from '../utils/processors';
 import * as binErrors from '../errors';
-import { promise } from '../../utils';
+import { promise, dirEmpty } from '../../utils';
 import config from '../../config';
 
 class CommandStart extends CommandPolykey {
@@ -33,12 +33,35 @@ class CommandStart extends CommandPolykey {
       options.clientPort =
         options.clientPort ?? config.defaults.networkConfig.clientPort;
       const { default: PolykeyAgent } = await import('../../PolykeyAgent');
-      // Password is necessary
-      // If recovery code is supplied, then this is the new password
-      const password = await binProcessors.processPassword(
-        options.passwordFile,
-        this.fs,
-      );
+      let password: string | undefined;
+      if (options.fresh) {
+        // If fresh, then get a new password
+        password = await binProcessors.processNewPassword(
+          options.passwordFile,
+          this.fs,
+        );
+      } else if (options.recoveryCodeFile != null) {
+        // If recovery code is supplied, then this is the new password
+        password = await binProcessors.processNewPassword(
+          options.passwordFile,
+          this.fs,
+        );
+      } else if (await dirEmpty(this.fs, options.nodePath)) {
+        // If the node path is empty, get a new password
+        password = await binProcessors.processNewPassword(
+          options.passwordFile,
+          this.fs,
+        );
+      } else {
+        // Otherwise this is the existing password
+        // however, the code is capable of doing partial bootstrapping
+        // so it's possible that this is also a new password
+        // if the root key isn't setup
+        password = await binProcessors.processPassword(
+          options.passwordFile,
+          this.fs,
+        );
+      }
       const recoveryCodeIn = await binProcessors.processRecoveryCode(
         options.recoveryCodeFile,
         this.fs,

diff --git a/src/bin/agent/CommandStatus.ts b/src/bin/agent/CommandStatus.ts
@@ -38,6 +38,10 @@ class CommandStatus extends CommandPolykey {
           }),
         );
       } else {
+        const meta = await binProcessors.processAuthentication(
+          options.passwordFile,
+          this.fs,
+        );
         let pkClient: PolykeyClient;
         this.exitHandlers.handlers.push(async () => {
           if (pkClient != null) await pkClient.stop();
@@ -51,10 +55,6 @@ class CommandStatus extends CommandPolykey {
             nodePath: options.nodePath,
             logger: this.logger.getChild(PolykeyClient.name),
           });
-          const meta = await binProcessors.processAuthentication(
-            options.passwordFile,
-            this.fs,
-          );
           const emptyMessage = new utilsPB.EmptyMessage();
           response = await binUtils.retryAuthentication(
             (auth) => pkClient.grpcClient.agentStatus(emptyMessage, auth),

diff --git a/src/bin/agent/CommandStop.ts b/src/bin/agent/CommandStop.ts
@@ -34,6 +34,10 @@ class CommandStop extends CommandPolykey {
       } else if (statusInfo?.status === 'STARTING') {
         throw new binErrors.ErrorCLIStatusStarting();
       }
+      const meta = await binProcessors.processAuthentication(
+        options.passwordFile,
+        this.fs,
+      );
       // Either the statusInfo is undefined or LIVE
       // Either way, the connection parameters now exist
       let pkClient: PolykeyClient;
@@ -48,14 +52,9 @@ class CommandStop extends CommandPolykey {
           port: clientStatus.clientPort!,
           logger: this.logger.getChild(PolykeyClient.name),
         });
-        const meta = await binProcessors.processAuthentication(
-          options.passwordFile,
-          this.fs,
-        );
         const emptyMessage = new utilsPB.EmptyMessage();
-        const grpcClient = pkClient.grpcClient;
         await binUtils.retryAuthentication(
-          (auth) => grpcClient.agentStop(emptyMessage, auth),
+          (auth) => pkClient.grpcClient.agentStop(emptyMessage, auth),
           meta,
         );
         this.logger.info('Stopping Agent');

diff --git a/src/bin/agent/CommandUnlock.ts b/src/bin/agent/CommandUnlock.ts
@@ -1,6 +1,3 @@
-import type { SessionToken } from '../../sessions/types';
-import type { Metadata } from '@grpc/grpc-js';
-
 import type PolykeyClient from '../../PolykeyClient';
 import CommandPolykey from '../CommandPolykey';
 import * as binUtils from '../utils';
@@ -17,10 +14,7 @@ class CommandUnlock extends CommandPolykey {
     this.addOption(binOptions.clientPort);
     this.action(async (options) => {
       const { default: PolykeyClient } = await import('../../PolykeyClient');
-      const sessionsPB = await import(
-        '../../proto/js/polykey/v1/sessions/sessions_pb'
-      );
-
+      const utilsPB = await import('../../proto/js/polykey/v1/utils/utils_pb');
       const clientOptions = await binProcessors.processClientOptions(
         options.nodePath,
         options.nodeId,
@@ -29,41 +23,29 @@ class CommandUnlock extends CommandPolykey {
         this.fs,
         this.logger.getChild(binProcessors.processClientOptions.name),
       );
-
-      let pkClient: PolykeyClient | undefined;
+      const meta = await binProcessors.processAuthentication(
+        options.passwordFile,
+        this.fs,
+      );
+      let pkClient: PolykeyClient;
       this.exitHandlers.handlers.push(async () => {
         if (pkClient != null) await pkClient.stop();
       });
       try {
         pkClient = await PolykeyClient.createPolykeyClient({
-          nodePath: options.nodePath,
           nodeId: clientOptions.nodeId,
           host: clientOptions.clientHost,
           port: clientOptions.clientPort,
+          nodePath: options.nodePath,
           logger: this.logger.getChild(PolykeyClient.name),
         });
-
-        const password = await binProcessors.processPassword(
-          options.passwordFile,
-          this.fs,
+        const emptyMessage = new utilsPB.EmptyMessage();
+        await binUtils.retryAuthentication(
+          (auth) => pkClient.grpcClient.sessionsUnlock(emptyMessage, auth),
+          meta,
         );
-        const grpcClient = pkClient.grpcClient;
-        const passwordMessage = new sessionsPB.Password();
-        passwordMessage.setPassword(password);
-        const responseMessage = await binUtils.retryAuthentication(
-          (metaRetried?: Metadata) => {
-            return metaRetried != null
-              ? grpcClient.sessionsUnlock(passwordMessage, metaRetried)
-              : grpcClient.sessionsUnlock(passwordMessage);
-          },
-        );
-        const token: SessionToken = responseMessage.getToken() as SessionToken;
-
-        // Write token to file
-        await pkClient.session.writeToken(token);
-        process.stdout.write('Client session started');
       } finally {
-        if (pkClient != null) await pkClient.stop();
+        if (pkClient! != null) await pkClient.stop();
       }
     });
   }