From c8c311a82f659d2c0629bb45bff1110e91cc3421 Mon Sep 17 00:00:00 2001
From: kkapsner <github@mail.kkapsner.de>
Date: Fri, 24 May 2019 18:30:57 +0200
Subject: [PATCH] Protect reopened documents.

---
 lib/iframeProtection.js | 47 ++++++++++++++++++++++++++++++++---------
 1 file changed, 37 insertions(+), 10 deletions(-)

diff --git a/lib/iframeProtection.js b/lib/iframeProtection.js
index 874178a..8b9b374 100644
--- a/lib/iframeProtection.js
+++ b/lib/iframeProtection.js
@@ -114,34 +114,61 @@
 			});
 		});
 		
-		// MutationObserver to intercept iframes while generating the DOM.
-		var observer = new MutationObserver(allCallback);
-		observer.observe(window.document.documentElement, {subtree: true, childList: true});
-		window.document.addEventListener("DOMContentLoaded", function(){
-			observer.disconnect();
-		});
+		// MutationObserver to intercept iFrames while generating the DOM.
+		const observe = function(){
+			var observer = new MutationObserver(allCallback);
+			var observing = false;
+			function observe(){
+				if (!observing){
+					observer.observe(window.document.documentElement, {subtree: true, childList: true});
+					observing = true;
+				}
+			}
+			observe();
+			window.document.addEventListener("DOMContentLoaded", function(){
+				if (observing){
+					observer.disconnect();
+					observing = false;
+				}
+			});
+			return observe;
+		}();
 		
 		// MutationObserver does not trigger fast enough when document.write is used
-		const documentWriteDescriptor = Object.getOwnPropertyDescriptor(wrappedWindow.HTMLDocument.prototype, "write");
+		const documentWriteDescriptor = Object.getOwnPropertyDescriptor(
+			wrappedWindow.HTMLDocument.prototype,
+			"write"
+		);
 		const documentWrite = documentWriteDescriptor.value;
 		documentWriteDescriptor.value = exportFunction(function write(str){
-			const parts = str.split(/(?=>)/);
+			const parts = str.split(/(?=<)/);
 			const length = parts.length;
+			const scripts = window.document.getElementsByTagName("script");
 			for (let i = 0; i < length; i += 1){
 				documentWrite.call(this, parts[i]);
 				allCallback();
+				if (scripts.length && scripts[scripts.length - 1].src){
+					observe();
+				}
 			}
 		}, window);
 		Object.defineProperty(wrappedWindow.HTMLDocument.prototype, "write", documentWriteDescriptor);
 		
-		const documentWritelnDescriptor = Object.getOwnPropertyDescriptor(wrappedWindow.HTMLDocument.prototype, "writeln");
+		const documentWritelnDescriptor = Object.getOwnPropertyDescriptor(
+			wrappedWindow.HTMLDocument.prototype,
+			"writeln"
+		);
 		const documentWriteln = documentWritelnDescriptor.value;
 		documentWritelnDescriptor.value = exportFunction(function writeln(str){
-			const parts = str.split(/(?=>)/);
+			const parts = str.split(/(?=<)/);
 			const length = parts.length - 1;
+			const scripts = window.document.getElementsByTagName("script");
 			for (let i = 0; i < length; i += 1){
 				documentWrite.call(this, parts[i]);
 				allCallback();
+				if (scripts.length && scripts[scripts.length - 1].src){
+					observe();
+				}
 			}
 			documentWriteln.call(this, parts[length]);
 		}, window);