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

/**
 * 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" }
  }

  /** Provides taint models for instances of `ServerConnection` objects passed to websocket handlers. */
  module ServerConnection {
    /**
     * A source of instances of `websockets.asyncio.ServerConnection` and `websockets.sync.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 `ServerConnection::instance()` to get references to instances of `websockets.asyncio.ServerConnection` and `websockets.sync.ServerConnection`.
     */
    abstract class InstanceSource extends DataFlow::LocalSourceNode { }

    /** Gets a reference to an instance of `websockets.asyncio.ServerConnection` or `websockets.sync.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.sync.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"] }
    }
  }
}