Tunnel factory can return undefined (#113232)

Fixes microsoft/vscode-internalbacklog#1680

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

@@ -31,7 +31,7 @@ export interface TunnelCreationOptions {
 }
 
 export interface ITunnelProvider {
-	forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<RemoteTunnel> | undefined;
+	forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<RemoteTunnel | undefined> | undefined;
 }
 
 export interface ITunnelService {
@@ -41,7 +41,7 @@ export interface ITunnelService {
 	readonly onTunnelOpened: Event<RemoteTunnel>;
 	readonly onTunnelClosed: Event<{ host: string, port: number }>;
 
-	openTunnel(addressProvider: IAddressProvider | undefined, remoteHost: string | undefined, remotePort: number, localPort?: number): Promise<RemoteTunnel> | undefined;
+	openTunnel(addressProvider: IAddressProvider | undefined, remoteHost: string | undefined, remotePort: number, localPort?: number): Promise<RemoteTunnel | undefined> | undefined;
 	closeTunnel(remoteHost: string, remotePort: number): Promise<void>;
 	setTunnelProvider(provider: ITunnelProvider | undefined): IDisposable;
 }
@@ -81,7 +81,7 @@ export abstract class AbstractTunnelService implements ITunnelService {
 	public onTunnelOpened: Event<RemoteTunnel> = this._onTunnelOpened.event;
 	private _onTunnelClosed: Emitter<{ host: string, port: number }> = new Emitter();
 	public onTunnelClosed: Event<{ host: string, port: number }> = this._onTunnelClosed.event;
-	protected readonly _tunnels = new Map</*host*/ string, Map</* port */ number, { refcount: number, readonly value: Promise<RemoteTunnel> }>>();
+	protected readonly _tunnels = new Map</*host*/ string, Map</* port */ number, { refcount: number, readonly value: Promise<RemoteTunnel | undefined> }>>();
 	protected _tunnelProvider: ITunnelProvider | undefined;
 
 	public constructor(
@@ -103,22 +103,33 @@ export abstract class AbstractTunnelService implements ITunnelService {
 	}
 
 	public get tunnels(): Promise<readonly RemoteTunnel[]> {
-		const promises: Promise<RemoteTunnel>[] = [];
-		Array.from(this._tunnels.values()).forEach(portMap => Array.from(portMap.values()).forEach(x => promises.push(x.value)));
-		return Promise.all(promises);
+		return new Promise(async (resolve) => {
+			const tunnels: RemoteTunnel[] = [];
+			const tunnelArray = Array.from(this._tunnels.values());
+			for (let portMap of tunnelArray) {
+				const portArray = Array.from(portMap.values());
+				for (let x of portArray) {
+					const tunnelValue = await x.value;
+					if (tunnelValue) {
+						tunnels.push(tunnelValue);
+					}
+				}
+			}
+			resolve(tunnels);
+		});
 	}
 
 	dispose(): void {
 		for (const portMap of this._tunnels.values()) {
 			for (const { value } of portMap.values()) {
-				value.then(tunnel => tunnel.dispose());
+				value.then(tunnel => tunnel?.dispose());
 			}
 			portMap.clear();
 		}
 		this._tunnels.clear();
 	}
 
-	openTunnel(addressProvider: IAddressProvider | undefined, remoteHost: string | undefined, remotePort: number, localPort: number): Promise<RemoteTunnel> | undefined {
+	openTunnel(addressProvider: IAddressProvider | undefined, remoteHost: string | undefined, remotePort: number, localPort: number): Promise<RemoteTunnel | undefined> | undefined {
 		if (!addressProvider) {
 			return undefined;
 		}
@@ -133,6 +144,10 @@ export abstract class AbstractTunnelService implements ITunnelService {
 		}
 
 		return resolvedTunnel.then(tunnel => {
+			if (!tunnel) {
+				this.removeEmptyTunnelFromMap(remoteHost!, remotePort);
+				return undefined;
+			}
 			const newTunnel = this.makeTunnel(tunnel);
 			if (tunnel.tunnelRemoteHost !== remoteHost || tunnel.tunnelRemotePort !== remotePort) {
 				this.logService.warn('Created tunnel does not match requirements of requested tunnel. Host or port mismatch.');
@@ -161,11 +176,13 @@ export abstract class AbstractTunnelService implements ITunnelService {
 		};
 	}
 
-	private async tryDisposeTunnel(remoteHost: string, remotePort: number, tunnel: { refcount: number, readonly value: Promise<RemoteTunnel> }): Promise<void> {
+	private async tryDisposeTunnel(remoteHost: string, remotePort: number, tunnel: { refcount: number, readonly value: Promise<RemoteTunnel | undefined> }): Promise<void> {
 		if (tunnel.refcount <= 0) {
 			const disposePromise: Promise<void> = tunnel.value.then(tunnel => {
-				tunnel.dispose(true);
-				this._onTunnelClosed.fire({ host: tunnel.tunnelRemoteHost, port: tunnel.tunnelRemotePort });
+				if (tunnel) {
+					tunnel.dispose(true);
+					this._onTunnelClosed.fire({ host: tunnel.tunnelRemoteHost, port: tunnel.tunnelRemotePort });
+				}
 			});
 			if (this._tunnels.has(remoteHost)) {
 				this._tunnels.get(remoteHost)!.delete(remotePort);
@@ -183,16 +200,30 @@ export abstract class AbstractTunnelService implements ITunnelService {
 		}
 	}
 
-	protected addTunnelToMap(remoteHost: string, remotePort: number, tunnel: Promise<RemoteTunnel>) {
+	protected addTunnelToMap(remoteHost: string, remotePort: number, tunnel: Promise<RemoteTunnel | undefined>) {
 		if (!this._tunnels.has(remoteHost)) {
 			this._tunnels.set(remoteHost, new Map());
 		}
 		this._tunnels.get(remoteHost)!.set(remotePort, { refcount: 1, value: tunnel });
 	}
 
-	protected getTunnelFromMap(remoteHost: string, remotePort: number): { refcount: number, readonly value: Promise<RemoteTunnel> } | undefined {
+	private async removeEmptyTunnelFromMap(remoteHost: string, remotePort: number) {
+		const hostMap = this._tunnels.get(remoteHost);
+		if (hostMap) {
+			const tunnel = hostMap.get(remotePort);
+			const tunnelResult = await tunnel;
+			if (!tunnelResult) {
+				hostMap.delete(remotePort);
+			}
+			if (hostMap.size === 0) {
+				this._tunnels.delete(remoteHost);
+			}
+		}
+	}
+
+	protected getTunnelFromMap(remoteHost: string, remotePort: number): { refcount: number, readonly value: Promise<RemoteTunnel | undefined> } | undefined {
 		const otherLocalhost = getOtherLocalhost(remoteHost);
-		let portMap: Map<number, { refcount: number, readonly value: Promise<RemoteTunnel> }> | undefined;
+		let portMap: Map<number, { refcount: number, readonly value: Promise<RemoteTunnel | undefined> }> | undefined;
 		if (otherLocalhost) {
 			const firstMap = this._tunnels.get(remoteHost);
 			const secondMap = this._tunnels.get(otherLocalhost);
@@ -207,7 +238,7 @@ export abstract class AbstractTunnelService implements ITunnelService {
 		return portMap ? portMap.get(remotePort) : undefined;
 	}
 
-	protected abstract retainOrCreateTunnel(addressProvider: IAddressProvider, remoteHost: string, remotePort: number, localPort?: number): Promise<RemoteTunnel> | undefined;
+	protected abstract retainOrCreateTunnel(addressProvider: IAddressProvider, remoteHost: string, remotePort: number, localPort?: number): Promise<RemoteTunnel | undefined> | undefined;
 
 	protected isPortPrivileged(port: number): boolean {
 		return port < 1024;
@@ -215,7 +246,7 @@ export abstract class AbstractTunnelService implements ITunnelService {
 }
 
 export class TunnelService extends AbstractTunnelService {
-	protected retainOrCreateTunnel(_addressProvider: IAddressProvider, remoteHost: string, remotePort: number, localPort?: number | undefined): Promise<RemoteTunnel> | undefined {
+	protected retainOrCreateTunnel(_addressProvider: IAddressProvider, remoteHost: string, remotePort: number, localPort?: number | undefined): Promise<RemoteTunnel | undefined> | undefined {
 		const existing = this.getTunnelFromMap(remoteHost, remotePort);
 		if (existing) {
 			++existing.refcount;

src/vs/platform/remote/node/tunnelService.ts

@@ -139,7 +139,7 @@ export class BaseTunnelService extends AbstractTunnelService {
 		super(logService);
 	}
 
-	protected retainOrCreateTunnel(addressProvider: IAddressProvider, remoteHost: string, remotePort: number, localPort?: number): Promise<RemoteTunnel> | undefined {
+	protected retainOrCreateTunnel(addressProvider: IAddressProvider, remoteHost: string, remotePort: number, localPort?: number): Promise<RemoteTunnel | undefined> | undefined {
 		const existing = this.getTunnelFromMap(remoteHost, remotePort);
 		if (existing) {
 			++existing.refcount;

src/vs/platform/webview/common/webviewPortMapping.ts

@@ -19,7 +19,7 @@ export interface IWebviewPortMapping {
  */
 export class WebviewPortMappingManager implements IDisposable {
 
-	private readonly _tunnels = new Map<number, Promise<RemoteTunnel>>();
+	private readonly _tunnels = new Map<number, RemoteTunnel>();
 
 	constructor(
 		private readonly _getExtensionLocation: () => URI | undefined,
@@ -62,17 +62,17 @@ export class WebviewPortMappingManager implements IDisposable {
 
 	dispose() {
 		for (const tunnel of this._tunnels.values()) {
-			tunnel.then(tunnel => tunnel.dispose());
+			tunnel.dispose();
 		}
 		this._tunnels.clear();
 	}
 
-	private getOrCreateTunnel(remoteAuthority: IAddress, remotePort: number): Promise<RemoteTunnel> | undefined {
+	private async getOrCreateTunnel(remoteAuthority: IAddress, remotePort: number): Promise<RemoteTunnel | undefined> {
 		const existing = this._tunnels.get(remotePort);
 		if (existing) {
 			return existing;
 		}
-		const tunnel = this.tunnelService.openTunnel({ getAddress: async () => remoteAuthority }, undefined, remotePort);
+		const tunnel = await this.tunnelService.openTunnel({ getAddress: async () => remoteAuthority }, undefined, remotePort);
 		if (tunnel) {
 			this._tunnels.set(remotePort, tunnel);
 		}

src/vs/workbench/api/browser/mainThreadTunnelService.ts

@@ -58,6 +58,9 @@ export class MainThreadTunnelService extends Disposable implements MainThreadTun
 				const forward = this._proxy.$forwardPort(tunnelOptions, tunnelCreationOptions);
 				if (forward) {
 					return forward.then(tunnel => {
+						if (!tunnel) {
+							return undefined;
+						}
 						return {
 							tunnelRemotePort: tunnel.remoteAddress.port,
 							tunnelRemoteHost: tunnel.remoteAddress.host,

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

@@ -1765,7 +1765,7 @@ export interface MainThreadThemingShape extends IDisposable {
 }
 
 export interface ExtHostTunnelServiceShape {
-	$forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<TunnelDto> | undefined;
+	$forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<TunnelDto | undefined>;
 	$closeTunnel(remote: { host: string, port: number }, silent?: boolean): Promise<void>;
 	$onDidTunnelsChange(): Promise<void>;
 	$registerCandidateFinder(): Promise<void>;

src/vs/workbench/api/common/extHostTunnelService.ts

@@ -59,7 +59,7 @@ export class ExtHostTunnelService implements IExtHostTunnelService {
 	async setTunnelExtensionFunctions(provider: vscode.RemoteAuthorityResolver | undefined): Promise<IDisposable> {
 		return { dispose: () => { } };
 	}
-	$forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<TunnelDto> | undefined { return undefined; }
+	async $forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<TunnelDto | undefined> { return undefined; }
 	async $closeTunnel(remote: { host: string, port: number }): Promise<void> { }
 	async $onDidTunnelsChange(): Promise<void> { }
 	async $registerCandidateFinder(): Promise<void> { }

src/vs/workbench/api/node/extHostTunnelService.ts

@@ -222,7 +222,7 @@ export class ExtHostTunnelService extends Disposable implements IExtHostTunnelSe
 		this._onDidChangeTunnels.fire();
 	}
 
-	$forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<TunnelDto> | undefined {
+	async $forwardPort(tunnelOptions: TunnelOptions, tunnelCreationOptions: TunnelCreationOptions): Promise<TunnelDto | undefined> {
 		if (this._forwardPortProvider) {
 			const providedPort = this._forwardPortProvider(tunnelOptions, tunnelCreationOptions);
 			if (providedPort !== undefined) {