Frame ancestry handler refactor (#1352)

* Validate source window before handling messages

* Add unregisterHandler to CrossFrameAPI

* Refactor the process FrameAncestryHandler uses to get ancestor frame IDs

* Store a mapping of child frame information

* Update getFrameAncestryInfo to only run once

2 files changed, 66 insertions, 61 deletions

diff --git a/ext/fg/js/frame-ancestry-handler.js b/ext/fg/js/frame-ancestry-handler.js
index d51467f8..31ee956b 100644
--- a/ext/fg/js/frame-ancestry-handler.js
+++ b/ext/fg/js/frame-ancestry-handler.js
@@ -33,7 +33,9 @@ class FrameAncestryHandler {
         this._frameId = frameId;
         this._isPrepared = false;
         this._requestMessageId = 'FrameAncestryHandler.requestFrameInfo';
-        this._responseMessageId = `${this._requestMessageId}.response`;
+        this._responseMessageIdBase = `${this._requestMessageId}.response.`;
+        this._getFrameAncestryInfoPromise = null;
+        this._childFrameMap = new Map();
     }
 
     /**
@@ -59,7 +61,16 @@ class FrameAncestryHandler {
      * @param timeout The maximum time to wait to receive a response to frame information requests.
      * @returns An array of frame IDs corresponding to the ancestors of the current frame.
      */
-    getFrameAncestryInfo(timeout=5000) {
+    async getFrameAncestryInfo() {
+        if (this._getFrameAncestryInfoPromise === null) {
+            this._getFrameAncestryInfoPromise = this._getFrameAncestryInfo(5000);
+        }
+        return await this._getFrameAncestryInfoPromise;
+    }
+
+    // Private
+
+    _getFrameAncestryInfo(timeout=5000) {
         return new Promise((resolve, reject) => {
             const targetWindow = window.parent;
             if (window === targetWindow) {
@@ -68,10 +79,9 @@ class FrameAncestryHandler {
             }
 
             const uniqueId = generateId(16);
-            const responseMessageId = this._responseMessageId;
+            let nonce = generateId(16);
+            const responseMessageId = `${this._responseMessageIdBase}${uniqueId}`;
             const results = [];
-            let resultsExpectedCount = null;
-            let resultsCount = 0;
             let timer = null;
 
             const cleanup = () => {
@@ -79,42 +89,24 @@ class FrameAncestryHandler {
                     clearTimeout(timer);
                     timer = null;
                 }
-                chrome.runtime.onMessage.removeListener(onMessage);
+                api.crossFrame.unregisterHandler(responseMessageId);
             };
-            const onMessage = (message, sender, sendResponse) => {
-                // Validate message
-                if (
-                    typeof message !== 'object' ||
-                    message === null ||
-                    message.action !== responseMessageId
-                ) {
-                    return;
-                }
-
-                const {params} = message;
-                if (params.uniqueId !== uniqueId) { return; } // Wrong ID
-
-                const {frameId, index, more} = params;
-                console.log({frameId, index, more});
-                if (typeof results[index] !== 'undefined') { return; } // Invalid repeat
+            const onMessage = (params) => {
+                if (params.nonce !== nonce) { return null; }
 
                 // Add result
-                results[index] = frameId;
-                ++resultsCount;
-                if (!more) {
-                    resultsExpectedCount = index + 1;
-                }
+                const {frameId, more} = params;
+                results.push(frameId);
+                nonce = generateId(16);
 
-                if (resultsExpectedCount !== null && resultsCount >= resultsExpectedCount) {
+                if (!more) {
                     // Cleanup
                     cleanup();
-                    sendResponse();
 
                     // Finish
                     resolve(results);
-                } else {
-                    resetTimeout();
                 }
+                return {nonce};
             };
             const onTimeout = () => {
                 timer = null;
@@ -127,61 +119,68 @@ class FrameAncestryHandler {
             };
 
             // Start
-            chrome.runtime.onMessage.addListener(onMessage);
+            api.crossFrame.registerHandlers([[responseMessageId, {async: false, handler: onMessage}]]);
             resetTimeout();
-            this._requestFrameInfo(targetWindow, uniqueId, this._frameId, 0);
+            const frameId = this._frameId;
+            this._requestFrameInfo(targetWindow, frameId, frameId, uniqueId, nonce);
         });
     }
 
-    // Private
-
     _onWindowMessage(event) {
+        const {source} = event;
+        if (source === window || source.parent !== window) { return; }
+
         const {data} = event;
         if (
             typeof data === 'object' &&
             data !== null &&
             data.action === this._requestMessageId
         ) {
-            try {
-                this._onRequestFrameInfo(data.params);
-            } catch (e) {
-                // NOP
-            }
+            this._onRequestFrameInfo(data.params, source);
         }
     }
 
-    _onRequestFrameInfo({uniqueId, originFrameId, index}) {
-        if (
-            typeof uniqueId !== 'string' ||
-            typeof originFrameId !== 'number' ||
-            !this._isNonNegativeInteger(index)
-        ) {
-            return;
-        }
+    async _onRequestFrameInfo(params, source) {
+        try {
+            let {originFrameId, childFrameId, uniqueId, nonce} = params;
+            if (
+                !this._isNonNegativeInteger(originFrameId) ||
+                typeof uniqueId !== 'string' ||
+                typeof nonce !== 'string'
+            ) {
+                return;
+            }
 
-        const {parent} = window;
-        const more = (window !== parent);
+            const frameId = this._frameId;
+            const {parent} = window;
+            const more = (window !== parent);
+            const responseParams = {frameId, nonce, more};
+            const responseMessageId = `${this._responseMessageIdBase}${uniqueId}`;
 
-        const responseParams = {uniqueId, frameId: this._frameId, index, more};
-        this._safeSendMessageToFrame(originFrameId, this._responseMessageId, responseParams);
+            try {
+                const response = await api.crossFrame.invoke(originFrameId, responseMessageId, responseParams);
+                if (response === null) { return; }
+                nonce = response.nonce;
+            } catch (e) {
+                return;
+            }
 
-        if (more) {
-            this._requestFrameInfo(parent, uniqueId, originFrameId, index + 1);
-        }
-    }
+            if (!this._childFrameMap.has(childFrameId)) {
+                this._childFrameMap.set(childFrameId, {window: source});
+            }
 
-    async _safeSendMessageToFrame(frameId, action, params) {
-        try {
-            await api.sendMessageToFrame(frameId, action, params);
+            if (more) {
+                this._requestFrameInfo(parent, originFrameId, frameId, uniqueId, nonce);
+            }
         } catch (e) {
             // NOP
         }
     }
 
-    _requestFrameInfo(targetWindow, uniqueId, originFrameId, index) {
+    _requestFrameInfo(targetWindow, originFrameId, childFrameId, uniqueId, nonce) {
         targetWindow.postMessage({
             action: this._requestMessageId,
-            params: {uniqueId, originFrameId, index}
+            params: {originFrameId, childFrameId, uniqueId, nonce}
         }, '*');
     }
 
diff --git a/ext/mixed/js/comm.js b/ext/mixed/js/comm.js
index ea450da6..997249c8 100644
--- a/ext/mixed/js/comm.js
+++ b/ext/mixed/js/comm.js
@@ -248,6 +248,12 @@ class CrossFrameAPI {
         }
     }
 
+    unregisterHandler(key) {
+        return this._messageHandlers.delete(key);
+    }
+
+    // Private
+
     _onConnect(port) {
         try {
             let details;