refactoring allowing to get the contents as soon as the frames are detected

lib/single-file/frame-tree/bg/frame-tree.js

@@ -32,12 +32,6 @@ this.FrameTree = (() => {
 				tabsData[sender.tab.id] = null;
 			}
 		}
-		if (message.method == "FrameTree.getDataRequest") {
-			browser.tabs.sendMessage(sender.tab.id, { method: "FrameTree.getDataRequest", windowId: message.windowId, sessionId: message.sessionId, tabId: sender.tab.id, options: message.options });
-		}
-		if (message.method == "FrameTree.getDataResponse") {
-			browser.tabs.sendMessage(message.tabId, { method: "FrameTree.getDataResponse", windowId: message.windowId, sessionId: message.sessionId, content: message.content, baseURI: message.baseURI, emptyStyleRulesText: message.emptyStyleRulesText });
-		}
 	});
 
 	return {

lib/single-file/frame-tree/content/frame-tree.js

@@ -24,7 +24,6 @@ this.FrameTree = this.FrameTree || (() => {
 
 	const MESSAGE_PREFIX = "__FrameTree__";
 	const TIMEOUT_INIT_REQUEST_MESSAGE = 500;
-	const TIMEOUT_DATA_RESPONSE_MESSAGE = 500;
 
 	const FrameTree = {
 		getFramesData
@@ -34,29 +33,12 @@ this.FrameTree = this.FrameTree || (() => {
 	if (window == top) {
 		browser.runtime.onMessage.addListener(onTopBackgroundMessage);
 	}
-	browser.runtime.onMessage.addListener(onBackgroundMessage);
 	addEventListener("message", onFrameWindowMessage, false);
 	return FrameTree;
 
 	async function getFramesData(options) {
 		const sessionId = options.sessionId;
 		const sessionFramesData = sessions.get(sessionId);
-		await Promise.all(sessionFramesData.frames.map(async frameData => {
-			return new Promise(resolve => {
-				sessionFramesData.dataRequestCallbacks.set(frameData.windowId, resolve);
-				if (frameData.sameDomain) {
-					top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "getDataRequest", windowId: frameData.windowId, sessionId, options: { removeHiddenElements: options.removeHiddenElements, compressHTML: options.compressHTML } }), "*");
-				} else {
-					browser.runtime.sendMessage({
-						method: "FrameTree.getDataRequest",
-						windowId: frameData.windowId,
-						sessionId,
-						options: { removeHiddenElements: options.removeHiddenElements, compressHTML: options.compressHTML }
-					});
-				}
-				frameData.getDataResponseTimeout = timeout.set(() => top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "getDataResponse", windowId: frameData.windowId, sessionId }), "*"), TIMEOUT_DATA_RESPONSE_MESSAGE);
-			});
-		}));
 		sessions.delete(sessionId);
 		return sessionFramesData.frames.sort((frame1, frame2) => frame2.windowId.split(".").length - frame1.windowId.split(".").length);
 	}
@@ -65,27 +47,6 @@ this.FrameTree = this.FrameTree || (() => {
 		if (message.method == "FrameTree.initRequest" && document.documentElement instanceof HTMLHtmlElement) {
 			initRequest(message);
 		}
-		if (message.method == "FrameTree.getDataResponse") {
-			getDataResponse(message);
-		}
-	}
-
-	function onBackgroundMessage(message) {
-		if (message.method == "FrameTree.getDataRequest" && FrameTree.windowId == message.windowId) {
-			const docData = docHelper.preProcessDoc(document, window, message.options);
-			browser.runtime.sendMessage({
-				method: "FrameTree.getDataResponse",
-				windowId: message.windowId,
-				sessionId: message.sessionId,
-				tabId: message.tabId,
-				content: docHelper.serialize(document),
-				emptyStyleRulesText: docData.emptyStyleRulesText,
-				canvasData: docData.canvasData,
-				baseURI: document.baseURI,
-				title: document.title
-			});
-			docHelper.postProcessDoc(document, message.options);
-		}
 	}
 
 	function onFrameWindowMessage(event) {
@@ -95,42 +56,52 @@ this.FrameTree = this.FrameTree || (() => {
 				initRequest(message);
 			} else if (message.method == "initResponse") {
 				initResponse(message);
-			} else if (message.method == "getDataResponse") {
-				getDataResponse(message);
 			}
 		}
 	}
 
 	function initRequest(message) {
-		FrameTree.windowId = message.windowId;
+		const windowId = message.windowId;
+		const sessionId = message.sessionId;
+		FrameTree.windowId = windowId;
 		const frameElements = document.querySelectorAll("iframe, frame, object[type=\"text/html\"][data]");
 		sessions.set(message.sessionId, {
 			frames: [],
 			dataRequestCallbacks: new Map()
 		});
-		if (frameElements.length) {
-			setFramesWinId(MESSAGE_PREFIX, frameElements, message.options, message.windowId, message.sessionId, window);
-		} else {
-			top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initResponse", framesData: [], windowId: message.windowId, sessionId: message.sessionId }), "*");
+		if (window != top) {
+			const docData = docHelper.preProcessDoc(document, window, message.options);
+			const content = docHelper.serialize(document);
+			docHelper.postProcessDoc(document, window, message.options);
+			top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initResponse", framesData: [{ windowId, content, baseURI: document.baseURI, title: document.title, emptyStyleRulesText: docData.emptyStyleRulesText, canvasData: docData.canvasData, processed: true }], sessionId }), "*");
 		}
+		processFrames(frameElements, message.options, windowId, sessionId, window);
 	}
 
 	function initResponse(message) {
 		if (window == top) {
 			if (message.framesData) {
-				message.framesData = message.framesData instanceof Array ? message.framesData : JSON.parse(message.framesData);
 				const sessionFramesData = sessions.get(message.sessionId);
-				if (sessionFramesData) {
-					sessionFramesData.frames = sessionFramesData.frames.concat(...message.framesData);
-					const frameData = sessionFramesData.frames.find(frameData => frameData.windowId == message.windowId);
-					if (message.windowId != "0") {
-						frameData.processed = true;
-					}
-					const pendingCount = sessionFramesData.frames.filter(frameData => !frameData.processed).length;
-					if (!pendingCount && !sessionFramesData.initResponseSent) {
-						sessionFramesData.initResponseSent = true;
-						browser.runtime.sendMessage({ method: "FrameTree.initResponse", sessionId: message.sessionId });
+				message.framesData.forEach(messageFrameData => {
+					let frameData = sessionFramesData.frames.find(frameData => messageFrameData.windowId == frameData.windowId);
+					if (!frameData) {
+						frameData = { windowId: messageFrameData.windowId };
+						sessionFramesData.frames.push(frameData);
 					}
+
+					frameData.content = messageFrameData.content;
+					frameData.baseURI = messageFrameData.baseURI;
+					frameData.title = messageFrameData.title;
+					frameData.emptyStyleRulesText = messageFrameData.emptyStyleRulesText;
+					frameData.canvasData = messageFrameData.canvasData;
+
+					frameData.processed = messageFrameData.processed;
+				});
+
+				const pendingCount = sessionFramesData.frames.filter(frameData => !frameData.processed).length;
+				if (!pendingCount && !sessionFramesData.initResponseSent) {
+					sessionFramesData.initResponseSent = true;
+					browser.runtime.sendMessage({ method: "FrameTree.initResponse", sessionId: message.sessionId });
 				}
 			}
 		} else {
@@ -138,71 +109,42 @@ this.FrameTree = this.FrameTree || (() => {
 		}
 	}
 
-	function getDataResponse(message) {
-		delete message.tabId;
-		delete message.method;
-		const sessionFramesData = sessions.get(message.sessionId);
-		const frameData = sessionFramesData.frames.find(frameData => frameData.windowId == message.windowId);
-		timeout.clear(frameData.getDataResponseTimeout);
-		frameData.content = message.content;
-		frameData.baseURI = message.baseURI;
-		frameData.title = message.title;
-		frameData.emptyStyleRulesText = message.emptyStyleRulesText;
-		frameData.canvasData = message.canvasData;
-		sessionFramesData.dataRequestCallbacks.get(message.windowId)(message);
-	}
-
-	function setFramesWinId(MESSAGE_PREFIX, frameElements, options, windowId, sessionId, win) {
-		const framesData = [];
+	function processFrames(frameElements, options, windowId, sessionId, win) {
 		if (win != top) {
 			win.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initResponse", windowId, sessionId }), "*");
 		}
+		let framesData = [];
 		frameElements.forEach((frameElement, frameIndex) => {
-			let src, sameDomain;
+			const frameWinId = windowId + "." + frameIndex;
+			frameElement.setAttribute(docHelper.WIN_ID_ATTRIBUTE_NAME, frameWinId);
+			framesData.push({ windowId: frameWinId });
 			try {
-				sameDomain = Boolean(frameElement.contentDocument && frameElement.contentWindow && top.addEventListener && top);
-				src = frameElement.src;
+				if (!frameElement.contentDocument) {
+					frameElement.contentWindow.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initRequest", windowId: frameWinId, sessionId, frameIndex, options }), "*");
+					timeout.set(() => top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initResponse", framesData: [], windowId: frameWinId, sessionId }), "*"), TIMEOUT_INIT_REQUEST_MESSAGE);
+				}
 			} catch (error) {
 				/* ignored */
 			}
-			framesData.push({ sameDomain, src, windowId: windowId + "." + frameIndex });
 		});
 		top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initResponse", framesData, windowId, sessionId }), "*");
+		framesData = [];
 		frameElements.forEach((frameElement, frameIndex) => {
 			const frameWinId = windowId + "." + frameIndex;
-			frameElement.setAttribute(docHelper.WIN_ID_ATTRIBUTE_NAME, frameWinId);
-			let frameDoc, frameWindow, topWindow;
-			let content, emptyStyleRulesText, canvasData;
+			const frameWindow = frameElement.contentWindow;
 			try {
-				frameDoc = frameElement.contentDocument;
-				frameWindow = frameElement.contentWindow;
-				topWindow = top.addEventListener && top;
+				const frameDoc = frameElement.contentDocument;
+				if (frameDoc) {
+					processFrames(frameDoc.querySelectorAll("iframe, frame, object[type=\"text/html\"][data]"), options, frameWinId, sessionId, frameWindow);
+					const docData = docHelper.preProcessDoc(frameDoc, frameWindow, options);
+					framesData.push({ windowId: frameWinId, content: docHelper.serialize(frameDoc), baseURI: frameDoc.baseURI, title: frameDoc.title, emptyStyleRulesText: docData.emptyStyleRulesText, canvasData: docData.canvasData, processed: true });
+					docHelper.postProcessDoc(frameDoc, frameWindow, options);
+				}
 			} catch (error) {
 				/* ignored */
 			}
-			if (frameWindow && frameDoc && topWindow) {
-				setFramesWinId(MESSAGE_PREFIX, frameDoc.querySelectorAll("iframe, frame, object[type=\"text/html\"][data]"), options, frameWinId, sessionId, frameWindow);
-				topWindow.addEventListener("message", onMessage, false);
-				const docData = docHelper.preProcessDoc(frameDoc, frameWindow, options);
-				content = docHelper.serialize(frameDoc);
-				emptyStyleRulesText = docData.emptyStyleRulesText;
-				canvasData = docData.canvasData;
-				docHelper.postProcessDoc(frameDoc, options);
-			} else if (frameWindow) {
-				frameWindow.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initRequest", windowId: frameWinId, sessionId, frameIndex, options }), "*");
-				timeout.set(() => top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initResponse", framesData: [], windowId: frameWinId, sessionId, frameIndex }), "*"), TIMEOUT_INIT_REQUEST_MESSAGE);
-			}
-
-			function onMessage(event) {
-				if (typeof event.data == "string" && event.data.startsWith(MESSAGE_PREFIX + "::")) {
-					const message = JSON.parse(event.data.substring(MESSAGE_PREFIX.length + 2));
-					if (message.method == "getDataRequest" && message.windowId == frameWinId) {
-						topWindow.removeEventListener("message", onMessage, false);
-						top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "getDataResponse", windowId: message.windowId, sessionId, content, baseURI: frameDoc.baseURI, title: document.title, emptyStyleRulesText, canvasData }), "*");
-					}
-				}
-			}
 		});
+		top.postMessage(MESSAGE_PREFIX + "::" + JSON.stringify({ method: "initResponse", framesData, windowId, sessionId }), "*");
 	}
 
 })();