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..baf04bdbc55fb544f70480d8b722f0a61622680b 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,133 @@ 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, preventing blocked WebSocket
+// connections from being opened.
+// [1] - https://bugs.chromium.org/p/chromium/issues/detail?id=129353
+function wrapWebSocket()
+{
+  if (typeof WebSocket == "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)
+    // As far as possible we must track everything we use that could be
+    // sabotaged by the website later in order to circumvent us.
+    var RealWebSocket = WebSocket;
+    var closeWebSocket = RealWebSocket.prototype.close;
+    var document = window.document;
+    var addEventListener = document.addEventListener;
+    var removeEventListener = document.removeEventListener;
+    var dispatchEvent = document.dispatchEvent;
+    var CustomEvent = window.CustomEvent;
+    var bind = Function.prototype.bind;
+    var call = Function.prototype.call;
+
+    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);
+        call.bind(removeEventListener)(document, incomingEventName, listener);
+      }
+      call.bind(addEventListener)(document, incomingEventName, listener);
+
+      call.bind(dispatchEvent)(document, new CustomEvent(eventName, {
+        detail: {url: url, protocols: protocols}
+      }));
     }
-  }
 
-  code.push("})();");
+    WebSocket = function(url, protocols)
+    {
+      var websocket = new RealWebSocket(url, protocols);
 
-  var script = document.createElement("script");
-  script.async = false;
-  script.textContent = code.join("\n");
-  document.documentElement.appendChild(script);
-  document.documentElement.removeChild(script);
+      checkRequest(url, protocols, function(blocked)
+      {
+        if (blocked)
+          call.bind(closeWebSocket)(websocket);
+      });
+
+      return websocket;
+    };
+    Object.defineProperties(WebSocket, {
+      CONNECTING: {value: 0, enumerable: true},
+      OPEN: {value: 1, enumerable: true},
+      CLOSING: {value: 2, enumerable: true},
+      CLOSED: {value: 3, enumerable: true}
+    });
+    WebSocket.prototype = RealWebSocket.prototype;
+    RealWebSocket.prototype.constructor = WebSocket;
+  }
+
+  injectJS(wrapper, eventName);
 }
 
 function init(document)
@@ -401,6 +486,8 @@ function init(document)
   var observer = null;
   var tracer = null;
 
+  wrapWebSocket();
+
   function getPropertyFilters(callback)
   {
     ext.backgroundPage.sendMessage({