From 0f46e3a093e7f0c07ad310d8c17e2582bdfd2741 Mon Sep 17 00:00:00 2001
From: toasted-nutbread <toasted-nutbread@users.noreply.github.com>
Date: Mon, 17 Feb 2020 11:02:21 -0500
Subject: Use a token to ensure that messages are coming from Yomichan

---
 ext/bg/js/backend.js |  9 ++++++++-
 ext/fg/js/float.js   | 50 ++++++++++++++++++++++++++++++++++++++++++++------
 ext/fg/js/popup.js   | 15 +++++++++++----
 ext/mixed/js/api.js  |  4 ++++
 4 files changed, 67 insertions(+), 11 deletions(-)

diff --git a/ext/bg/js/backend.js b/ext/bg/js/backend.js
index 458ea483..a4d085e0 100644
--- a/ext/bg/js/backend.js
+++ b/ext/bg/js/backend.js
@@ -46,6 +46,8 @@ class Backend {
         this.popupWindow = null;
 
         this.apiForwarder = new BackendApiForwarder();
+
+        this.messageToken = yomichan.generateId(16);
     }
 
     async prepare() {
@@ -614,6 +616,10 @@ class Backend {
         });
     }
 
+    async _onApiGetMessageToken() {
+        return this.messageToken;
+    }
+
     // Command handlers
 
     async _onCommandSearch(params) {
@@ -875,7 +881,8 @@ Backend._messageHandlers = new Map([
     ['clipboardGet', (self, ...args) => self._onApiClipboardGet(...args)],
     ['getDisplayTemplatesHtml', (self, ...args) => self._onApiGetDisplayTemplatesHtml(...args)],
     ['getQueryParserTemplatesHtml', (self, ...args) => self._onApiGetQueryParserTemplatesHtml(...args)],
-    ['getZoom', (self, ...args) => self._onApiGetZoom(...args)]
+    ['getZoom', (self, ...args) => self._onApiGetZoom(...args)],
+    ['getMessageToken', (self, ...args) => self._onApiGetMessageToken(...args)]
 ]);
 
 Backend._commandHandlers = new Map([
diff --git a/ext/fg/js/float.js b/ext/fg/js/float.js
index 440a9731..8f21a9c5 100644
--- a/ext/fg/js/float.js
+++ b/ext/fg/js/float.js
@@ -16,7 +16,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-/*global popupNestedInitialize, apiForward, Display*/
+/*global popupNestedInitialize, apiForward, apiGetMessageToken, Display*/
 
 class DisplayFloat extends Display {
     constructor() {
@@ -30,6 +30,8 @@ class DisplayFloat extends Display {
 
         this._orphaned = false;
         this._prepareInvoked = false;
+        this._messageToken = null;
+        this._messageTokenPromise = null;
 
         yomichan.on('orphaned', () => this.onOrphaned());
         window.addEventListener('message', (e) => this.onMessage(e), false);
@@ -75,11 +77,23 @@ class DisplayFloat extends Display {
     }
 
     onMessage(e) {
-        const {action, params} = e.data;
-        const handler = DisplayFloat._messageHandlers.get(action);
-        if (typeof handler !== 'function') { return; }
-
-        handler(this, params);
+        const data = e.data;
+        if (typeof data !== 'object' || data === null) { return; } // Invalid data
+
+        const token = data.token;
+        if (typeof token !== 'string') { return; } // Invalid data
+
+        if (this._messageToken === null) {
+            // Async
+            this.getMessageToken()
+                .then(
+                    () => { this.handleAction(token, data); },
+                    () => {}
+                );
+        } else {
+            // Sync
+            this.handleAction(token, data);
+        }
     }
 
     onKeyDown(e) {
@@ -94,6 +108,30 @@ class DisplayFloat extends Display {
         return super.onKeyDown(e);
     }
 
+    async getMessageToken() {
+        // this._messageTokenPromise is used to ensure that only one call to apiGetMessageToken is made.
+        if (this._messageTokenPromise === null) {
+            this._messageTokenPromise = apiGetMessageToken();
+        }
+        const messageToken = await this._messageTokenPromise;
+        if (this._messageToken === null) {
+            this._messageToken = messageToken;
+        }
+        this._messageTokenPromise = null;
+    }
+
+    handleAction(token, {action, params}) {
+        if (token !== this._messageToken) {
+            // Invalid token
+            return;
+        }
+
+        const handler = DisplayFloat._messageHandlers.get(action);
+        if (typeof handler !== 'function') { return; }
+
+        handler(this, params);
+    }
+
     getOptionsContext() {
         return this.optionsContext;
     }
diff --git a/ext/fg/js/popup.js b/ext/fg/js/popup.js
index 900e7325..4927f4bd 100644
--- a/ext/fg/js/popup.js
+++ b/ext/fg/js/popup.js
@@ -16,7 +16,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
-/*global apiInjectStylesheet*/
+/*global apiInjectStylesheet, apiGetMessageToken*/
 
 class Popup {
     constructor(id, depth, frameIdPromise) {
@@ -34,6 +34,7 @@ class Popup {
         this._contentScale = 1.0;
         this._containerSizeContentScale = null;
         this._targetOrigin = chrome.runtime.getURL('/').replace(/\/$/, '');
+        this._messageToken = null;
 
         this._container = document.createElement('iframe');
         this._container.className = 'yomichan-float';
@@ -198,6 +199,10 @@ class Popup {
             // NOP
         }
 
+        if (this._messageToken === null) {
+            this._messageToken = await apiGetMessageToken();
+        }
+
         return new Promise((resolve) => {
             const parentFrameId = (typeof this._frameId === 'number' ? this._frameId : null);
             this._container.setAttribute('src', chrome.runtime.getURL('/fg/float.html'));
@@ -349,9 +354,11 @@ class Popup {
     }
 
     _invokeApi(action, params={}) {
-        if (this._container.contentWindow) {
-            this._container.contentWindow.postMessage({action, params}, this._targetOrigin);
-        }
+        const token = this._messageToken;
+        const contentWindow = this._container.contentWindow;
+        if (token === null || contentWindow === null) { return; }
+
+        contentWindow.postMessage({action, params, token}, this._targetOrigin);
     }
 
     static _getFullscreenElement() {
diff --git a/ext/mixed/js/api.js b/ext/mixed/js/api.js
index 14900ecf..7ea68d59 100644
--- a/ext/mixed/js/api.js
+++ b/ext/mixed/js/api.js
@@ -113,6 +113,10 @@ function apiGetZoom() {
     return _apiInvoke('getZoom');
 }
 
+function apiGetMessageToken() {
+    return _apiInvoke('getMessageToken');
+}
+
 function _apiInvoke(action, params={}) {
     const data = {action, params};
     return new Promise((resolve, reject) => {
-- 
cgit v1.2.3