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..dba04846fd04654e841a6b901e23928a642348b5 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,221 @@ 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]));
+      }
+    }
+  };
+
+  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;
 
-  var disableables = ["style", "style.sheet"];
-  for (var i = 0; i < disableables.length; i++)
+  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;
+
+    // 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 can't create a WebSocket until we know it shouldn't be blocked, but we
+    // can only check this asynchronously. We queue up any actions that are
+    // performed on a WebSocket wrapper instance before its real WebSocket
+    // exists. If not blocked we perform queued actions on the WebSocket after
+    // it has been created.
+    function processQueue(queue, websocket)
     {
-      var origin = "CSSStyleSheet.prototype." + method;
-      code.push("  var " + method + " = " + origin + ";",
-                "  " + origin + " = function(index)",
-                "  {",
-                "    if (this != style.sheet)",
-                "      " + method + ".call(this, index);",
-                "  }");
+      for (var i = 0; i < queue.length; i += 1)
+      {
+        var action = queue[i][0];
+        var key = queue[i][1];
+        var value = queue[i][2];
+        switch (action)
+        {
+          case "set":
+            websocket[key] = value;
+            break;
+          case "call":
+            websocket[key].apply(websocket, value);
+            break;
+        }
+      }
     }
-  }
 
-  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, protocol)
+    {
+      var storage = {
+        url: url,
+        protocol: protocol,
+        queue: [],
+        blocked: false,
+        websocket: null
+      };
+      instanceStorage.set(this, storage);
+
+      var incomingEventName = eventName + "-" + url;
+      function listener(event)
+      {
+        storage.blocked = event.detail;
+        if (!storage.blocked)
+        {
+          storage.websocket = new RealWebSocket(url, protocol);
+          processQueue(storage.queue, storage.websocket);
+        }
+        delete storage.queue; // FIXME onerror!?

[kzar, 2016/07/06 16:39:20]

So far I don't fire the error / close event when a WebSocket connection is
blocked.

I could do that here but I'm not sure how to reproduce the event properly, for
example isTrusted! Closest I can think to do is to create a new WebSocket with a
bogus URL, providing their onerror listener. Of course they could then check the
URL...

[kzar, 2016/07/07 07:51:31]

Note: My earlier implementation which always created a WebSocket but then called
.close() if blocked didn't have this limitation. A real close event was fired.

If you really insist that we don't just do that again we could still do
something similar. We could create a WebSocket when the block/allow result comes
back, but then immediately close it if blocked. That would be even less likely
to leak connections than the earlier implementation. (And I never witnessed the
earlier implementation leaking connections anyway.)

+
+        document.removeEventListener(incomingEventName, listener);
+      }
+      document.addEventListener(incomingEventName, listener);
+
+      document.dispatchEvent(new CustomEvent(eventName, {
+        detail: {url: url, protocol: protocol}
+      }));
+    };
+
+    function proxyProperties(original)
+    {
+      var properties = Object.create(null);
+
+      Object.keys(original).map(function(key)
+      {
+        if (key == "prototype")

[kzar, 2016/07/06 16:39:20]

(For Safari)

+          return;
+
+        var descriptor = Object.getOwnPropertyDescriptor(original, key);
+
+        if (typeof descriptor.value == "function")
+        {
+          descriptor.value = function()
+          {
+            var storage = instanceStorage.get(this);
+            if (!storage)

[kzar, 2016/07/06 16:39:20]

So that the "Uncaught TypeError: Illegal invocation(…)" exception is thrown when
doing something like `WebSocket.prototype.send(...)`. (Similar below too.)

+              original[key].apply(original, arguments);
+            if (storage.websocket)
+              storage.websocket[key].apply(storage.websocket, arguments);
+            else if (!storage.blocked)
+              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 (storage.blocked && key == "readyState")
+              return RealWebSocket.CLOSED;
+            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 if (!storage.blocked)
+              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 +574,8 @@ function init(document)
   var observer = null;
   var tracer = null;
 
+  wrapWebSocket();
+
   function getPropertyFilters(callback)
   {
     ext.backgroundPage.sendMessage({