Move the remote connection data to the IRemoteAuthorityResolverService

src/vs/platform/remote/browser/remoteAuthorityResolverService.ts

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ResolvedAuthority, IRemoteAuthorityResolverService, ResolverResult } from 'vs/platform/remote/common/remoteAuthorityResolver';
+import { ResolvedAuthority, IRemoteAuthorityResolverService, ResolverResult, IRemoteConnectionData } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import { RemoteAuthorities } from 'vs/base/common/network';
 import { URI } from 'vs/base/common/uri';
 
@@ -11,11 +11,13 @@ export class RemoteAuthorityResolverService implements IRemoteAuthorityResolverS
 
 	declare readonly _serviceBrand: undefined;
 	private readonly _cache: Map<string, ResolverResult>;
+	private readonly _connectionTokens: Map<string, string>;
 
 	constructor(
 		resourceUriProvider: ((uri: URI) => URI) | undefined
 	) {
 		this._cache = new Map<string, ResolverResult>();
+		this._connectionTokens = new Map<string, string>();
 		if (resourceUriProvider) {
 			RemoteAuthorities.setDelegate(resourceUriProvider);
 		}
@@ -30,6 +32,19 @@ export class RemoteAuthorityResolverService implements IRemoteAuthorityResolverS
 		return this._cache.get(authority)!;
 	}
 
+	getConnectionData(authority: string): IRemoteConnectionData | null {
+		if (!this._cache.has(authority)) {
+			return null;
+		}
+		const resolverResult = this._cache.get(authority)!;
+		const connectionToken = this._connectionTokens.get(authority);
+		return {
+			host: resolverResult.authority.host,
+			port: resolverResult.authority.port,
+			connectionToken: connectionToken
+		};
+	}
+
 	private _doResolveAuthority(authority: string): ResolverResult {
 		if (authority.indexOf(':') >= 0) {
 			const pieces = authority.split(':');
@@ -38,12 +53,17 @@ export class RemoteAuthorityResolverService implements IRemoteAuthorityResolverS
 		return { authority: { authority, host: authority, port: 80 } };
 	}
 
-	clearResolvedAuthority(authority: string): void {
+	_clearResolvedAuthority(authority: string): void {
+	}
+
+	_setResolvedAuthority(resolvedAuthority: ResolvedAuthority) {
 	}
 
-	setResolvedAuthority(resolvedAuthority: ResolvedAuthority) {
+	_setResolvedAuthorityError(authority: string, err: any): void {
 	}
 
-	setResolvedAuthorityError(authority: string, err: any): void {
+	_setAuthorityConnectionToken(authority: string, connectionToken: string): void {
+		this._connectionTokens.set(authority, connectionToken);
+		RemoteAuthorities.setConnectionToken(authority, connectionToken);
 	}
 }

src/vs/platform/remote/common/remoteAuthorityResolver.ts

@@ -31,6 +31,12 @@ export interface ResolverResult {
 	tunnelInformation?: TunnelInformation;
 }
 
+export interface IRemoteConnectionData {
+	host: string;
+	port: number;
+	connectionToken: string | undefined;
+}
+
 export enum RemoteAuthorityResolverErrorCode {
 	Unknown = 'Unknown',
 	NotAvailable = 'NotAvailable',
@@ -80,8 +86,10 @@ export interface IRemoteAuthorityResolverService {
 	readonly _serviceBrand: undefined;
 
 	resolveAuthority(authority: string): Promise<ResolverResult>;
+	getConnectionData(authority: string): IRemoteConnectionData | null;
 
-	clearResolvedAuthority(authority: string): void;
-	setResolvedAuthority(resolvedAuthority: ResolvedAuthority, resolvedOptions?: ResolvedOptions): void;
-	setResolvedAuthorityError(authority: string, err: any): void;
+	_clearResolvedAuthority(authority: string): void;
+	_setResolvedAuthority(resolvedAuthority: ResolvedAuthority, resolvedOptions?: ResolvedOptions): void;
+	_setResolvedAuthorityError(authority: string, err: any): void;
+	_setAuthorityConnectionToken(authority: string, connectionToken: string): void;
 }

src/vs/platform/remote/electron-browser/remoteAuthorityResolverService.ts

@@ -3,16 +3,29 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ResolvedAuthority, IRemoteAuthorityResolverService, ResolverResult, ResolvedOptions } from 'vs/platform/remote/common/remoteAuthorityResolver';
+import { ResolvedAuthority, IRemoteAuthorityResolverService, ResolverResult, ResolvedOptions, IRemoteConnectionData } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import * as errors from 'vs/base/common/errors';
 import { RemoteAuthorities } from 'vs/base/common/network';
 
 class PendingResolveAuthorityRequest {
+
+	public value: ResolverResult | null;
+
 	constructor(
-		public readonly resolve: (value: ResolverResult) => void,
-		public readonly reject: (err: any) => void,
+		private readonly _resolve: (value: ResolverResult) => void,
+		private readonly _reject: (err: any) => void,
 		public readonly promise: Promise<ResolverResult>,
 	) {
+		this.value = null;
+	}
+
+	resolve(value: ResolverResult): void {
+		this.value = value;
+		this._resolve(this.value);
+	}
+
+	reject(err: any): void {
+		this._reject(err);
 	}
 }
 
@@ -20,44 +33,67 @@ export class RemoteAuthorityResolverService implements IRemoteAuthorityResolverS
 
 	declare readonly _serviceBrand: undefined;
 
-	private _resolveAuthorityRequests: { [authority: string]: PendingResolveAuthorityRequest; };
+	private readonly _resolveAuthorityRequests: Map<string, PendingResolveAuthorityRequest>;
+	private readonly _connectionTokens: Map<string, string>;
 
 	constructor() {
-		this._resolveAuthorityRequests = Object.create(null);
+		this._resolveAuthorityRequests = new Map<string, PendingResolveAuthorityRequest>();
+		this._connectionTokens = new Map<string, string>();
 	}
 
 	resolveAuthority(authority: string): Promise<ResolverResult> {
-		if (!this._resolveAuthorityRequests[authority]) {
+		if (!this._resolveAuthorityRequests.has(authority)) {
 			let resolve: (value: ResolverResult) => void;
 			let reject: (err: any) => void;
 			const promise = new Promise<ResolverResult>((_resolve, _reject) => {
 				resolve = _resolve;
 				reject = _reject;
 			});
-			this._resolveAuthorityRequests[authority] = new PendingResolveAuthorityRequest(resolve!, reject!, promise);
+			this._resolveAuthorityRequests.set(authority, new PendingResolveAuthorityRequest(resolve!, reject!, promise));
 		}
-		return this._resolveAuthorityRequests[authority].promise;
+		return this._resolveAuthorityRequests.get(authority)!.promise;
 	}
 
-	clearResolvedAuthority(authority: string): void {
-		if (this._resolveAuthorityRequests[authority]) {
-			this._resolveAuthorityRequests[authority].reject(errors.canceled());
-			delete this._resolveAuthorityRequests[authority];
+	getConnectionData(authority: string): IRemoteConnectionData | null {
+		if (!this._resolveAuthorityRequests.has(authority)) {
+			return null;
+		}
+		const request = this._resolveAuthorityRequests.get(authority)!;
+		if (!request.value) {
+			return null;
+		}
+		const connectionToken = this._connectionTokens.get(authority);
+		return {
+			host: request.value.authority.host,
+			port: request.value.authority.port,
+			connectionToken: connectionToken
+		};
+	}
+
+	_clearResolvedAuthority(authority: string): void {
+		if (this._resolveAuthorityRequests.has(authority)) {
+			this._resolveAuthorityRequests.get(authority)!.reject(errors.canceled());
+			this._resolveAuthorityRequests.delete(authority);
 		}
 	}
 
-	setResolvedAuthority(resolvedAuthority: ResolvedAuthority, options?: ResolvedOptions) {
-		if (this._resolveAuthorityRequests[resolvedAuthority.authority]) {
-			const request = this._resolveAuthorityRequests[resolvedAuthority.authority];
+	_setResolvedAuthority(resolvedAuthority: ResolvedAuthority, options?: ResolvedOptions): void {
+		if (this._resolveAuthorityRequests.has(resolvedAuthority.authority)) {
+			const request = this._resolveAuthorityRequests.get(resolvedAuthority.authority)!;
 			RemoteAuthorities.set(resolvedAuthority.authority, resolvedAuthority.host, resolvedAuthority.port);
 			request.resolve({ authority: resolvedAuthority, options });
 		}
 	}
 
-	setResolvedAuthorityError(authority: string, err: any): void {
-		if (this._resolveAuthorityRequests[authority]) {
-			const request = this._resolveAuthorityRequests[authority];
+	_setResolvedAuthorityError(authority: string, err: any): void {
+		if (this._resolveAuthorityRequests.has(authority)) {
+			const request = this._resolveAuthorityRequests.get(authority)!;
 			request.reject(err);
 		}
 	}
+
+	_setAuthorityConnectionToken(authority: string, connectionToken: string): void {
+		this._connectionTokens.set(authority, connectionToken);
+		RemoteAuthorities.setConnectionToken(authority, connectionToken);
+	}
 }

src/vs/workbench/api/common/extHost.protocol.ts

@@ -31,7 +31,7 @@ import { LogLevel } from 'vs/platform/log/common/log';
 import { IMarkerData } from 'vs/platform/markers/common/markers';
 import { IProgressOptions, IProgressStep } from 'vs/platform/progress/common/progress';
 import * as quickInput from 'vs/platform/quickinput/common/quickInput';
-import { RemoteAuthorityResolverErrorCode, ResolverResult, TunnelDescription } from 'vs/platform/remote/common/remoteAuthorityResolver';
+import { RemoteAuthorityResolverErrorCode, ResolverResult, TunnelDescription, IRemoteConnectionData } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import * as statusbar from 'vs/workbench/services/statusbar/common/statusbar';
 import { ClassifiedEvent, GDPRClassification, StrictPropertyCheck } from 'vs/platform/telemetry/common/gdprTypings';
 import { ITelemetryInfo } from 'vs/platform/telemetry/common/telemetry';
@@ -85,12 +85,6 @@ export interface IWorkspaceData extends IStaticWorkspaceData {
 	folders: { uri: UriComponents, name: string, index: number; }[];
 }
 
-export interface IRemoteConnectionData {
-	host: string;
-	port: number;
-	connectionToken: string;
-}
-
 export interface IInitData {
 	version: string;
 	commit?: string;

src/vs/workbench/services/extensions/browser/extensionService.ts

@@ -78,16 +78,9 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 			remoteAuthority: remoteAuthority,
 			getInitData: async () => {
 				await this.whenInstalledExtensionsRegistered();
-				const resolverResult = await this._remoteAuthorityResolverService.resolveAuthority(remoteAuthority);
+				const connectionData = this._remoteAuthorityResolverService.getConnectionData(remoteAuthority) || undefined;
 				const remoteEnvironment = this._remoteExtensionsEnvironmentData!;
-				return {
-					connectionData: {
-						host: resolverResult.authority.host,
-						port: resolverResult.authority.port,
-						connectionToken: remoteEnvironment.connectionToken
-					},
-					remoteEnvironment: remoteEnvironment
-				};
+				return { connectionData, remoteEnvironment };
 			}
 		};
 	}

src/vs/workbench/services/extensions/common/remoteExtensionHostClient.ts

@@ -11,12 +11,12 @@ import { ILogService } from 'vs/platform/log/common/log';
 import { connectRemoteAgentExtensionHost, IRemoteExtensionHostStartParams, IConnectionOptions, ISocketFactory } from 'vs/platform/remote/common/remoteAgentConnection';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { IWorkspaceContextService, WorkbenchState } from 'vs/platform/workspace/common/workspace';
-import { IInitData, UIKind, IRemoteConnectionData } from 'vs/workbench/api/common/extHost.protocol';
+import { IInitData, UIKind } from 'vs/workbench/api/common/extHost.protocol';
 import { MessageType, createMessageOfType, isMessageOfType } from 'vs/workbench/services/extensions/common/extensionHostProtocol';
 import { IExtensionHostStarter, ExtensionHostLogFileName } from 'vs/workbench/services/extensions/common/extensions';
 import { parseExtensionDevOptions } from 'vs/workbench/services/extensions/common/extensionDevOptions';
 import { IRemoteAgentEnvironment } from 'vs/platform/remote/common/remoteAgentEnvironment';
-import { IRemoteAuthorityResolverService } from 'vs/platform/remote/common/remoteAuthorityResolver';
+import { IRemoteAuthorityResolverService, IRemoteConnectionData } from 'vs/platform/remote/common/remoteAuthorityResolver';
 import * as platform from 'vs/base/common/platform';
 import { Schemas } from 'vs/base/common/network';
 import { Disposable } from 'vs/base/common/lifecycle';
@@ -34,7 +34,7 @@ import { IOutputChannelRegistry, Extensions } from 'vs/workbench/services/output
 import { localize } from 'vs/nls';
 
 export interface IRemoteInitData {
-	readonly connectionData: IRemoteConnectionData;
+	readonly connectionData: IRemoteConnectionData | undefined;
 	readonly remoteEnvironment: IRemoteAgentEnvironment;
 }
 

src/vs/workbench/services/extensions/electron-browser/extensionService.ts

@@ -53,7 +53,6 @@ class DeltaExtensionsQueueItem {
 export class ExtensionService extends AbstractExtensionService implements IExtensionService {
 
 	private readonly _remoteEnvironment: Map<string, IRemoteAgentEnvironment>;
-	private readonly _remoteResolvedAuthority: Map<string, ResolverResult>;
 
 	private readonly _extensionScanner: CachedExtensionScanner;
 	private _deltaExtensionsQueue: DeltaExtensionsQueueItem[];
@@ -97,7 +96,6 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 		}
 
 		this._remoteEnvironment = new Map<string, IRemoteAgentEnvironment>();
-		this._remoteResolvedAuthority = new Map<string, ResolverResult>();
 
 		this._extensionScanner = instantiationService.createInstance(CachedExtensionScanner);
 		this._deltaExtensionsQueue = [];
@@ -348,16 +346,9 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 			remoteAuthority: remoteAuthority,
 			getInitData: async () => {
 				await this.whenInstalledExtensionsRegistered();
+				const connectionData = this._remoteAuthorityResolverService.getConnectionData(remoteAuthority) || undefined;
 				const remoteEnvironment = this._remoteEnvironment.get(remoteAuthority)!;
-				const resolvedAuthority = this._remoteResolvedAuthority.get(remoteAuthority)!.authority;
-				return {
-					connectionData: {
-						host: resolvedAuthority.host,
-						port: resolvedAuthority.port,
-						connectionToken: remoteEnvironment.connectionToken
-					},
-					remoteEnvironment: remoteEnvironment
-				};
+				return { connectionData, remoteEnvironment };
 			}
 		};
 	}
@@ -443,12 +434,12 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 		}
 
 		const extensionHost = this._extensionHostProcessManagers[0];
-		this._remoteAuthorityResolverService.clearResolvedAuthority(remoteAuthority);
+		this._remoteAuthorityResolverService._clearResolvedAuthority(remoteAuthority);
 		try {
 			const result = await extensionHost.resolveAuthority(remoteAuthority);
-			this._remoteAuthorityResolverService.setResolvedAuthority(result.authority, result.options);
+			this._remoteAuthorityResolverService._setResolvedAuthority(result.authority, result.options);
 		} catch (err) {
-			this._remoteAuthorityResolverService.setResolvedAuthorityError(remoteAuthority, err);
+			this._remoteAuthorityResolverService._setResolvedAuthorityError(remoteAuthority, err);
 		}
 	}
 
@@ -481,7 +472,7 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 						console.log(`Error handled: Not showing a notification for the error`);
 					}
 				}
-				this._remoteAuthorityResolverService.setResolvedAuthorityError(remoteAuthority, err);
+				this._remoteAuthorityResolverService._setResolvedAuthorityError(remoteAuthority, err);
 
 				// Proceed with the local extension host
 				await this._startLocalExtensionHost(extensionHost, localExtensions, localExtensions.map(extension => extension.identifier));
@@ -489,7 +480,7 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 			}
 
 			// set the resolved authority
-			this._remoteAuthorityResolverService.setResolvedAuthority(resolvedAuthority.authority, resolvedAuthority.options);
+			this._remoteAuthorityResolverService._setResolvedAuthority(resolvedAuthority.authority, resolvedAuthority.options);
 			this._remoteExplorerService.setTunnelInformation(resolvedAuthority.tunnelInformation);
 
 			// monitor for breakage
@@ -501,7 +492,7 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 						return;
 					}
 					if (e.type === PersistentConnectionEventType.ConnectionLost) {
-						this._remoteAuthorityResolverService.clearResolvedAuthority(remoteAuthority);
+						this._remoteAuthorityResolverService._clearResolvedAuthority(remoteAuthority);
 					}
 				});
 				connection.onReconnecting(() => this._resolveAuthorityAgain());
@@ -558,7 +549,6 @@ export class ExtensionService extends AbstractExtensionService implements IExten
 
 			// save for remote extension's init data
 			this._remoteEnvironment.set(remoteAuthority, remoteEnv);
-			this._remoteResolvedAuthority.set(remoteAuthority, resolvedAuthority);
 
 			await this._startLocalExtensionHost(extensionHost, remoteEnv.extensions.concat(localExtensions), localExtensions.map(extension => extension.identifier));
 		} else {

src/vs/workbench/services/remote/browser/remoteAgentServiceImpl.ts

@@ -27,7 +27,7 @@ export class RemoteAgentService extends AbstractRemoteAgentService implements IR
 		@ISignService signService: ISignService,
 		@ILogService logService: ILogService
 	) {
-		super(environmentService);
+		super(environmentService, remoteAuthorityResolverService);
 
 		this.socketFactory = new BrowserSocketFactory(webSocketFactory);
 		const remoteAuthority = environmentService.configuration.remoteAuthority;

src/vs/workbench/services/remote/common/abstractRemoteAgentService.ts

@@ -30,7 +30,8 @@ export abstract class AbstractRemoteAgentService extends Disposable {
 	private _environment: Promise<IRemoteAgentEnvironment | null> | null;
 
 	constructor(
-		@IEnvironmentService protected readonly _environmentService: IEnvironmentService
+		@IEnvironmentService protected readonly _environmentService: IEnvironmentService,
+		@IRemoteAuthorityResolverService private readonly _remoteAuthorityResolverService: IRemoteAuthorityResolverService
 	) {
 		super();
 		this._environment = null;
@@ -41,7 +42,11 @@ export abstract class AbstractRemoteAgentService extends Disposable {
 	getEnvironment(bail?: boolean): Promise<IRemoteAgentEnvironment | null> {
 		if (!this._environment) {
 			this._environment = this._withChannel(
-				(channel, connection) => RemoteExtensionEnvironmentChannelClient.getEnvironmentData(channel, connection.remoteAuthority, this._environmentService.extensionDevelopmentLocationURI),
+				async (channel, connection) => {
+					const env = await RemoteExtensionEnvironmentChannelClient.getEnvironmentData(channel, connection.remoteAuthority, this._environmentService.extensionDevelopmentLocationURI);
+					this._remoteAuthorityResolverService._setAuthorityConnectionToken(connection.remoteAuthority, env.connectionToken);
+					return env;
+				},
 				null
 			);
 		}

src/vs/workbench/services/remote/common/remoteAgentEnvironmentChannel.ts

@@ -9,7 +9,6 @@ import { IChannel } from 'vs/base/parts/ipc/common/ipc';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { IRemoteAgentEnvironment } from 'vs/platform/remote/common/remoteAgentEnvironment';
 import { IDiagnosticInfoOptions, IDiagnosticInfo } from 'vs/platform/diagnostics/common/diagnostics';
-import { RemoteAuthorities } from 'vs/base/common/network';
 import { ITelemetryData } from 'vs/platform/telemetry/common/telemetry';
 
 export interface IGetEnvironmentDataArguments {
@@ -44,8 +43,6 @@ export class RemoteExtensionEnvironmentChannelClient {
 
 		const data = await channel.call<IRemoteAgentEnvironmentDTO>('getEnvironmentData', args);
 
-		RemoteAuthorities.setConnectionToken(remoteAuthority, data.connectionToken);
-
 		return {
 			pid: data.pid,
 			connectionToken: data.connectionToken,

src/vs/workbench/services/remote/electron-browser/remoteAgentServiceImpl.ts

@@ -26,7 +26,7 @@ export class RemoteAgentService extends AbstractRemoteAgentService implements IR
 		@ILogService logService: ILogService,
 		@IProductService productService: IProductService
 	) {
-		super(environmentService);
+		super(environmentService, remoteAuthorityResolverService);
 		this.socketFactory = nodeSocketFactory;
 		if (environmentService.configuration.remoteAuthority) {
 			this._connection = this._register(new RemoteAgentConnection(environmentService.configuration.remoteAuthority, productService.commit, nodeSocketFactory, remoteAuthorityResolverService, signService, logService));