Caching user login preferences (#32)

src/controllers/ExtensionController.ts

@@ -327,7 +327,8 @@ export class ExtensionController implements Disposable {
       this._coreCredentialsCache,
       this._dheClientCache,
       this._dheCredentialsCache,
-      this._dheJsApiCache
+      this._dheJsApiCache,
+      this._toaster
     );
 
     this._dheServiceCache = new DheServiceCache(this._dheServiceFactory);

src/controllers/UserLoginController.ts

@@ -1,13 +1,20 @@
-import * as vscode from 'vscode';
 import type { SecretService, URLMap } from '../services';
 import { ControllerBase } from './ControllerBase';
 import type { LoginCredentials as DheLoginCredentials } from '@deephaven-enterprise/jsapi-types';
 import {
   GENERATE_DHE_KEY_PAIR_CMD,
   REQUEST_DHE_USER_CREDENTIALS_CMD,
 } from '../common';
-import { formatDHPublicKey, generateBase64KeyPair, Logger } from '../util';
-import type { ServerState } from '../types';
+import {
+  createAuthenticationMethodQuickPick,
+  formatDHPublicKey,
+  generateBase64KeyPair,
+  Logger,
+  promptForOperateAs,
+  promptForPassword,
+  promptForUsername,
+} from '../util';
+import type { LoginWorkflowType, ServerState, Username } from '../types';
 
 const logger = new Logger('UserLoginController');
 
@@ -45,7 +52,7 @@ export class UserLoginController extends ControllerBase {
     serverState: ServerState
   ): Promise<void> => {
     const serverUrl = serverState.url;
-    await this.onDidRequestDheUserCredentials(serverUrl);
+    await this.onDidRequestDheUserCredentials(serverUrl, 'generatePrivateKey');
 
     const credentials = this.dheCredentialsCache.get(serverUrl);
     if (credentials?.username == null) {
@@ -61,8 +68,7 @@ export class UserLoginController extends ControllerBase {
     );
 
     // Get existing server keys or create a new object
-    const serverKeys =
-      (await this.secretService.getServerKeys(serverUrl)) ?? {};
+    const serverKeys = await this.secretService.getServerKeys(serverUrl);
 
     // Store the new private key for the user
     await this.secretService.storeServerKeys(serverUrl, {
@@ -80,31 +86,90 @@ export class UserLoginController extends ControllerBase {
    * Handle the request for DHE user credentials. If credentials are provided,
    * they will be stored in the credentials cache.
    * @param serverUrl The server URL to request credentials for.
+   * @param isGeneratePrivateKeyWorkflow Whether the request is part of a private key generation workflow.
    * @returns A promise that resolves when the credentials have been provided or declined.
    */
-  onDidRequestDheUserCredentials = async (serverUrl: URL): Promise<void> => {
+  onDidRequestDheUserCredentials = async (
+    serverUrl: URL,
+    workflowType: LoginWorkflowType = 'login'
+  ): Promise<void> => {
     // Remove any existing credentials for the server
     this.dheCredentialsCache.delete(serverUrl);
 
-    const username = await vscode.window.showInputBox({
-      placeHolder: 'Username',
-      prompt: 'Enter your Deephaven username',
-    });
+    const title =
+      workflowType === 'generatePrivateKey' ? 'Generate Private Key' : 'Login';
+
+    const secretKeys = await this.secretService.getServerKeys(serverUrl);
+    const userLoginPreferences =
+      await this.secretService.getUserLoginPreferences(serverUrl);
+
+    const privateKeyUserNames = Object.keys(secretKeys) as Username[];
+
+    if (workflowType === 'login' && privateKeyUserNames.length > 0) {
+      const authenticationMethod = await createAuthenticationMethodQuickPick(
+        title,
+        privateKeyUserNames
+      );
+
+      if (authenticationMethod?.type === 'privateKey') {
+        const username = authenticationMethod.label;
+
+        // Operate As
+        const operateAs = await promptForOperateAs(
+          title,
+          userLoginPreferences.operateAsUser[username] ?? username
+        );
+        if (operateAs == null) {
+          return;
+        }
+
+        await this.secretService.storeUserLoginPreferences(serverUrl, {
+          lastLogin: username,
+          operateAsUser: {
+            ...userLoginPreferences.operateAsUser,
+            [username]: operateAs,
+          },
+        });
+
+        // TODO: login with private key
+        logger.debug('Login with private key:', authenticationMethod.label);
+
+        return;
+      }
+    }
 
+    // Username
+    const username = await promptForUsername(
+      title,
+      userLoginPreferences.lastLogin
+    );
     if (username == null) {
       return;
     }
 
-    const token = await vscode.window.showInputBox({
-      placeHolder: 'Password',
-      prompt: 'Enter your Deephaven password',
-      password: true,
-    });
-
+    // Password
+    const token = await promptForPassword(title);
     if (token == null) {
       return;
     }
 
+    // Operate As
+    const operateAs = await promptForOperateAs(
+      title,
+      userLoginPreferences.operateAsUser[username] ?? username
+    );
+    if (operateAs == null) {
+      return;
+    }
+
+    await this.secretService.storeUserLoginPreferences(serverUrl, {
+      lastLogin: username,
+      operateAsUser: {
+        ...userLoginPreferences.operateAsUser,
+        [username]: operateAs,
+      },
+    });
+
     const dheCredentials: DheLoginCredentials = {
       username,
       token,

src/services/DheService.ts

@@ -12,6 +12,7 @@ import {
   type IConfigService,
   type IDheService,
   type IDheServiceFactory,
+  type IToastService,
   type Lazy,
   type QuerySerial,
   type UniqueID,
@@ -48,7 +49,8 @@ export class DheService implements IDheService {
     coreCredentialsCache: URLMap<Lazy<DhcType.LoginCredentials>>,
     dheClientCache: IAsyncCacheService<URL, EnterpriseClient>,
     dheCredentialsCache: URLMap<DheLoginCredentials>,
-    dheJsApiCache: IAsyncCacheService<URL, DheType>
+    dheJsApiCache: IAsyncCacheService<URL, DheType>,
+    toaster: IToastService
   ): IDheServiceFactory => {
     return {
       create: (serverUrl: URL): IDheService =>
@@ -58,7 +60,8 @@ export class DheService implements IDheService {
           coreCredentialsCache,
           dheClientCache,
           dheCredentialsCache,
-          dheJsApiCache
+          dheJsApiCache,
+          toaster
         ),
     };
   };
@@ -73,7 +76,8 @@ export class DheService implements IDheService {
     coreCredentialsCache: URLMap<Lazy<DhcType.LoginCredentials>>,
     dheClientCache: IAsyncCacheService<URL, EnterpriseClient>,
     dheCredentialsCache: URLMap<DheLoginCredentials>,
-    dheJsApiCache: IAsyncCacheService<URL, DheType>
+    dheJsApiCache: IAsyncCacheService<URL, DheType>,
+    toaster: IToastService
   ) {
     this.serverUrl = serverUrl;
     this._config = configService;
@@ -82,6 +86,7 @@ export class DheService implements IDheService {
     this._dheCredentialsCache = dheCredentialsCache;
     this._dheJsApiCache = dheJsApiCache;
     this._querySerialSet = new Set<QuerySerial>();
+    this._toaster = toaster;
     this._workerInfoMap = new URLMap<WorkerInfo, WorkerURL>();
   }
 
@@ -95,6 +100,7 @@ export class DheService implements IDheService {
   private readonly _dheCredentialsCache: URLMap<DheLoginCredentials>;
   private readonly _dheJsApiCache: IAsyncCacheService<URL, DheType>;
   private readonly _querySerialSet: Set<QuerySerial>;
+  private readonly _toaster: IToastService;
   private readonly _workerInfoMap: URLMap<WorkerInfo, WorkerURL>;
 
   readonly serverUrl: URL;
@@ -111,7 +117,7 @@ export class DheService implements IDheService {
    * @returns DHE client or null if initialization failed.
    */
   private _initClient = async (): Promise<EnterpriseClient | null> => {
-    const dheClient = await this._dheClientCache.get(this.serverUrl);
+    const dheClientPromise = this._dheClientCache.get(this.serverUrl);
 
     if (!this._dheCredentialsCache.has(this.serverUrl)) {
       await vscode.commands.executeCommand(
@@ -128,12 +134,15 @@ export class DheService implements IDheService {
       }
     }
 
+    const dheClient = await dheClientPromise;
     const dheCredentials = this._dheCredentialsCache.get(this.serverUrl)!;
 
     try {
       await dheClient.login(dheCredentials);
     } catch (err) {
+      this._dheCredentialsCache.delete(this.serverUrl);
       logger.error('An error occurred while connecting to DHE server:', err);
+      this._toaster.error(`Login failed to '${this.serverUrl.toString()}'`);
       return null;
     }
 

src/services/SecretService.ts

@@ -1,5 +1,5 @@
 import type { SecretStorage } from 'vscode';
-import type { OperateAsUserStored, ServerSecretKeysStored } from '../types';
+import type { ServerSecretKeys, UserLoginPreferences } from '../types';
 
 class Key {
   static operateAsUser(serverUrl: URL): string {
@@ -53,38 +53,42 @@ export class SecretService {
   };
 
   /**
-   * Get a map of user -> operatas users for a given server.
+   * Get user login preferences for a given server.
    * @param serverUrl The server URL to get the map for.
-   * @returns The map of user -> operate as user or null.
+   * @returns The user login preferences for the server.
    */
-  getOperateAsUser = async (
+  getUserLoginPreferences = async (
     serverUrl: URL
-  ): Promise<OperateAsUserStored | null> => {
-    return this._getJson(Key.operateAsUser(serverUrl));
+  ): Promise<UserLoginPreferences> => {
+    const preferences = await this._getJson<UserLoginPreferences>(
+      Key.operateAsUser(serverUrl)
+    );
+    return preferences ?? { operateAsUser: {} };
   };
 
   /**
-   * Store a map of user -> operate as user for a given server.
+   * Store user login preferences for a given server.
    * @param serverUrl The server URL to store the map for.
-   * @param operateAsUser The map of user -> operate as user.
+   * @param preferences The user login preferences to store.
    */
-  storeOperateAsUser = async (
+  storeUserLoginPreferences = async (
     serverUrl: URL,
-    operateAsUser: string
+    preferences: UserLoginPreferences
   ): Promise<void> => {
     const key = Key.operateAsUser(serverUrl);
-    await this._storeJson(key, operateAsUser);
+    await this._storeJson(key, 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<ServerSecretKeysStored | null> => {
-    return this._getJson(Key.serverKeys(serverUrl));
+  getServerKeys = async (serverUrl: URL): Promise<ServerSecretKeys> => {
+    const maybeServerKeys = await this._getJson<ServerSecretKeys>(
+      Key.serverKeys(serverUrl)
+    );
+    return maybeServerKeys ?? {};
   };
 
   /**
@@ -94,7 +98,7 @@ export class SecretService {
    */
   storeServerKeys = async (
     serverUrl: URL,
-    serverKeys: ServerSecretKeysStored
+    serverKeys: ServerSecretKeys
   ): Promise<void> => {
     const key = Key.serverKeys(serverUrl);
     await this._storeJson(key, serverKeys);

src/types/commonTypes.d.ts

@@ -37,12 +37,19 @@ export interface EnterpriseConnectionConfig {
   experimentalWorkerConfig?: WorkerConfig;
 }
 
+export type AuthenticationMethod = 'password' | 'privateKey';
+export type LoginWorkflowType = 'login' | 'generatePrivateKey';
+export type Username = Brand<'Username', string>;
+export type OperateAsUsername = Brand<'OperateAsUsername', string>;
 export type Base64PrivateKey = Brand<'Base64PrivateKey', string>;
 export type Base64PublicKey = Brand<'Base64PublicKey', string>;
 export type DHPrivateKey = Brand<'DHPrivateKey', string>;
 export type DHPublicKey = Brand<'DHPublicKey', string>;
-export type OperateAsUserStored = Record<string, string>;
-export type ServerSecretKeysStored = Record<string, Base64PrivateKey>;
+export type ServerSecretKeys = Record<string, Base64PrivateKey>;
+export type UserLoginPreferences = {
+  lastLogin?: Username;
+  operateAsUser: Record<Username, OperateAsUsername>;
+};
 
 export type ServerConnectionConfig =
   | CoreConnectionConfig

src/types/uiTypes.d.ts

@@ -1,11 +1,28 @@
 import * as vscode from 'vscode';
-import type { ConnectionState, ServerState } from './commonTypes';
+import type {
+  AuthenticationMethod,
+  ConnectionState,
+  ServerState,
+  Username,
+} from './commonTypes';
 
 export type SeparatorPickItem = {
   label: string;
   kind: vscode.QuickPickItemKind.Separator;
 };
 
+export type AuthenticationMethodPickItem =
+  | {
+      label: 'Username / Password';
+      type: 'password';
+      iconPath: vscode.ThemeIcon;
+    }
+  | {
+      label: Username;
+      type: 'privateKey';
+      iconPath: vscode.ThemeIcon;
+    };
+
 export type ConnectionPickItem<TType, TData> = vscode.QuickPickItem & {
   type: TType;
   data: TData;