Model tornado websockets

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
   // ---------------------------------------------------------------------------