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..6aec1ea1e72cfa3a0597d1b2f460fd3013ac19a9 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,297 @@ 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. We go to some lengths to avoid breaking code
+// using WebSockets, circumvention and as far as possible detection.
+// [1] - https://bugs.chromium.org/p/chromium/issues/detail?id=129353
+function wrapWebSocket()
+{
+  if (typeof WebSocket == "undefined" ||
+      typeof WeakMap == "undefined" ||
+      typeof Proxy == "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;
+
+    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}
+      }));
     }
-  }
 
-  code.push("})();");
+    // We need to store state for our wrapped WebSocket instances, that webpages
+    // can't access. We use a WeakMap to avoid leaking memory in the case that
+    // all other references to a WebSocket instance have been deleted.
+    var instanceStorage = new WeakMap();
 
-  var script = document.createElement("script");
-  script.async = false;
-  script.textContent = code.join("\n");
-  document.documentElement.appendChild(script);
-  document.documentElement.removeChild(script);
+    var eventNames = ["close", "open", "message", "error"];
+    var eventAttrNames = ["onclose", "onopen", "onmessage", "onerror"];
+
+    function addRemoveEventListener(storage, key, type, listener)

[kzar, 2016/07/12 11:34:32]

Should we care about the useCapture parameter for addEventListener? So far I
just ignore it.

+    {
+      if (typeof listener == "object")
+        listener = listener.handleEvent;
+
+      if (!(eventNames.indexOf(type) > -1 && typeof listener == "function"))
+        return;
+
+      var listeners = storage.listeners[type];
+      var listenerIndex = listeners.indexOf(listener);
+
+      if (key == "addEventListener")
+      {
+        if (listenerIndex == -1)
+          listeners.push(listener);
+      }
+      else if (listenerIndex > -1)
+        listeners.splice(listenerIndex, 1);
+    }
+
+    // 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 result, we create the WebSocket (if allowed) and perform
+    // the queued actions.
+    function processQueue(storage)

[kzar, 2016/07/12 11:34:33]

Queuing up assignments and method calls seems rather pointless now that we
manage events ourselves. In my testing I couldn't find any website which use
this, it's usual to wait for a connection to be opened before attempting to send
for example. I'd like to remove storage.queue and all related logic, any
opinions?

+    {
+      for (var i = 0; i < storage.queue.length; i += 1)
+      {
+        var action = storage.queue[i][0];
+        var key = storage.queue[i][1];
+        var value = storage.queue[i][2];
+
+        if (action == "set")
+          storage.websocket[key] = value;
+        else if (action == "call")
+          storage.websocket[key].apply(storage.websocket, value);
+      }
+    }
+
+    var defaults = {
+      readyState: RealWebSocket.CONNECTING,
+      bufferedAmount: 0,
+      extensions: "",
+      binaryType: "blob"
+    };
+
+    // We cannot dispatch WebSocket events directly to their listeners since
+    // event.target would give a way back to the original WebSocket constructor.
+    // Instead we must listen for events ourselves and pass them on, taking care
+    // to spoof the event target and isTrusted flag.
+    function wrappedEventListener(name, me)
+    {
+      var storage = instanceStorage.get(me);
+      return function(event)
+      {
+        var eventProxy = new Proxy(event, {
+          get: function(target, key)
+          {
+            if (key == "isTrusted" && "isTrusted" in target)
+              return true;
+            if (key == "target" || key == "srcElement" || key == "currentTarget")
+              return me;
+            return target[key];
+          }
+        });
+
+        var listeners = storage.listeners[name];
+        for (var i = 0; i < listeners.length; i += 1)
+          listeners[i].call(me, eventProxy);
+        var listener = storage.listeners["on" + name];
+        if (typeof listener == "function")
+          listener.call(me, eventProxy);
+      };
+    }
+
+    WebSocket = function(url, protocols)

[kzar, 2016/07/12 11:34:33]

So far I don't intercept WebSocket.toString() so it could be used to detect us.
I'm inclined to wrap WebSocket in a Proxy to fix that. On the other hand this
code is already too complex for my taste. Opinions?

[lainverse, 2016/07/12 12:22:20]

If we going all the way to avoid detection then we have to. However, make sure
to handle chained toString calls and calls on functions, like
ws.toString.toString(), ws.close.toString(). Also, there is toLocaleString
function with the same output (not sure if some specificl locale parameters
provided, but should be the same).

+    {
+      var me = this;
+      var storage = {
+        url: url,
+        protocol: protocols || "",
+        queue: [],
+        websocket: null,
+        blocked: false,
+        listeners: {
+          onclose: null,
+          onopen: null,
+          onmessage: null,
+          onerror: null,
+          close: [],
+          open: [],
+          message: [],
+          error: []
+        }
+      };
+      instanceStorage.set(me, storage);
+
+      checkRequest(url, protocols, function(blocked)
+      {
+        if (blocked)
+        {
+          storage.blocked = true;
+          wrappedEventListener("error", me)(new Error("error"));
+        }
+        else
+        {
+          storage.websocket = new RealWebSocket(url, protocols);
+          for (var i = 0; i < eventNames.length; i += 1)

[kzar, 2016/07/12 11:34:32]

If a website creates a WebSocket (which manages to connect), doesn't add a close
event listener and then deletes all references to it the WebSocket should be
garbage collected.

As we always add a close event listener here the (actual, not wrapper of)
WebSocket wouldn't be garbage collected until the connection was closed.
Personally I don't think this is a big deal, and certainly not worth the extra
logic to fix. Any opinions?

+            storage.websocket.addEventListener(
+              eventNames[i],
+              wrappedEventListener(eventNames[i], me)
+            );
+          processQueue(storage);
+        }
+        delete storage.queue;
+      });
+    };
+    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 = new Proxy(RealWebSocket.prototype,{
+      set: function(target, key, value, me)
+      {
+        var storage = instanceStorage.get(me);
+
+        if (!storage)
+          target[key] = value;
+        else if (eventAttrNames.indexOf(key) > -1)
+          storage.listeners[key] = value;
+        else if (storage.websocket)
+          storage.websocket[key] = value;
+        else if (!storage.blocked)
+          storage.queue.push(["set", key, value]);
+
+        return true;
+      },
+      get: function(target, key, me)
+      {
+        if (key == "__proto__" && me != WebSocket.prototype)
+          return WebSocket.prototype;
+
+        if (key == "constructor")
+          return WebSocket;
+
+        var storage = instanceStorage.get(me);
+        if (!storage)
+          return target[key];
+
+        if (key == "addEventListener" || key == "removeEventListener")
+          return function()
+          {
+            if (arguments.length > 1)
+              addRemoveEventListener(storage, key, arguments[0], arguments[1]);
+          };
+
+        if (eventAttrNames.indexOf(key) > -1)
+          return storage.listeners[key];
+
+        var desc = Object.getOwnPropertyDescriptor(target, key);
+        if (desc && typeof desc.value == "function")
+          return function()
+          {
+            if (storage.websocket)
+              storage.websocket[key].apply(storage.websocket, arguments);
+            else if (!storage.blocked)
+              storage.queue.push(["call", key, arguments]);
+          };
+
+        if (storage.websocket)
+          return storage.websocket[key];
+        if (key == "url" || key == "protocol")
+          return storage[key];
+        if (storage.blocked && key == "readyState")
+          return WebSocket.CLOSED;
+        if (key in defaults)
+          return defaults[key];
+        return undefined;
+      }
+    });
+  }
+
+  injectJS(wrapper, eventName);
 }
 
 function init(document)
@@ -401,6 +650,8 @@ function init(document)
   var observer = null;
   var tracer = null;
 
+  wrapWebSocket();
+
   function getPropertyFilters(callback)
   {
     ext.backgroundPage.sendMessage({