add SockJS to the existing WebSocket model

javascript/ql/src/semmle/javascript/frameworks/WebSocket.qll

@@ -1,5 +1,5 @@
 /**
- * Provides classes for working with [WebSocket](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket) and [ws](https://github.com/websockets/ws).
+ * Provides classes for working with [WebSocket](https://developer.mozilla.org/en-US/docs/Web/API/WebSocket), [ws](https://github.com/websockets/ws), and [SockJS](http://sockjs.org).
  *
  * The model is based on the EventEmitter model, and there is therefore a
  * data-flow step from where a WebSocket event is sent to where the message
@@ -18,26 +18,58 @@ import javascript
  */
 private string channelName() { result = "message" }
 
+/**
+ * The names of the libraries modelled in this file.
+ */
+private module LibraryNames {
+  string sockjs() { result = "SockJS" }
+
+  string websocket() { result = "WebSocket" }
+
+  string ws() { result = "ws" }
+}
+
+/**
+ * Holds if the websocket library named `client` can send a message to the library named `server`.
+ * Both `client` and `server` are library names defined in `LibraryNames`.
+ */
+private predicate areLibrariesCompatible(string client, string server) {
+  // sockjs is a WebSocket emulating library, but not actually an implementation of WebSockets.
+  client = LibraryNames::sockjs() and server = LibraryNames::sockjs()
+  or
+  server = LibraryNames::ws() and
+  (client = LibraryNames::ws() or client = LibraryNames::websocket())
+}
+
 /**
  * Provides classes that model WebSockets clients.
  */
 module ClientWebSocket {
+  private import LibraryNames
+
   /**
    * A class that can be used to instantiate a WebSocket instance.
    */
   class SocketClass extends DataFlow::SourceNode {
-    boolean isNode;
+    string library; // the name of the WebSocket library. Can be one of the libraries defined in `LibraryNames`.
 
     SocketClass() {
-      this = DataFlow::globalVarRef("WebSocket") and isNode = false
+      this = DataFlow::globalVarRef("WebSocket") and library = websocket()
       or
-      this = DataFlow::moduleImport("ws") and isNode = true
+      this = DataFlow::moduleImport("ws") and library = ws()
+      or
+      // the sockjs-client library:https://www.npmjs.com/package/sockjs-client
+      library = sockjs() and
+      (
+        this = DataFlow::moduleImport("sockjs-client") or
+        this = DataFlow::globalVarRef("SockJS")
+      )
     }
 
     /**
-     * Holds if this class is an import of the "ws" module.
+     * Gets the WebSocket library name.
      */
-    predicate isNode() { isNode = true }
+    string getLibrary() { result = library }
   }
 
   /**
@@ -49,11 +81,9 @@ module ClientWebSocket {
     ClientSocket() { this = socketClass.getAnInstantiation() }
 
     /**
-     * Holds if this ClientSocket is created from the "ws" module.
-     *
-     * The predicate is used to differentiate where the behavior of the "ws" module differs from the native WebSocket in browsers.
+     * Gets the WebSocket library name.
      */
-    predicate isNode() { socketClass.isNode() }
+    string getLibrary() { result = socketClass.getLibrary() }
   }
 
   /**
@@ -68,7 +98,10 @@ module ClientWebSocket {
 
     override DataFlow::Node getSentItem(int i) { i = 0 and result = this.getArgument(0) }
 
-    override ServerWebSocket::ReceiveNode getAReceiver() { any() }
+    override ServerWebSocket::ReceiveNode getAReceiver() {
+      areLibrariesCompatible(emitter.getLibrary(),
+        result.getEmitter().(ServerWebSocket::ServerSocket).getLibrary())
+    }
   }
 
   /**
@@ -116,7 +149,7 @@ module ClientWebSocket {
    */
   private class WSReceiveNode extends ClientWebSocket::ReceiveNode {
     WSReceiveNode() {
-      emitter.isNode() and
+      emitter.getLibrary() = ws() and
       this = getAMessageHandler(emitter, EventEmitter::on())
     }
 
@@ -128,21 +161,38 @@ module ClientWebSocket {
  * Provides classes that model WebSocket servers.
  */
 module ServerWebSocket {
+  private import LibraryNames
+
+  /**
+   * Gets a server created by a library named `library`.
+   */
+  DataFlow::SourceNode getAServer(string library) {
+    library = ws() and
+    result = DataFlow::moduleImport("ws").getAConstructorInvocation("Server")
+    or
+    library = sockjs() and
+    result = DataFlow::moduleImport("sockjs").getAMemberCall("createServer")
+  }
+
   /**
    * A server WebSocket instance.
    */
   class ServerSocket extends EventEmitter::Range, DataFlow::SourceNode {
+    string library;
+
     ServerSocket() {
       exists(DataFlow::CallNode onCall |
-        onCall =
-          DataFlow::moduleImport("ws")
-              .getAConstructorInvocation("Server")
-              .getAMemberCall(EventEmitter::on()) and
+        onCall = getAServer(library).getAMemberCall(EventEmitter::on()) and
         onCall.getArgument(0).mayHaveStringValue("connection")
       |
         this = onCall.getCallback(1).getParameter(0)
       )
     }
+
+    /**
+     * Gets the name of the library that created this server socket.
+     */
+    string getLibrary() { result = library }
   }
 
   /**
@@ -151,7 +201,13 @@ module ServerWebSocket {
   class SendNode extends EventDispatch::Range, DataFlow::CallNode {
     override ServerSocket emitter;
 
-    SendNode() { this = emitter.getAMemberCall("send") }
+    SendNode() {
+      emitter.getLibrary() = ws() and
+      this = emitter.getAMemberCall("send")
+      or
+      emitter.getLibrary() = sockjs() and
+      this = emitter.getAMemberCall("write")
+    }
 
     override string getChannel() { result = channelName() }
 
@@ -160,7 +216,10 @@ module ServerWebSocket {
       result = getArgument(0)
     }
 
-    override ClientWebSocket::ReceiveNode getAReceiver() { any() }
+    override ClientWebSocket::ReceiveNode getAReceiver() {
+      areLibrariesCompatible(result.getEmitter().(ClientWebSocket::ClientSocket).getLibrary(),
+        emitter.getLibrary())
+    }
   }
 
   /**
@@ -170,8 +229,14 @@ module ServerWebSocket {
     override ServerSocket emitter;
 
     ReceiveNode() {
-      this = emitter.getAMemberCall(EventEmitter::on()) and
-      this.getArgument(0).mayHaveStringValue("message")
+      exists(string eventName |
+        emitter.getLibrary() = ws() and eventName = "message"
+        or
+        emitter.getLibrary() = sockjs() and eventName = "data"
+      |
+        this = emitter.getAMemberCall(EventEmitter::on()) and
+        this.getArgument(0).mayHaveStringValue(eventName)
+      )
     }
 
     override string getChannel() { result = channelName() }

javascript/ql/test/library-tests/frameworks/WebSocket/browser.js

@@ -12,4 +12,21 @@
 	socket.onmessage = function (event) {
 		console.log("Message from server 2", event.data)
 	};
-})();
+})();
+
+
+(function () {
+	var sock = new SockJS('http://0.0.0.0:9999/echo');
+	sock.onopen = function () {
+		sock.send('test');
+	};
+	
+	sock.onmessage = function (e) {
+		console.log('message', e.data);
+		sock.close();
+	};
+	
+	sock.addEventListener('message', function (event) {
+		console.log('Using addEventListener ', event.data);
+	});
+})

javascript/ql/test/library-tests/frameworks/WebSocket/test.expected

@@ -1,22 +1,32 @@
 clientSocket
 | browser.js:2:17:2:52 | new Web ... :8080') |
+| browser.js:19:13:19:50 | new Soc ... /echo') |
 | client.js:4:13:4:45 | new Web ... e.org') |
 clientSend
 | browser.js:5:3:5:33 | socket. ... wser!') |
+| browser.js:21:3:21:19 | sock.send('test') |
 | client.js:7:3:7:28 | ws.send ... ient!') |
 clientReceive
 | browser.js:8:37:10:2 | functio ... ta);\\n\\t} |
 | browser.js:12:21:14:2 | functio ... ata)\\n\\t} |
+| browser.js:24:19:27:2 | functio ... e();\\n\\t} |
+| browser.js:29:35:31:2 | functio ... ta);\\n\\t} |
 | client.js:10:19:12:2 | functio ... ta);\\n\\t} |
 serverSocket
 | server.js:6:43:6:44 | ws |
+| sockjs.js:8:40:8:43 | conn |
 serverSend
 | server.js:11:3:11:28 | ws.send ... rver!') |
+| sockjs.js:11:9:11:51 | conn.wr ... test))) |
 serverReceive
 | server.js:7:3:9:4 | ws.on(' ... );\\n\\t\\t}) |
+| sockjs.js:9:5:12:6 | conn.on ... \\n    }) |
 taintStep
 | browser.js:5:15:5:32 | 'Hi from browser!' | server.js:7:38:7:44 | message |
+| browser.js:21:13:21:18 | 'test' | sockjs.js:9:31:9:37 | message |
 | client.js:7:11:7:27 | 'Hi from client!' | server.js:7:38:7:44 | message |
 | server.js:11:11:11:27 | 'Hi from server!' | browser.js:9:39:9:48 | event.data |
 | server.js:11:11:11:27 | 'Hi from server!' | browser.js:13:40:13:49 | event.data |
 | server.js:11:11:11:27 | 'Hi from server!' | client.js:10:37:10:40 | data |
+| sockjs.js:11:20:11:50 | JSON.st ... .test)) | browser.js:25:26:25:31 | e.data |
+| sockjs.js:11:20:11:50 | JSON.st ... .test)) | browser.js:30:42:30:51 | event.data |