Python: Model response_class attribute of Flask class

RasmusWL · RasmusWL · commit 44ba3469db28 · 2020-10-23T14:31:34.000+02:00
diff --git a/python/ql/src/experimental/semmle/python/frameworks/Flask.qll b/python/ql/src/experimental/semmle/python/frameworks/Flask.qll
@@ -171,6 +171,21 @@ private module FlaskModel {
       // completely disallowed in QL. I added an underscore to move thing forwards for
       // now :(
       DataFlow::Node make_response_() { result = instance_attr("make_response") }
+
+      /** Gets a reference to the `response_class` attribute on the `flask.Flask` class or an instance. */
+      private DataFlow::Node response_class(DataFlow::TypeTracker t) {
+        t.startInAttr("response_class") and
+        result in [classRef(), instance()]
+        or
+        exists(DataFlow::TypeTracker t2 | result = response_class(t2).track(t2, t))
+      }
+
+      /**
+       * Gets a reference to the `response_class` attribute on the `flask.Flask` class or an instance.
+       *
+       * See https://flask.palletsprojects.com/en/1.1.x/api/#flask.Flask.response_class
+       */
+      DataFlow::Node response_class() { result = response_class(DataFlow::TypeTracker::end()) }
     }
   }
 
@@ -183,7 +198,7 @@ private module FlaskModel {
     /** Gets a reference to the `flask.Response` class. */
     private DataFlow::Node classRef(DataFlow::TypeTracker t) {
       t.start() and
-      result = flask_attr("Response")
+      result in [flask_attr("Response"), flask::Flask::response_class()]
       or
       exists(DataFlow::TypeTracker t2 | result = classRef(t2).track(t2, t))
     }
diff --git a/python/ql/test/experimental/library-tests/frameworks/flask/response_test.py b/python/ql/test/experimental/library-tests/frameworks/flask/response_test.py
@@ -39,15 +39,15 @@ def html4():  # $routeHandler
 def html5():  # $routeHandler
     # note: flask.Flask.response_class is set to `flask.Response` by default.
     # it can be overridden, but we don't try to handle that right now.
-    resp = Flask.response_class("<h1>hello</h1>")  # $f-:HttpResponse $f-:mimetype=text/html $f-:responseBody="<h1>hello</h1>"
+    resp = Flask.response_class("<h1>hello</h1>")  # $HttpResponse $mimetype=text/html $responseBody="<h1>hello</h1>"
     return resp
 
 
 @app.route("/html6")  # $routeSetup="/html6"
 def html6():  # $routeHandler
     # note: app.response_class (flask.Flask.response_class) is set to `flask.Response` by default.
     # it can be overridden, but we don't try to handle that right now.
-    resp = app.response_class("<h1>hello</h1>")  # $f-:HttpResponse $f-:mimetype=text/html $f-:responseBody="<h1>hello</h1>"
+    resp = app.response_class("<h1>hello</h1>")  # $HttpResponse $mimetype=text/html $responseBody="<h1>hello</h1>"
     return resp
 
 
@@ -127,15 +127,15 @@ def Response6():  # $routeHandler
 def Flask_response_class():  # $routeHandler
     # note: flask.Flask.response_class is set to `flask.Response` by default.
     # it can be overridden, but we don't try to handle that right now.
-    resp = Flask.response_class("<h1>hello</h1>", mimetype="text/plain")  # $f-:HttpResponse $f-:mimetype=text/plain $f-:responseBody="<h1>hello</h1>"
+    resp = Flask.response_class("<h1>hello</h1>", mimetype="text/plain")  # $HttpResponse $mimetype=text/plain $responseBody="<h1>hello</h1>"
     return resp
 
 
 @app.route("/content-type/app-response-class")  # $routeSetup="/content-type/app-response-class"
 def app_response_class():  # $routeHandler
     # note: app.response_class (flask.Flask.response_class) is set to `flask.Response` by default.
     # it can be overridden, but we don't try to handle that right now.
-    resp = app.response_class("<h1>hello</h1>", mimetype="text/plain")  # $f-:HttpResponse $f-:mimetype=text/plain $f-:responseBody="<h1>hello</h1>"
+    resp = app.response_class("<h1>hello</h1>", mimetype="text/plain")  # $HttpResponse $mimetype=text/plain $responseBody="<h1>hello</h1>"
     return resp
 
 
