Issue 4335 - Adblock Plus should forget pop-ups after the initial load

lib/child/contentPolicy.js

Index: lib/child/contentPolicy.js
===================================================================
--- a/lib/child/contentPolicy.js
+++ b/lib/child/contentPolicy.js
@@ -242,21 +242,21 @@ var PolicyImplementation =
 
     let registrar = Components.manager.QueryInterface(Ci.nsIComponentRegistrar);
     registrar.registerFactory(this.classID, this.classDescription, this.contractID, this);
 
     let catMan = Utils.categoryManager;
     for (let category of this.xpcom_categories)
       catMan.addCategoryEntry(category, this.contractID, this.contractID, false, true);
 
-    Services.obs.addObserver(this, "content-document-global-created", true);
+    Services.obs.addObserver(this, "document-element-inserted", true);
 
     onShutdown.add(() =>
     {
-      Services.obs.removeObserver(this, "content-document-global-created");
+      Services.obs.removeObserver(this, "document-element-inserted");
 
       for (let category of this.xpcom_categories)
         catMan.deleteCategoryEntry(category, this.contractID, false);
 
       registrar.unregisterFactory(this.classID, this);
     });
   },
 
@@ -307,60 +307,84 @@ var PolicyImplementation =
   {
     return Ci.nsIContentPolicy.ACCEPT;
   },
 
   //
   // nsIObserver interface implementation
   //
   _openers: new WeakMap(),
+  _alreadyLoaded: {},
 
   observe: function(subject, topic, data, uri)
   {
     switch (topic)
     {
-      case "content-document-global-created":
+      case "document-element-inserted":
       {
-        let opener = this._openers.get(subject);
-        if (opener && Components.utils.isDeadWrapper(opener))
+        let window = subject.defaultView;
+        if (!window)
+          return;
+
+        let type = window.QueryInterface(Ci.nsIInterfaceRequestor)
+                         .getInterface(Ci.nsIWebNavigation)
+                         .QueryInterface(Ci.nsIDocShellTreeItem)
+                         .itemType;
+        if (type != Ci.nsIDocShellTreeItem.typeContent)
+          return;
+
+        let opener = this._openers.get(window);
+        if (opener == this._alreadyLoaded)
+        {
+          // This window has loaded already, ignore it regardless of whether
+          // window.opener is still set.
+          return;
+        }
+
+        if (opener && Cu.isDeadWrapper(opener))
           opener = null;
 
         if (!opener)
         {
           // We don't know the opener for this window yet, try to find it
-          if (subject instanceof Ci.nsIDOMWindow)
-            opener = subject.opener;
-
+          opener = window.opener;
           if (!opener)
             return;
 
           // The opener might be an intermediate window, get the real one
           while (opener.location == "about:blank" && opener.opener)
             opener = opener.opener;
 
-          this._openers.set(subject, opener);
+          this._openers.set(window, opener);
+
+          let forgetPopup = event =>
+          {
+            subject.removeEventListener("DOMContentLoaded", forgetPopup);
+            this._openers.set(window, this._alreadyLoaded);
+          };
+          subject.addEventListener("DOMContentLoaded", forgetPopup);
         }
 
-        if (!uri && subject instanceof Ci.nsIDOMWindow)
-          uri = subject.location.href;
+        if (!uri)
+          uri = window.location.href;
         if (!shouldAllow(opener, opener.document, "POPUP", uri))
         {
-          subject.stop();
-          Utils.runAsync(() => subject.close());
+          window.stop();
+          Utils.runAsync(() => window.close());
         }
         else if (uri == "about:blank")
         {
           // An about:blank pop-up most likely means that a load will be
           // initiated asynchronously. Wait for that.
           Utils.runAsync(() =>
           {
-            let channel = subject.QueryInterface(Ci.nsIInterfaceRequestor)
-                                 .getInterface(Ci.nsIDocShell)
-                                 .QueryInterface(Ci.nsIDocumentLoader)
-                                 .documentChannel;
+            let channel = window.QueryInterface(Ci.nsIInterfaceRequestor)
+                                .getInterface(Ci.nsIDocShell)
+                                .QueryInterface(Ci.nsIDocumentLoader)
+                                .documentChannel;
             if (channel)
               this.observe(subject, topic, data, channel.URI.spec);
           });
         }
         break;
       }
     }
   },
@@ -389,18 +413,18 @@ var PolicyImplementation =
       if (contentType == Ci.nsIContentPolicy.TYPE_DOCUMENT)
       {
         if (wnd.history.length <= 1 && wnd.opener)
         {
           // Special treatment for pop-up windows - this will close the window
           // rather than preventing the redirect. Note that we might not have
           // seen the original channel yet because the redirect happened before
           // the async code in observe() had a chance to run.
-          this.observe(wnd, "content-document-global-created", null, oldChannel.URI.spec);
-          this.observe(wnd, "content-document-global-created", null, newChannel.URI.spec);
+          this.observe(wnd, "document-element-inserted", null, oldChannel.URI.spec);
+          this.observe(wnd, "document-element-inserted", null, newChannel.URI.spec);
         }
         return;
       }
 
       shouldAllowAsync(wnd, wnd.document, types.get(contentType), newChannel.URI.spec, function(allow)
       {
         callback.onRedirectVerifyCallback(allow ? Cr.NS_OK : Cr.NS_BINDING_ABORTED);
       });