Refactored secret storage, added clear storage cmd, and wired up more…

… auth (32)

package.json

@@ -160,6 +160,10 @@
         "command": "vscode-deephaven.downloadLogs",
         "title": "Deephaven: Download Logs"
       },
+      {
+        "command": "vscode-deephaven.clearSecretStorage",
+        "title": "Deephaven: Clear Secrets"
+      },
       {
         "command": "vscode-deephaven.connectToServer",
         "title": "Deephaven: Connect to Server",
@@ -637,6 +641,10 @@
           "command": "vscode-deephaven.selectConnection",
           "when": "true"
         },
+        {
+          "command": "vscode-deephaven.clearSecretStorage",
+          "when": "true"
+        },
         {
           "command": "vscode-deephaven.downloadLogs",
           "when": "true"

src/common/commands.ts

@@ -8,6 +8,7 @@ function cmd<T extends string>(cmd: T): `${typeof EXTENSION_ID}.${T}` {
   return `${EXTENSION_ID}.${cmd}`;
 }
 
+export const CLEAR_SECRET_STORAGE_CMD = cmd('clearSecretStorage');
 export const CONNECT_TO_SERVER_CMD = cmd('connectToServer');
 export const CREATE_NEW_TEXT_DOC_CMD = cmd('createNewTextDoc');
 export const DISCONNECT_EDITOR_CMD = cmd('disconnectEditor');

src/controllers/ExtensionController.ts

@@ -6,6 +6,7 @@ import type {
   LoginCredentials as DheLoginCredentials,
 } from '@deephaven-enterprise/jsapi-types';
 import {
+  CLEAR_SECRET_STORAGE_CMD,
   CONNECT_TO_SERVER_CMD,
   CREATE_NEW_TEXT_DOC_CMD,
   DISCONNECT_EDITOR_CMD,
@@ -389,6 +390,9 @@ export class ExtensionController implements Disposable {
   initializeCommands = (): void => {
     assertDefined(this._connectionController, 'connectionController');
 
+    /** Clear secret storage */
+    this.registerCommand(CLEAR_SECRET_STORAGE_CMD, this.onClearSecretStorage);
+
     /** Create server connection */
     this.registerCommand(CONNECT_TO_SERVER_CMD, this.onConnectToServer);
 
@@ -537,6 +541,14 @@ export class ExtensionController implements Disposable {
     this._serverManager?.updateStatus();
   };
 
+  /**
+   * Handle clearing secret storage.
+   */
+  onClearSecretStorage = async (): Promise<void> => {
+    await this._secretService?.clearStorage();
+    this._toaster?.info('Stored secrets have been removed.');
+  };
+
   /**
    * Handle connecting to a server
    */

src/controllers/UserLoginController.ts

@@ -61,12 +61,6 @@ export class UserLoginController extends ControllerBase {
     serverState: ServerState
   ): Promise<void> => {
     const serverUrl = serverState.url;
-    // await this.onDidRequestDheUserCredentials(serverUrl, 'generatePrivateKey');
-
-    // const credentials = await this.dheCredentialsCache.get(serverUrl)?.();
-    // if (credentials?.username == null) {
-    //   return;
-    // }
 
     const title = 'Generate Private Key';
 
@@ -93,32 +87,9 @@ export class UserLoginController extends ControllerBase {
     const keyPair = generateBase64KeyPair();
     const { type, publicKey } = keyPair;
 
-    let dheClient = await this.dheClientCache.get(serverUrl);
-
-    const uploadKeyResult = await uploadPublicKey(
-      dheClient,
-      dheCredentials,
-      publicKey,
-      type
-    );
-    logger.debug('uploadKeyResult:', uploadKeyResult.status);
-
-    // TODO: Need to move the login logic to lazy credentials call
-
-    // Have to use a new client to login with the private key
-    this.dheClientCache.invalidate(serverUrl);
-    dheClient = await this.dheClientCache.get(serverUrl);
-
-    await authWithPrivateKey({
-      dheClient,
-      keyPair,
-      username,
-      operateAs: username,
-    });
+    const dheClient = await this.dheClientCache.get(serverUrl);
 
-    return;
-
-    // TODO: Need to store public key + algorithm in the server keys
+    await uploadPublicKey(dheClient, dheCredentials, publicKey, type);
 
     // Get existing server keys or create a new object
     const serverKeys = await this.secretService.getServerKeys(serverUrl);
@@ -128,11 +99,6 @@ export class UserLoginController extends ControllerBase {
       ...serverKeys,
       [dheCredentials.username]: keyPair,
     });
-
-    // Remove credentials from cache since presumably a valid key pair was
-    // generated and we'll want the user to authenticate with the private key
-    // instead.
-    this.dheCredentialsCache.delete(serverUrl);
   };
 
   /**
@@ -183,12 +149,31 @@ export class UserLoginController extends ControllerBase {
           },
         });
 
-        this.dheCredentialsCache.set(serverUrl, async () => {
-          logger.debug('Login with private key:', authenticationMethod.label);
-          // TODO: login with private key
-          throw new Error('Login with private key not implemented');
+        logger.debug('Login with private key:', authenticationMethod.label);
+        // Have to use a new client to login with the private key
+        this.dheClientCache.invalidate(serverUrl);
+        const dheClient = await this.dheClientCache.get(serverUrl);
+
+        const keyPair = (await this.secretService.getServerKeys(serverUrl))?.[
+          username
+        ];
+
+        await authWithPrivateKey({
+          dheClient,
+          keyPair,
+          username,
+          operateAs: username,
         });
 
+        this.dheCredentialsCache.set(
+          serverUrl,
+          async (): Promise<DheLoginCredentials> => {
+            // TODO: Need to figure out how to instantiate client + login at
+            // connection time
+            throw new Error('Not implemented');
+          }
+        );
+
         return;
       }
     }

src/services/SecretService.ts

@@ -1,18 +1,14 @@
 import type { SecretStorage } from 'vscode';
-import type { ServerSecretKeys, UserLoginPreferences } from '../types';
+import type { UserKeyPairs, UserLoginPreferences } from '../types';
 
-class Key {
-  static operateAsUser(serverUrl: URL): string {
-    return `operateAsUser.${serverUrl.toString()}`;
-  }
-
-  static serverKeys(serverUrl: URL): string {
-    return `serverKeys.${serverUrl.toString()}`;
-  }
-}
+const OPERATE_AS_USER_KEY = 'operateAsUser' as const;
+const SERVER_KEYS_KEY = 'serverKeys' as const;
 
 /**
- * Wrapper around `vscode.SecretStorage` for storing and retrieving secrets.
+ * Wrapper around `vscode.SecretStorage` for storing and retrieving secrets. We
+ * are storing everything as known keys so that we can easily find and delete
+ * them. There don't appear to be any apis to delete all secrets at once or to
+ * determine which keys exist.
  * NOTE: For debugging, the secret store contents can be dumped to devtools
  * console via:
  * > Developer: Log Storage Database Contents
@@ -43,6 +39,14 @@ export class SecretService {
     }
   };
 
+  /**
+   * Clear all stored secrets.
+   */
+  clearStorage = async (): Promise<void> => {
+    await this._secrets.delete(OPERATE_AS_USER_KEY);
+    await this._secrets.delete(SERVER_KEYS_KEY);
+  };
+
   /**
    * Store a JSON-serializable value.
    * @param key Secret storage key
@@ -60,10 +64,12 @@ export class SecretService {
   getUserLoginPreferences = async (
     serverUrl: URL
   ): Promise<UserLoginPreferences> => {
-    const preferences = await this._getJson<UserLoginPreferences>(
-      Key.operateAsUser(serverUrl)
-    );
-    return preferences ?? { operateAsUser: {} };
+    const preferences =
+      await this._getJson<Record<string, UserLoginPreferences>>(
+        SERVER_KEYS_KEY
+      );
+
+    return preferences?.[serverUrl.toString()] ?? { operateAsUser: {} };
   };
 
   /**
@@ -75,20 +81,27 @@ export class SecretService {
     serverUrl: URL,
     preferences: UserLoginPreferences
   ): Promise<void> => {
-    const key = Key.operateAsUser(serverUrl);
-    await this._storeJson(key, preferences);
+    const existingPreferences =
+      await this._getJson<Record<string, UserLoginPreferences>>(
+        SERVER_KEYS_KEY
+      );
+
+    await this._storeJson(OPERATE_AS_USER_KEY, {
+      ...existingPreferences,
+      [serverUrl.toString()]: preferences,
+    });
   };
 
   /**
    * Get a map of user -> private keys for a given server.
    * @param serverUrl
    * @returns The map of user -> private key or null.
    */
-  getServerKeys = async (serverUrl: URL): Promise<ServerSecretKeys> => {
-    const maybeServerKeys = await this._getJson<ServerSecretKeys>(
-      Key.serverKeys(serverUrl)
-    );
-    return maybeServerKeys ?? {};
+  getServerKeys = async (serverUrl: URL): Promise<UserKeyPairs> => {
+    const maybeServerKeys =
+      await this._getJson<Record<string, UserKeyPairs>>(SERVER_KEYS_KEY);
+
+    return maybeServerKeys?.[serverUrl.toString()] ?? {};
   };
 
   /**
@@ -98,9 +111,14 @@ export class SecretService {
    */
   storeServerKeys = async (
     serverUrl: URL,
-    serverKeys: ServerSecretKeys
+    serverKeys: UserKeyPairs
   ): Promise<void> => {
-    const key = Key.serverKeys(serverUrl);
-    await this._storeJson(key, serverKeys);
+    const existingKeys =
+      await this._getJson<Record<string, UserKeyPairs>>(SERVER_KEYS_KEY);
+
+    await this._storeJson(SERVER_KEYS_KEY, {
+      ...existingKeys,
+      [serverUrl.toString()]: serverKeys,
+    });
   };
 }

src/types/commonTypes.d.ts

@@ -51,7 +51,7 @@ export type Base64KeyPair = {
   publicKey: Base64PublicKey;
   privateKey: Base64PrivateKey;
 };
-export type ServerSecretKeys = Record<string, Base64KeyPair>;
+export type UserKeyPairs = Record<Username, Base64KeyPair>;
 export type UserLoginPreferences = {
   lastLogin?: Username;
   operateAsUser: Record<Username, OperateAsUsername>;