1 files changed, 72 insertions, 63 deletions

diff --git a/ext/fg/js/float.js b/ext/fg/js/float.js
index 77e8edd8..845bf7f6 100644
--- a/ext/fg/js/float.js
+++ b/ext/fg/js/float.js
@@ -18,7 +18,7 @@
 /* global
  * Display
  * apiBroadcastTab
- * apiGetMessageToken
+ * apiSendMessageToFrame
  * popupNestedInitialize
  */
 
@@ -27,12 +27,11 @@ class DisplayFloat extends Display {
         super(document.querySelector('#spinner'), document.querySelector('#definitions'));
         this.autoPlayAudioTimer = null;
 
-        this._popupId = null;
+        this._secret = yomichan.generateId(16);
+        this._token = null;
 
         this._orphaned = false;
-        this._prepareInvoked = false;
-        this._messageToken = null;
-        this._messageTokenPromise = null;
+        this._initializedNestedPopups = false;
 
         this._onKeyDownHandlers = new Map([
             ['C', (e) => {
@@ -46,38 +45,23 @@ class DisplayFloat extends Display {
         ]);
 
         this._windowMessageHandlers = new Map([
-            ['setOptionsContext', ({optionsContext}) => this.setOptionsContext(optionsContext)],
-            ['setContent', ({type, details}) => this.setContent(type, details)],
-            ['clearAutoPlayTimer', () => this.clearAutoPlayTimer()],
-            ['setCustomCss', ({css}) => this.setCustomCss(css)],
-            ['prepare', ({popupInfo, optionsContext, childrenSupported, scale}) => this.prepare(popupInfo, optionsContext, childrenSupported, scale)],
-            ['setContentScale', ({scale}) => this.setContentScale(scale)]
+            ['initialize', {handler: this._initialize.bind(this), authenticate: false}],
+            ['configure', {handler: this._configure.bind(this)}],
+            ['setOptionsContext', {handler: ({optionsContext}) => this.setOptionsContext(optionsContext)}],
+            ['setContent', {handler: ({type, details}) => this.setContent(type, details)}],
+            ['clearAutoPlayTimer', {handler: () => this.clearAutoPlayTimer()}],
+            ['setCustomCss', {handler: ({css}) => this.setCustomCss(css)}],
+            ['setContentScale', {handler: ({scale}) => this.setContentScale(scale)}]
         ]);
-
-        yomichan.on('orphaned', this.onOrphaned.bind(this));
-        window.addEventListener('message', this.onMessage.bind(this), false);
     }
 
-    async prepare(popupInfo, optionsContext, childrenSupported, scale) {
-        if (this._prepareInvoked) { return; }
-        this._prepareInvoked = true;
-
-        const {id, parentFrameId} = popupInfo;
-        this._popupId = id;
-
-        this.optionsContext = optionsContext;
-
+    async prepare() {
         await super.prepare();
-        await this.updateOptions();
-
-        if (childrenSupported) {
-            const {depth, url} = optionsContext;
-            popupNestedInitialize(id, depth, parentFrameId, url);
-        }
 
-        this.setContentScale(scale);
+        yomichan.on('orphaned', this.onOrphaned.bind(this));
+        window.addEventListener('message', this.onMessage.bind(this), false);
 
-        apiBroadcastTab('popupPrepareCompleted', {targetPopupId: this._popupId});
+        apiBroadcastTab('popupPrepared', {secret: this._secret});
     }
 
     onError(error) {
@@ -102,46 +86,30 @@ class DisplayFloat extends Display {
 
     onMessage(e) {
         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);
+        if (typeof data !== 'object' || data === null) {
+            this._logMessageError(e, 'Invalid data');
+            return;
         }
-    }
 
-    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;
+        const action = data.action;
+        if (typeof action !== 'string') {
+            this._logMessageError(e, 'Invalid data');
+            return;
         }
-        this._messageTokenPromise = null;
-    }
 
-    handleAction(token, {action, params}) {
-        if (token !== this._messageToken) {
-            // Invalid token
+        const handlerInfo = this._windowMessageHandlers.get(action);
+        if (typeof handlerInfo === 'undefined') {
+            this._logMessageError(e, `Invalid action: ${JSON.stringify(action)}`);
             return;
         }
 
-        const handler = this._windowMessageHandlers.get(action);
-        if (typeof handler !== 'function') { return; }
+        if (handlerInfo.authenticate !== false && !this._isMessageAuthenticated(data)) {
+            this._logMessageError(e, 'Invalid authentication');
+            return;
+        }
 
-        handler(params);
+        const handler = handlerInfo.handler;
+        handler(data.params);
     }
 
     autoPlayAudio() {
@@ -193,4 +161,45 @@ class DisplayFloat extends Display {
             return '';
         }
     }
+
+    _logMessageError(event, type) {
+        yomichan.logWarning(new Error(`Popup received invalid message from origin ${JSON.stringify(event.origin)}: ${type}`));
+    }
+
+    _initialize(params) {
+        if (this._token !== null) { return; } // Already initialized
+        if (!isObject(params)) { return; } // Invalid data
+
+        const secret = params.secret;
+        if (secret !== this._secret) { return; } // Invalid authentication
+
+        const {token, frameId} = params;
+        this._token = token;
+
+        apiSendMessageToFrame(frameId, 'popupInitialized', {secret, token});
+    }
+
+    async _configure({messageId, frameId, popupId, optionsContext, childrenSupported, scale}) {
+        this.optionsContext = optionsContext;
+
+        await this.updateOptions();
+
+        if (childrenSupported && !this._initializedNestedPopups) {
+            const {depth, url} = optionsContext;
+            popupNestedInitialize(popupId, depth, frameId, url);
+            this._initializedNestedPopups = true;
+        }
+
+        this.setContentScale(scale);
+
+        apiSendMessageToFrame(frameId, 'popupConfigured', {messageId});
+    }
+
+    _isMessageAuthenticated(message) {
+        return (
+            this._token !== null &&
+            this._token === message.token &&
+            this._secret === message.secret
+        );
+    }
 }