Merge pull request #7228 from RasmusWL/fastapi-improvements

yoff · web-flow · commit f10f053c36b7 · 2021-12-02T12:58:53.000+01:00
Python: FastAPI improvements
diff --git a/python/change-notes/2021-11-24-FastAPI-Custom-APIRouter-Subclass.md b/python/change-notes/2021-11-24-FastAPI-Custom-APIRouter-Subclass.md
@@ -0,0 +1,2 @@
+lgtm,codescanning
+* Extended the modeling of FastAPI such that custom subclasses of `fastapi.APIRouter` are recognized.
diff --git a/python/change-notes/2021-11-24-FastAPI-FileResponse-FileSystemAccess copy.md b/python/change-notes/2021-11-24-FastAPI-FileResponse-FileSystemAccess copy.md
@@ -0,0 +1,2 @@
+lgtm,codescanning
+* Extended the modeling of FastAPI such that `fastapi.responses.FileResponse` are considered `FileSystemAccess`, making them sinks for the _Uncontrolled data used in path expression_ (`py/path-injection`) query.
diff --git a/python/ql/lib/semmle/python/frameworks/FastApi.qll b/python/ql/lib/semmle/python/frameworks/FastApi.qll
@@ -33,7 +33,7 @@ private module FastApi {
   module APIRouter {
     /** Gets a reference to an instance of `fastapi.APIRouter`. */
     API::Node instance() {
-      result = API::moduleImport("fastapi").getMember("APIRouter").getReturn()
+      result = API::moduleImport("fastapi").getMember("APIRouter").getASubclass*().getReturn()
     }
   }
 
@@ -226,6 +226,17 @@ private module FastApi {
       }
     }
 
+    /**
+     * A direct instantiation of a FileResponse.
+     */
+    private class FileResponseInstantiation extends ResponseInstantiation, FileSystemAccess::Range {
+      FileResponseInstantiation() { baseApiNode = getModeledResponseClass("FileResponse") }
+
+      override DataFlow::Node getAPathArgument() {
+        result in [this.getArg(0), this.getArgByName("path")]
+      }
+    }
+
     /**
      * An implicit response from a return of FastAPI request handler.
      */
@@ -256,7 +267,8 @@ private module FastApi {
      * An implicit response from a return of FastAPI request handler, that has
      * `response_class` set to a `FileResponse`.
      */
-    private class FastApiRequestHandlerFileResponseReturn extends FastApiRequestHandlerReturn {
+    private class FastApiRequestHandlerFileResponseReturn extends FastApiRequestHandlerReturn,
+      FileSystemAccess::Range {
       FastApiRequestHandlerFileResponseReturn() {
         exists(API::Node responseClass |
           responseClass.getAUse() = routeSetup.getResponseClassArg() and
@@ -265,6 +277,8 @@ private module FastApi {
       }
 
       override DataFlow::Node getBody() { none() }
+
+      override DataFlow::Node getAPathArgument() { result = this }
     }
 
     /**
diff --git a/python/ql/test/library-tests/frameworks/fastapi/response_test.py b/python/ql/test/library-tests/frameworks/fastapi/response_test.py
@@ -136,10 +136,10 @@ async def file_response(): # $ requestHandler
 
     # We don't really have any good QL modeling of passing a file-path, whose content
     # will be returned as part of the response... so will leave this as a TODO for now.
-    resp = fastapi.responses.FileResponse(__file__) # $ HttpResponse
+    resp = fastapi.responses.FileResponse(__file__) # $ HttpResponse getAPathArgument=__file__
     return resp # $ SPURIOUS: HttpResponse mimetype=application/json responseBody=resp
 
 
 @app.get("/file_response2", response_class=fastapi.responses.FileResponse) # $ routeSetup="/file_response2"
 async def file_response2(): # $ requestHandler
-    return __file__ # $ HttpResponse
+    return __file__ # $ HttpResponse getAPathArgument=__file__
diff --git a/python/ql/test/library-tests/frameworks/fastapi/router.py b/python/ql/test/library-tests/frameworks/fastapi/router.py
@@ -1,5 +1,6 @@
 # like blueprints in Flask
 # see https://fastapi.tiangolo.com/tutorial/bigger-applications/
+# see basic.py for instructions for how to run this code.
 
 from fastapi import APIRouter, FastAPI
 
@@ -30,4 +31,23 @@ async def items(): # $ requestHandler
 app.include_router(outer_router, prefix="/outer")
 app.include_router(items_router)
 
-# see basic.py for instructions for how to run this code.
+# Using a custom router
+
+class MyCustomRouter(APIRouter):
+    """
+    Which automatically removes trailing slashes
+    """
+    def api_route(self, path: str, **kwargs):
+        path = path.rstrip("/")
+        return super().api_route(path, **kwargs)
+
+
+custom_router = MyCustomRouter()
+
+
+@custom_router.get("/bar/") # $ routeSetup="/bar/"
+async def items(): # $ requestHandler
+    return {"msg": "custom_router /bar/"} # $ HttpResponse
+
+
+app.include_router(custom_router)

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+lgtm,codescanning`
	`2`	+* Extended the modeling of FastAPI such that custom subclasses of `fastapi.APIRouter` are recognized.