From cfa27fdb54c1cce60f14734ff7df306719601b20 Mon Sep 17 00:00:00 2001 From: Brian Ingles Date: Fri, 20 Sep 2024 09:26:59 -0500 Subject: [PATCH] First pass at DHE connections (#79) --- packages/require-jsapi/src/dhc.ts | 3 ++ packages/require-jsapi/src/dhe.ts | 72 ++++++++++++++++++++++++++ src/controllers/ExtensionController.ts | 24 ++++++++- src/services/CacheByUrlService.ts | 31 +++++++++++ src/services/DhService.ts | 4 ++ src/services/DhcService.ts | 50 +++++++++--------- src/services/DhcServiceFactory.ts | 10 ++-- src/services/ServerManager.ts | 30 ++++++++--- src/services/index.ts | 1 + src/types/serviceTypes.d.ts | 10 ++-- 10 files changed, 191 insertions(+), 44 deletions(-) create mode 100644 src/services/CacheByUrlService.ts diff --git a/packages/require-jsapi/src/dhc.ts b/packages/require-jsapi/src/dhc.ts index 036e66c9..f3d3d649 100644 --- a/packages/require-jsapi/src/dhc.ts +++ b/packages/require-jsapi/src/dhc.ts @@ -4,6 +4,9 @@ import type { dh as DhType } from '@deephaven/jsapi-types'; import { polyfillDh } from './polyfill'; import { downloadFromURL, hasStatusCode } from './serverUtils'; +export const AUTH_HANDLER_TYPE_DHE = + 'io.deephaven.enterprise.dnd.authentication.DheAuthenticationHandler'; + /** * Check if a given server is running by checking if the `dh-core.js` file is * accessible. diff --git a/packages/require-jsapi/src/dhe.ts b/packages/require-jsapi/src/dhe.ts index 091b716f..4cd736dc 100644 --- a/packages/require-jsapi/src/dhe.ts +++ b/packages/require-jsapi/src/dhe.ts @@ -1,4 +1,10 @@ import { hasStatusCode } from './serverUtils'; +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import type { + EnterpriseDhType as DheType, + EnterpriseClient, +} from '@deephaven-enterprise/jsapi-types'; /** * Check if a given server is running by checking if the `irisapi/irisapi.nocache.js` @@ -15,3 +21,69 @@ export async function isDheServerRunning(serverUrl: URL): Promise { return false; } } + +export async function createDheClient( + dhe: DheType, + serverUrl: URL +): Promise { + const dheClient = new dhe.Client(serverUrl.toString()); + + return new Promise(resolve => { + const unsubscribe = dheClient.addEventListener( + dhe.Client.EVENT_CONNECT, + () => { + unsubscribe(); + resolve(dheClient); + } + ); + }); +} + +export async function initDheApi(serverUrl: URL): Promise { + polyfillDh(); + return getDhe(serverUrl, true); +} + +declare global { + export const iris: DheType; +} + +export async function getDhe( + serverUrl: URL, + download: boolean +): Promise { + const tmpDir = getTempDir(false, urlToDirectoryName(serverUrl)); + const dheFilePath = path.join(tmpDir, 'irisapi.nocache.js'); + + if (download) { + const dhe = await downloadFromURL( + path.join(serverUrl.toString(), 'irisapi/irisapi.nocache.js') + ); + + fs.writeFileSync(dheFilePath, dhe); + } + + require(dheFilePath); + + return iris; +} + +export async function getDheAuthToken( + client: EnterpriseClient +): Promise<{ type: string; token: string }> { + const token = await client.createAuthToken('RemoteQueryProcessor'); + return { + type: 'io.deephaven.proto.auth.Token', + token, + }; +} + +export function getWsUrl(serverUrl: URL): URL { + const url = new URL('/socket', serverUrl); + if (url.protocol === 'http:') { + url.protocol = 'ws:'; + } else { + url.protocol = 'wss:'; + } + return url; +} diff --git a/src/controllers/ExtensionController.ts b/src/controllers/ExtensionController.ts index 1c1f6d58..d908e2f0 100644 --- a/src/controllers/ExtensionController.ts +++ b/src/controllers/ExtensionController.ts @@ -1,4 +1,6 @@ import * as vscode from 'vscode'; +import type { dh as DhcType } from '@deephaven/jsapi-types'; +import type { EnterpriseDhType as DheType } from '@deephaven-enterprise/jsapi-types'; import { CONNECT_TO_SERVER_CMD, CREATE_NEW_TEXT_DOC_CMD, @@ -32,9 +34,16 @@ import { ServerConnectionPanelTreeProvider, runSelectedLinesHoverProvider, } from '../providers'; -import { DhcServiceFactory, PanelService, ServerManager } from '../services'; +import { + CacheByUrlService, + DhcServiceFactory, + PanelService, + ServerManager, + URLMap, +} from '../services'; import type { Disposable, + ICacheService, IConfigService, IDhService, IDhServiceFactory, @@ -51,6 +60,7 @@ import { ServerConnectionTreeDragAndDropController } from './ServerConnectionTre import { ConnectionController } from './ConnectionController'; import { PipServerController } from './PipServerController'; import { PanelController } from './PanelController'; +import { initDheApi } from '../dh/dhe'; const logger = new Logger('ExtensionController'); @@ -84,10 +94,12 @@ export class ExtensionController implements Disposable { readonly _config: IConfigService; private _connectionController: ConnectionController | null = null; + private _credentialsCache: URLMap | null = null; private _panelController: PanelController | null = null; private _panelService: IPanelService | null = null; private _pipServerController: PipServerController | null = null; private _dhcServiceFactory: IDhServiceFactory | null = null; + private _dheJsApiCache: ICacheService | null = null; private _serverManager: IServerManager | null = null; // Tree providers @@ -243,19 +255,27 @@ export class ExtensionController implements Disposable { assertDefined(this._outputChannel, 'outputChannel'); assertDefined(this._toaster, 'toaster'); + this._credentialsCache = new URLMap(); + this._panelService = new PanelService(); this._context.subscriptions.push(this._panelService); this._dhcServiceFactory = new DhcServiceFactory( + this._credentialsCache, this._panelService, this._pythonDiagnostics, this._outputChannel, this._toaster ); + this._dheJsApiCache = new CacheByUrlService(initDheApi); + this._context.subscriptions.push(this._dheJsApiCache); + this._serverManager = new ServerManager( this._config, - this._dhcServiceFactory + this._credentialsCache, + this._dhcServiceFactory, + this._dheJsApiCache ); this._context.subscriptions.push(this._serverManager); diff --git a/src/services/CacheByUrlService.ts b/src/services/CacheByUrlService.ts new file mode 100644 index 00000000..401c906a --- /dev/null +++ b/src/services/CacheByUrlService.ts @@ -0,0 +1,31 @@ +import type { ICacheService } from '../types'; +import { URLMap } from './URLMap'; + +/** + * Cache service that stores values by URL. + */ +export class CacheByUrlService implements ICacheService { + constructor(loader: (url: URL) => Promise) { + this._loader = loader; + this._promiseMap = new URLMap>(); + } + + private readonly _loader: (url: URL) => Promise; + private readonly _promiseMap = new URLMap>(); + + dispose = async (): Promise => { + this._promiseMap.clear(); + }; + + get = async (url: URL): Promise => { + if (!this._promiseMap.has(url)) { + this._promiseMap.set(url, this._loader(url)); + } + + return this._promiseMap.get(url)!; + }; + + invalidate = (url: URL): void => { + this._promiseMap.delete(url); + }; +} diff --git a/src/services/DhService.ts b/src/services/DhService.ts index d3d46da0..cbb43a71 100644 --- a/src/services/DhService.ts +++ b/src/services/DhService.ts @@ -19,6 +19,7 @@ import { } from '../common'; import type { ConnectionAndSession } from '../dh/dhc'; import { NoConsoleTypesError, parseServerError } from '../dh/errorUtils'; +import type { URLMap } from './URLMap'; const logger = new Logger('DhService'); @@ -27,11 +28,13 @@ export abstract class DhService { constructor( serverUrl: URL, + credentialsCache: URLMap, panelService: IPanelService, diagnosticsCollection: vscode.DiagnosticCollection, outputChannel: vscode.OutputChannel, toaster: IToastService ) { + this.credentialsCache = credentialsCache; this.serverUrl = serverUrl; this.panelService = panelService; this.diagnosticsCollection = diagnosticsCollection; @@ -45,6 +48,7 @@ export abstract class DhService public readonly serverUrl: URL; protected readonly subscriptions: (() => void)[] = []; + protected readonly credentialsCache: URLMap; protected readonly outputChannel: vscode.OutputChannel; protected readonly toaster: IToastService; private readonly panelService: IPanelService; diff --git a/src/services/DhcService.ts b/src/services/DhcService.ts index ee044382..09aae08f 100644 --- a/src/services/DhcService.ts +++ b/src/services/DhcService.ts @@ -13,14 +13,14 @@ import { const logger = new Logger('DhcService'); export class DhcService extends DhService { - private _psk?: string; + getPsk(): string | null { + const credentials = this.credentialsCache.get(this.serverUrl); - getPsk(): string | undefined { - return this._psk; - } + if (credentials?.type !== AUTH_HANDLER_TYPE_PSK) { + return null; + } - setPsk(psk: string): void { - this._psk = psk; + return credentials.token ?? null; } protected async initApi(): Promise { @@ -45,29 +45,29 @@ export class DhcService extends DhService { dh: typeof DhcType, client: DhcType.CoreClient ): Promise> { - const authConfig = new Set( - (await client.getAuthConfigValues()).map(([, value]) => value) - ); + if (!this.credentialsCache.has(this.serverUrl)) { + const authConfig = new Set( + (await client.getAuthConfigValues()).map(([, value]) => value) + ); - if (authConfig.has(AUTH_HANDLER_TYPE_ANONYMOUS)) { - return initDhcSession(client, { - type: dh.CoreClient.LOGIN_TYPE_ANONYMOUS, - }); - } else if (authConfig.has(AUTH_HANDLER_TYPE_PSK)) { - if (this._psk == null) { - this._psk = await vscode.window.showInputBox({ - placeHolder: 'Pre-Shared Key', - prompt: 'Enter your Deephaven pre-shared key', - password: true, + if (authConfig.has(AUTH_HANDLER_TYPE_ANONYMOUS)) { + this.credentialsCache.set(this.serverUrl, { + type: dh.CoreClient.LOGIN_TYPE_ANONYMOUS, + }); + } else if (authConfig.has(AUTH_HANDLER_TYPE_PSK)) { + this.credentialsCache.set(this.serverUrl, { + type: AUTH_HANDLER_TYPE_PSK, + token: await vscode.window.showInputBox({ + placeHolder: 'Pre-Shared Key', + prompt: 'Enter your Deephaven pre-shared key', + password: true, + }), }); } + } - const connectionAndSession = await initDhcSession(client, { - type: 'io.deephaven.authentication.psk.PskAuthenticationHandler', - token: this._psk, - }); - - return connectionAndSession; + if (this.credentialsCache.has(this.serverUrl)) { + return initDhcSession(client, this.credentialsCache.get(this.serverUrl)!); } throw new Error('No supported authentication methods found.'); diff --git a/src/services/DhcServiceFactory.ts b/src/services/DhcServiceFactory.ts index de121508..8ca09f3d 100644 --- a/src/services/DhcServiceFactory.ts +++ b/src/services/DhcServiceFactory.ts @@ -1,31 +1,31 @@ import * as vscode from 'vscode'; +import type { dh as DhcType } from '@deephaven/jsapi-types'; import { DhcService } from './DhcService'; import type { IDhServiceFactory, IPanelService, IToastService } from '../types'; +import type { URLMap } from './URLMap'; /** * Factory for creating DhcService instances. */ export class DhcServiceFactory implements IDhServiceFactory { constructor( + private credentialsCache: URLMap, private panelService: IPanelService, private diagnosticsCollection: vscode.DiagnosticCollection, private outputChannel: vscode.OutputChannel, private toaster: IToastService ) {} - create = (serverUrl: URL, psk?: string): DhcService => { + create = (serverUrl: URL): DhcService => { const dhService = new DhcService( serverUrl, + this.credentialsCache, this.panelService, this.diagnosticsCollection, this.outputChannel, this.toaster ); - if (psk != null) { - dhService.setPsk(psk); - } - return dhService; }; } diff --git a/src/services/ServerManager.ts b/src/services/ServerManager.ts index e460cf50..64ebd8c7 100644 --- a/src/services/ServerManager.ts +++ b/src/services/ServerManager.ts @@ -1,12 +1,16 @@ import * as vscode from 'vscode'; import { randomUUID } from 'node:crypto'; +import type { dh as DhcType } from '@deephaven/jsapi-types'; +import type { EnterpriseDhType as DheType } from '@deephaven-enterprise/jsapi-types'; import { + AUTH_HANDLER_TYPE_PSK, isDhcServerRunning, isDheServerRunning, } from '@deephaven/require-jsapi'; import { UnsupportedConsoleTypeError } from '../common'; import type { ConsoleType, + ICacheService, IConfigService, IDhService, IDhServiceFactory, @@ -22,11 +26,14 @@ const logger = new Logger('ServerManager'); export class ServerManager implements IServerManager { constructor( configService: IConfigService, - dhcServiceFactory: IDhServiceFactory + credentialsCache: URLMap, + dhcServiceFactory: IDhServiceFactory, + dheJsApiCache: ICacheService ) { this._configService = configService; + this._credentialsCache = credentialsCache; this._dhcServiceFactory = dhcServiceFactory; - + this._dheJsApiCache = dheJsApiCache; this._serverMap = new URLMap(); this._connectionMap = new URLMap(); this._uriConnectionsMap = new URIMap(); @@ -38,7 +45,9 @@ export class ServerManager implements IServerManager { private readonly _configService: IConfigService; private readonly _connectionMap: URLMap; + private readonly _credentialsCache: URLMap; private readonly _dhcServiceFactory: IDhServiceFactory; + private readonly _dheJsApiCache: ICacheService; private readonly _uriConnectionsMap: URIMap; private _serverMap: URLMap; @@ -99,15 +108,20 @@ export class ServerManager implements IServerManager { const serverState = this._serverMap.get(serverUrl); - // TODO: implement DHE #76 - if (serverState == null || serverState.type !== 'DHC') { + if (serverState == null) { return null; } - const connection = this._dhcServiceFactory.create( - serverUrl, - serverState.isManaged ? serverState.psk : undefined - ); + if (serverState.isManaged) { + this._credentialsCache.set(serverUrl, { + type: AUTH_HANDLER_TYPE_PSK, + token: serverState.psk, + }); + } else if (serverState.type === 'DHE') { + return null; + } + + const connection = this._dhcServiceFactory.create(serverUrl); this._connectionMap.set(serverUrl, connection); this._onDidUpdate.fire(); diff --git a/src/services/index.ts b/src/services/index.ts index 298d0575..14e92959 100644 --- a/src/services/index.ts +++ b/src/services/index.ts @@ -1,3 +1,4 @@ +export * from './CacheByUrlService'; export * from './ConfigService'; export * from './DhService'; export * from './DhcService'; diff --git a/src/types/serviceTypes.d.ts b/src/types/serviceTypes.d.ts index 18431007..8ebc3aa2 100644 --- a/src/types/serviceTypes.d.ts +++ b/src/types/serviceTypes.d.ts @@ -13,6 +13,11 @@ import type { VariableID, } from '../types/commonTypes'; +export interface ICacheService extends Disposable { + get: (key: TKey) => Promise; + invalidate: (key: TKey) => void; +} + /** * Configuration service interface. */ @@ -61,10 +66,7 @@ export interface IFactory { /** * Factory for creating IDhService instances. */ -export type IDhServiceFactory = IFactory< - IDhService, - [serverUrl: URL, psk?: string] ->; +export type IDhServiceFactory = IFactory; export interface IPanelService extends Disposable { readonly onDidUpdate: vscode.Event;