Issue 5650 - Apply emulation filters to shadow DOMs

lib/content/elemHideEmulation.js

Index: lib/content/elemHideEmulation.js
===================================================================
--- a/lib/content/elemHideEmulation.js
+++ b/lib/content/elemHideEmulation.js
@@ -31,34 +31,42 @@
 {
   let {children} = node.parentNode;
   for (let i = 0; i < children.length; i++)
     if (children[i] == node)
       return i + 1;
   return 0;
 }
 
+function isShadowRoot(node)
+{
+  return typeof ShadowRoot != "undefined" && node instanceof ShadowRoot;
+}
+
 function makeSelector(node, selector)
 {
   if (node == null)
     return null;
-  if (!node.parentElement)
+
+  // If this is the topmost element in a shadow DOM, climb up one more level
+  // and then use a ":host" prefix.
+  if (!node.parentElement && !isShadowRoot(node.parentNode))
   {
-    let newSelector = ":root";
+    let newSelector = isShadowRoot(node) ? ":host" : ":root";
     if (selector)
       newSelector += " > " + selector;
     return newSelector;
   }
   let idx = positionInParent(node);
   if (idx > 0)
   {
     let newSelector = `${node.tagName}:nth-child(${idx})`;
     if (selector)
       newSelector += " > " + selector;
-    return makeSelector(node.parentElement, newSelector);
+    return makeSelector(node.parentElement || node.parentNode, newSelector);
   }
 
   return selector;
 }
 
 function parseSelectorContent(content, startIndex)
 {
   let parens = 1;
@@ -138,16 +146,68 @@
       yield* evaluate(chain, index + 1, selector, element, styles);
   }
   // Just in case the getSelectors() generator above had to run some heavy
   // document.querySelectorAll() call which didn't produce any results, make
   // sure there is at least one point where execution can pause.
   yield null;
 }
 
+function isDescendantOf(node, subtrees)
+{
+  return subtrees.some(subtree => subtree.contains(node));
+}
+
+function* extractAddedSubtrees(mutations)
+{
+  let knownSubtrees = [];
+
+  for (let mutation of mutations)
+  {
+    for (let node of mutation.addedNodes)
+    {
+      if (node instanceof Element && !isDescendantOf(node, knownSubtrees))
+      {
+        knownSubtrees.push(node);
+        yield node;
+      }
+    }
+  }
+}
+
+function* traverse(nodes)
+{
+  for (let node of nodes)
+  {
+    yield* traverse(node.children);
+    yield node;
+  }
+}
+
+function niceLoop(iterator, callback)
+{
+  let loop = () =>
+  {
+    let cycleStart = performance.now();
+
+    for (let next = iterator.next(); !next.done; next = iterator.next())
+    {
+      callback(next.value);
+
+      if (performance.now() - cycleStart > MAX_SYNCHRONOUS_PROCESSING_TIME)
+      {
+        setTimeout(loop, 0);
+        return;
+      }
+    }
+  };
+
+  loop();
+}
+
 function PlainSelector(selector)
 {
   this._selector = selector;
 }
 
 PlainSelector.prototype = {
   /**
    * Generator function returning a pair of selector
@@ -297,22 +357,38 @@
 };
 
 function isSelectorHidingOnlyPattern(pattern)
 {
   return pattern.selectors.some(s => s.preferHideWithSelector) &&
     !pattern.selectors.some(s => s.requiresHiding);
 }
 
-function ElemHideEmulation(addSelectorsFunc, hideElemsFunc)
+function ElemHideEmulation(document, root, addSelectorsFunc, hideElemsFunc,
+                           shadowAttachedEventType)
 {
   this.document = document;
+  this.root = root || document;
   this.addSelectorsFunc = addSelectorsFunc;
   this.hideElemsFunc = hideElemsFunc;
+  this.shadowAttachedEventType = shadowAttachedEventType;
+  this.patterns = [];
   this.observer = new MutationObserver(this.observe.bind(this));
+  this.shadowEmulations = new WeakMap();
+
+  if (shadowAttachedEventType)
+  {
+    this.root.addEventListener(shadowAttachedEventType,
+                               this.onShadowAttached.bind(this), true);
+  }
+
+  if (this.root == this.document)
+    this.document.addEventListener("load", this.onLoad.bind(this), true);
+  else
+    this.findShadowRoots(this.root.children);
 }
 
 ElemHideEmulation.prototype = {
   isSameOrigin(stylesheet)
   {
     try
     {
       return new URL(stylesheet.href).origin == this.document.location.origin;
@@ -404,17 +480,17 @@
 
     let elements = [];
     let elementFilters = [];
 
     let cssStyles = [];
 
     let stylesheetOnlyChange = !!stylesheets;
     if (!stylesheets)
-      stylesheets = this.document.styleSheets;
+      stylesheets = this.root.styleSheets;
 
     for (let stylesheet of stylesheets)
     {
       // Explicitly ignore third-party stylesheets to ensure consistent behavior
       // between Firefox and Chrome.
       if (!this.isSameOrigin(stylesheet))
         continue;
 
@@ -454,30 +530,30 @@
 
         if (stylesheetOnlyChange &&
             !pattern.selectors.some(selector => selector.dependsOnStyles))
         {
           pattern = null;
           return processPatterns();
         }
         generator = evaluate(pattern.selectors, 0, "",
-                             this.document, cssStyles);
+                             this.root, cssStyles);
       }
       for (let selector of generator)
       {
         if (selector != null)
         {
           if (isSelectorHidingOnlyPattern(pattern))
           {
             selectors.push(selector);
             selectorFilters.push(pattern.text);
           }
           else
           {
-            for (let element of this.document.querySelectorAll(selector))
+            for (let element of this.root.querySelectorAll(selector))
             {
               elements.push(element);
               elementFilters.push(pattern.text);
             }
           }
         }
         if (performance.now() - cycleStart > MAX_SYNCHRONOUS_PROCESSING_TIME)
         {
@@ -555,46 +631,102 @@
     {
       this._filteringInProgress = true;
       this._addSelectors(stylesheets, completion);
     }
   },
 
   onLoad(event)
   {
+    this.findShadowRoots(this.root.children);
+
+    if (this.patterns.length == 0)
+      return;
+
     let stylesheet = event.target.sheet;
     if (stylesheet)
       this.queueFiltering([stylesheet]);
   },
 
+  onShadowAttached(event)
+  {
+    event.stopImmediatePropagation();
+
+    if (this.patterns.length == 0)
+      return;
+
+    // The shadow root may not be available if it's a closed shadow root.
+    let shadowRoot = event.target.shadowRoot;
+    if (!shadowRoot)
+      return;
+
+    this.addShadowRoot(shadowRoot);
+  },
+
+  addShadowRoot(shadowRoot)
+  {
+    if (!this.shadowEmulations.has(shadowRoot))
+    {
+      let emulation = new ElemHideEmulation(this.document,
+                                            shadowRoot,
+                                            this.addSelectorsFunc,
+                                            this.hideElemsFunc,
+                                            this.shadowAttachedEventType);
+      this.shadowEmulations.set(shadowRoot, emulation);
+      emulation.apply(this.patterns, true);
+    }
+  },
+
+  findShadowRoots(nodes)
+  {
+    niceLoop(traverse(nodes), node =>
+    {
+      let shadowRoot = node.shadowRoot;
+      if (shadowRoot)
+        this.addShadowRoot(shadowRoot);
+    });
+  },
+
   observe(mutations)
   {
+    if (typeof ShadowRoot != "undefined")
+    {
+      // Find any preattached shadows.
+      this.findShadowRoots(extractAddedSubtrees(mutations));
+    }
+
     this.queueFiltering();
   },
 
-  apply(patterns)
+  apply(patterns, parsed)
   {
-    this.patterns = [];
-    for (let pattern of patterns)
+    if (parsed)
+    {
+      this.patterns = patterns;
+    }
+    else
     {
-      let selectors = this.parseSelector(pattern.selector);
-      if (selectors != null && selectors.length > 0)
-        this.patterns.push({selectors, text: pattern.text});
+      this.patterns = [];
+      for (let pattern of patterns)
+      {
+        let selectors = this.parseSelector(pattern.selector);
+        if (selectors != null && selectors.length > 0)
+          this.patterns.push({selectors, text: pattern.text});
+      }
     }
 
     if (this.patterns.length > 0)
     {
       this.queueFiltering();
       this.observer.observe(
-        this.document,
+        this.root,
         {
           childList: true,
           attributes: true,
           characterData: true,
           subtree: true
         }
       );
-      this.document.addEventListener("load", this.onLoad.bind(this), true);
     }
   }
 };
 
 exports.ElemHideEmulation = ElemHideEmulation;