Use the new lastCommittedOrigin for remote port forwarding in webviews

Fixes microsoft#102449

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

@@ -14,7 +14,7 @@ export const IWebviewManagerService = createDecorator<IWebviewManagerService>('w
 export interface IWebviewManagerService {
 	_serviceBrand: unknown;
 
-	registerWebview(id: string, webContentsId: number | undefined, windowId: number, metadata: RegisterWebviewMetadata): Promise<void>;
+	registerWebview(id: string, windowId: number, metadata: RegisterWebviewMetadata): Promise<void>;
 	unregisterWebview(id: string): Promise<void>;
 	updateWebviewMetadata(id: string, metadataDelta: Partial<RegisterWebviewMetadata>): Promise<void>;
 

src/vs/platform/webview/electron-main/webviewMainService.ts

@@ -33,7 +33,7 @@ export class WebviewMainService extends Disposable implements IWebviewManagerSer
 		this.portMappingProvider = this._register(new WebviewPortMappingProvider(tunnelService));
 	}
 
-	public async registerWebview(id: string, webContentsId: number | undefined, windowId: number, metadata: RegisterWebviewMetadata): Promise<void> {
+	public async registerWebview(id: string, windowId: number, metadata: RegisterWebviewMetadata): Promise<void> {
 		const extensionLocation = metadata.extensionLocation ? URI.from(metadata.extensionLocation) : undefined;
 
 		this.protocolProvider.registerWebview(id, {
@@ -43,7 +43,7 @@ export class WebviewMainService extends Disposable implements IWebviewManagerSer
 			localResourceRoots: metadata.localResourceRoots.map(x => URI.from(x))
 		});
 
-		this.portMappingProvider.registerWebview(id, webContentsId, {
+		this.portMappingProvider.registerWebview(id, {
 			extensionLocation,
 			mappings: metadata.portMappings,
 			resolvedAuthority: metadata.remoteConnectionData,

src/vs/platform/webview/electron-main/webviewPortMappingProvider.ts

@@ -3,14 +3,18 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { session } from 'electron';
+import { OnBeforeRequestListenerDetails, session } from 'electron';
 import { Disposable } from 'vs/base/common/lifecycle';
 import { URI } from 'vs/base/common/uri';
 import { IAddress } from 'vs/platform/remote/common/remoteAgentConnection';
 import { ITunnelService } from 'vs/platform/remote/common/tunnel';
 import { webviewPartitionId } from 'vs/platform/webview/common/resourceLoader';
 import { IWebviewPortMapping, WebviewPortMappingManager } from 'vs/platform/webview/common/webviewPortMapping';
 
+interface OnBeforeRequestListenerDetails_Extended extends OnBeforeRequestListenerDetails {
+	readonly lastCommittedOrigin?: string;
+}
+
 interface PortMappingData {
 	readonly extensionLocation: URI | undefined;
 	readonly mappings: readonly IWebviewPortMapping[];
@@ -20,13 +24,10 @@ interface PortMappingData {
 export class WebviewPortMappingProvider extends Disposable {
 
 	private readonly _webviewData = new Map<string, {
-		readonly webContentsId: number | undefined;
 		readonly manager: WebviewPortMappingManager;
 		metadata: PortMappingData;
 	}>();
 
-	private _webContentsIdsToWebviewIds = new Map<number, /* id */ string>();
-
 	constructor(
 		@ITunnelService private readonly _tunnelService: ITunnelService,
 	) {
@@ -40,12 +41,15 @@ export class WebviewPortMappingProvider extends Disposable {
 				'*://127.0.0.1:*/*',
 				'*://0.0.0.0:*/*',
 			]
-		}, async (details, callback) => {
-			const webviewId = details.webContentsId && this._webContentsIdsToWebviewIds.get(details.webContentsId);
-			if (!webviewId) {
+		}, async (details: OnBeforeRequestListenerDetails_Extended, callback) => {
+			let origin: URI;
+			try {
+				origin = URI.parse(details.lastCommittedOrigin!);
+			} catch {
 				return callback({});
 			}
 
+			const webviewId = origin.authority;
 			const entry = this._webviewData.get(webviewId);
 			if (!entry) {
 				return callback({});
@@ -56,26 +60,20 @@ export class WebviewPortMappingProvider extends Disposable {
 		});
 	}
 
-	public async registerWebview(id: string, webContentsId: number | undefined, metadata: PortMappingData): Promise<void> {
+	public async registerWebview(id: string, metadata: PortMappingData): Promise<void> {
 		const manager = new WebviewPortMappingManager(
 			() => this._webviewData.get(id)?.metadata.extensionLocation,
 			() => this._webviewData.get(id)?.metadata.mappings || [],
 			this._tunnelService);
 
-		this._webviewData.set(id, { webContentsId, metadata, manager });
-		if (typeof webContentsId === 'number') {
-			this._webContentsIdsToWebviewIds.set(webContentsId, id);
-		}
+		this._webviewData.set(id, { metadata, manager });
 	}
 
 	public unregisterWebview(id: string): void {
 		const existing = this._webviewData.get(id);
 		if (existing) {
 			existing.manager.dispose();
 			this._webviewData.delete(id);
-			if (typeof existing.webContentsId === 'number') {
-				this._webContentsIdsToWebviewIds.delete(existing.webContentsId);
-			}
 		}
 	}
 

src/vs/workbench/contrib/webview/electron-browser/iframeWebviewElement.ts

@@ -52,7 +52,7 @@ export class ElectronIframeWebview extends IFrameWebview {
 		super(id, options, contentOptions, extension, webviewThemeDataProvider,
 			noficationService, tunnelService, fileService, requestService, telemetryService, environmentService, _workbenchEnvironmentService, _remoteAuthorityResolverService, logService);
 
-		this._resourceRequestManager = this._register(instantiationService.createInstance(WebviewResourceRequestManager, id, extension, this.content.options, Promise.resolve(undefined)));
+		this._resourceRequestManager = this._register(instantiationService.createInstance(WebviewResourceRequestManager, id, extension, this.content.options));
 	}
 
 	protected createElement(options: WebviewOptions, contentOptions: WebviewContentOptions) {

src/vs/workbench/contrib/webview/electron-browser/resourceLoading.ts

@@ -58,7 +58,6 @@ export class WebviewResourceRequestManager extends Disposable {
 		private readonly id: string,
 		private readonly extension: WebviewExtensionDescription | undefined,
 		initialContentOptions: WebviewContentOptions,
-		getWebContentsId: Promise<number | undefined>,
 		@ILogService private readonly _logService: ILogService,
 		@IRemoteAuthorityResolverService remoteAuthorityResolverService: IRemoteAuthorityResolverService,
 		@IWorkbenchEnvironmentService environmentService: IWorkbenchEnvironmentService,
@@ -79,15 +78,13 @@ export class WebviewResourceRequestManager extends Disposable {
 		const remoteAuthority = environmentService.configuration.remoteAuthority;
 		const remoteConnectionData = remoteAuthority ? remoteAuthorityResolverService.getConnectionData(remoteAuthority) : null;
 
-		this._ready = getWebContentsId.then(async (webContentsId) => {
-			this._logService.debug(`WebviewResourceRequestManager(${this.id}): did-start-loading`);
-			await this._webviewManagerService.registerWebview(this.id, webContentsId, electronService.windowId, {
-				extensionLocation: this.extension?.location.toJSON(),
-				localResourceRoots: this._localResourceRoots.map(x => x.toJSON()),
-				remoteConnectionData: remoteConnectionData,
-				portMappings: this._portMappings,
-			});
-
+		this._logService.debug(`WebviewResourceRequestManager(${this.id}): did-start-loading`);
+		this._ready = this._webviewManagerService.registerWebview(this.id, electronService.windowId, {
+			extensionLocation: this.extension?.location.toJSON(),
+			localResourceRoots: this._localResourceRoots.map(x => x.toJSON()),
+			remoteConnectionData: remoteConnectionData,
+			portMappings: this._portMappings,
+		}).then(() => {
 			this._logService.debug(`WebviewResourceRequestManager(${this.id}): did register`);
 		});
 

src/vs/workbench/contrib/webview/electron-browser/webviewElement.ts

@@ -142,17 +142,7 @@ export class ElectronWebviewBasedWebview extends BaseWebview<WebviewTag> impleme
 
 		this._myLogService.debug(`Webview(${this.id}): init`);
 
-		const webviewId = new Promise<number | undefined>((resolve, reject) => {
-			const sub = this._register(addDisposableListener(this.element!, 'dom-ready', once(() => {
-				if (!this.element) {
-					reject();
-					throw new Error('No element');
-				}
-				resolve(this.element.getWebContentsId());
-				sub.dispose();
-			})));
-		});
-		this._resourceRequestManager = this._register(instantiationService.createInstance(WebviewResourceRequestManager, id, extension, this.content.options, webviewId));
+		this._resourceRequestManager = this._register(instantiationService.createInstance(WebviewResourceRequestManager, id, extension, this.content.options));
 
 		this._register(addDisposableListener(this.element!, 'dom-ready', once(() => {
 			this._register(ElectronWebviewBasedWebview.getWebviewKeyboardHandler(configurationService, mainProcessService).add(this.element!));