Fixed worker logins (151)

src/controllers/ExtensionController.ts

@@ -62,7 +62,6 @@ import type {
   ISecretService,
   IServerManager,
   IToastService,
-  Lazy,
   ServerConnectionPanelNode,
   ServerConnectionPanelTreeView,
   ServerConnectionTreeView,
@@ -120,8 +119,6 @@ export class ExtensionController implements Disposable {
   private _connectionController: ConnectionController | null = null;
   private _coreClientCache: URLMap<CoreAuthenticatedClient> | null = null;
   private _coreClientFactory: ICoreClientFactory | null = null;
-  private _coreCredentialsCache: URLMap<Lazy<DhcType.LoginCredentials>> | null =
-    null;
   private _coreJsApiCache: IAsyncCacheService<URL, typeof DhcType> | null =
     null;
   private _dheClientCache: URLMap<DheAuthenticatedClient> | null = null;
@@ -136,6 +133,7 @@ export class ExtensionController implements Disposable {
   private _secretService: ISecretService | null = null;
   private _serverManager: IServerManager | null = null;
   private _userLoginController: UserLoginController | null = null;
+  private _workerURLToServerURLMap: URLMap<URL> | null = null;
 
   // Tree providers
   private _serverTreeProvider: ServerTreeProvider | null = null;
@@ -214,11 +212,13 @@ export class ExtensionController implements Disposable {
     assertDefined(this._dheClientCache, 'dheClientCache');
     assertDefined(this._panelService, 'panelService');
     assertDefined(this._serverManager, 'serverManager');
+    assertDefined(this._workerURLToServerURLMap, 'workerURLToServerURLMap');
 
     this._panelController = new PanelController(
       this._dheClientCache,
       this._serverManager,
-      this._panelService
+      this._panelService,
+      this._workerURLToServerURLMap
     );
 
     this._context.subscriptions.push(this._panelController);
@@ -251,6 +251,7 @@ export class ExtensionController implements Disposable {
     assertDefined(this._dheClientFactory, 'dheClientFactory');
     assertDefined(this._secretService, 'secretService');
     assertDefined(this._toaster, 'toaster');
+    assertDefined(this._workerURLToServerURLMap, 'workerURLToServerURLMap');
 
     this._userLoginController = new UserLoginController(
       this._coreClientCache,
@@ -259,7 +260,8 @@ export class ExtensionController implements Disposable {
       this._dheClientCache,
       this._dheClientFactory,
       this._secretService,
-      this._toaster
+      this._toaster,
+      this._workerURLToServerURLMap
     );
 
     this._context.subscriptions.push(this._userLoginController);
@@ -324,8 +326,6 @@ export class ExtensionController implements Disposable {
     assertDefined(this._secretService, 'secretService');
     assertDefined(this._toaster, 'toaster');
 
-    this._coreCredentialsCache = new URLMap<Lazy<DhcType.LoginCredentials>>();
-
     this._coreJsApiCache = new CoreJsApiCache();
     this._context.subscriptions.push(this._coreJsApiCache);
 
@@ -365,7 +365,6 @@ export class ExtensionController implements Disposable {
 
     this._dheServiceFactory = DheService.factory(
       this._config,
-      this._coreCredentialsCache,
       this._dheClientCache,
       this._dheJsApiCache,
       this._toaster
@@ -374,6 +373,8 @@ export class ExtensionController implements Disposable {
     this._dheServiceCache = new DheServiceCache(this._dheServiceFactory);
     this._context.subscriptions.push(this._dheServiceCache);
 
+    this._workerURLToServerURLMap = new URLMap();
+
     this._serverManager = new ServerManager(
       this._config,
       this._coreClientCache,
@@ -382,7 +383,8 @@ export class ExtensionController implements Disposable {
       this._dheServiceCache,
       this._outputChannel,
       this._secretService,
-      this._toaster
+      this._toaster,
+      this._workerURLToServerURLMap
     );
     this._context.subscriptions.push(this._serverManager);
 
@@ -667,8 +669,6 @@ export class ExtensionController implements Disposable {
 
     // DHC ServerState
     if (serverOrConnectionState.type === 'DHC') {
-      this._coreCredentialsCache?.delete(serverOrConnectionState.url);
-
       await this._serverManager?.disconnectFromServer(
         serverOrConnectionState.url
       );

src/controllers/PanelController.ts

@@ -32,13 +32,15 @@ export class PanelController extends ControllerBase {
   constructor(
     dheClientCache: URLMap<DheAuthenticatedClient>,
     serverManager: IServerManager,
-    panelService: IPanelService
+    panelService: IPanelService,
+    workerURLToServerURLMap: URLMap<URL>
   ) {
     super();
 
     this._dheClientCache = dheClientCache;
     this._panelService = panelService;
     this._serverManager = serverManager;
+    this._workerURLToServerURLMap = workerURLToServerURLMap;
 
     this.registerCommand(OPEN_VARIABLE_PANELS_CMD, this._onOpenPanels);
     this.registerCommand(
@@ -56,6 +58,7 @@ export class PanelController extends ControllerBase {
   private readonly _dheClientCache: URLMap<DheAuthenticatedClient>;
   private readonly _panelService: IPanelService;
   private readonly _serverManager: IServerManager;
+  private readonly _workerURLToServerURLMap: URLMap<URL>;
 
   /**
    * Handle `postMessage` messages from the panel.
@@ -81,8 +84,9 @@ export class PanelController extends ControllerBase {
 
     // Respond to login credentials request from DH iframe
     if (message === DEEPHAVEN_POST_MSG.loginOptionsRequest) {
-      const dheClient = this._dheClientCache.get(serverOrWorkerUrl);
-
+      const dheServerUrl = this._workerURLToServerURLMap.get(serverOrWorkerUrl);
+      const dheClient =
+        dheServerUrl == null ? null : this._dheClientCache.get(dheServerUrl);
       const credentials =
         dheClient == null ? null : await getWorkerCredentials(dheClient);
 

src/controllers/UserLoginController.ts

@@ -26,9 +26,10 @@ import type {
   IToastService,
   ServerState,
 } from '../types';
-import { hasInteractivePermission } from '../dh/dhe';
+import { getWorkerCredentials, hasInteractivePermission } from '../dh/dhe';
 import {
   AUTH_HANDLER_TYPE_ANONYMOUS,
+  AUTH_HANDLER_TYPE_DHE,
   AUTH_HANDLER_TYPE_PSK,
   loginClient,
 } from '../dh/dhc';
@@ -46,7 +47,8 @@ export class UserLoginController extends ControllerBase {
     dheClientCache: URLMap<DheAuthenticatedClient>,
     dheClientFactory: IDheClientFactory,
     secretService: ISecretService,
-    toastService: IToastService
+    toastService: IToastService,
+    workerURLToServerURLMap: URLMap<URL>
   ) {
     super();
 
@@ -57,6 +59,7 @@ export class UserLoginController extends ControllerBase {
     this.dheClientFactory = dheClientFactory;
     this.secretService = secretService;
     this.toast = toastService;
+    this.workerURLToServerURLMap = workerURLToServerURLMap;
 
     this.registerCommand(
       GENERATE_DHE_KEY_PAIR_CMD,
@@ -81,6 +84,7 @@ export class UserLoginController extends ControllerBase {
   private readonly dheClientFactory: IDheClientFactory;
   private readonly secretService: ISecretService;
   private readonly toast: IToastService;
+  private readonly workerURLToServerURLMap: URLMap<URL>;
 
   /**
    * Login with a given key pair and remove the public key from the server.
@@ -218,6 +222,16 @@ export class UserLoginController extends ControllerBase {
       };
 
       this.secretService.storePsk(serverUrl, token);
+    } else if (authConfig.has(AUTH_HANDLER_TYPE_DHE)) {
+      const dheServerUrl = this.workerURLToServerURLMap.get(serverUrl);
+      const dheClient =
+        dheServerUrl == null ? null : this.dheClientCache.get(dheServerUrl);
+
+      if (dheClient == null) {
+        throw new Error(`No DHE server client found for worker '${serverUrl}'`);
+      }
+
+      credentials = await getWorkerCredentials(dheClient);
     } else {
       throw new Error('No supported authentication methods found.');
     }

src/services/DheService.ts

@@ -1,5 +1,4 @@
 import * as vscode from 'vscode';
-import type { dh as DhcType } from '@deephaven/jsapi-types';
 import type { AuthenticatedClient as DheAuthenticatedClient } from '@deephaven-enterprise/auth-nodejs';
 import type { EnterpriseDhType as DheType } from '@deephaven-enterprise/jsapi-types';
 import {
@@ -10,7 +9,6 @@ import {
   type IDheService,
   type IDheServiceFactory,
   type IToastService,
-  type Lazy,
   type QuerySerial,
   type UniqueID,
   type WorkerConfig,
@@ -21,7 +19,6 @@ import { Logger } from '../util';
 import {
   createInteractiveConsoleQuery,
   deleteQueries,
-  getWorkerCredentials,
   getWorkerInfoFromQuery,
 } from '../dh/dhe';
 import { CREATE_DHE_AUTHENTICATED_CLIENT_CMD } from '../common';
@@ -35,15 +32,13 @@ export class DheService implements IDheService {
   /**
    * Creates a factory function that can be used to create DheService instances.
    * @param configService Configuration service.
-   * @param coreCredentialsCache Core credentials cache.
    * @param dheClientCache DHE client cache.
    * @param dheJsApiCache DHE JS API cache.
    * @param toaster Toast service for notifications.
    * @returns A factory function that can be used to create DheService instances.
    */
   static factory = (
     configService: IConfigService,
-    coreCredentialsCache: URLMap<Lazy<DhcType.LoginCredentials>>,
     dheClientCache: URLMap<DheAuthenticatedClient>,
     dheJsApiCache: IAsyncCacheService<URL, DheType>,
     toaster: IToastService
@@ -53,7 +48,6 @@ export class DheService implements IDheService {
         new DheService(
           serverUrl,
           configService,
-          coreCredentialsCache,
           dheClientCache,
           dheJsApiCache,
           toaster
@@ -68,14 +62,12 @@ export class DheService implements IDheService {
   private constructor(
     serverUrl: URL,
     configService: IConfigService,
-    coreCredentialsCache: URLMap<Lazy<DhcType.LoginCredentials>>,
     dheClientCache: URLMap<DheAuthenticatedClient>,
     dheJsApiCache: IAsyncCacheService<URL, DheType>,
     toaster: IToastService
   ) {
     this.serverUrl = serverUrl;
     this._config = configService;
-    this._coreCredentialsCache = coreCredentialsCache;
     this._dheClientCache = dheClientCache;
     this._dheJsApiCache = dheJsApiCache;
     this._querySerialSet = new Set<QuerySerial>();
@@ -88,9 +80,6 @@ export class DheService implements IDheService {
   private _clientPromise: Promise<DheAuthenticatedClient | null> | null = null;
   private _isConnected: boolean = false;
   private readonly _config: IConfigService;
-  private readonly _coreCredentialsCache: URLMap<
-    Lazy<DhcType.LoginCredentials>
-  >;
   private readonly _dheClientCache: URLMap<DheAuthenticatedClient>;
   private readonly _dheJsApiCache: IAsyncCacheService<URL, DheType>;
   private readonly _querySerialSet: Set<QuerySerial>;
@@ -241,10 +230,6 @@ export class DheService implements IDheService {
     if (workerInfo == null) {
       throw new Error('Failed to create worker.');
     }
-    const workerUrl = new URL(workerInfo.grpcUrl);
-    this._coreCredentialsCache.set(workerUrl, () =>
-      getWorkerCredentials(dheClient)
-    );
 
     this._workerInfoMap.set(workerInfo.grpcUrl, workerInfo);
 

src/services/ServerManager.ts

@@ -45,7 +45,8 @@ export class ServerManager implements IServerManager {
     dheServiceCache: IAsyncCacheService<URL, IDheService>,
     outputChannel: vscode.OutputChannel,
     secretService: ISecretService,
-    toaster: IToastService
+    toaster: IToastService,
+    workerURLToServerURLMap: URLMap<URL>
   ) {
     this._configService = configService;
     this._connectionMap = new URLMap<ConnectionState>();
@@ -58,7 +59,7 @@ export class ServerManager implements IServerManager {
     this._serverMap = new URLMap<ServerState>();
     this._toaster = toaster;
     this._uriConnectionsMap = new URIMap<ConnectionState>();
-    this._workerURLToServerURLMap = new URLMap<URL>();
+    this._workerURLToServerURLMap = workerURLToServerURLMap;
 
     this.canStartServer = false;