Model tornado websockets

Author: joefarebrother

python/ql/lib/semmle/python/frameworks/Tornado.qll

@@ -135,6 +135,8 @@ module Tornado {
         API::Node subclassRef() {
           result = web().getMember("RequestHandler").getASubclass*()
           or
+          result = WebSocket::WebSocketHandler::subclassRef()
+          or
           result = ModelOutput::getATypeNode("tornado.web.RequestHandler~Subclass").getASubclass*()
         }
 
@@ -428,6 +430,42 @@ module Tornado {
         }
       }
     }
+
+    // ---------------------------------------------------------------------------
+    // tornado.websocket
+    // ---------------------------------------------------------------------------
+    /** Gets a reference to the `tornado.websocket` module. */
+    API::Node websocket() { result = Tornado::tornado().getMember("websocket") }
+
+    module WebSocket {
+      module WebSocketHandler {
+        /** Gets a reference to the `tornado.websocket.WebSocketHandler` class or any subclass. */
+        API::Node subclassRef() {
+          result = websocket().getMember("WebSocketHandler").getASubclass*()
+          or
+          result =
+            ModelOutput::getATypeNode("tornado.websocket.WebSocketHandler~Subclass").getASubclass*()
+        }
+
+        class WebSocketHandlerClass extends Web::RequestHandler::RequestHandlerClass {
+          WebSocketHandlerClass() { this.getParent() = subclassRef().asSource().asExpr() }
+
+          override Function getARequestHandler() {
+            result = super.getARequestHandler()
+            or
+            result = this.getAMethod() and
+            result.getName() = "open"
+          }
+
+          /** Gets a function that could handle incoming websocket events, if any. */
+          Function getAWebSocketEventHandler() {
+            result = this.getAMethod() and
+            result.getName() =
+              ["on_message", "on_close", "on_ping", "on_pong", "select_subprotocol", "check_origin"]
+          }
+        }
+      }
+    }
   }
 
   // ---------------------------------------------------------------------------
@@ -542,6 +580,27 @@ module Tornado {
     override string getFramework() { result = "Tornado" }
   }
 
+  /** A request handler for WebSocket events */
+  private class TornadoWebSocketEventHandler extends Http::Server::RequestHandler::Range {
+    TornadoWebSocketEventHandler() {
+      exists(TornadoModule::WebSocket::WebSocketHandler::WebSocketHandlerClass cls |
+        cls.getAWebSocketEventHandler() = this
+      )
+    }
+
+    override Parameter getARoutedParameter() {
+      // The `open` method is handled as a normal request handler in `TornadoRouteSetup` or `TornadoRequestHandlerWithoutKnownRoute`.
+      // For other event handlers (such as `on_message`), all parameters should be remote flow sources, as they are not affected by routing.
+      result in [
+          this.getArg(_), this.getArgByName(_), this.getVararg().(Parameter),
+          this.getKwarg().(Parameter)
+        ] and
+      not result = this.getArg(0)
+    }
+
+    override string getFramework() { result = "Tornado" }
+  }
+
   // ---------------------------------------------------------------------------
   // Response modeling
   // ---------------------------------------------------------------------------

python/ql/test/library-tests/frameworks/tornado/routing_test.py

@@ -54,6 +54,27 @@ class PossiblyNotRouted(tornado.web.RequestHandler):
     def get(self):  # $ requestHandler
         self.write("NotRouted") # $ HttpResponse
 
+class WebSocket(tornado.websocket.WebSocketHandler):
+    def open(self, x): # $ requestHandler routedParameter=x
+        self.write_message("WebSocket open {}".format(x)) 
+
+    def on_message(self, data): # $ requestHandler routedParameter=data
+        self.write_message("WebSocket on_message {}".format(data))
+
+    def on_ping(self, data): # $ requestHandler routedParameter=data
+        print("ping", data)
+
+    def on_pong(self, data): # $ requestHandler routedParameter=data
+        print("pong", data)
+
+    def select_subprotocol(self, subs): # $ requestHandler routedParameter=subs
+        print("select_subprotocol", subs)
+
+    def check_origin(self, origin): # $ requestHandler routedParameter=origin
+        print("check_origin", origin)
+        return True
+
+    
 
 def make_app():
     # see https://www.tornadoweb.org/en/stable/routing.html for even more examples
@@ -74,6 +95,7 @@ def make_app():
             (tornado.routing.HostMatches(r"(localhost|127\.0\.0\.1)"), [
                 ("/only-localhost", OnlyLocalhost)  # $ routeSetup="/only-localhost"
             ]),
+            (r"/websocket/([0-9]+)", WebSocket),  # $ routeSetup="/websocket/([0-9]+)"
 
         ],
         debug=True,