Python: Add starlette.websocket model

tausbn · RasmusWL · commit 83e6e51e95f4 · 2023-12-08T11:27:50.000+01:00
diff --git a/python/ql/lib/semmle/python/frameworks/Starlette.qll b/python/ql/lib/semmle/python/frameworks/Starlette.qll
@@ -13,6 +13,7 @@ private import semmle.python.Concepts
 private import semmle.python.ApiGraphs
 private import semmle.python.frameworks.internal.InstanceTaintStepsHelper
 private import semmle.python.frameworks.Stdlib
+private import semmle.python.frameworks.data.ModelsAsData
 
 /**
  * INTERNAL: Do not use.
@@ -35,6 +36,8 @@ module Starlette {
       result = API::moduleImport("starlette").getMember("websockets").getMember("WebSocket")
       or
       result = API::moduleImport("fastapi").getMember("WebSocket")
+      or
+      result = ModelOutput::getATypeNode("starlette.websockets.WebSocket~Subclass").getASubclass*()
     }
 
     /**
diff --git a/python/ql/src/meta/ClassHierarchy/Find.ql b/python/ql/src/meta/ClassHierarchy/Find.ql
@@ -14,6 +14,7 @@ private import semmle.python.frameworks.Django
 private import semmle.python.frameworks.Tornado
 private import semmle.python.frameworks.Stdlib
 private import semmle.python.frameworks.Requests
+private import semmle.python.frameworks.Starlette
 import semmle.python.frameworks.data.internal.ApiGraphModelsExtensions as Extensions
 
 class FlaskViewClasses extends FindSubclassesSpec {
@@ -227,6 +228,12 @@ class HttpClientHttpResponse extends FindSubclassesSpec {
   override API::Node getAlreadyModeledClass() { result = StdlibPrivate::HttpResponse::classRef() }
 }
 
+class StarletteWebsocket extends FindSubclassesSpec {
+  StarletteWebsocket() { this = "starlette.websockets.WebSocket~Subclass" }
+
+  override API::Node getAlreadyModeledClass() { result = Starlette::WebSocket::classRef() }
+}
+
 bindingset[fullyQualified]
 predicate fullyQualifiedToYamlFormat(string fullyQualified, string type2, string path) {
   exists(int firstDot | firstDot = fullyQualified.indexOf(".", 0, 0) |
