Issue 1727 - Prevent circumvention via WebSocket

include.preload.js

Index: include.preload.js
diff --git a/include.preload.js b/include.preload.js
index 949d039663f9130843b4d8d0a3ddfcb434522066..598155518ec78036f4c8a163446a6117612a5801 100644
--- a/include.preload.js
+++ b/include.preload.js
@@ -17,6 +17,7 @@
 
 var MutationObserver = window.MutationObserver || window.WebKitMutationObserver;
 var SELECTOR_GROUP_SIZE = 200;
+var id = Math.random().toString(36).substr(2);
 
 var typeMap = {
   "img": "IMAGE",
@@ -349,49 +350,225 @@ function reinjectStyleSheetWhenRemoved(document, style)
   return observer;
 }
 
+function injectJS(f)
+{
+  var args = JSON.stringify(Array.prototype.slice.call(arguments, 1));
+  args = args.substring(1, args.length - 1);
+  var codeString = "(" + f.toString() + ")(" + args + ");";
+
+  var script = document.createElement("script");
+  script.async = false;
+  script.textContent = codeString;
+  document.documentElement.appendChild(script);
+  document.documentElement.removeChild(script);
+}
+
 function protectStyleSheet(document, style)
 {
-  var id = Math.random().toString(36).substr(2)
   style.id = id;
 
-  var code = [
-    "(function()",
-    "{",
-    '  var style = document.getElementById("' + id + '") ||',
-    '              document.documentElement.shadowRoot.getElementById("' + id + '");',
-    '  style.removeAttribute("id");'
-  ];
+  var protector = function(id)
+  {
+    var style = document.getElementById(id) ||
+                document.documentElement.shadowRoot.getElementById(id);
+    style.removeAttribute("id");
+
+    var i;
+    var disableables = [style, style.sheet];
+    for (i = 0; i < disableables.length; i += 1)
+      Object.defineProperty(disableables[i], "disabled",
+                            {value: false, enumerable: true});
+
+    var methods = ["deleteRule", "removeRule"];
+    for (i = 0; i < methods.length; i += 1)
+    {
+      if (methods[i] in CSSStyleSheet.prototype)
+      {
+        (function(method)
+        {
+          var original = CSSStyleSheet.prototype[method];
+          CSSStyleSheet.prototype[method] = function(index)
+          {
+            if (this != style.sheet)
+              original.call(this, index);
+          };
+        }(methods[i]));
+      }
+    }
+  };
 
-  var disableables = ["style", "style.sheet"];
-  for (var i = 0; i < disableables.length; i++)
+  injectJS(protector, id);
+}
+
+// Neither Chrome[1] nor Safari allow us to intercept WebSockets, and therefore
+// some ad networks are misusing them as a way to serve adverts and circumvent
+// us. As a workaround we wrap WebSocket, closing connections that would have
+// otherwise been blocked.
+// [1] - https://bugs.chromium.org/p/chromium/issues/detail?id=129353
+function wrapWebSocket()
+{
+  if (typeof WebSocket == "undefined" || typeof WeakMap == "undefined")
+    return;
+
+  var eventName = "abpws-" + id;
+
+  document.addEventListener(eventName, function(event)
   {
-    code.push("  Object.defineProperty(" + disableables[i] + ', "disabled", '
-                                         + "{value: false, enumerable: true});");
-  }
+    ext.backgroundPage.sendMessage({
+      type: "websocket-request",
+      url: event.detail.url
+    }, function (block)
+    {
+      document.dispatchEvent(
+        new CustomEvent(eventName + "-" + event.detail.url, {detail: block})
+      );
+    });
+  });
 
-  var methods = ["deleteRule", "removeRule"];
-  for (var j = 0; j < methods.length; j++)
+  function wrapper(eventName)
   {
-    var method = methods[j];
-    if (method in CSSStyleSheet.prototype)
+    var RealWebSocket = WebSocket;
+
+    function checkRequest(url, protocols, callback)
     {
-      var origin = "CSSStyleSheet.prototype." + method;
-      code.push("  var " + method + " = " + origin + ";",
-                "  " + origin + " = function(index)",
-                "  {",
-                "    if (this != style.sheet)",
-                "      " + method + ".call(this, index);",
-                "  }");
+      var incomingEventName = eventName + "-" + url;
+      function listener(event)
+      {
+        callback(event.detail);
+        document.removeEventListener(incomingEventName, listener);
+      }
+      document.addEventListener(incomingEventName, listener);
+
+      document.dispatchEvent(new CustomEvent(eventName, {
+        detail: {url: url, protocols: protocols}
+      }));
+    }
+
+    // To avoid detection we make our WebSocket wrapper as similar as possible.
+    // Most WebSocket properties and methods actually belong to
+    // WebSocket.prototype, but they require access to the instance state. To
+    // acheive this without adding extra properties we need to manage our own
+    // storage here. To avoid leaking memory in the case that all other
+    // references to a WebSocket have been deleted we must use a WeakMap.
+    var instanceStorage = new WeakMap();
+
+    // We check if a WebSocket should be blocked before actually creating it. As
+    // this is done asynchonously we must queue up any actions (method calls and
+    // assignments) that happen in the mean time.
+    // Once we have a block/allow result we create the WebSocket and perform the
+    // queued actions. (For blocked requests we only need to manage error and
+    // close listeners, other actions are all ignored. This is so real close and
+    // error events are fired, with the minimal risk of leaking messages.)
+    function processQueue(queue, websocket, blocked)
+    {
+      for (var i = 0; i < queue.length; i += 1)
+      {
+        var action = queue[i][0];
+        var key = queue[i][1];
+        var value = queue[i][2];
+
+        if (action == "set" &&
+            (!blocked || key == "onclose" || key == "onerror"))
+          websocket[key] = value;
+        else if (action == "call" &&
+                 (!blocked ||
+                  ((key == "addEventListener" || key == "removeEventListener") &&
+                   (value[0] == "error" || value[0] == "close"))))
+          websocket[key].apply(websocket, value);
+      }
     }
-  }
 
-  code.push("})();");
+    var defaults = {
+      readyState: RealWebSocket.CONNECTING,
+      bufferedAmount: 0,
+      extensions: "",
+      binaryType: "blob"
+    };
 
-  var script = document.createElement("script");
-  script.async = false;
-  script.textContent = code.join("\n");
-  document.documentElement.appendChild(script);
-  document.documentElement.removeChild(script);
+    WebSocket = function(url, protocols)
+    {
+      var storage = {
+        url: url,
+        protocol: protocols,
+        queue: [],
+        websocket: null
+      };
+      instanceStorage.set(this, storage);
+
+      checkRequest(url, protocols, function(blocked)
+      {
+        storage.websocket = new RealWebSocket(url, protocols);
+        processQueue(storage.queue, storage.websocket, blocked);
+        if (blocked)
+          storage.websocket.close();
+        delete storage.queue;
+      });
+    };
+
+    function proxyProperties(original)
+    {
+      var properties = Object.create(null);
+
+      Object.keys(original).map(function(key)
+      {
+        if (key == "prototype")
+          return;
+
+        var descriptor = Object.getOwnPropertyDescriptor(original, key);
+
+        if (typeof descriptor.value == "function")
+        {
+          descriptor.value = function()
+          {
+            var storage = instanceStorage.get(this);
+            if (!storage)
+              original[key].apply(original, arguments);
+            else if (storage.websocket)
+              storage.websocket[key].apply(storage.websocket, arguments);
+            else
+              storage.queue.push(["call", key, arguments]);
+          };
+        }
+        else if (typeof descriptor.value == "undefined")
+        {
+          descriptor.get = function()
+          {
+            var storage = instanceStorage.get(this);
+            if (!storage)
+              return original[key];
+            if (storage.websocket)
+              return storage.websocket[key];
+            if (key == "url" || key == "protocol")
+              return storage[key];
+            if (key in defaults)
+              return defaults[key];
+            return null;
+          };
+          descriptor.set = function(value)
+          {
+            var storage = instanceStorage.get(this);
+            if (!storage)
+              original[key] = value;
+            else if (storage.websocket)
+              storage.websocket[key] = value;
+            else
+              storage.queue.push(["set", key, value]);
+            return value;
+          };
+        }
+        properties[key] = descriptor;
+      });
+      return properties;
+    }
+
+    Object.defineProperties(WebSocket, proxyProperties(RealWebSocket));
+    WebSocket.prototype = Object.create(
+      RealWebSocket.prototype.__proto__,
+      proxyProperties(RealWebSocket.prototype)
+    );
+  }
+
+  injectJS(wrapper, eventName);
 }
 
 function init(document)
@@ -401,6 +578,8 @@ function init(document)
   var observer = null;
   var tracer = null;
 
+  wrapWebSocket();
+
   function getPropertyFilters(callback)
   {
     ext.backgroundPage.sendMessage({