Implement websockets models

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

@@ -0,0 +1,83 @@
+/**
+ * Provides definitions and modeling for the `websockets` PyPI package.
+ *
+ * See https://websockets.readthedocs.io/en/stable/
+ */
+
+private import python
+private import semmle.python.dataflow.new.RemoteFlowSources
+private import semmle.python.Concepts
+private import semmle.python.ApiGraphs
+private import semmle.python.frameworks.internal.PoorMansFunctionResolution
+private import semmle.python.frameworks.internal.InstanceTaintStepsHelper
+
+/**
+ * Provides models for the `websockets` PyPI package.
+ * See https://websockets.readthedocs.io/en/stable/
+ */
+module Websockets {
+  private class HandlerArg extends DataFlow::Node {
+    HandlerArg() {
+      exists(DataFlow::CallCfgNode c |
+        c =
+          API::moduleImport("websockets")
+              .getMember(["asyncio", "sync"])
+              .getMember("server")
+              .getMember(["serve", "unix_serve"])
+              .getACall()
+      |
+        (this = c.getArg(0) or this = c.getArgByName("handler"))
+      )
+    }
+  }
+
+  /** A websocket handler that is passed to `serve`. */
+  // TODO: handlers defined via route maps, e.g. through `websockets.asyncio.router.route`, are more complex to handle.
+  class WebSocketHandler extends Http::Server::RequestHandler::Range {
+    WebSocketHandler() { poorMansFunctionTracker(this) = any(HandlerArg a) }
+
+    override Parameter getARoutedParameter() { result = this.getAnArg() }
+
+    override string getFramework() { result = "websockets" }
+  }
+
+  module ServerConnection {
+    /**
+     * A source of instances of `websockets.asyncio.ServerConnection` and `websockets.threading.ServerConnection`, extend this class to model new instances.
+     *
+     * This can include instantiations of the class, return values from function
+     * calls, or a special parameter that will be set when functions are called by an external
+     * library.
+     *
+     * Use the predicate `WebSocket::instance()` to get references to instances of `websockets.asyncio.ServerConnection` and `websockets.threading.ServerConnection`.
+     */
+    abstract class InstanceSource extends DataFlow::LocalSourceNode { }
+
+    /** Gets a reference to an instance of `websockets.asyncio.ServerConnection` or `websockets.threading.ServerConnection`. */
+    private DataFlow::TypeTrackingNode instance(DataFlow::TypeTracker t) {
+      t.start() and
+      result instanceof InstanceSource
+      or
+      exists(DataFlow::TypeTracker t2 | result = instance(t2).track(t2, t))
+    }
+
+    /** Gets a reference to an instance of `websockets.asyncio.ServerConnection` or `websockets.threading.ServerConnection`. */
+    DataFlow::Node instance() { instance(DataFlow::TypeTracker::end()).flowsTo(result) }
+
+    private class HandlerParam extends DataFlow::Node, InstanceSource {
+      HandlerParam() { exists(WebSocketHandler h | this = DataFlow::parameterNode(h.getArg(0))) }
+    }
+
+    private class InstanceTaintSteps extends InstanceTaintStepsHelper {
+      InstanceTaintSteps() { this = "websockets.asyncio.ServerConnection" }
+
+      override DataFlow::Node getInstance() { result = instance() }
+
+      override string getAttributeName() { none() }
+
+      override string getAsyncMethodName() { result = ["recv", "recv_streaming"] }
+
+      override string getMethodName() { result = ["recv", "recv_streaming"] }
+    }
+  }
+}