3 files changed, 108 insertions, 136 deletions

diff --git a/ext/fg/js/frame-ancestry-handler.js b/ext/fg/js/frame-ancestry-handler.js
index 31ee956b..b1ed7114 100644
--- a/ext/fg/js/frame-ancestry-handler.js
+++ b/ext/fg/js/frame-ancestry-handler.js
@@ -23,6 +23,7 @@
  * This class is used to return the ancestor frame IDs for the current frame.
  * This is a workaround to using the `webNavigation.getAllFrames` API, which
  * would require an additional permission that is otherwise unnecessary.
+ * It is also used to track the correlation between child frame elements and their IDs.
  */
 class FrameAncestryHandler {
     /**
@@ -55,6 +56,14 @@ class FrameAncestryHandler {
     }
 
     /**
+     * Returns whether or not this frame is the root frame in the tab.
+     * @returns `true` if it is the root, otherwise `false`.
+     */
+    isRootFrame() {
+        return (window === window.parent);
+    }
+
+    /**
      * Gets the frame ancestry information for the current frame. If the frame is the
      * root frame, an empty array is returned. Otherwise, an array of frame IDs is returned,
      * starting from the nearest ancestor.
@@ -68,6 +77,26 @@ class FrameAncestryHandler {
         return await this._getFrameAncestryInfoPromise;
     }
 
+    /**
+     * Gets the frame element of a child frame given a frame ID.
+     * For this function to work, the `getFrameAncestryInfo` function needs to have
+     * been invoked previously.
+     * @param frameId The frame ID of the child frame to get.
+     * @returns The element corresponding to the frame with ID `frameId`, otherwise `null`.
+     */
+    getChildFrameElement(frameId) {
+        const frameInfo = this._childFrameMap.get(frameId);
+        if (typeof frameInfo === 'undefined') { return null; }
+
+        let {frameElement} = frameInfo;
+        if (typeof frameElement === 'undefined') {
+            frameElement = this._findFrameElementWithContentWindow(frameInfo.window);
+            frameInfo.frameElement = frameElement;
+        }
+
+        return frameElement;
+    }
+
     // Private
 
     _getFrameAncestryInfo(timeout=5000) {
@@ -166,7 +195,7 @@ class FrameAncestryHandler {
             }
 
             if (!this._childFrameMap.has(childFrameId)) {
-                this._childFrameMap.set(childFrameId, {window: source});
+                this._childFrameMap.set(childFrameId, {window: source, frameElement: void 0});
             }
 
             if (more) {
@@ -192,4 +221,49 @@ class FrameAncestryHandler {
             Math.floor(value) === value
         );
     }
+
+    _findFrameElementWithContentWindow(contentWindow) {
+        // Check frameElement, for non-null same-origin frames
+        try {
+            const {frameElement} = contentWindow;
+            if (frameElement !== null) { return frameElement; }
+        } catch (e) {
+            // NOP
+        }
+
+        // Check frames
+        const frameTypes = ['iframe', 'frame', 'embed'];
+        for (const frameType of frameTypes) {
+            for (const frame of document.getElementsByTagName(frameType)) {
+                if (frame.contentWindow === contentWindow) {
+                    return frame;
+                }
+            }
+        }
+
+        // Check for shadow roots
+        const rootElements = [document.documentElement];
+        while (rootElements.length > 0) {
+            const rootElement = rootElements.shift();
+            const walker = document.createTreeWalker(rootElement, NodeFilter.SHOW_ELEMENT);
+            while (walker.nextNode()) {
+                const element = walker.currentNode;
+
+                if (element.contentWindow === contentWindow) {
+                    return element;
+                }
+
+                const shadowRoot = (
+                    element.shadowRoot ||
+                    element.openOrClosedShadowRoot // Available to Firefox 63+ for WebExtensions
+                );
+                if (shadowRoot) {
+                    rootElements.push(shadowRoot);
+                }
+            }
+        }
+
+        // Not found
+        return null;
+    }
 }
diff --git a/ext/fg/js/frame-offset-forwarder.js b/ext/fg/js/frame-offset-forwarder.js
index b5d5424c..0a0b4a18 100644
--- a/ext/fg/js/frame-offset-forwarder.js
+++ b/ext/fg/js/frame-offset-forwarder.js
@@ -16,162 +16,55 @@
  */
 
 /* global
+ * FrameAncestryHandler
  * api
  */
 
 class FrameOffsetForwarder {
     constructor(frameId) {
         this._frameId = frameId;
-        this._isPrepared = false;
-        this._cacheMaxSize = 1000;
-        this._frameCache = new Set();
-        this._unreachableContentWindowCache = new Set();
-        this._windowMessageHandlers = new Map([
-            ['getFrameOffset', this._onMessageGetFrameOffset.bind(this)]
-        ]);
+        this._frameAncestryHandler = new FrameAncestryHandler(frameId);
     }
 
     prepare() {
-        if (this._isPrepared) { return; }
-        window.addEventListener('message', this._onMessage.bind(this), false);
-        this._isPrepared = true;
+        this._frameAncestryHandler.prepare();
+        api.crossFrame.registerHandlers([
+            ['FrameOffsetForwarder.getChildFrameRect', {async: false, handler: this._onMessageGetChildFrameRect.bind(this)}]
+        ]);
     }
 
     async getOffset() {
-        if (window === window.parent) {
+        if (this._frameAncestryHandler.isRootFrame()) {
             return [0, 0];
         }
 
-        const uniqueId = generateId(16);
-
-        const frameOffsetPromise = yomichan.getTemporaryListenerResult(
-            chrome.runtime.onMessage,
-            ({action, params}, {resolve}) => {
-                if (action === 'frameOffset' && isObject(params) && params.uniqueId === uniqueId) {
-                    resolve(params);
-                }
-            },
-            5000
-        );
-
-        this._getFrameOffsetParent([0, 0], uniqueId, this._frameId);
-
-        const {offset} = await frameOffsetPromise;
-        return offset;
-    }
-
-    // Private
-
-    _onMessage(event) {
-        const data = event.data;
-        if (data === null || typeof data !== 'object') { return; }
-
-        try {
-            const {action, params} = event.data;
-            const handler = this._windowMessageHandlers.get(action);
-            if (typeof handler !== 'function') { return; }
-            handler(params, event);
-        } catch (e) {
-            // NOP
-        }
-    }
-
-    _onMessageGetFrameOffset({offset, uniqueId, frameId}, e) {
-        let sourceFrame = null;
-        if (!this._unreachableContentWindowCache.has(e.source)) {
-            sourceFrame = this._findFrameWithContentWindow(e.source);
-        }
-        if (sourceFrame === null) {
-            // closed shadow root etc.
-            this._addToCache(this._unreachableContentWindowCache, e.source);
-            this._replyFrameOffset(null, uniqueId, frameId);
-            return;
-        }
-
-        const [forwardedX, forwardedY] = offset;
-        const {x, y} = sourceFrame.getBoundingClientRect();
-        offset = [forwardedX + x, forwardedY + y];
+        const ancestorFrameIds = await this._frameAncestryHandler.getFrameAncestryInfo();
 
-        if (window === window.parent) {
-            this._replyFrameOffset(offset, uniqueId, frameId);
-        } else {
-            this._getFrameOffsetParent(offset, uniqueId, frameId);
+        let childFrameId = this._frameId;
+        const promises = [];
+        for (const frameId of ancestorFrameIds) {
+            promises.push(api.crossFrame.invoke(frameId, 'FrameOffsetForwarder.getChildFrameRect', {frameId: childFrameId}));
+            childFrameId = frameId;
         }
-    }
 
-    _findFrameWithContentWindow(contentWindow) {
-        const ELEMENT_NODE = Node.ELEMENT_NODE;
-        for (const elements of this._getFrameElementSources()) {
-            while (elements.length > 0) {
-                const element = elements.shift();
-                if (element.contentWindow === contentWindow) {
-                    this._addToCache(this._frameCache, element);
-                    return element;
-                }
+        const results = await Promise.all(promises);
 
-                const shadowRoot = (
-                    element.shadowRoot ||
-                    element.openOrClosedShadowRoot // Available to Firefox 63+ for WebExtensions
-                );
-                if (shadowRoot) {
-                    for (const child of shadowRoot.children) {
-                        if (child.nodeType === ELEMENT_NODE) {
-                            elements.push(child);
-                        }
-                    }
-                }
-
-                for (const child of element.children) {
-                    if (child.nodeType === ELEMENT_NODE) {
-                        elements.push(child);
-                    }
-                }
-            }
+        let xOffset = 0;
+        let yOffset = 0;
+        for (const {x, y} of results) {
+            xOffset += x;
+            yOffset += y;
         }
-
-        return null;
+        return [xOffset, yOffset];
     }
 
-    *_getFrameElementSources() {
-        const frameCache = [];
-        for (const frame of this._frameCache) {
-            // removed from DOM
-            if (!frame.isConnected) {
-                this._frameCache.delete(frame);
-                continue;
-            }
-            frameCache.push(frame);
-        }
-        yield frameCache;
-        // will contain duplicates, but frame elements are cheap to handle
-        yield [...document.querySelectorAll('frame,iframe')];
-        yield [document.documentElement];
-    }
-
-    _addToCache(cache, value) {
-        let freeSlots = this._cacheMaxSize - cache.size;
-        if (freeSlots <= 0) {
-            for (const cachedValue of cache) {
-                cache.delete(cachedValue);
-                ++freeSlots;
-                if (freeSlots > 0) { break; }
-            }
-        }
-        cache.add(value);
-    }
+    // Private
 
-    _getFrameOffsetParent(offset, uniqueId, frameId) {
-        window.parent.postMessage({
-            action: 'getFrameOffset',
-            params: {
-                offset,
-                uniqueId,
-                frameId
-            }
-        }, '*');
-    }
+    _onMessageGetChildFrameRect({frameId}) {
+        const frameElement = this._frameAncestryHandler.getChildFrameElement(frameId);
+        if (frameElement === null) { return null; }
 
-    _replyFrameOffset(offset, uniqueId, frameId) {
-        api.sendMessageToFrame(frameId, 'frameOffset', {offset, uniqueId});
+        const {x, y, width, height} = frameElement.getBoundingClientRect();
+        return {x, y, width, height};
     }
 }
diff --git a/ext/fg/js/popup-proxy.js b/ext/fg/js/popup-proxy.js
index 6d6c7fb9..bb037705 100644
--- a/ext/fg/js/popup-proxy.js
+++ b/ext/fg/js/popup-proxy.js
@@ -34,7 +34,7 @@ class PopupProxy extends EventDispatcher {
         this._ownerFrameId = ownerFrameId;
         this._frameOffsetForwarder = frameOffsetForwarder;
 
-        this._frameOffset = null;
+        this._frameOffset = [0, 0];
         this._frameOffsetPromise = null;
         this._frameOffsetUpdatedAt = null;
         this._frameOffsetExpireTimeout = 1000;
@@ -194,7 +194,12 @@ class PopupProxy extends EventDispatcher {
     async _updateFrameOffsetInner(now) {
         this._frameOffsetPromise = this._frameOffsetForwarder.getOffset();
         try {
-            const offset = await this._frameOffsetPromise;
+            let offset = null;
+            try {
+                offset = await this._frameOffsetPromise;
+            } catch (e) {
+                // NOP
+            }
             this._frameOffset = offset !== null ? offset : [0, 0];
             if (offset === null) {
                 this.trigger('offsetNotFound');