Allow extensions to provide a port forwarding implementation.

Part of #81388
microsoft · Dec 16, 2019 · 92a83fb · 92a83fb
1 parent c11a0c1
commit 92a83fb
Show file tree

Hide file tree

Showing 9 changed files with 183 additions and 34 deletions.
diff --git a/src/vs/platform/remote/common/tunnel.ts b/src/vs/platform/remote/common/tunnel.ts
@@ -6,17 +6,28 @@
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import { URI } from 'vs/base/common/uri';
 import { Event } from 'vs/base/common/event';
+import { IDisposable } from 'vs/base/common/lifecycle';
 
 export const ITunnelService = createDecorator<ITunnelService>('tunnelService');
 
 export interface RemoteTunnel {
 	readonly tunnelRemotePort: number;
 	readonly tunnelRemoteHost: string;
-	readonly tunnelLocalPort: number;
+	readonly tunnelLocalPort?: number;
 	readonly localAddress: string;
 	dispose(): void;
 }
 
+export interface TunnelOptions {
+	remote: { port: number, host: string };
+	localPort?: number;
+	name?: string;
+}
+
+export interface ITunnelProvider {
+	forwardPort(tunnelOptions: TunnelOptions): Promise<RemoteTunnel> | undefined;
+}
+
 export interface ITunnelService {
 	_serviceBrand: undefined;
 
@@ -26,6 +37,7 @@ export interface ITunnelService {
 
 	openTunnel(remotePort: number, localPort?: number): Promise<RemoteTunnel> | undefined;
 	closeTunnel(remotePort: number): Promise<void>;
+	setTunnelProvider(provider: ITunnelProvider | undefined): IDisposable;
 }
 
 export function extractLocalHostUriMetaDataForPortMapping(uri: URI): { address: string, port: number } | undefined {

diff --git a/src/vs/platform/remote/common/tunnelService.ts b/src/vs/platform/remote/common/tunnelService.ts
@@ -3,8 +3,9 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ITunnelService, RemoteTunnel } from 'vs/platform/remote/common/tunnel';
+import { ITunnelService, RemoteTunnel, ITunnelProvider } from 'vs/platform/remote/common/tunnel';
 import { Event, Emitter } from 'vs/base/common/event';
+import { IDisposable } from 'vs/base/common/lifecycle';
 
 export class NoOpTunnelService implements ITunnelService {
 	_serviceBrand: undefined;
@@ -19,4 +20,7 @@ export class NoOpTunnelService implements ITunnelService {
 	}
 	async closeTunnel(_remotePort: number): Promise<void> {
 	}
+	setTunnelProvider(provider: ITunnelProvider | undefined): IDisposable {
+		throw new Error('Method not implemented.');
+	}
 }
diff --git a/src/vs/vscode.proposed.d.ts b/src/vs/vscode.proposed.d.ts
@@ -39,9 +39,11 @@ declare module 'vscode' {
 		name?: string;
 	}
 
-	export interface Tunnel extends Disposable {
+	export interface Tunnel {
 		remote: { port: number, host: string };
 		localAddress: string;
+		onDispose: Event<void>;
+		dispose(): void;
 	}
 
 	/**
@@ -72,7 +74,7 @@ declare module 'vscode' {
 		 * When not implemented, the core will use its default forwarding logic.
 		 * When implemented, the core will use this to forward ports.
 		 */
-		forwardPort?(tunnelOptions: TunnelOptions): Thenable<Tunnel | undefined>;
+		forwardPort?(tunnelOptions: TunnelOptions): Thenable<Tunnel> | undefined;
 	}
 
 	export namespace workspace {

diff --git a/src/vs/workbench/api/browser/mainThreadTunnelService.ts b/src/vs/workbench/api/browser/mainThreadTunnelService.ts
@@ -7,22 +7,24 @@ import { MainThreadTunnelServiceShape, IExtHostContext, MainContext, ExtHostCont
 import { TunnelOptions, TunnelDto } from 'vs/workbench/api/common/extHostTunnelService';
 import { extHostNamedCustomer } from 'vs/workbench/api/common/extHostCustomers';
 import { IRemoteExplorerService } from 'vs/workbench/services/remote/common/remoteExplorerService';
+import { ITunnelProvider, ITunnelService } from 'vs/platform/remote/common/tunnel';
 
 @extHostNamedCustomer(MainContext.MainThreadTunnelService)
 export class MainThreadTunnelService implements MainThreadTunnelServiceShape {
 	private readonly _proxy: ExtHostTunnelServiceShape;
 
 	constructor(
 		extHostContext: IExtHostContext,
-		@IRemoteExplorerService private readonly remoteExplorerService: IRemoteExplorerService
+		@IRemoteExplorerService private readonly remoteExplorerService: IRemoteExplorerService,
+		@ITunnelService private readonly tunnelService: ITunnelService
 	) {
 		this._proxy = extHostContext.getProxy(ExtHostContext.ExtHostTunnelService);
 	}
 
 	async $openTunnel(tunnelOptions: TunnelOptions): Promise<TunnelDto | undefined> {
 		const tunnel = await this.remoteExplorerService.forward(tunnelOptions.remote.port, tunnelOptions.localPort, tunnelOptions.name);
 		if (tunnel) {
-			return { remote: { host: tunnel.tunnelRemoteHost, port: tunnel.tunnelRemotePort }, localAddress: tunnel.localAddress };
+			return TunnelDto.fromServiceTunnel(tunnel);
 		}
 		return undefined;
 	}
@@ -35,6 +37,28 @@ export class MainThreadTunnelService implements MainThreadTunnelServiceShape {
 		this.remoteExplorerService.registerCandidateFinder(() => this._proxy.$findCandidatePorts());
 	}
 
+	async $setTunnelProvider(): Promise<void> {
+		const tunnelProvider: ITunnelProvider = {
+			forwardPort: (tunnelOptions: TunnelOptions) => {
+				const forward = this._proxy.$forwardPort(tunnelOptions);
+				if (forward) {
+					return forward.then(tunnel => {
+						return {
+							tunnelRemotePort: tunnel.remote.port,
+							tunnelRemoteHost: tunnel.remote.host,
+							localAddress: tunnel.localAddress,
+							dispose: () => {
+								this._proxy.$closeTunnel({ host: tunnel.remote.host, port: tunnel.remote.port });
+							}
+						};
+					});
+				}
+				return undefined;
+			}
+		};
+		this.tunnelService.setTunnelProvider(tunnelProvider);
+	}
+
 	dispose(): void {
 		//
 	}

diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts
@@ -777,6 +777,7 @@ export interface MainThreadTunnelServiceShape extends IDisposable {
 	$openTunnel(tunnelOptions: TunnelOptions): Promise<TunnelDto | undefined>;
 	$closeTunnel(remotePort: number): Promise<void>;
 	$registerCandidateFinder(): Promise<void>;
+	$setTunnelProvider(): Promise<void>;
 }
 
 // -- extension host
@@ -1396,6 +1397,8 @@ export interface ExtHostStorageShape {
 
 export interface ExtHostTunnelServiceShape {
 	$findCandidatePorts(): Promise<{ port: number, detail: string }[]>;
+	$forwardPort(tunnelOptions: TunnelOptions): Promise<TunnelDto> | undefined;
+	$closeTunnel(remote: { host: string, port: number }): Promise<void>;
 }
 
 // --- proxy identifiers

diff --git a/src/vs/workbench/api/common/extHostExtensionService.ts b/src/vs/workbench/api/common/extHostExtensionService.ts
@@ -32,6 +32,7 @@ import { IExtHostInitDataService } from 'vs/workbench/api/common/extHostInitData
 import { IExtensionStoragePaths } from 'vs/workbench/api/common/extHostStoragePaths';
 import { IExtHostRpcService } from 'vs/workbench/api/common/extHostRpcService';
 import { ServiceCollection } from 'vs/platform/instantiation/common/serviceCollection';
+import { IExtHostTunnelService } from 'vs/workbench/api/common/extHostTunnelService';
 
 interface ITestRunner {
 	/** Old test runner API, as exported from `vscode/lib/testrunner` */
@@ -76,6 +77,7 @@ export abstract class AbstractExtHostExtensionService implements ExtHostExtensio
 	protected readonly _extHostWorkspace: ExtHostWorkspace;
 	protected readonly _extHostConfiguration: ExtHostConfiguration;
 	protected readonly _logService: ILogService;
+	protected readonly _extHostTunnelService: IExtHostTunnelService;
 
 	protected readonly _mainThreadWorkspaceProxy: MainThreadWorkspaceShape;
 	protected readonly _mainThreadTelemetryProxy: MainThreadTelemetryShape;
@@ -104,7 +106,8 @@ export abstract class AbstractExtHostExtensionService implements ExtHostExtensio
 		@IExtHostConfiguration extHostConfiguration: IExtHostConfiguration,
 		@ILogService logService: ILogService,
 		@IExtHostInitDataService initData: IExtHostInitDataService,
-		@IExtensionStoragePaths storagePath: IExtensionStoragePaths
+		@IExtensionStoragePaths storagePath: IExtensionStoragePaths,
+		@IExtHostTunnelService extHostTunnelService: IExtHostTunnelService
 	) {
 		this._hostUtils = hostUtils;
 		this._extHostContext = extHostContext;
@@ -113,6 +116,7 @@ export abstract class AbstractExtHostExtensionService implements ExtHostExtensio
 		this._extHostWorkspace = extHostWorkspace;
 		this._extHostConfiguration = extHostConfiguration;
 		this._logService = logService;
+		this._extHostTunnelService = extHostTunnelService;
 		this._disposables = new DisposableStore();
 
 		this._mainThreadWorkspaceProxy = this._extHostContext.getProxy(MainContext.MainThreadWorkspace);
@@ -641,6 +645,7 @@ export abstract class AbstractExtHostExtensionService implements ExtHostExtensio
 
 		try {
 			const result = await resolver.resolve(remoteAuthority, { resolveAttempt });
+			this._disposables.add(await this._extHostTunnelService.setForwardPortProvider(resolver));
 
 			// Split merged API result into separate authority/options
 			const authority: ResolvedAuthority = {

diff --git a/src/vs/workbench/api/common/extHostTunnelService.ts b/src/vs/workbench/api/common/extHostTunnelService.ts
@@ -6,6 +6,8 @@
 import { ExtHostTunnelServiceShape } from 'vs/workbench/api/common/extHost.protocol';
 import { createDecorator } from 'vs/platform/instantiation/common/instantiation';
 import * as vscode from 'vscode';
+import { RemoteTunnel } from 'vs/platform/remote/common/tunnel';
+import { IDisposable } from 'vs/base/common/lifecycle';
 
 export interface TunnelOptions {
 	remote: { port: number, host: string };
@@ -19,9 +21,24 @@ export interface TunnelDto {
 	localAddress: string;
 }
 
+export namespace TunnelDto {
+	export function fromApiTunnel(tunnel: vscode.Tunnel): TunnelDto {
+		return { remote: tunnel.remote, localAddress: tunnel.localAddress };
+	}
+	export function fromServiceTunnel(tunnel: RemoteTunnel): TunnelDto {
+		return { remote: { host: tunnel.tunnelRemoteHost, port: tunnel.tunnelRemotePort }, localAddress: tunnel.localAddress };
+	}
+}
+
+export interface Tunnel extends vscode.Disposable {
+	remote: { port: number, host: string };
+	localAddress: string;
+}
+
 export interface IExtHostTunnelService extends ExtHostTunnelServiceShape {
 	readonly _serviceBrand: undefined;
 	makeTunnel(forward: TunnelOptions): Promise<vscode.Tunnel | undefined>;
+	setForwardPortProvider(provider: vscode.RemoteAuthorityResolver | undefined): Promise<IDisposable>;
 }
 
 export const IExtHostTunnelService = createDecorator<IExtHostTunnelService>('IExtHostTunnelService');
@@ -34,4 +51,8 @@ export class ExtHostTunnelService implements IExtHostTunnelService {
 	async $findCandidatePorts(): Promise<{ port: number; detail: string; }[]> {
 		return [];
 	}
+	async setForwardPortProvider(provider: vscode.RemoteAuthorityResolver | undefined): Promise<IDisposable> { return { dispose: () => { } }; }
+	$forwardPort(tunnelOptions: TunnelOptions): Promise<TunnelDto> | undefined { return undefined; }
+	async $closeTunnel(remote: { host: string, port: number }): Promise<void> { }
+
 }
diff --git a/src/vs/workbench/api/node/extHostTunnelService.ts b/src/vs/workbench/api/node/extHostTunnelService.ts
@@ -6,18 +6,37 @@
 import { MainThreadTunnelServiceShape, MainContext } from 'vs/workbench/api/common/extHost.protocol';
 import { IExtHostRpcService } from 'vs/workbench/api/common/extHostRpcService';
 import * as vscode from 'vscode';
-import { Disposable } from 'vs/base/common/lifecycle';
+import { Disposable, IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { IExtHostInitDataService } from 'vs/workbench/api/common/extHostInitDataService';
 import { URI } from 'vs/base/common/uri';
 import { exec } from 'child_process';
 import * as resources from 'vs/base/common/resources';
 import * as fs from 'fs';
 import { isLinux } from 'vs/base/common/platform';
-import { IExtHostTunnelService, TunnelOptions } from 'vs/workbench/api/common/extHostTunnelService';
+import { IExtHostTunnelService, TunnelOptions, TunnelDto } from 'vs/workbench/api/common/extHostTunnelService';
+import { asPromise } from 'vs/base/common/async';
+import { Event, Emitter } from 'vs/base/common/event';
+
+class ExtensionTunnel implements vscode.Tunnel {
+	private _onDispose: Emitter<void> = new Emitter();
+	onDispose: Event<void> = this._onDispose.event;
+
+	constructor(
+		public readonly remote: { port: number; host: string; },
+		public readonly localAddress: string,
+		private readonly _dispose: () => void) { }
+
+	dispose(): void {
+		this._onDispose.fire();
+		this._dispose();
+	}
+}
 
 export class ExtHostTunnelService extends Disposable implements IExtHostTunnelService {
 	readonly _serviceBrand: undefined;
 	private readonly _proxy: MainThreadTunnelServiceShape;
+	private _forwardPortProvider: ((tunnelOptions: TunnelOptions) => Thenable<vscode.Tunnel> | undefined) | undefined;
+	private _extensionTunnels: Map<string, Map<number, vscode.Tunnel>> = new Map();
 
 	constructor(
 		@IExtHostRpcService extHostRpc: IExtHostRpcService,
@@ -32,13 +51,9 @@ export class ExtHostTunnelService extends Disposable implements IExtHostTunnelSe
 	async makeTunnel(forward: TunnelOptions): Promise<vscode.Tunnel | undefined> {
 		const tunnel = await this._proxy.$openTunnel(forward);
 		if (tunnel) {
-			const disposableTunnel: vscode.Tunnel = {
-				remote: tunnel.remote,
-				localAddress: tunnel.localAddress,
-				dispose: () => {
-					return this._proxy.$closeTunnel(tunnel.remote.port);
-				}
-			};
+			const disposableTunnel: vscode.Tunnel = new ExtensionTunnel(tunnel.remote, tunnel.localAddress, () => {
+				return this._proxy.$closeTunnel(tunnel.remote.port);
+			});
 			this._register(disposableTunnel);
 			return disposableTunnel;
 		}
@@ -49,6 +64,46 @@ export class ExtHostTunnelService extends Disposable implements IExtHostTunnelSe
 		return this._proxy.$registerCandidateFinder();
 	}
 
+	async setForwardPortProvider(provider: vscode.RemoteAuthorityResolver | undefined): Promise<IDisposable> {
+		if (provider && provider.forwardPort) {
+			this._forwardPortProvider = provider.forwardPort;
+			await this._proxy.$setTunnelProvider();
+		} else {
+			this._forwardPortProvider = undefined;
+		}
+		return toDisposable(() => {
+			this._forwardPortProvider = undefined;
+		});
+	}
+
+	async $closeTunnel(remote: { host: string, port: number }): Promise<void> {
+		if (this._extensionTunnels.has(remote.host)) {
+			const hostMap = this._extensionTunnels.get(remote.host)!;
+			if (hostMap.has(remote.port)) {
+				hostMap.get(remote.port)!.dispose();
+				hostMap.delete(remote.port);
+			}
+		}
+	}
+
+	$forwardPort(tunnelOptions: TunnelOptions): Promise<TunnelDto> | undefined {
+		if (this._forwardPortProvider) {
+			const providedPort = this._forwardPortProvider!(tunnelOptions);
+			if (providedPort !== undefined) {
+				return asPromise(() => providedPort).then(tunnel => {
+					if (!this._extensionTunnels.has(tunnelOptions.remote.host)) {
+						this._extensionTunnels.set(tunnelOptions.remote.host, new Map());
+					}
+					this._extensionTunnels.get(tunnelOptions.remote.host)!.set(tunnelOptions.remote.port, tunnel);
+					this._register(tunnel.onDispose(() => this._proxy.$closeTunnel(tunnel.remote.port)));
+					return Promise.resolve(TunnelDto.fromApiTunnel(tunnel));
+				});
+			}
+		}
+		return undefined;
+	}
+
+
 	async $findCandidatePorts(): Promise<{ port: number, detail: string }[]> {
 		if (!isLinux) {
 			return [];