JS: Refactor WebSocket to use API graphs

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

@@ -47,6 +47,20 @@ private predicate areLibrariesCompatible(
   (client = LibraryNames::ws() or client = LibraryNames::websocket())
 }
 
+/** Treats `WebSocket` as an entry point for API graphs. */
+private class WebSocketEntryPoint extends API::EntryPoint {
+  WebSocketEntryPoint() { this = "global.WebSocket" }
+
+  override DataFlow::SourceNode getASource() { result = DataFlow::globalVarRef("WebSocket") }
+}
+
+/** Treats `SockJS` as an entry point for API graphs. */
+private class SockJSEntryPoint extends API::EntryPoint {
+  SockJSEntryPoint() { this = "global.SockJS" }
+
+  override DataFlow::SourceNode getASource() { result = DataFlow::globalVarRef("SockJS") }
+}
+
 /**
  * Provides classes that model WebSockets clients.
  */
@@ -56,19 +70,19 @@ module ClientWebSocket {
   /**
    * A class that can be used to instantiate a WebSocket instance.
    */
-  class SocketClass extends DataFlow::SourceNode {
+  class SocketClass extends API::Node {
     LibraryName library; // the name of the WebSocket library. Can be one of the libraries defined in `LibraryNames`.
 
     SocketClass() {
-      this = DataFlow::globalVarRef("WebSocket") and library = websocket()
+      this = any(WebSocketEntryPoint e).getANode() and library = websocket()
       or
-      this = DataFlow::moduleImport("ws") and library = ws()
+      this = API::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")
+        this = API::moduleImport("sockjs-client") or
+        this = any(SockJSEntryPoint e).getANode()
       )
     }
 
@@ -81,10 +95,10 @@ module ClientWebSocket {
   /**
    * A client WebSocket instance.
    */
-  class ClientSocket extends EventEmitter::Range, DataFlow::NewNode, ClientRequest::Range {
+  class ClientSocket extends EventEmitter::Range, API::NewNode, ClientRequest::Range {
     SocketClass socketClass;
 
-    ClientSocket() { this = socketClass.getAnInstantiation() }
+    ClientSocket() { this = socketClass.getAnInvocation() }
 
     /**
      * Gets the WebSocket library name.
@@ -115,10 +129,10 @@ module ClientWebSocket {
   /**
    * A message sent from a WebSocket client.
    */
-  class SendNode extends EventDispatch::Range, DataFlow::CallNode {
+  class SendNode extends EventDispatch::Range, API::CallNode {
     override ClientSocket emitter;
 
-    SendNode() { this = emitter.getAMemberCall("send") }
+    SendNode() { this = emitter.getReturn().getMember("send").getACall() }
 
     override string getChannel() { result = channelName() }
 
@@ -145,8 +159,8 @@ module ClientWebSocket {
   private DataFlow::FunctionNode getAMessageHandler(
     ClientWebSocket::ClientSocket emitter, string methodName
   ) {
-    exists(DataFlow::CallNode call |
-      call = emitter.getAMemberCall(methodName) and
+    exists(API::CallNode call |
+      call = emitter.getReturn().getMember(methodName).getACall() and
       call.getArgument(0).mayHaveStringValue("message") and
       result = call.getCallback(1)
     )
@@ -161,7 +175,13 @@ module ClientWebSocket {
     WebSocketReceiveNode() {
       this = getAMessageHandler(emitter, "addEventListener")
       or
-      this = emitter.getAPropertyWrite("onmessage").getRhs()
+      this = emitter.getReturn().getMember("onmessage").getAValueReachingSink()
+      or
+      exists(DataFlow::MethodCallNode bindCall |
+        bindCall = emitter.getReturn().getMember("onmessage").getAValueReachingSink() and
+        bindCall.getMethodName() = "bind" and
+        this = bindCall.getReceiver().getAFunctionValue()
+      )
     }
 
     override DataFlow::Node getReceivedItem(int i) {
@@ -192,19 +212,19 @@ module ServerWebSocket {
   /**
    * Gets a server created by a library named `library`.
    */
-  DataFlow::SourceNode getAServer(LibraryName library) {
+  API::InvokeNode getAServer(LibraryName library) {
     library = ws() and
-    result = DataFlow::moduleImport("ws").getAConstructorInvocation("Server")
+    result = API::moduleImport("ws").getMember("Server").getAnInvocation()
     or
     library = sockjs() and
-    result = DataFlow::moduleImport("sockjs").getAMemberCall("createServer")
+    result = API::moduleImport("sockjs").getMember("createServer").getAnInvocation()
   }
 
   /**
    * Gets a `socket.on("connection", (msg, req) => {})` call.
    */
   private DataFlow::CallNode getAConnectionCall(LibraryName library) {
-    result = getAServer(library).getAMemberCall(EventEmitter::on()) and
+    result = getAServer(library).getReturn().getMember(EventEmitter::on()).getACall() and
     result.getArgument(0).mayHaveStringValue("connection")
   }
 
@@ -324,15 +344,18 @@ module ServerWebSocket {
       result = this.getCallback(1).getParameter(0)
     }
   }
+}
 
-  /**
-   * A data flow node representing data received from a client, viewed as remote user input.
-   */
-  private class ReceivedItemAsRemoteFlow extends RemoteFlowSource {
-    ReceivedItemAsRemoteFlow() { this = any(ReceiveNode rercv).getReceivedItem(_) }
+/**
+ * A data flow node representing data received from a client or server, viewed as remote user input.
+ */
+private class ReceivedItemAsRemoteFlow extends RemoteFlowSource {
+  ReceivedItemAsRemoteFlow() {
+    this = any(ClientWebSocket::ReceiveNode rercv).getReceivedItem(_) or
+    this = any(ServerWebSocket::ReceiveNode rercv).getReceivedItem(_)
+  }
 
-    override string getSourceType() { result = "WebSocket client data" }
+  override string getSourceType() { result = "WebSocket transmitted data" }
 
-    override predicate isUserControlledObject() { any() }
-  }
+  override predicate isUserControlledObject() { any() }
 }

javascript/ql/src/Security/trest/test.ql

@@ -0,0 +1,10 @@
+import javascript
+
+API::NewNode getAWebSocketInstance() { result instanceof ClientWebSocket::ClientSocket }
+
+from DataFlow::Node handler
+where
+  handler = getAWebSocketInstance().getReturn().getMember("onmessage").asSource()
+  or
+  handler = getAWebSocketInstance().getAPropertyWrite("onmessage").getRhs()
+select handler, "This is a WebSocket onmessage handler."

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

@@ -0,0 +1,74 @@
+import { MyWebSocket, MySockJS, myWebSocketInstance, mySockJSInstance } from './browser.js';
+
+(function () {
+	const socket = new MyWebSocket('ws://localhost:9080'); // $ clientSocket
+
+	socket.addEventListener('open', function (event) {
+		socket.send('Hi from browser!'); // $ clientSend
+	});
+
+	socket.addEventListener('message', function (event) {
+		console.log('Message from server ', event.data); // $ remoteFlow
+	}); // $ clientReceive
+
+	socket.onmessage = function (event) {
+		console.log("Message from server 2", event.data); // $ remoteFlow
+	}; // $ clientReceive
+})();
+
+
+(function () {
+	var sock = new MySockJS('http://0.0.0.0:9999/echo'); // $ clientSocket
+	sock.onopen = function () {
+		sock.send('test'); // $ clientSend
+	};
+
+	sock.onmessage = function (e) {
+		console.log('message', e.data); // $ remoteFlow
+		sock.close();
+	}; // $ clientReceive
+
+	sock.addEventListener('message', function (event) {
+		console.log('Using addEventListener ', event.data); // $ remoteFlow
+	}); // $ clientReceive
+})();
+
+
+(function () {
+    myWebSocketInstance.addEventListener('open', function (event) {
+        myWebSocketInstance.send('Hi from browser!'); // $ clientSend
+    });
+
+    myWebSocketInstance.addEventListener('message', function (event) {
+        console.log('Message from server ', event.data); // $ remoteFlow
+    }); // $ clientReceive
+
+    myWebSocketInstance.onmessage = function (event) {
+        console.log("Message from server 2", event.data); // $ remoteFlow
+    }; // $ clientReceive
+})();
+
+
+(function () {
+    mySockJSInstance.onopen = function () {
+        mySockJSInstance.send('test'); // $ clientSend
+    };
+
+    mySockJSInstance.onmessage = function (e) {
+        console.log('message', e.data); // $ remoteFlow
+        mySockJSInstance.close();
+    }; // $ clientReceive
+
+    mySockJSInstance.addEventListener('message', function (event) {
+        console.log('Using addEventListener ', event.data); // $ remoteFlow
+    }); // $ clientReceive
+})();
+
+
+const recv_message = function (e) {
+    console.log('Received message:', e.data); // $ remoteFlow
+}; // $ clientReceive
+
+(function () {
+    myWebSocketInstance.onmessage = recv_message.bind(this);
+})();

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

@@ -1,32 +1,37 @@
 (function () {
-	const socket = new WebSocket('ws://localhost:8080');
+	const socket = new WebSocket('ws://localhost:8080'); // $clientSocket
 
 	socket.addEventListener('open', function (event) {
-		socket.send('Hi from browser!');
+		socket.send('Hi from browser!'); // $clientSend
 	});
 
 	socket.addEventListener('message', function (event) {
-		console.log('Message from server ', event.data);
-	});
+		console.log('Message from server ', event.data); // $ remoteFlow
+	}); // $clientReceive
 
 	socket.onmessage = function (event) {
-		console.log("Message from server 2", event.data)
-	};
+		console.log("Message from server 2", event.data); // $ remoteFlow
+	}; // $clientReceive
 })();
 
 
 (function () {
-	var sock = new SockJS('http://0.0.0.0:9999/echo');
+	var sock = new SockJS('http://0.0.0.0:9999/echo'); // $clientSocket
 	sock.onopen = function () {
-		sock.send('test');
+		sock.send('test'); // $clientSend
 	};
 
 	sock.onmessage = function (e) {
-		console.log('message', e.data);
+		console.log('message', e.data); // $ remoteFlow
 		sock.close();
-	};
+	}; // $clientReceive
 
 	sock.addEventListener('message', function (event) {
-		console.log('Using addEventListener ', event.data);
-	});
-})
+		console.log('Using addEventListener ', event.data); // $ remoteFlow
+	}); // $clientReceive
+})();
+
+export const MyWebSocket = WebSocket;
+export const MySockJS = SockJS;
+export const myWebSocketInstance = new WebSocket('ws://localhost:8080'); // $ clientSocket
+export const mySockJSInstance = new SockJS('http://0.0.0.0:9999/echo'); // $ clientSocket

javascript/ql/test/library-tests/frameworks/WebSocket/client-custom.js

@@ -0,0 +1,23 @@
+const { MyWebSocketWS, myWebSocketWSInstance } = require('./client.js');
+
+(function () {
+	const ws = new MyWebSocketWS('ws://example.org'); // $ clientSocket
+
+	ws.on('open', function open() {
+		ws.send('Hi from client!'); // $ clientSend
+	});
+
+	ws.on('message', function incoming(data) { // $ remoteFlow
+		console.log(data);
+	}); // $ clientReceive
+})();
+
+(function () {
+	myWebSocketWSInstance.on('open', function open() {
+		myWebSocketWSInstance.send('Hi from client!'); // $ clientSend
+	});
+
+	myWebSocketWSInstance.on('message', function incoming(data) { // $ remoteFlow
+		console.log(data);
+	}); // $ clientReceive
+})();

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

@@ -1,13 +1,16 @@
-(function () {
-	const WebSocket = require('ws');
+const WebSocket = require('ws');
 
-	const ws = new WebSocket('ws://example.org');
+(function () {
+	const ws = new WebSocket('ws://example.org'); // $clientSocket
 
 	ws.on('open', function open() {
-		ws.send('Hi from client!');
+		ws.send('Hi from client!'); // $clientSend
 	});
 
-	ws.on('message', function incoming(data) {
+	ws.on('message', function incoming(data) { // $ remoteFlow
 		console.log(data);
-	});
-})();
+	}); // $clientReceive
+})();
+
+module.exports.MyWebSocketWS = require('ws');
+module.exports.myWebSocketWSInstance = new WebSocket('ws://example.org'); // $ clientSocket

javascript/ql/test/library-tests/frameworks/WebSocket/server-custom.js

@@ -0,0 +1,23 @@
+const { MyWebSocketServer, myWebSocketServerInstance } = require('./server.js');
+
+(function () {
+	const wss = new MyWebSocketServer({ port: 8080 });
+
+	wss.on('connection', function connection(ws) { // $ serverSocket
+		ws.on('message', function incoming(message) { // $ remoteFlow
+			console.log('received: %s', message);
+		}); // $ serverReceive
+
+		ws.send('Hi from server!'); // $ serverSend
+	});
+})();
+
+(function () {
+	myWebSocketServerInstance.on('connection', function connection(ws) { // $ serverSocket
+		ws.on('message', function incoming(message) { // $ remoteFlow
+			console.log('received: %s', message);
+		}); // $ serverReceive
+
+		ws.send('Hi from server!'); // $ serverSend
+	});
+})();

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

@@ -1,13 +1,16 @@
-(function () {
-	const WebSocket = require('ws');
+const WebSocket = require('ws');
 
+(function () {
 	const wss = new WebSocket.Server({ port: 8080 });
 
-	wss.on('connection', function connection(ws) {
-		ws.on('message', function incoming(message) {
+	wss.on('connection', function connection(ws) { // $serverSocket
+		ws.on('message', function incoming(message) { // $remoteFlow
 			console.log('received: %s', message);
-		});
+		}); // $serverReceive
 
-		ws.send('Hi from server!');
+		ws.send('Hi from server!'); // $serverSend
 	});
-})();
+})();
+
+module.exports.MyWebSocketServer = require('ws').Server;
+module.exports.myWebSocketServerInstance = new WebSocket.Server({ port: 8080 });