Python: FastAPI: Proper modeling of implicit returns

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

@@ -73,6 +73,9 @@ private module FastApi {
     override Function getARequestHandler() { result.getADecorator().getAFlowNode() = node }
 
     override string getFramework() { result = "FastAPI" }
+
+    /** Gets the argument specifying the response class to use, if any. */
+    DataFlow::Node getResponseClassArg() { result = this.getArgByName("response_class") }
   }
 
   // ---------------------------------------------------------------------------
@@ -199,6 +202,65 @@ private module FastApi {
       }
     }
 
+    /**
+     * An implicit response from a return of FastAPI request handler.
+     */
+    private class FastApiRequestHandlerReturn extends HTTP::Server::HttpResponse::Range,
+      DataFlow::CfgNode {
+      FastApiRouteSetup routeSetup;
+
+      FastApiRequestHandlerReturn() {
+        node = routeSetup.getARequestHandler().getAReturnValueFlowNode()
+      }
+
+      override DataFlow::Node getBody() { result = this }
+
+      override DataFlow::Node getMimetypeOrContentTypeArg() { none() }
+
+      override string getMimetypeDefault() {
+        exists(API::Node responseClass |
+          responseClass.getAUse() = routeSetup.getResponseClassArg() and
+          result = getDefaultMimeType(responseClass)
+        )
+        or
+        not exists(routeSetup.getResponseClassArg()) and
+        result = "application/json"
+      }
+    }
+
+    /**
+     * An implicit response from a return of FastAPI request handler, that has
+     * `response_class` set to a `FileResponse`.
+     */
+    private class FastApiRequestHandlerFileResponseReturn extends FastApiRequestHandlerReturn {
+      FastApiRequestHandlerFileResponseReturn() {
+        exists(API::Node responseClass |
+          responseClass.getAUse() = routeSetup.getResponseClassArg() and
+          responseClass = getModeledResponseClass("FileResponse").getASubclass*()
+        )
+      }
+
+      override DataFlow::Node getBody() { none() }
+    }
+
+    /**
+     * An implicit response from a return of FastAPI request handler, that has
+     * `response_class` set to a `RedirectResponse`.
+     */
+    private class FastApiRequestHandlerRedirectReturn extends FastApiRequestHandlerReturn,
+      HTTP::Server::HttpRedirectResponse::Range {
+      FastApiRequestHandlerRedirectReturn() {
+        exists(API::Node responseClass |
+          responseClass.getAUse() = routeSetup.getResponseClassArg() and
+          responseClass = getModeledResponseClass("RedirectResponse").getASubclass*()
+        )
+      }
+
+      override DataFlow::Node getBody() { none() }
+
+      override DataFlow::Node getRedirectLocation() { result = this }
+    }
+
     /**
      * INTERNAL: Do not use.
      *
@@ -263,23 +325,4 @@ private module FastApi {
       override DataFlow::Node getValueArg() { none() }
     }
   }
-
-  /**
-   * Implicit response from returns of FastAPI request handlers
-   */
-  private class FastApiRequestHandlerReturn extends HTTP::Server::HttpResponse::Range,
-    DataFlow::CfgNode {
-    FastApiRequestHandlerReturn() {
-      exists(Function requestHandler |
-        requestHandler = any(FastApiRouteSetup rs).getARequestHandler() and
-        node = requestHandler.getAReturnValueFlowNode()
-      )
-    }
-
-    override DataFlow::Node getBody() { result = this }
-
-    override DataFlow::Node getMimetypeOrContentTypeArg() { none() }
-
-    override string getMimetypeDefault() { result = "application/json" }
-  }
 }

python/ql/test/library-tests/frameworks/fastapi/response_test.py

@@ -47,7 +47,7 @@ async def response_parameter_custom_type(response: MyXmlResponse): # $ requestHa
     response.set_cookie("key", "value") # $ CookieWrite CookieName="key" CookieValue="value"
     response.headers["Custom-Response-Type"] = "yes, but only after function has run"
     xml_data = "<foo>FOO</foo>"
-    return xml_data # $ HttpResponse responseBody=xml_data SPURIOUS: mimetype=application/json MISSING: mimetype=application/xml
+    return xml_data # $ HttpResponse responseBody=xml_data mimetype=application/xml
 
 
 # Direct response construction
@@ -68,7 +68,7 @@ async def direct_response(): # $ requestHandler
 @app.get("/direct_response2", response_class=fastapi.responses.Response) # $ routeSetup="/direct_response2"
 async def direct_response2(): # $ requestHandler
     xml_data = "<foo>FOO</foo>"
-    return xml_data # $ HttpResponse responseBody=xml_data SPURIOUS: mimetype=application/json
+    return xml_data # $ HttpResponse responseBody=xml_data
 
 
 @app.get("/my_xml_response") # $ routeSetup="/my_xml_response"
@@ -81,7 +81,7 @@ async def my_xml_response(): # $ requestHandler
 @app.get("/my_xml_response2", response_class=MyXmlResponse) # $ routeSetup="/my_xml_response2"
 async def my_xml_response2(): # $ requestHandler
     xml_data = "<foo>FOO</foo>"
-    return xml_data # $ HttpResponse responseBody=xml_data SPURIOUS: mimetype=application/json MISSING: mimetype=application/xml
+    return xml_data # $ HttpResponse responseBody=xml_data mimetype=application/xml
 
 
 @app.get("/html_response") # $ routeSetup="/html_response"
@@ -94,7 +94,7 @@ async def html_response(): # $ requestHandler
 @app.get("/html_response2", response_class=fastapi.responses.HTMLResponse) # $ routeSetup="/html_response2"
 async def html_response2(): # $ requestHandler
     hello_world = "<h1>Hello World!</h1>"
-    return hello_world # $ HttpResponse responseBody=hello_world SPURIOUS: mimetype=application/json MISSING: mimetype=text/html
+    return hello_world # $ HttpResponse responseBody=hello_world mimetype=text/html
 
 
 @app.get("/redirect") # $ routeSetup="/redirect"
@@ -107,7 +107,7 @@ async def redirect(): # $ requestHandler
 @app.get("/redirect2", response_class=fastapi.responses.RedirectResponse) # $ routeSetup="/redirect2"
 async def redirect2(): # $ requestHandler
     next = "https://www.example.com"
-    return next # $ HttpResponse SPURIOUS: mimetype=application/json responseBody=next MISSING: HttpRedirectResponse redirectLocation=next
+    return next # $ HttpResponse HttpRedirectResponse redirectLocation=next
 
 
 @app.get("/streaming_response") # $ routeSetup="/streaming_response"
@@ -142,4 +142,4 @@ async def file_response(): # $ requestHandler
 
 @app.get("/file_response2", response_class=fastapi.responses.FileResponse) # $ routeSetup="/file_response2"
 async def file_response2(): # $ requestHandler
-    return __file__ # $ HttpResponse SPURIOUS: mimetype=application/json responseBody=__file__
+    return __file__ # $ HttpResponse