Popup init update (#497)

* Add API function to send a message to a specific frameId in a tab

* Update _windowMessageHandlers to support additional info per handler

* Remove message token

* Add new authorization check

* Set up new initialization handler

* Update initialization

* Remove message token

* Replace 'prepare' with 'configure'

* Create new prepare function

* Change configure guard

* Log errors in onMessage

* Improve popup initialize function

* Clear secret/token in _resetFrame

* Remove backend message token

* Clear src and srcdoc attributes before loading

* Don't treat about:blank unloads as load events

3 files changed, 230 insertions, 92 deletions

diff --git a/ext/fg/js/float-main.js b/ext/fg/js/float-main.js
index e7e50a54..20771910 100644
--- a/ext/fg/js/float-main.js
+++ b/ext/fg/js/float-main.js
@@ -56,5 +56,6 @@ async function popupNestedInitialize(id, depth, parentFrameId, url) {
 
 (async () => {
     apiForwardLogsToBackend();
-    new DisplayFloat();
+    const display = new DisplayFloat();
+    await display.prepare();
 })();
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
+        );
+    }
 }
diff --git a/ext/fg/js/popup.js b/ext/fg/js/popup.js
index f5cb6f77..7db53f0d 100644
--- a/ext/fg/js/popup.js
+++ b/ext/fg/js/popup.js
@@ -17,7 +17,6 @@
 
 /* global
  * DOM
- * apiGetMessageToken
  * apiInjectStylesheet
  * apiOptionsGet
  */
@@ -39,8 +38,9 @@ class Popup {
         this._contentScale = 1.0;
         this._containerSizeContentScale = null;
         this._targetOrigin = chrome.runtime.getURL('/').replace(/\/$/, '');
-        this._messageToken = null;
         this._previousOptionsContextSource = null;
+        this._containerSecret = null;
+        this._containerToken = null;
 
         this._container = document.createElement('iframe');
         this._container.className = 'yomichan-float';
@@ -216,40 +216,154 @@ class Popup {
         return injectPromise;
     }
 
+    _initializeFrame(frame, targetOrigin, frameId, setupFrame, timeout=10000) {
+        return new Promise((resolve, reject) => {
+            const tokenMap = new Map();
+            let timer = null;
+            let containerLoadedResolve = null;
+            let containerLoadedReject = null;
+            const containerLoaded = new Promise((resolve2, reject2) => {
+                containerLoadedResolve = resolve2;
+                containerLoadedReject = reject2;
+            });
+
+            const postMessage = (action, params) => {
+                const contentWindow = frame.contentWindow;
+                if (contentWindow === null) { throw new Error('Frame missing content window'); }
+
+                let validOrigin = true;
+                try {
+                    validOrigin = (contentWindow.location.origin === targetOrigin);
+                } catch (e) {
+                    // NOP
+                }
+                if (!validOrigin) { throw new Error('Unexpected frame origin'); }
+
+                contentWindow.postMessage({action, params}, targetOrigin);
+            };
+
+            const onMessage = (message) => {
+                onMessageInner(message);
+                return false;
+            };
+
+            const onMessageInner = async (message) => {
+                try {
+                    if (!isObject(message)) { return; }
+                    const {action, params} = message;
+                    if (!isObject(params)) { return; }
+                    await containerLoaded;
+                    if (timer === null) { return; } // Done
+
+                    switch (action) {
+                        case 'popupPrepared':
+                            {
+                                const {secret} = params;
+                                const token = yomichan.generateId(16);
+                                tokenMap.set(secret, token);
+                                postMessage('initialize', {secret, token, frameId});
+                            }
+                            break;
+                        case 'popupInitialized':
+                            {
+                                const {secret, token} = params;
+                                const token2 = tokenMap.get(secret);
+                                if (typeof token2 !== 'undefined' && token === token2) {
+                                    cleanup();
+                                    resolve({secret, token});
+                                }
+                            }
+                            break;
+                    }
+                } catch (e) {
+                    cleanup();
+                    reject(e);
+                }
+            };
+
+            const onLoad = () => {
+                if (containerLoadedResolve === null) {
+                    cleanup();
+                    reject(new Error('Unexpected load event'));
+                    return;
+                }
+
+                if (Popup.isFrameAboutBlank(frame)) {
+                    return;
+                }
+
+                containerLoadedResolve();
+                containerLoadedResolve = null;
+                containerLoadedReject = null;
+            };
+
+            const cleanup = () => {
+                if (timer === null) { return; } // Done
+                clearTimeout(timer);
+                timer = null;
+
+                containerLoadedResolve = null;
+                if (containerLoadedReject !== null) {
+                    containerLoadedReject(new Error('Terminated'));
+                    containerLoadedReject = null;
+                }
+
+                chrome.runtime.onMessage.removeListener(onMessage);
+                frame.removeEventListener('load', onLoad);
+            };
+
+            // Start
+            timer = setTimeout(() => {
+                cleanup();
+                reject(new Error('Timeout'));
+            }, timeout);
+
+            chrome.runtime.onMessage.addListener(onMessage);
+            frame.addEventListener('load', onLoad);
+
+            // Prevent unhandled rejections
+            containerLoaded.catch(() => {}); // NOP
+
+            setupFrame(frame);
+        });
+    }
+
     async _createInjectPromise() {
-        if (this._messageToken === null) {
-            this._messageToken = await apiGetMessageToken();
-        }
+        this._injectStyles();
+
+        const {secret, token} = await this._initializeFrame(this._container, this._targetOrigin, this._frameId, (frame) => {
+            frame.removeAttribute('src');
+            frame.removeAttribute('srcdoc');
+            frame.setAttribute('src', chrome.runtime.getURL('/fg/float.html'));
+            this._observeFullscreen(true);
+            this._onFullscreenChanged();
+        });
+        this._containerSecret = secret;
+        this._containerToken = token;
 
+        // Configure
+        const messageId = yomichan.generateId(16);
         const popupPreparedPromise = yomichan.getTemporaryListenerResult(
             chrome.runtime.onMessage,
-            ({action, params}, {resolve}) => {
+            (message, {resolve}) => {
                 if (
-                    action === 'popupPrepareCompleted' &&
-                    isObject(params) &&
-                    params.targetPopupId === this._id
+                    isObject(message) &&
+                    message.action === 'popupConfigured' &&
+                    isObject(message.params) &&
+                    message.params.messageId === messageId
                 ) {
                     resolve();
                 }
             }
         );
-
-        const parentFrameId = (typeof this._frameId === 'number' ? this._frameId : null);
-        this._container.setAttribute('src', chrome.runtime.getURL('/fg/float.html'));
-        this._container.addEventListener('load', () => {
-            this._invokeApi('prepare', {
-                popupInfo: {
-                    id: this._id,
-                    parentFrameId
-                },
-                optionsContext: this._optionsContext,
-                childrenSupported: this._childrenSupported,
-                scale: this._contentScale
-            });
+        this._invokeApi('configure', {
+            messageId,
+            frameId: this._frameId,
+            popupId: this._id,
+            optionsContext: this._optionsContext,
+            childrenSupported: this._childrenSupported,
+            scale: this._contentScale
         });
-        this._observeFullscreen(true);
-        this._onFullscreenChanged();
-        this._injectStyles();
 
         return popupPreparedPromise;
     }
@@ -267,6 +381,8 @@ class Popup {
         this._container.removeAttribute('src');
         this._container.removeAttribute('srcdoc');
 
+        this._containerSecret = null;
+        this._containerToken = null;
         this._injectPromise = null;
         this._injectPromiseComplete = false;
     }
@@ -401,11 +517,12 @@ class Popup {
     }
 
     _invokeApi(action, params={}) {
-        const token = this._messageToken;
+        const secret = this._containerSecret;
+        const token = this._containerToken;
         const contentWindow = this._container.contentWindow;
-        if (token === null || contentWindow === null) { return; }
+        if (secret === null || token === null || contentWindow === null) { return; }
 
-        contentWindow.postMessage({action, params, token}, this._targetOrigin);
+        contentWindow.postMessage({action, params, secret, token}, this._targetOrigin);
     }
 
     _getFrameParentElement() {
@@ -653,6 +770,17 @@ class Popup {
         injectedStylesheets.set(id, styleNode);
         return styleNode;
     }
+
+    static isFrameAboutBlank(frame) {
+        try {
+            const contentDocument = frame.contentDocument;
+            if (contentDocument === null) { return false; }
+            const url = contentDocument.location.href;
+            return /^about:blank(?:[#?]|$)/.test(url);
+        } catch (e) {
+            return false;
+        }
+    }
 }
 
 Popup._injectedStylesheets = new Map();