Make sure image views (and custom editors) work properly on the web

This fixes an issue where webviews for custom editors did not have any associated extension information, which caused them try reading `file:` uri resources instead of remote uri resources

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

@@ -90,11 +90,11 @@ export class MainThreadEditorInsets implements MainThreadEditorInsetsShape {
 
 		const webview = this._webviewService.createWebview('' + handle, {
 			enableFindWidget: false,
-			extension: { id: extensionId, location: URI.revive(extensionLocation) }
 		}, {
 			allowScripts: options.enableScripts,
 			localResourceRoots: options.localResourceRoots ? options.localResourceRoots.map(uri => URI.revive(uri)) : undefined
 		});
+		webview.extension = { id: extensionId, location: URI.revive(extensionLocation) };
 
 		const webviewZone = new EditorWebviewZone(editor, line, height, webview);
 

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

@@ -8,7 +8,6 @@ import { Disposable, IDisposable } from 'vs/base/common/lifecycle';
 import { isWeb } from 'vs/base/common/platform';
 import { startsWith } from 'vs/base/common/strings';
 import { URI, UriComponents } from 'vs/base/common/uri';
-import { generateUuid } from 'vs/base/common/uuid';
 import * as modes from 'vs/editor/common/modes';
 import { localize } from 'vs/nls';
 import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
@@ -173,6 +172,11 @@ export class MainThreadWebviews extends Disposable implements MainThreadWebviews
 		webview.webview.contentOptions = reviveWebviewOptions(options as any /*todo@mat */);
 	}
 
+	public $setExtension(handle: WebviewPanelHandle, extensionId: ExtensionIdentifier, extensionLocation: UriComponents): void {
+		const webview = this.getWebviewEditorInput(handle);
+		webview.webview.extension = { id: extensionId, location: URI.revive(extensionLocation) };
+	}
+
 	public $reveal(handle: WebviewPanelHandle, showOptions: WebviewPanelShowOptions): void {
 		const webview = this.getWebviewEditorInput(handle);
 		if (webview.isDisposed()) {

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

@@ -532,7 +532,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
 			},
 			registerWebviewEditorProvider: (viewType: string, provider: vscode.WebviewEditorProvider) => {
 				checkProposedApiEnabled(extension);
-				return extHostWebviews.registerWebviewEditorProvider(viewType, provider);
+				return extHostWebviews.registerWebviewEditorProvider(extension, viewType, provider);
 			},
 			registerDecorationProvider(provider: vscode.DecorationProvider) {
 				checkProposedApiEnabled(extension);

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

@@ -550,6 +550,8 @@ export interface MainThreadWebviewsShape extends IDisposable {
 
 	$setHtml(handle: WebviewPanelHandle, value: string): void;
 	$setOptions(handle: WebviewPanelHandle, options: modes.IWebviewOptions): void;
+	$setExtension(handle: WebviewPanelHandle, extensionId: ExtensionIdentifier, extensionLocation: UriComponents): void;
+
 	$postMessage(handle: WebviewPanelHandle, value: any): Promise<boolean>;
 
 	$registerSerializer(viewType: string): void;

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

@@ -264,7 +264,7 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 	private readonly _proxy: MainThreadWebviewsShape;
 	private readonly _webviewPanels = new Map<WebviewPanelHandle, ExtHostWebviewEditor>();
 	private readonly _serializers = new Map<string, vscode.WebviewPanelSerializer>();
-	private readonly _editorProviders = new Map<string, vscode.WebviewEditorProvider>();
+	private readonly _editorProviders = new Map<string, { readonly provider: vscode.WebviewEditorProvider, readonly extension: IExtensionDescription }>();
 
 	constructor(
 		mainContext: IMainContext,
@@ -313,14 +313,15 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 	}
 
 	public registerWebviewEditorProvider(
+		extension: IExtensionDescription,
 		viewType: string,
-		provider: vscode.WebviewEditorProvider
+		provider: vscode.WebviewEditorProvider,
 	): vscode.Disposable {
 		if (this._editorProviders.has(viewType)) {
 			throw new Error(`Editor provider for '${viewType}' already registered`);
 		}
 
-		this._editorProviders.set(viewType, provider);
+		this._editorProviders.set(viewType, { extension, provider, });
 		this._proxy.$registerEditorProvider(viewType);
 
 		return new Disposable(() => {
@@ -415,21 +416,24 @@ export class ExtHostWebviews implements ExtHostWebviewsShape {
 
 	async $resolveWebviewEditor(
 		resource: UriComponents,
-		webviewHandle: WebviewPanelHandle,
+		handle: WebviewPanelHandle,
 		viewType: string,
 		title: string,
 		state: any,
 		position: EditorViewColumn,
 		options: modes.IWebviewOptions & modes.IWebviewPanelOptions
 	): Promise<void> {
-		const provider = this._editorProviders.get(viewType);
-		if (!provider) {
+		const entry = this._editorProviders.get(viewType);
+		if (!entry) {
 			return Promise.reject(new Error(`No provider found for '${viewType}'`));
 		}
+		const { provider, extension } = entry;
 
-		const webview = new ExtHostWebview(webviewHandle, this._proxy, options, this.initData);
-		const revivedPanel = new ExtHostWebviewEditor(webviewHandle, this._proxy, viewType, title, typeof position === 'number' && position >= 0 ? typeConverters.ViewColumn.to(position) : undefined, options, webview);
-		this._webviewPanels.set(webviewHandle, revivedPanel);
+		this._proxy.$setExtension(handle, extension.identifier, extension.extensionLocation);
+
+		const webview = new ExtHostWebview(handle, this._proxy, options, this.initData);
+		const revivedPanel = new ExtHostWebviewEditor(handle, this._proxy, viewType, title, typeof position === 'number' && position >= 0 ? typeConverters.ViewColumn.to(position) : undefined, options, webview);
+		this._webviewPanels.set(handle, revivedPanel);
 		return Promise.resolve(provider.resolveWebviewEditor(URI.revive(resource), revivedPanel));
 	}
 

src/vs/workbench/contrib/customEditor/browser/customEditorInput.ts

@@ -38,7 +38,7 @@ export class CustomFileEditorInput extends WebviewInput {
 		@IExtensionService private readonly _extensionService: IExtensionService,
 		@IDialogService private readonly dialogService: IDialogService,
 	) {
-		super(id, viewType, '', undefined, webview);
+		super(id, viewType, '', webview);
 		this._editorResource = resource;
 	}
 

src/vs/workbench/contrib/webview/browser/dynamicWebviewEditorOverlay.ts

@@ -24,6 +24,11 @@ export class DynamicWebviewEditorOverlay extends Disposable implements WebviewEd
 	private _html: string = '';
 	private _initialScrollProgress: number = 0;
 	private _state: string | undefined = undefined;
+	private _extension: {
+		readonly location: URI;
+		readonly id?: ExtensionIdentifier;
+	} | undefined;
+
 	private _owner: any = undefined;
 
 	public constructor(
@@ -82,6 +87,7 @@ export class DynamicWebviewEditorOverlay extends Disposable implements WebviewEd
 			this._webview.value = webview;
 			webview.state = this._state;
 			webview.html = this._html;
+			webview.extension = this._extension;
 			if (this.options.tryRestoreScrollPosition) {
 				webview.initialScrollProgress = this._initialScrollProgress;
 			}
@@ -133,6 +139,12 @@ export class DynamicWebviewEditorOverlay extends Disposable implements WebviewEd
 		this.withWebview(webview => webview.contentOptions = value);
 	}
 
+	public get extension() { return this._extension; }
+	public set extension(value) {
+		this._extension = value;
+		this.withWebview(webview => webview.extension = value);
+	}
+
 	private readonly _onDidFocus = this._register(new Emitter<void>());
 	public readonly onDidFocus: Event<void> = this._onDidFocus.event;
 

src/vs/workbench/contrib/webview/browser/webview.ts

@@ -43,10 +43,6 @@ export interface IWebviewService {
 export const WebviewResourceScheme = 'vscode-resource';
 
 export interface WebviewOptions {
-	readonly extension?: {
-		readonly location: URI;
-		readonly id?: ExtensionIdentifier;
-	};
 	readonly enableFindWidget?: boolean;
 	readonly tryRestoreScrollPosition?: boolean;
 	readonly retainContextWhenHidden?: boolean;
@@ -63,6 +59,10 @@ export interface Webview extends IDisposable {
 
 	html: string;
 	contentOptions: WebviewContentOptions;
+	extension: {
+		readonly location: URI;
+		readonly id?: ExtensionIdentifier;
+	} | undefined;
 	initialScrollProgress: number;
 	state: string | undefined;
 

src/vs/workbench/contrib/webview/browser/webviewEditorInput.ts

@@ -6,7 +6,6 @@ import * as dom from 'vs/base/browser/dom';
 import { memoize } from 'vs/base/common/decorators';
 import { URI } from 'vs/base/common/uri';
 import { IEditorModel } from 'vs/platform/editor/common/editor';
-import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions';
 import { EditorInput, EditorModel, GroupIdentifier, IEditorInput, Verbosity } from 'vs/workbench/common/editor';
 import { WebviewEditorOverlay } from 'vs/workbench/contrib/webview/browser/webview';
 import { UnownedDisposable as Unowned } from 'vs/base/common/lifecycle';
@@ -67,16 +66,11 @@ export class WebviewInput extends EditorInput {
 		public readonly id: string,
 		public readonly viewType: string,
 		name: string,
-		public readonly extension: undefined | {
-			readonly location: URI;
-			readonly id: ExtensionIdentifier;
-		},
 		webview: Unowned<WebviewEditorOverlay>
 	) {
 		super();
 
 		this._name = name;
-		this.extension = extension;
 
 		this._webview = this._register(webview.acquire()); // The input owns this webview
 	}
@@ -113,6 +107,10 @@ export class WebviewInput extends EditorInput {
 		return this._webview;
 	}
 
+	public get extension() {
+		return this._webview.extension;
+	}
+
 	public get iconPath() {
 		return this._iconPath;
 	}
@@ -150,14 +148,10 @@ export class RevivedWebviewEditorInput extends WebviewInput {
 		id: string,
 		viewType: string,
 		name: string,
-		extension: undefined | {
-			readonly location: URI;
-			readonly id: ExtensionIdentifier
-		},
 		private readonly reviver: (input: WebviewInput) => Promise<void>,
 		webview: Unowned<WebviewEditorOverlay>
 	) {
-		super(id, viewType, name, extension, webview);
+		super(id, viewType, name, webview);
 	}
 
 	public async resolve(): Promise<IEditorModel> {

src/vs/workbench/contrib/webview/browser/webviewEditorService.ts

@@ -150,7 +150,7 @@ export class WebviewEditorService implements IWebviewEditorService {
 	): WebviewInput {
 		const webview = this.createWebiew(id, extension, options);
 
-		const webviewInput = this._instantiationService.createInstance(WebviewInput, id, viewType, title, extension, new UnownedDisposable(webview), undefined);
+		const webviewInput = this._instantiationService.createInstance(WebviewInput, id, viewType, title, new UnownedDisposable(webview), undefined);
 		this._editorService.openEditor(webviewInput, {
 			pinned: true,
 			preserveFocus: showOptions.preserveFocus,
@@ -197,7 +197,7 @@ export class WebviewEditorService implements IWebviewEditorService {
 		const webview = this.createWebiew(id, extension, options);
 		webview.state = state;
 
-		const webviewInput = new RevivedWebviewEditorInput(id, viewType, title, extension, async (webview: WebviewInput): Promise<void> => {
+		const webviewInput = new RevivedWebviewEditorInput(id, viewType, title, async (webview: WebviewInput): Promise<void> => {
 			const didRevive = await this.tryRevive(webview);
 			if (didRevive) {
 				return Promise.resolve(undefined);
@@ -263,14 +263,15 @@ export class WebviewEditorService implements IWebviewEditorService {
 	}
 
 	private createWebiew(id: string, extension: { location: URI; id: ExtensionIdentifier; } | undefined, options: WebviewInputOptions) {
-		return this._webviewService.createWebviewEditorOverlay(id, {
-			extension: extension,
+		const webview = this._webviewService.createWebviewEditorOverlay(id, {
 			enableFindWidget: options.enableFindWidget,
 			retainContextWhenHidden: options.retainContextWhenHidden
 		}, {
 			...options,
 			localResourceRoots: options.localResourceRoots || this.getDefaultLocalResourceRoots(extension),
 		});
+		webview.extension = extension;
+		return webview;
 	}
 
 	private getDefaultLocalResourceRoots(extension: undefined | {

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

@@ -36,9 +36,14 @@ export class IFrameWebview extends Disposable implements Webview {
 
 	private readonly _portMappingManager: WebviewPortMappingManager;
 
+	public extension: {
+		readonly location: URI;
+		readonly id?: ExtensionIdentifier;
+	} | undefined;
+
 	constructor(
 		private readonly id: string,
-		private _options: WebviewOptions,
+		_options: WebviewOptions,
 		contentOptions: WebviewContentOptions,
 		@IThemeService themeService: IThemeService,
 		@ITunnelService tunnelService: ITunnelService,
@@ -52,7 +57,7 @@ export class IFrameWebview extends Disposable implements Webview {
 		}
 
 		this._portMappingManager = this._register(new WebviewPortMappingManager(
-			this._options.extension ? this._options.extension.location : undefined,
+			() => this.extension ? this.extension.location : undefined,
 			() => this.content.options.portMapping || [],
 			tunnelService
 		));
@@ -307,7 +312,7 @@ export class IFrameWebview extends Disposable implements Webview {
 
 	private async loadResource(requestPath: string, uri: URI) {
 		try {
-			const result = await loadLocalResource(uri, this.fileService, this._options.extension ? this._options.extension.location : undefined,
+			const result = await loadLocalResource(uri, this.fileService, this.extension ? this.extension.location : undefined,
 				() => (this.content.options.localResourceRoots || []));
 
 			if (result.type === 'success') {

src/vs/workbench/contrib/webview/common/portMapping.ts

@@ -14,7 +14,7 @@ export class WebviewPortMappingManager extends Disposable {
 	private readonly _tunnels = new Map<number, Promise<RemoteTunnel>>();
 
 	constructor(
-		private readonly extensionLocation: URI | undefined,
+		private readonly getExtensionLocation: () => URI | undefined,
 		private readonly mappings: () => ReadonlyArray<modes.IWebviewPortMapping>,
 		private readonly tunnelService: ITunnelService
 	) {
@@ -30,7 +30,8 @@ export class WebviewPortMappingManager extends Disposable {
 
 		for (const mapping of this.mappings()) {
 			if (mapping.webviewPort === requestLocalHostInfo.port) {
-				if (this.extensionLocation && this.extensionLocation.scheme === REMOTE_HOST_SCHEME) {
+				const extensionLocation = this.getExtensionLocation();
+				if (extensionLocation && extensionLocation.scheme === REMOTE_HOST_SCHEME) {
 					const tunnel = await this.getOrCreateTunnel(mapping.extensionHostPort);
 					if (tunnel) {
 						return encodeURI(uri.with({

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

@@ -92,7 +92,7 @@ class WebviewSession extends Disposable {
 class WebviewProtocolProvider extends Disposable {
 	constructor(
 		webview: WebviewTag,
-		private readonly _extensionLocation: URI | undefined,
+		private readonly _getExtensionLocation: () => URI | undefined,
 		private readonly _getLocalResourceRoots: () => ReadonlyArray<URI>,
 		private readonly _fileService: IFileService,
 	) {
@@ -111,7 +111,7 @@ class WebviewProtocolProvider extends Disposable {
 			return;
 		}
 
-		registerFileProtocol(contents, WebviewResourceScheme, this._fileService, this._extensionLocation, () =>
+		registerFileProtocol(contents, WebviewResourceScheme, this._fileService, this._getExtensionLocation(), () =>
 			this._getLocalResourceRoots()
 		);
 	}
@@ -123,12 +123,12 @@ class WebviewPortMappingProvider extends Disposable {
 
 	constructor(
 		session: WebviewSession,
-		extensionLocation: URI | undefined,
+		getExtensionLocation: () => URI | undefined,
 		mappings: () => ReadonlyArray<modes.IWebviewPortMapping>,
 		tunnelService: ITunnelService,
 	) {
 		super();
-		this._manager = this._register(new WebviewPortMappingManager(extensionLocation, mappings, tunnelService));
+		this._manager = this._register(new WebviewPortMappingManager(getExtensionLocation, mappings, tunnelService));
 
 		session.onBeforeRequest(async details => {
 			const redirect = await this._manager.getRedirect(details.url);
@@ -233,8 +233,13 @@ export class ElectronWebviewBasedWebview extends Disposable implements Webview,
 	private readonly _onDidFocus = this._register(new Emitter<void>());
 	public readonly onDidFocus: Event<void> = this._onDidFocus.event;
 
+	public extension: {
+		readonly location: URI;
+		readonly id?: ExtensionIdentifier;
+	} | undefined;
+
 	constructor(
-		private readonly _options: WebviewOptions,
+		options: WebviewOptions,
 		contentOptions: WebviewContentOptions,
 		@IInstantiationService instantiationService: IInstantiationService,
 		@IThemeService themeService: IThemeService,
@@ -279,13 +284,13 @@ export class ElectronWebviewBasedWebview extends Disposable implements Webview,
 
 		this._register(new WebviewProtocolProvider(
 			this._webview,
-			this._options.extension ? this._options.extension.location : undefined,
+			() => this.extension ? this.extension.location : undefined,
 			() => (this.content.options.localResourceRoots || []),
 			fileService));
 
 		this._register(new WebviewPortMappingProvider(
 			session,
-			_options.extension ? _options.extension.location : undefined,
+			() => this.extension ? this.extension.location : undefined,
 			() => (this.content.options.portMapping || []),
 			tunnelService,
 		));
@@ -382,7 +387,7 @@ export class ElectronWebviewBasedWebview extends Disposable implements Webview,
 			this._send('devtools-opened');
 		}));
 
-		if (_options.enableFindWidget) {
+		if (options.enableFindWidget) {
 			this._webviewFindWidget = this._register(instantiationService.createInstance(WebviewFindWidget, this));
 
 			this._register(addDisposableListener(this._webview, 'found-in-page', e => {
@@ -529,9 +534,9 @@ export class ElectronWebviewBasedWebview extends Disposable implements Webview,
 		}
 		this._hasAlertedAboutMissingCsp = true;
 
-		if (this._options.extension && this._options.extension.id) {
+		if (this.extension && this.extension.id) {
 			if (this._environementService.isExtensionDevelopment) {
-				this._onMissingCsp.fire(this._options.extension.id);
+				this._onMissingCsp.fire(this.extension.id);
 			}
 
 			type TelemetryClassification = {
@@ -542,7 +547,7 @@ export class ElectronWebviewBasedWebview extends Disposable implements Webview,
 			};
 
 			this._telemetryService.publicLog2<TelemetryData, TelemetryClassification>('webviewMissingCsp', {
-				extension: this._options.extension.id.value
+				extension: this.extension.id.value
 			});
 		}
 	}