From 355b5692ff64a3604ddad9ea4a8076f6fb9545cd Mon Sep 17 00:00:00 2001
From: Matt Bierner <matb@microsoft.com>
Date: Thu, 20 Jun 2019 14:36:48 -0700
Subject: [PATCH] Add version check to service worker

Try to make sure our page is talking to the expected version of the service worker
---
 .../contrib/webview/browser/pre/index.html    | 99 ++++++++++++++-----
 .../contrib/webview/browser/pre/main.js       | 28 +-----
 .../webview/browser/pre/service-worker.js     | 16 ++-
 3 files changed, 94 insertions(+), 49 deletions(-)

diff --git a/src/vs/workbench/contrib/webview/browser/pre/index.html b/src/vs/workbench/contrib/webview/browser/pre/index.html
index 9727c148c3f..4128a65be4a 100644
--- a/src/vs/workbench/contrib/webview/browser/pre/index.html
+++ b/src/vs/workbench/contrib/webview/browser/pre/index.html
@@ -16,38 +16,91 @@
 	<script>
 		(function () {
 			const id = document.location.search.match(/\bid=([\w\-]+)/)[1];
-			const handlers = new Map();
 
-			const postMessageToVsCode = (channel, data) => {
-				window.parent.postMessage({ target: id, channel, data }, '*');
-			};
+			const hostMessaging = new class HostMessaging {
+				constructor() {
+					this.handlers = new Map();
+					window.addEventListener('message', (e) => {
+						if (e.data && (e.data.command === 'onmessage' || e.data.command === 'do-update-state')) {
+							// Came from inner iframe
+							this.postMessage(e.data.command, e.data.data);
+							return;
+						}
 
-			window.addEventListener('message', (e) => {
-				if (e.data && (e.data.command === 'onmessage' || e.data.command === 'do-update-state')) {
-					// Came from inner iframe
-					postMessageToVsCode(e.data.command, e.data.data);
-					return;
+						const channel = e.data.channel;
+						const handler = this.handlers.get(channel);
+						if (handler) {
+							handler(e, e.data.args);
+						} else {
+							console.log('no handler for ', e);
+						}
+					});
 				}
 
-				const channel = e.data.channel;
-				const handler = handlers.get(channel);
-				if (handler) {
-					handler(e, e.data.args);
-				} else {
-					console.log('no handler for ', e);
+				postMessage(channel, data) {
+					window.parent.postMessage({ target: id, channel, data }, '*');
 				}
+
+				onMessage(channel, handler) {
+					this.handlers.set(channel, handler);
+				}
+			}();
+
+			const workerReady = new Promise(async (resolveWorkerReady) => {
+				if (!navigator.serviceWorker) {
+					resolveWorkerReady();
+				}
+
+				const expectedWorkerVersion = 1;
+
+				navigator.serviceWorker.register('service-worker.js').then(async registration => {
+					await navigator.serviceWorker.ready;
+
+					const versionHandler = (event) => {
+						if (event.data.channel !== 'version') {
+							return;
+						}
+
+						navigator.serviceWorker.removeEventListener('message', versionHandler);
+						if (event.data.version === expectedWorkerVersion) {
+							return resolveWorkerReady();
+						} else {
+							console.log('Wrong version')
+							// If we have the wrong version, try once to unregister and re-register
+							return registration.unregister()
+								.then(() => navigator.serviceWorker.register('service-worker.js'))
+								.then(navigator.serviceWorker.ready)
+								.finally(resolveWorkerReady);
+						}
+					};
+					navigator.serviceWorker.addEventListener('message', versionHandler);
+					registration.active.postMessage({ channel: 'version' });
+				});
+
+				const forwardFromHostToWorker = (channel) => {
+					hostMessaging.onMessage(channel, event => {
+						navigator.serviceWorker.ready.then(registration => {
+							registration.active.postMessage({ channel: channel, data: event.data.args });
+						});
+					});
+				}
+				forwardFromHostToWorker('did-load-resource');
+				forwardFromHostToWorker('did-load-localhost');
+
+				navigator.serviceWorker.addEventListener('message', event => {
+					if (['load-resource', 'load-localhost'].includes(event.data.channel)) {
+						hostMessaging.postMessage(event.data.channel, event.data);
+					}
+				});
 			});
 
 			createWebviewManager({
-				origin: document.location.href,
-				postMessage: (channel, data) => {
-					postMessageToVsCode(channel, data);
-				},
-				onMessage: (channel, handler) => {
-					handlers.set(channel, handler);
-				}
+				postMessage: hostMessaging.postMessage.bind(hostMessaging),
+				onMessage: hostMessaging.onMessage.bind(hostMessaging),
+				ready: workerReady,
 			});
 		}());
 	</script>
 </body>
-</html>
+
+</html>
\ No newline at end of file
diff --git a/src/vs/workbench/contrib/webview/browser/pre/main.js b/src/vs/workbench/contrib/webview/browser/pre/main.js
index caef053b16f..c30f0c79592 100644
--- a/src/vs/workbench/contrib/webview/browser/pre/main.js
+++ b/src/vs/workbench/contrib/webview/browser/pre/main.js
@@ -99,7 +99,8 @@
 	 *   postMessage: (channel: string, data?: any) => void,
 	 *   onMessage: (channel: string, handler: any) => void,
 	 *   injectHtml?: (document: HTMLDocument) => void,
-	 *   focusIframeOnCreate?: boolean
+	 *   focusIframeOnCreate?: boolean,
+	 *   ready?: Promise<void>
 	 * }} HostCommunications
 	 */
 
@@ -120,29 +121,6 @@
 		// Service worker for resource loading
 		const FAKE_LOAD = !!navigator.serviceWorker;
 
-		const workerReady = new Promise(resolve => {
-			if (!navigator.serviceWorker) {
-				resolve();
-			}
-			navigator.serviceWorker.register('service-worker.js').finally(resolve);
-
-			function forwardFromHostToWorker(channel) {
-				host.onMessage(channel, event => {
-					navigator.serviceWorker.ready.then(registration => {
-						registration.active.postMessage({ channel: channel, data: event.data.args });
-					});
-				});
-			}
-			forwardFromHostToWorker('did-load-resource');
-			forwardFromHostToWorker('did-load-localhost');
-
-			navigator.serviceWorker.addEventListener('message', event => {
-				if (['load-resource', 'load-localhost'].includes(event.data.channel)) {
-					host.postMessage(event.data.channel, event.data);
-				}
-			});
-		});
-
 		/**
 		 * @param {HTMLDocument?} document
 		 * @param {HTMLElement?} body
@@ -269,7 +247,7 @@
 			let updateId = 0;
 			host.onMessage('content', async (_event, data) => {
 				const currentUpdateId = ++updateId;
-				await workerReady;
+				await host.ready;
 				if (currentUpdateId !== updateId) {
 					return;
 				}
diff --git a/src/vs/workbench/contrib/webview/browser/pre/service-worker.js b/src/vs/workbench/contrib/webview/browser/pre/service-worker.js
index 09157ea5121..7e200dc4e98 100644
--- a/src/vs/workbench/contrib/webview/browser/pre/service-worker.js
+++ b/src/vs/workbench/contrib/webview/browser/pre/service-worker.js
@@ -2,6 +2,8 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
+const VERSION = 1;
+
 /**
  * Root path for resources
  */
@@ -106,8 +108,20 @@ const notFoundResponse = new Response('Not Found', {
 	status: 404,
 });
 
-self.addEventListener('message', (event) => {
+self.addEventListener('message', async (event) => {
 	switch (event.data.channel) {
+		case 'version':
+			{
+				self.clients.get(event.source.id).then(client => {
+					if (client) {
+						client.postMessage({
+							channel: 'version',
+							version: VERSION
+						});
+					}
+				});
+				return;
+			}
 		case 'did-load-resource':
 			{
 				const webviewId = getWebviewIdForClient(event.source);
-- 
GitLab