Python: Allow custom fastapi.APIRouter subclasses

Author: RasmusWL

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.

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()
     }
   }
 

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)